Add MXFP4 training support via a dequantize-on-the-fly fallback - #47464
Add MXFP4 training support via a dequantize-on-the-fly fallback#47464rexminnis wants to merge 2 commits into
Conversation
The triton MXFP4 kernels are forward-only, so lifting the is_trainable gate alone trains on silently biased residual-only gradients (expert outputs are autograd constants). Instead, mlp_forward now branches when gradients are required: the eager reference math (top-k router + expert loop) runs over weights dequantized on the fly per expert, keeping the packed MXFP4 storage untouched and the kernel path for inference. With gradient checkpointing the dequantized transients are bounded to ~one decoder layer, which puts gpt-oss-20b LoRA fine-tuning at ~15GB peak on a single 24GB GPU (validated end to end on an L4: cos=0.999992 vs kernel forward, true gradient flow, 5-step LoRA run + adapter round-trip through PEFT and vLLM serving).
|
cc @Rocketknight1 @salmanmohammadi — you both reviewed #40180 (native MXFP4 training, deferred pending backward kernels). This PR takes the complementary route: weights stay packed MXFP4 (15 GB peak for gpt-oss-20b LoRA on a single 24 GB L4, vs 42 GB for bf16 weights alone), and the backward runs eager reference math over per-expert on-the-fly dequantization — so gradients through the expert path are real (cosine 0.999992 vs the Triton forward), avoiding the silent-gradient issue. Scoped to LoRA/adapter training, where the frozen base never needs |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: mxfp4 |
CI recapDashboard: View test results in Grafana |
|
Sorry, no code agent PRs from first-time contributors! |
2 participants
Add MXFP4 training support via a dequantize-on-the-fly fallback
Fixes #40170 (partially — see "Relation to backward kernels" below), closes the gap behind #40236.
What this does
Enables LoRA/adapter fine-tuning of MXFP4-quantized models (gpt-oss-20b/120b) with the quantized weights kept packed, on hardware where the dequantized-bf16 recipe cannot fit:
Mxfp4HfQuantizer.is_trainable→True(with awarning_onceexplaining the fallback + recommending gradient checkpointing).mlp_forwardgains a grad-aware branch: whentorch.is_grad_enabled()and gradients are required, the MoE block runs the eager reference math (top-k router +GptOssExperts-style expert loop) over weights dequantized on the fly; otherwise the existing triton kernel path runs unchanged. Inference performance is untouched.Mxfp4GptOssExpertsgains_dequantize_weights()(layout-safe:convert_layout(..., StridedLayout)beforeupcast_from_mxfp_torch, per-expert chunking to bound the fp32 padding intermediates) andforward_dequantized()(the autograd-visible expert forward, mirroring the eagerGptOssExperts.forwardmath).The packed MXFP4 storage is never modified. The dequantized tensors are transients owned by the autograd graph; with gradient checkpointing enabled, at most ~one decoder layer's worth (~1.7 GB bf16 for gpt-oss-20b) is alive at a time.
Why (the silent-gradient trap)
The triton MXFP4 kernels implement forward only. If a user lifts the
is_trainablegate today (one-line patch, circulating in gists), training appears to work — losses go down, grad norms look healthy — but the expert outputs havegrad_fn=None, so gradients flow through residual connections alone. Every LoRA trained that way learns from biased gradients with no error or warning. We verified this empirically before writing this PR:This PR replaces that trap with a correct (if slower) path.
Validation (gpt-oss-20b, single 24 GB NVIDIA L4, SM89)
The exact code in this PR (injected over a transformers 5.14.1 install) was validated end to end, not a prototype variant:
grad_fn; gradient reaches the embedding output (norm ≈ 2.86) with all base parameters frozen.lora_Babs-sum 702.7 ≠ 0) loads onto the quantized base viaPeftModel.from_pretrained, greedy generation visibly shifts vs base; adapter also serves via vLLM (merged into checkpoint, and dynamically on vLLM ≥ 0.25.1).Hardware note: validated on SM89 where the triton weight layout is already
StridedLayout;_dequantize_weightscallsconvert_layout(..., StridedLayout)explicitly so swizzled layouts (SM90+) should round-trip, but we could not validate on Hopper/Blackwell — reviewer confirmation welcome.What this unlocks for users
Relation to the backward triton kernels (#40170) and to #40180
This is deliberately a pure-Python fallback, not the fused MXFP4 backward kernels the issue ultimately wants: ~2–4× slower per step than a fused implementation would be, since the backward recomputes dequantization. It is intended to (a) make training correct and possible now, and (b) serve as the numerical reference implementation to validate the real backward kernels against when they land — at which point the grad-aware branch can dispatch to them instead.
To be explicit about scope, in the terms raised on #40180 ("isn't this MXFP4 weight storage and computation in a fallback dtype?"): yes — by design, and that split is the contract. The storage half is where the memory win lives and it is fully preserved — the packed weights are never materialized in bf16 (that's the 14.9 GiB peak above, vs ~42 GB for dequantized-bf16 weights alone). The compute half is an eager fallback that is correct, which is sufficient for the LoRA/adapter use case: the frozen base weights never need
grad_weight— gradients only need to flow through the expert math to the adapters, which the autograd-visible dequant path provides. Native MXFP4 compute remains #40170's fused kernels; #40180 was deferred pending those kernels, and this PR intentionally does not presume them — it is the reference they can be validated against.Tests
test_mxfp4_expert_output_requires_grad— quantized load, forward under grad; asserts expert output hasgrad_fnand backward reaches inputs (GPU + triton kernels required).test_mxfp4_dequant_forward_matches_kernel— cosine similarity ≥ 0.999 between the two paths on random inputs.