Investigate gfx1201 (RDNA4) training: Flash Linear Attention + Gemma4-31B QLoRA OOM

[ulises-c]

Currently broken on stable release — investigate nightly releases or other workarounds for Flash Linear Attention (FLA) on gfx1201 (RDNA4 / AMD R9700).

Status: investigated (brainstorm). Findings summarized below; full detail in the comments. Nothing committed to the tree yet.

TL;DR — two kernel stacks, two blockers

The original FLA question turned out to be one of two distinct gfx1201 problems, split by which model we're training:

	Qwen targets (Qwen3.6-27B teacher · Qwen3.5-0.8B-Base classifier)	Gemma4-31B teacher
Attention arch	hybrid Mamba + attention (linear-attention layers)	softmax (sliding-window + global)
Needs FLA?	YES	NO (Makefile:342)
Real blocker	FLA Triton kernel deadlock	OOM at load (CPU RAM)

1. FLA blocker (the Qwen targets)

flash-linear-attention==0.5.0 JIT-deadlocks at training step 0 on gfx1201. Root cause: Triton 3.6.0's tritonamdgpu-pipeline pass has a use-after-free when num_stages>=2 (gfx12xx isn't in Triton's verified target list). Same bug that crashes SageAttention.

• Fix: make patch-fla-rocm (rewrites num_stages→1, num_warps→4) — confirmed working at seq 768.
• Nightly status: not fixed upstream in any nightly; every fix in the wild is a downstream num_stages=1 workaround.
• Alternatives: TRITON_DISABLE_AUTOTUNING=1 (cheapest, try first) · pin Triton 3.5.1 (avoids the 3.6.0 regression at the source) · TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1.

2. Gemma4-31B SFT blocker — NOT FLA, it's load-time CPU RAM

