feat(mxfp4): switch MXFP4 linears to FP8 at a preset iteration - #1010
feat(mxfp4): switch MXFP4 linears to FP8 at a preset iteration#1010jasainio wants to merge 5 commits into
Conversation
The Q/K RMSNorm can return a wider dtype than it was given, leaving q/k in fp32 while v stays bf16. Every dense flash-attention backend rejects a mixed-precision (q, k, v) triple rather than casting for us, so the MXFP4 Flux config could not train at all on a Primus-Turbo new enough to expose the merged FlashAttnFunc. Reference Flux performs this same cast inside QKNorm; the Megatron port had dropped it. Applied to both streams of JointSelfAttention and to FluxSingleAttention.
Primus-Turbo PR #447 changed get_gemm_backend() from returning a bare BackendType to returning a BackendChoice carrying its own auto_tune flag. Comparing a BackendChoice against a BackendType silently evaluates False, which turned every MXFP4 run into a preshuffle-contract failure on Turbo 0.4.1. Read through either shape so the same code works on both the newer Turbo the Flux local spec requires and the older pinned ones still in use elsewhere.
…ices torchrun was always launched with --local-ranks-filter 0, so only local rank 0 reached the console and a crash on any other rank produced no error message at all. Diagnosing a GPU fault on rank 2 took seven runs because its "illegal memory access" was being thrown away every time. PRIMUS_LOCAL_RANKS_FILTER now overrides the default, with "all" keeping every rank. HIP_VISIBLE_DEVICES was also overwritten unconditionally, silently clobbering a value passed in from outside. A caller-supplied value now wins, which is what allows a specific device set or ordering to be pinned -- needed to tell a rank-specific software fault apart from a failing GPU. Both default to the previous behaviour when the variables are unset.
Flips every MXFP4 linear to dynamic tensorwise FP8 at a configured iteration, with no checkpoint, no weight conversion and no optimizer state remapping. Neither precision stores a quantized weight -- both inherit a plain BF16 nn.Parameter from Megatron and quantize inside forward -- so the switch is purely a change of which autograd Function _forward_impl dispatches to. There is nothing for DDP's bucket groups or the optimizer's identity-keyed state to notice, which is what keeps it within the MLPerf rules. The train_step wrapper derives its state solely from the iteration number, so every rank reaches the same decision without communicating; determinism, idempotence and order-independence are covered by tests. The pre-warm is hosted inside mlperf_warmup because that is the only place that runs a real grad-enabled step at the production micro_batch_size. With automatic_dynamic_shapes on, pre-warming at any other shape would mark dims dynamic and change the compiled graph for the MXFP4 path too. It doubles as the only runtime proof the switch is not a silent no-op: it asserts Dynamo traced a new graph, and an empty layer plan raises rather than letting a BF16 fallback masquerade as a successful switch. Validated on MI355X: 57/57 layers flip at the boundary, the whole model needs 3 additional Dynamo graphs, and loss continues without a spike. Saved activations grow 0.531 -> 1.005 bytes/element, matching the projection.
Formatting only, no behaviour change. Satisfies the code-lint gate, which runs black with --line-length=110.
Note:
|
2 participants
Summary
Flips every MXFP4 linear to dynamic tensorwise FP8 at a configured iteration, with no checkpoint, no weight conversion and no optimizer state remapping.
Neither precision stores a quantized weight — both inherit a plain BF16
nn.Parameterfrom Megatron and quantize inside forward — so the switch is purely a change of which autograd Function_forward_impldispatches to. There is nothing for DDP's bucket groups or the optimizer's identity-keyed state to notice, which is what keeps it within the MLPerf rules.Two independent bug fixes come first and are worth reviewing on their own. Both repair the baseline MXFP4 Flux path on Primus-Turbo 0.4.1, so they matter to anyone who does not care about the switch:
cb360974fix(flux)— cast q/k back to value's dtype after QK norm7f1b047dfix(mxfp4)— acceptBackendChoicefromget_gemm_backend3db3604efix(runner)— stop discarding non-zero rank output; let caller pin devicesf00e4bc8feat(mxfp4)— the switch itselfThe runner commit is a standalone bug fix:
torchrunwas always launched with--local-ranks-filter 0, so a crash on any rank other than zero produced no error message at all. That cost seven runs of debugging during this work, because the failing rank had been printingillegal memory accessthe whole time and it was being thrown away.Design notes for reviewers
train_stepwrapper derives its state solely from the iteration number, so every rank reaches the same decision without communicating. Determinism, idempotence and order-independence are covered by tests.mlperf_warmupbecause that is the only place that runs a real grad-enabled step at the productionmicro_batch_size. Withautomatic_dynamic_shapeson, pre-warming at any other shape would mark dims dynamic and change the compiled graph for the MXFP4 path too._switch_-prefixed attributes are deliberately distinct from the existing_fp8_bwd_dtype/_fp8_gran_value/_fp8_backend_valuetriple, which is passed on every call including the pure-MXFP4 path; reusing it would change traced constants on a path this switch must leave bit-identical.Test plan
unique_graphs 8 -> 11)Known gaps
No full-length convergence run has completed. Every 8-GPU attempt at MBS 64 was killed by a hardware fault on one GPU of the test node (
hipErrorIllegalAddress,SQC (inst)fetch at address zero, following physical GPU 2 across rank permutations and disappearing when that GPU is excluded). The MBS 64 validation above was therefore run with that GPU excluded. This is unrelated to the change set, but it means production-config convergence and post-switch steady-state performance numbers are still outstanding.