Trained a 32B FLUX.2 LoRA on a 24GB AMD 7900 XTX, native ROCm on Windows — full guide + patches

TL;DR: Everyone says QLoRA past ~13B is dead on a 24GB card. I got the full 32B FLUX.2 dev transformer QLoRA-training resident on the GPU on a 7900 XTX under native ROCm on Windows (no ZLUDA, no CUDA shim) — on a box with only 32GB of system RAM (the 64GB base + 48GB text encoder load through a big pagefile) — ~9-10 s/it sustained at 448/bs2 once the (wild) stall issue below is handled. It took a dozen distinct failures to get there. Config + fixes below. Setup GPU: RX 7900 XTX 24GB (gfx1100). No FP8/FP4 hardware, so uint4 weight-only quant (optimum-quanto). Host: 32GB RAM + ~100GB pagefile (you need the virtual headroom for the one-time bf16 loads). torch 2.12.0+rocm7.15, trainer: ai-toolkit AMD ROCm fork (cupertinomiranda/ai-toolkit-amd-rocm-support). Base: official FLUX.2-dev — the repo-root single-file flux2-dev.safetensors (64GB bf16), NOT the diffusers transformer/ subdir (different keys). Text encoder: Mistral-Small-3.1-24B (yes, FLUX.2 uses a 24B LLM as its TE). The walls, in order (each one blocks the next) Load & quantize the 64GB base without dying: 1. safetensors mmap load_file crashes natively on the 64GB file (no traceback, process just dies; fine at 33GB). → Manual non-mmap loader: read the header, then per-tensor seek/read/frombuffer. 2. Transformer OOMs at ~38GB before quantizing — the trainer moves the full bf16 to GPU before packing. → Quantize on CPU; only the ~20GB uint4 result touches the card. 3. 0xC0000005 while loading the text encoder — the 64GB bf16 is still referenced when Mistral's 48GB loads on top. → del transformer_state_dict; gc.collect() right after load_state_dict. 4. Mistral OOMs the GPU (c10 abort) — same as #2 for the TE. → Quantize Mistral on CPU first, then .to(device). Make it train on the GPU, not the CPU: 5. Block-swap (layer_offloading) deadlocks the HIP driver (hangs at sampling AND first step, needs two kill passes). → layer_offloading: false, keep the base resident. 6. In-training sampling deadlocks + uint4 previews are black frames. → disable_sampling: true, evaluate in ComfyUI instead. 7. uint4→GPU move fragments/OOMs. → Launch with PYTORCH_HIP_ALLOC_CONF=expandable_segments:True. 8. Re-quantizing every launch costs ~8 min. → Save the quanto state-dict once as a .pt; training torch.loads it in seconds. 9. Base won't stay resident (looks like CPU training) — low_vram parks it on CPU during TE-caching and never brings it back. → After the TE caches + unloads, move the base back to GPU; gate the load-path's transformer→GPU line on low_vram so it doesn't collide with the resident TE. The two that cost me a whole night: 10. "It's training on CPU" — except it wasn't. A separate process reading VRAM via torch.cuda.mem_get_info() lies on ROCm/Windows — reported 0.2GB while the process actually held 20GB. Combined with "1 busy CPU core" (which is normal for GPU training) it looked exactly like CPU. I killed several working runs over this. → Trust an in-process VRAM print, the Windows GPU Engine(*engtype_compute)Utilization counter, and the drop in system RAM when the base moves off CPU. Never trust a cross-process VRAM read here. 11. Resident but crawling at 200 s/step. The 20GB base leaves no headroom, so activations spill to host RAM over PCIe (expandable_segments lets it overflow instead of OOMing → thrash). → Cut resolution until the spill is small. Evaluate it: 12. In-training previews are useless, so render checkpoints in ComfyUI + ComfyUI-GGUF: Q3_K_M GGUF unet + Mistral Q5 GGUF (CLIPLoaderGGUF type=flux2) + flux2 VAE. The ai-toolkit LoRA keys (diffusion_model…lora_A/lora_B) load with zero conversion. Resolution is the speed knob (measured, batch 1, grad-checkpointing on) Max res Host spill Step time 1024 2.27 GB ~204 s 768 0.82 GB ~79 s 512 0.83 GB ~20–40 s 768 and 512 spill the same ~0.8GB — that part's fixed overhead, not activations (the allocator won't touch the last ~0.6GB of VRAM). The 768→512 gain is just less compute. Identity trains fine at 512. ⚠ Caveat discovered later: these step times were measured on STALLED runs (see Part 2) — the real, saturated cost is ~6-10× lower. The spill relationship holds; the absolute times were the stall talking. I now train at 448/bs2. Config that works yaml model: arch: "flux2" quantize: true qtype: "uint4" # quanto; also drives the TE quant in this fork quantize_te: true low_vram: true # park during TE-cache, move back resident to train layer_offloading: false # block-swap DEADLOCKS on ROCm model_kwargs: use_uint4_cache: true # load the pre-quantized .pt in seconds datasets: - resolution: [ 512 ] cache_latents_to_disk: true cache_text_embeddings: true train: gradient_checkpointing: true disable_sampling: true Launch: set PYTORCH_HIP_ALLOC_CONF=expandable_segments:True (+ the CUDA alias) and run python -u run.py config.yaml directly — a detached Start-Process -RedirectStandardOutput silently eats early output if the child dies during import. vcvars64 is NOT needed. Part 2 — the week after (this is the part you actually want) 13. The step rate is a LIE, and POWER is the diagnostic. My runs swung 6.9→45→67→118 s/it with clock, temp, and VRAM-spill all flat. Turns out this card has a failure state where a lone training context runs at ~1/8 speed: high clock (~3100 MHz), 100% "GPU load"… and only ~230 W draw with the memory controller at 2-7%. Spinning, not working. Saturated looks like lower clock (~2500) at ~385 W. Once you know the tell, one glance at wattage tells you which state you're in. (Root cause is somewhere in the driver/scheduler — invisible from Windows.) 14. The fix is absurd and reproducible: run a SECOND process doing heavy GEMMs for ~25 s. The stalled trainer flips to saturated — 10× on demand — and stays saturated after the rescuer exits. Two catches, both measured: it must be a fresh process (a long-lived idle context is itself degraded, ~7 TFLOP/s on a 77 TFLOP/s card, and lifts nothing), and fresh processes are born degraded ~half the time — check the burst's own TFLOP/s and just respawn until one runs fast. I ended up with a watchdog daemon that reads the power telemetry and fires bursts automatically; my last 2000-step run needed 8 unattended rescues and finished at ~9-10 s/it average. 15. Stalls cluster at predictable moments — process start (every launch/resume I measured) and right after checkpoint saves — so the daemon also fires a prophylactic burst ~60 s after those events. Most stalls now never establish at all. 16. batch_size 2 is ~1.45×/sample — but only when saturated. Stalled, it's a net LOSS (the stall tax scales with work per step). The two levers are coupled: fix the stall first, then batch 2 is free money. bs2 fits with ~150 MB to spare at 448; bs3 does not fit. Scale LR accordingly (I used sqrt: 1e-4 → 1.41e-4). 17. Lossless pause/resume for mid-run previews. ai-toolkit resumes cleanly (checkpoint + optimizer.pt), so I patched two flag files into the train loop: SAVE_NOW (checkpoint at the current step, keep going) and STOP_NOW (checkpoint + clean exit — zero steps lost). Pause, render the checkpoint in ComfyUI, relaunch, it resumes at the exact step. Mid-run previews every 500 steps cost ~10 min each. 18. Renders hit the same stall (a 20-step render swung 200 s ↔ 800 s). Same power tell, same burst fix — teach your watchdog to cover render contexts too. Results, final Three finished identity LoRAs so far (rank 16, 448 res, 2000 steps @ bs2 ≈ 6.5 h each on this one card), subject-verified likeness — the people they depict sign off on them, which is the only metric that matters. Face geometry converges late: checkpoints look "recognizable" by 1000 and keep visibly truing up until 2000; don't early-stop at "looks close." Full config + all the patches (copy-paste ready): https://github.com/drhawktopus/flux2-32b-qlora-rocm-windows Happy to answer questions — hope this saves someone the week it cost me. submitted by /u/elderon_echar [link] [comments]

Source: r/StableDiffusion | 2026-07-23

Related