Gemma4 uses softmax attention (no FLA). The actual wall was OOM during model load: from_pretrained(load_in_4bit=True) stages the full ~62 GB BF16 model in CPU RAM before NF4 quantization (a bitsandbytes detail; low_cpu_mem_usage/offload don't prevent it). The 62 GB R9700 box thrashes into zram (~20 h just to load).

• Fix: load a pre-quantized bnb checkpoint — unsloth/gemma-4-31B-it-unsloth-bnb-4bit (already NF4, ~19 GB; embedded quant config byte-matches our pipeline; it's the instruction-tuned variant). Drops load RAM from ~62 GB to ~16 GB.

The one gate common to both: bitsandbytes 4-bit runtime on gfx1201

Officially listed in bnb's RDNA arch table; verify with make test-gpu-stack (steps 4/5). May need a source build: cmake -DCOMPUTE_BACKEND=hip -DBNB_ROCM_ARCH="gfx1201".

Hardware estimate — Gemma4-31B QLoRA (seq 1280, batch 1×16, grad-ckpt on, eval off)

Resource	Pre-quant checkpoint (recommended)	Naive load_in_4bit from BF16
GPU VRAM (peak, train)	~22–27 GB	~22–27 GB
System RAM (peak, load)	~8–16 GB	~62–70 GB ← the blocker
Download	~19 GB	~62 GB
Disk working set	~30 GB	~75 GB

Minimum viable: GPU ≥ 32 GB (R9700 / 5090) · 24 GB (3090) marginal · 20 GB (RTX 4000 Ada) won't fit. System RAM ~16 GB with the pre-quant checkpoint. Disk ~30 GB. Eval stays disabled (TRAIN_EVAL_STRATEGY=no) — its 256K-vocab logits.float() copy adds an ~8.6 GB spike that OOMs even 32 GB.

Corrections logged during the investigation

• HSA_OVERRIDE_GFX_VERSION: our hang came from 11.0.0 (the RDNA3 value). The correct gfx1201 value is 12.0.1 — but on ROCm 7.2 no override should be needed at all.
• expandable_segments:True: HANDOFF.md (05-25) says "unsupported on gfx1201"; the PR feat(stage2-sft): format fix, loaders, configs, DPO pair builder #79 session (05-26) re-added it as a working fix (releases freed BF16 shards from VRAM). The Makefile is correct; HANDOFF.md Issue 2 is stale.

Detailed notes (comments below)

1. Deep dive — FLA vs FA2 on gfx1201 (root cause, nightly status, reference map)
2. AMD Vision Mamba — hipify path validates the build mechanism, CDNA-only
3. Brainstorm — scope + untried FLA levers (from PR feat(stage2-sft): format fix, loaders, configs, DPO pair builder #79: autotuning, Triton 3.5.1, AOTriton experimental)
4. The real Gemma4-31B blocker — CPU-RAM staging at load
5. Pre-quantized bnb candidate + hardware estimate

Open / next

• Confirm bnb 4-bit runtime on gfx1201 (make test-gpu-stack; source build if needed) — gates both paths.
• Spike TRITON_DISABLE_AUTOTUNING=1, then Triton 3.5.1, as cleaner FLA fixes than the per-uv sync patch.
• Validate Gemma4-31B SFT end-to-end on the R9700 via the pre-quant checkpoint (CUDA-validated only so far).

Refs: PR #79 · apollo-mg RDNA4 gfx1201 gists · scripts/patch_fla_rocm.sh · docs/GPU_SUPPORT.md · HANDOFF.md.

[ulises-c]

Deep dive — Flash Linear Attention (FLA) on gfx1201 (RDNA4 / R9700)

Investigation for this issue. Bottom line: the FLA hang is a Triton 3.6.0 compiler bug that is still unfixed upstream as of late May 2026 — no nightly resolves it. Our num_stages=1 patch (make patch-fla-rocm) and a Triton 3.5.1 downgrade are the only two working paths. Details, root cause, the related FA2 landscape, and corrections to some widely-repeated community advice below.

Sourcing note

The three r/ROCm links (refs 1–3) are hard-blocked to automated fetch (old./www./bare reddit.com all refused). Those posts are apollo-mg's own threads pointing at his gists, which I read directly as the primary source:

• Master guide: AI on AMD RDNA4 w/ ROCm 7.1
• FLA "ASCENDED" / RDNA4 config
• The "NVCC trap" — causal-conv1d on RDNA4

Much of the older web material (HSA override tricks, "gfx1201 unsupported") predates ROCm 7.2 and is stale — our own HANDOFF.md empirical results (Apr–May 2026) are weighted as ground truth where they conflict.

---

1. The key distinction: two kernel stacks, two blockers, two fixes

Qwen3.5 (Qwen3_5ForConditionalGeneration) interleaves 48 linear_attention (gated-delta / DeltaNet) layers + 16 full_attention layers (HANDOFF.md:50). These ride different kernel paths, and almost every reference attached to this issue is actually about the second one:

Path	Used by	Library	gfx1201 blocker	Fix
Flash Linear Attention (FLA) — this issue	48 linear/DeltaNet layers	flash-linear-attention (Triton) + causal-conv1d	Triton 3.6.0 tritonamdgpu-pipeline use-after-free at num_stages>=2 → JIT deadlock at step 0	make patch-fla-rocm (num_stages→1) or downgrade Triton to 3.5.1
Flash Attention 2 (FA2) — the other refs	16 full-attn layers; serving	flash-attn / AOTriton / rocWMMA	CK backend assumes Wave64 (CDNA); RDNA3/4 are Wave32 → assembly/compile failure	FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE (swap CK→Triton backend)

Reference → path map: refs 2, 4, 5, 8 + ThunderKittens (refs 6,7) are all FA2; ref 1 is FA2 in llama.cpp (rocWMMA); ref 3 is serving/MTP. Only the apollo-mg gists cover FLA proper. So: the FA2 material is the related landscape (and the full-attention path), but our actual training blocker is FLA, covered in §2–4.

---

2. FLA blocker + fix (confirmed, ground truth)

flash-linear-attention==0.5.0 installs and imports fine; on the first training step it JIT-compiles ~120 Triton kernels, then one kernel hangs forever — process alive at 0% CPU, VRAM full, no error, progress bar stuck at step 0 (HANDOFF.md:53-60).

Root cause: Triton 3.6.0's AMD software-pipelining pass (tritonamdgpu-pipeline) has a use-after-free when num_stages >= 2 on gfx12xx. gfx12xx is not in Triton 3.6.0's verified target list. Same compiler bug, two manifestations:

• SageAttention path: hard crash — RuntimeError: PassManager::run failed 'tt.load' op operation destroyed but still has uses
• FLA path: silent deadlock (the hang we see)

Fix (ours): scripts/patch_fla_rocm.sh rewrites num_stages=[2-9]→1 (and num_warps>4→4 for Wave32) in the installed FLA wheel, then clears the Triton + FLA caches. Re-run after every uv sync. This mirrors the upstream SageAttention fix pattern, thu-ml/SageAttention#365 / kijai/ComfyUI-WanVideoWrapper#2007, which gates num_stages=1 on torch.version.hip is not None.

---

3. Nightly / upstream status — the direct answer to this issue

This issue asks to "look into nightly releases or other workarounds." Resolved:

No nightly fixes the FLA hang. The bug lives in Triton 3.6.0's tritonamdgpu-pipeline pass, and as of late May 2026 gfx12xx is still not in Triton's verified target list. Every "fix" in the wild is a downstream num_stages=1 workaround in the calling library (SageAttention PR #365, our patch_fla_rocm.sh) — not a Triton compiler fix. Upgrading torch/Triton nightly does not help and risks re-introducing the bug; pin and patch instead.

---

4. Alternative worth trying: downgrade Triton to 3.5.1 (avoids the bug entirely)

The bug is specific to 3.6.0's pipeline pass. A prebuilt triton-3.5.1+rocm7.2.1 wheel exists (triton-3.5.1+rocm7.2.1.gita272dfa8-cp312-...whl), and apollo-mg's gist calls Triton 3.5.1 + PyTorch 2.9.1 the "stable native pairing for gfx1201." If 3.5.1 predates the pipeline regression, downgrading would make the FLA fast-path work without re-patching the wheel on every uv sync.

⚠️ Must verify before adopting: we currently run torch 2.11.0+rocm7.2 / Triton 3.6.0 (HANDOFF.md:25). The documented 3.5.1 pairing is with torch 2.9.1 — unconfirmed whether 3.5.1 is ABI-compatible with our torch 2.11 build. Two candidate stacks to A/B:

• (current) Keep torch 2.11 + Triton 3.6.0 + patch-fla-rocm — confirmed working at seq 768.
• (try) Pin Triton 3.5.1 (test against torch 2.11 first; fall back to torch 2.9.1 if ABI breaks).

---

5. The FA2 / full-attention landscape (refs 1,4,5,6,7,8)

For the 16 full_attention layers and for serving, the FA2 story on RDNA4:

• CK vs Triton: stock flash-attn defaults to the Composable Kernel backend (Wave64/CDNA) and fails to compile on Wave32 RDNA. Set FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE and pip install flash-attn --no-build-isolation to JIT a gfx1201 Triton kernel at runtime (fp16/bf16/fp32, fwd+bwd, causal, varlen). This is exactly what scripts/amd_r9700_setup.sh already does (builds the ROCm/flash-attention fork, Triton backend).
• comfyui-amd-go-fast (ref 5): a --use-pytorch-cross-attention + sdpa monkeypatch over the ROCm/flash-attention@howiejay/navi_support branch (noted possibly-broken now) — Navi 3 only. Tuning knobs worth borrowing: PYTORCH_TUNABLEOP_ENABLED=1 MIOPEN_FIND_MODE=FAST (~up to 20% on ROCm).
• rocWMMA FA2 (ref 8): the AOTriton feat: R9700 + Mac Mini M4 two-PC LAN evaluation setup #16 thread's actual payload is Repeerc/flash-attention-v2-RDNA3-minimal — a rocWMMA FA2 kernel, but RDNA3 (gfx1100) only, no backward pass, weak BF16 → not viable for our training.
• ThunderKittens (refs 6,7): NVIDIA-only (Hopper/Blackwell). Its AMD sibling HipKittens supports CDNA3/CDNA4 (MI300/MI350) only — gfx1201 is NOT supported. Neither is usable on the R9700. Note for the record so nobody re-investigates.
• SDPA may already give us flash on the full-attn layers: train_sft.py:185 uses attn_implementation="sdpa". On ROCm, torch SDPA can route the full-attention layers to the AOTriton flash kernel if this torch build's AOTriton supports gfx1201; otherwise it silently falls back to the math/mem-efficient path. Worth verifying with the SDPA backend context manager — but this is a marginal optimization for 16 layers, not a blocker.

---

6. causal-conv1d + nvcc — the other candidate root cause (two theories, reconciled)

There are two competing claims about what actually blocks the linear-attention path on gfx1201, and which one bites you depends on how you install causal-conv1d:

• Theory A — "the NVCC trap" (apollo-mg gist 2 + the post that prompted this): the linear-attention path never even builds. causal-conv1d's setup.py assumes CUDA — it shells out to nvcc -V / torch.cuda.get_device_capability() (→ FileNotFoundError, then a NameError when the handler references an undefined version var on native ROCm), and PyTorch's C++ extension builder injects -gencode arch=compute_80, which amdclang++ rejects. csrc/ also #include <cuda_runtime.h> directly. Upstream tracking: Dao-AILab/causal-conv1d#99.
• Theory B — our experience (HANDOFF.md): causal-conv1d 1.6.2.post1 builds and runs (fwd+bwd verified on gfx1201), and the next wall is the FLA Triton hang (§2).

These reconcile cleanly — same library, two install paths:

Install method	Outcome
pip install causal-conv1d (isolated build)	❌ NVCC trap — isolated build env pulls CUDA torch, hits the nvcc/compute_80 failure (Theory A)
uv pip install --no-build-isolation causal-conv1d (+ wheel)	✅ builds against the ROCm torch already in .venv → 1.6.2.post1, works (Theory B — HANDOFF.md:89)
EmbeddedLLM/causal-conv1d-rocm (rocm-port)	✅ alt — setup.py patched to gate on torch.version.hip (clean fork)

So the NVCC trap is real, but for us it's an install-method artifact, not the blocker — --no-build-isolation is what makes Theory A not apply here. The reason that one flag fixes it: the isolated build spins up a fresh env and resolves torch from PyPI (=CUDA), which is what drags nvcc/compute_80 into the build; --no-build-isolation reuses our ROCm torch so the HIP path is taken. (Don't ever run plain pip install causal-conv1d here.)

Also worth confirming (⏳): whether causal-conv1d is even required. Qwen3.5 DeltaNet uses it only for the short causal conv in the linear path; if absent, transformers/FLA should fall back to a pure-PyTorch conv (slower, but no nvcc, no build). If true, dropping causal-conv1d entirely is a valid "make it run at all" fallback that sidesteps Theory A completely — verify the fallback exists before relying on it.

---

7. Corrections to widely-repeated community advice (stale post-ROCm-7.2)

• DO NOT HSA_OVERRIDE_GFX_VERSION=11.0.0. Half the guides call this "the critical fix." It was only ever for pre-7.2 ROCm that didn't recognize gfx1201 — it lies to ROCm that the card is RDNA3 (11.0.0). ROCm 7.2 supports gfx1201 natively, so forcing the override now hangs the GPU (confirmed, HANDOFF.md:112). Leave it unset.
• FP8 silently falls back to FP32 on gfx1201 — gfx1201 is missing from AITER's arch table (ROCm/TransformerEngine#520); throughput ~halves. Two-line fix ('gfx1201': 'MI350X') exists, unmerged. Affects serving/FP8, not LoRA training.
• vLLM still doesn't recognize gfx1201 upstream (vllm#28649) — serve via HF Transformers (serve_teacher_local.sh), as docs/GPU_SUPPORT.md already says.

---

8. Confirmed version matrix

Component	Confirmed-working (ours)	apollo-mg gist	Notes
ROCm	7.2 (torch.version.hip 7.2.26015)	7.1 / 7.2.0	7.2 = native gfx1201
PyTorch	2.11.0+rocm7.2	2.9.1 (stable) / 2.12 nightly	nightly not needed for FLA
Triton	3.6.0 + num_stages=1 patch	3.5.1 (no patch)	3.6.0 has the UAF; 3.5.1 = candidate alt
flash-linear-attention	0.5.0 (patched)	bare-metal Triton	hangs unpatched
causal-conv1d	1.6.2.post1 (--no-build-isolation)	NVCC-trap workaround	fwd+bwd verified

Also confirmed gfx1201-specific (HANDOFF.md): TORCH_USE_HIPBLASLT=0 (RoPE hipBLASLt crash, ~75–90s/step on rocBLAS); PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8 only (expandable_segments:True ignored, max_split_size_mb blocks large allocs); always uv run --no-sync.

---

9. Recommendation / open items

1. Keep make patch-fla-rocm as the supported path (current, confirmed at seq 768). ✅
2. Spike: pin triton==3.5.1+rocm7.2.1, test FLA fast-path against our torch 2.11 — if it works, drop the per-uv sync patch entirely. ⏳
3. Verify whether SDPA already routes the 16 full-attn layers through AOTriton flash on gfx1201 (backend ctx manager) — marginal. ⏳
4. Watch upstream: Triton adding gfx12xx to its verified target list + fixing the tritonamdgpu-pipeline UAF closes this issue for real. Until then, num_stages=1 is the only fix anywhere. 🔭

[ulises-c]

Addendum — AMD's official Mamba/causal-conv1d-on-ROCm recipe (and why it doesn't close gfx1201)

Ref: ROCm Blogs — Vision Mamba.

This is AMD's own walkthrough for the SSM/causal-conv1d path, and it pins down the mechanism behind §6's Theory-A/B reconciliation:

• AMD builds it from source — pip install -e ./causal-conv1d / ./mamba-1p1p1 — and it works because PyTorch's C++ extension builder runs hipify_torch to translate the CUDA source to HIP at build time. No nvcc, no Triton when the build sees a ROCm torch. So --no-build-isolation "fixing" the NVCC trap isn't a hack — it's what routes the build down this official hipify path (it reuses our ROCm torch instead of pulling CUDA torch into an isolated env). Confirms Theory B's mechanism.

Why this still doesn't resolve this issue, though:

	AMD Vision Mamba blog	Our gfx1201 reality
GPU	MI210 / MI300X — CDNA2/3, Wave64	R9700 — RDNA4 (gfx1201), Wave32
ROCm / torch	6.2 / 2.1.2 (~2024)	7.2 / 2.11
Kernel path	hipified C++ (selective_scan)	FLA Triton kernels

The blog validates the build on Wave64, never on gfx1201/Wave32, and on a 2024 stack. The hipified C++ kernels are CDNA-tuned — no guarantee they run correctly/fast on RDNA4 Wave32 (the same Wave64-vs-Wave32 mismatch that breaks stock FA2 CK kernels, §1/§5). That gap is exactly why the gfx1201 community routes the linear path through FLA's Triton kernels (JIT'd per-arch) rather than this hipified C++ path.

Net: this gives us a third candidate for the 48 linear layers on gfx1201 — the hipified C++ causal_conv1d/selective_scan path — as an alternative to "FLA Triton + num_stages=1 patch." Worth a spike, but expect Wave64-tuning risk, and it's a 2024-era recipe. The Triton+patch path (§2) remains the confirmed one.

[ulises-c]

Brainstorm notes — FLA on gfx1201, scoped to our actual training targets (from PR #79 + session)

Cross-referencing PR #79's "AMD ROCm / gfx1201 known issues" comment (the provenance for scripts/patch_fla_rocm.sh — num_stages=1 + num_warps→4) with this session's findings. Brainstorm only; nothing committed.

Scope refinement — which of our models FLA actually gates

Model	Role	Arch	FLA / this issue?
Qwen3.6-27B	teacher (SFT)	48 linear_attention (gated delta) + 16 full	gated — FLA Triton hang
Qwen3.5-0.8B-Base	34-way state classifier (this branch)	hybrid Mamba+attention — LoRA targets incl. in_proj_qkv/out_proj (scripts/train_state_classifier_34way.py:40)	gated — same FLA path
Gemma4-31B-IT	teacher (SFT)	softmax (sliding-window + global)	NOT gated — Makefile:342 "no FLA needed"; SDPA stable on gfx1201

So #100 is on the critical path for both Qwen targets (incl. the 0.8B classifier this branch trains) but is irrelevant to the Gemma4-31B SFT loop. If the near-term goal is a working SFT loop on the R9700, Gemma4-31B is the path that sidesteps this issue entirely (only gate there is bitsandbytes-4bit, which officially lists gfx1201).

Stakes (why we can't just stay on the seq=512 fallback)

Per PR #79: the torch-fallback at seq=512 is stable but truncates ~50% of dialogues (p50=599 tok) and the 27B run ETA is ~143 h / 3 epochs. FLA fast-path is what unlocks seq≥768 and a sane ETA. So "make FLA work" is load-bearing, not a nice-to-have, for the Qwen targets.

Untried levers to brainstorm (ordered by ROI), beyond our current num_stages=1 patch

The current patch (3.6.0 + rewrite num_stages→1, num_warps→4) works but must be re-applied after every uv sync. Alternatives that may be cleaner or faster:

• A — Freeze the "known-good" stack: ROCm 7.2 + torch 2.9.1 + Triton 3.5.1. Multiple gfx1201 reports + apollo-mg's Poachers Reproduction Guide gist report FLA working here with no patch — the tritonamdgpu-pipeline UAF is a 3.6.0 regression, so 3.5.1 avoids it at the source. Tradeoff: downgrade from our current torch 2.11; verify ABI before adopting. Highest-confidence fix.
• B — TRITON_DISABLE_AUTOTUNING=1. The hang is in one autotuned kernel config; disabling autotune means the broken candidate is never generated/selected. Cheap, stack-agnostic, reversible — worth trying first before any downgrade.
• C — TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. Enables AOTriton's experimental flash paths on RDNA4; reported to help gfx1201 attention. Stackable with B.
• D — Conservative kernel configs (superset of current patch): BLOCK_M=64, BLOCK_N=64, num_warps=2, num_stages=1. If A/B/C don't land, hand-cap the FLA @triton.autotune configs harder.
• E — Docker kernel extraction (Poachers method). Pull rocm/vllm-dev:rocm7.2_navi, extract AMD's pre-optimized Triton kernels, run bare-metal. Heavier/hackier but "verified working" with FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE. Last resort.

Corrections / refinements to earlier notes

• HSA override — nuance, not a flat "never." Our GPU hang was from HSA_OVERRIDE_GFX_VERSION=11.0.0 (that's the RDNA3 value — wrong arch). The Poachers gist uses HSA_OVERRIDE_GFX_VERSION=12.0.1, the correct gfx1201 identifier. On native ROCm 7.2 no override should be needed at all; if one ever is, use 12.0.1, never 11.0.0.
• Deadlock signature corroborated. Community reports describe the same failure as "hsa_signal_wait polling a DMA-completion flag that never flips" — exactly our symptom (120 kernels compile, one hangs at 0% CPU, no error). Confirms it's a single bad kernel specialization, not a general ROCm failure → fixable via stack/kernel-path changes (A–E), not abandonment.

Suggested experiment order on the R9700

1. TRITON_DISABLE_AUTOTUNING=1 (+ optionally TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1) on the current stack — cheapest test, no reinstall.
2. If still hanging: pin Triton 3.5.1 (test against torch 2.11 first; fall back to torch 2.9.1 if ABI breaks) — drops the per-uv sync patch entirely.
3. Keep patch-fla-rocm (3.6.0 + num_stages=1) as the confirmed fallback.

New primary sources this round: PR #79 thread; apollo-mg Poachers gist (ROCm 7.2 / torch 2.9.1 / Triton 3.5.1, kernels from rocm/vllm-dev:rocm7.2_navi).

[ulises-c]

Finding — the real Gemma4-31B blocker on R9700 is CPU-RAM staging at load, not FLA and not VRAM

Logging here for completeness (cross-ref PR #79 last comment, commit cef2e13). This blocker is orthogonal to this issue's FLA topic — Gemma 4 uses softmax attention and never touches flash-linear-attention. It's a bitsandbytes loading issue and would hit any GPU paired with ≤62 GB system RAM. Recorded here because it's the thing that actually stopped Gemma4-31B QLoRA on the R9700.

Root cause

from_pretrained(..., load_in_4bit=True) materializes the full ~62 GB BF16 model in CPU anon-RSS before NF4 quantization — regardless of low_cpu_mem_usage=True, offload_state_dict=True, offload_folder, GC threshold, or malloc tuning. bitsandbytes quantizes only after all weights are resident. So the napkin math was right after all: 4-bit QLoRA still needs ~62 GB of CPU RAM transiently at load (the 16 GB figure is the post-quant resident size, not the load peak).

The R9700 box has 62 GB RAM + 62 GB zram. With swap it survives past shard 820/1188, but zram thrashing slows loading to ~200 s/shard (~20 h just to load) → not viable.

Correction to an earlier note in this thread

The earlier deep dive flagged expandable_segments:True as a likely gfx1201 bug (per HANDOFF.md 2026-05-25 "unsupported, silently ignored"). The PR #79 session (2026-05-26, one day later) re-added it deliberately and found it works and helps — the HIP caching allocator was holding freed BF16 shards in VRAM, and expandable_segments:True releases them correctly. So the Makefile targets are correct and HANDOFF.md Issue 2 is the stale entry. Retracting the "remove it" suggestion.

Fix — use the pre-quantization path (already built)

The clean fix is already wired in the repo and drops load-time RAM from ~62 GB to ~6 GB:

1. Pre-quantize once on a high-RAM machine: scripts/prequant_gemma4.py → make prequant-gemma4-31b-l40s (needs the ~62 GB BF16 staging there, where RAM is plentiful), persist the NF4 checkpoint.
2. Transfer the ~16 GB NF4 weights to the R9700: make transfer-gemma4-31b-nf4 HOST=....
3. Train from pre-quantized: make train-gemma4-31b-stage2-preq (sets TRAIN_PREQ=true → train_sft.py:169 "loading pre-quantized NF4 checkpoint — skipping BF16 staging").

Cheaper alternative worth trying first

Download a community pre-quantized 4-bit checkpoint (e.g. an Unsloth-style *-bnb-4bit Gemma4-31B) directly to the R9700 — the weights arrive already NF4 (~16 GB), so there's no 62 GB BF16 staging at all and no high-RAM machine needed. Caveats to verify: the pre-quant must match our config (NF4 + double-quant + bf16 compute), and bnb 4-bit runtime kernels must work on gfx1201 (officially listed; may need a source build -DBNB_ROCM_ARCH="gfx1201").

Options that DON'T fix it (from PR #79)

• 8-bit (load_in_8bit) — still stages 62 GB BF16 in CPU at load; same RAM wall. ❌
• Accept the ~20 h load — technically possible, painful. ❌
• Run on NVIDIA (3090/5090) — valid Track B, but doesn't unblock the R9700. ↔
• Smaller model (Qwen2.5-14B QLoRA) — known to load+train fine; fallback. ↔

Net

On the R9700, Gemma4-31B QLoRA is gated by load-time CPU RAM, not VRAM and not FLA. Recommended order: (1) try a community 4-bit checkpoint (no staging, no second machine) → (2) fall back to the built-in offline prequant path. The 32 GB VRAM budget itself is fine once loaded (~22–30 GB peak with TRAIN_EVAL_STRATEGY=no + grad-ckpt, both already set).

[ulises-c]

Finding — pre-quantized 4-bit Gemma4-31B candidate (fixes the CPU-RAM load blocker)

Follow-up to the OOM finding above. Searched HF Hub for a pre-quantized bitsandbytes checkpoint so the R9700 can skip the ~62 GB BF16 CPU staging at load.

Recommended: unsloth/gemma-4-31B-it-unsloth-bnb-4bit (↓583k)

• 19 GB on disk, already NF4 (4 safetensors shards) → loads straight to 4-bit, no full-BF16 CPU materialization. Directly removes the ~62 GB load-RAM wall.
• Embedded quantization_config byte-matches our pipeline: bnb_4bit_quant_type=nf4, bnb_4bit_use_double_quant=true, bnb_4bit_compute_dtype=bfloat16 — identical to the BitsAndBytesConfig in train_sft.py:161. Loads via the existing TRAIN_PREQ=true path (bnb_config=None, reads config from the repo).
• Arch + LoRA targets intact: Gemma4ForConditionalGeneration, decoder at model.language_model.layers.* → our anchored LoRA regex still matches; decoder linears are 4-bit so PEFT injects adapters correctly.
• It's the -it (instruction-tuned) variant → continue-SFT on an already-IT model, as intended.

Caveats

• Unsloth dynamic 4-bit: llm_int8_skip_modules keeps embeddings/lm_head/vision_tower/audio_tower/projectors/router/per-layer gates in 16-bit. Better for fine-tune accuracy; resident VRAM slightly above pure NF4 but fits 32 GB. No plain uniform *-bnb-4bit exists — only this dynamic one (other hits are low-download derivatives).
• bnb 4-bit runtime on gfx1201 remains the one true gate — dequant kernels run at train time regardless of weight provenance. Pre-quant fixes loading, not this. Verify with make test-gpu-stack (steps 4/5); may need source build -DBNB_ROCM_ARCH="gfx1201".

Rejected (4-bit but unusable by the bnb QLoRA path)

Repo(s)	Format	Why not
cyankiwi / QuantTrio / ebircak / lcu0312 *-AWQ*	AWQ	needs autoawq; PEFT-QLoRA expects bnb Linear4bit
ebircak *-GPTQ	GPTQ	not bnb
nvidia / RedHatAI / LilaRest *-NVFP4*	FP4	vLLM/TensorRT inference, not a bnb training path
Intel *-int4-AutoRound	AutoRound	different runtime
mlx-community / lmstudio *-MLX-4bit	MLX	Apple Metal only
unsloth / bartowski / lmstudio *-GGUF	GGUF	llama.cpp inference, not trainable
*-abliterated / *-heretic	—	modified behavior, wrong for a teacher baseline

Usage (env overrides the config default; mirrors make train-gemma4-31b-stage2-preq)

nohup env TORCH_USE_HIPBLASLT=0 \
  PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8,expandable_segments:True \
  TRAIN_BASE_MODEL=unsloth/gemma-4-31B-it-unsloth-bnb-4bit TRAIN_PREQ=true \
  uv run --no-sync python scripts/train_sft.py \
  --config configs/train-sft-stage2-gemma4-31b.env \
  > outputs/sft-stage2-gemma4-31b/train.log 2>&1 &

Hardware requirements estimate (Gemma4-31B QLoRA, seq=1280, batch=1×16 grad-accum, grad-ckpt on, eval off)

Resource	Path B — pre-quant checkpoint (recommended)	Path A — naive load_in_4bit from BF16 base
GPU VRAM (peak, train)	~22–27 GB	~22–27 GB (same once loaded)
System RAM (peak, load)	~8–16 GB (weights stream in already-4-bit)	~62–70 GB (full BF16 staged in CPU before quant ← the blocker)
Download	~19 GB (NF4 shards)	~62 GB (BF16 base)
Disk working set	~30 GB (19 GB weights + ~11 GB adapter checkpoints over the run)	~75 GB (62 GB base + offload dir + checkpoints)

VRAM breakdown (Path B): 4-bit base resident ~16–18 GB (after vision_tower/audio_tower drop) · LoRA+grads+optimizer (adapters only) ~1–2 GB · activations w/ grad-ckpt ~3–5 GB · logits transient (256K vocab, fp32 loss) ~1.3 GB. Eval is disabled (TRAIN_EVAL_STRATEGY=no) because the eval logits.float() copy adds an ~8.6 GB spike that OOMs even 32 GB.

Minimum viable hardware:

• GPU: ≥ 32 GB VRAM (R9700 / RTX 5090). 24 GB (RTX 3090) is marginal — fits only with reduced seq len. 20 GB (RTX 4000 Ada) will not fit — 16–18 GB base leaves too little headroom; OOMs in the training step regardless of tuning.
• System RAM: ~16 GB is plenty with the pre-quant checkpoint (vs ~64 GB for the naive path — exactly why the 62 GB R9700 box thrashed into zram).
• Disk: ~30 GB free.

Net: the pre-quant checkpoint is the cheapest unblock for Gemma4-31B SFT on the R9700 — drops load-RAM from ~62 GB to ~16 GB, no second high-RAM machine, no offline prequant step. Only open question is bnb-4bit gfx1201 runtime, which the offline-prequant path would also have to clear.

[ulises-c]

Reproduction environment — gfx1201 (RDNA4)

Posting system specs for anyone hitting the same issues via search.

Hardware

Component	Spec
CPU	AMD Ryzen 9 5900X (12-core / 24-thread, up to 4.95 GHz)
GPU	AMD Radeon AI PRO R9700 — gfx1201 (RDNA4, 64 CUs, 32 GB VRAM, 2350 MHz)
System RAM	64 GB

Software

Layer	Version
OS	CachyOS (Arch-based)
Kernel	7.0.5-2-cachyos
ROCm	7.2 (HIP 7.2.26015, ROCm Runtime 1.18, rocm-smi-lib 7.8.0)
PyTorch	2.11.0+rocm7.2
transformers	5.9.0
peft	0.19.1
trl	1.4.0
bitsandbytes	0.49.2

Search terms that led here

gfx1201 · RDNA4 · R9700 · ROCm 7.2 · caching_allocator_warmup · Found no NVIDIA driver · bitsandbytes CPU dispatch · flash-linear-attention deadlock · num_stages Triton gfx12 · Gemma 4 QLoRA OOM load

[ulises-c]

make test-gpu-stack results — all 13/13 passed

Ran against the environment above. Posting results to close out the open gate from the issue.

Step	Check	Result
1	ROCm driver + GPU visibility	✅ ROCm 7.2.3, gfx1201
2	torch GPU tensor op	✅ 2.11.0+rocm7.2, HIP 7.2.26015, 31 GB visible
3	bitsandbytes import + GPU detection	✅ 0.49.2, ROCm/HIP backend confirmed
4	bitsandbytes 8-bit Linear (LLM.int8)	✅ forward pass OK
5	bitsandbytes 4-bit NF4 (QLoRA)	✅ QLoRA quantization supported on gfx1201
6	transformers BitsAndBytesConfig	✅ 5.9.0, 4-bit NF4 + double-quant config OK
7	PEFT LoraConfig	✅ 0.19.1
8	TRL SFTConfig	✅ 1.4.0
9	SDPA / flash attention	✅ flash=True, mem_efficient=True, math=True
10	llama-server binary + ROCm flags	✅ AMDGPU_TARGETS=gfx1201, GGML_HIP=ON
11	llama-server flash attention	✅ GGML_CUDA_FA=ON
12	Vulkan backend	✅ GGML_VULKAN=ON, RADV GFX1201 detected
13	Serve script settings	✅ ubatch/KV quant/parallel slots verified

Step 5 is the gate from this issue — bnb 4-bit NF4 runtime on gfx1201 is confirmed working without a source build (stock bitsandbytes==0.49.2 from PyPI).

Note: flash_attn PyPI wheel fails on gfx1201 (CK-tile ISA mismatch). PyTorch SDPA covers the same path on ROCm and is sufficient for SFT.