feat(mxfp6): MXFP6 (E2M3) Flux linears, a fused MLP, and the MLPerf recipe - #1064
Draft
jasainio wants to merge 19 commits into
Draft
feat(mxfp6): MXFP6 (E2M3) Flux linears, a fused MLP, and the MLPerf recipe#1064jasainio wants to merge 19 commits into
jasainio wants to merge 19 commits into
Conversation
Wires Primus-Turbo's gemm_fp6 into Flux through a new local spec provider, mirroring the MXFP4 path but without its configuration surface: MXFP6 has one backend (AITER A6W6), a mandatory Hadamard rotation folded into the packer, and no preshuffle choice, so the preshuffle contract and recipe flags have nothing to guard and are dropped rather than retargeted. fp6 is a separate config field rather than an fp4: "mxfp6" overload. Overloading fp4 would let an MXFP6 config select the MXFP4 provider while skipping its preshuffle check, producing a plausible-looking run in the wrong precision. Disjoint fields make that unrepresentable, at the cost that Megatron never sees fp6 -- so BaseDiffusionConfig rejects fp6 with fp4, with fp8, and with mxfp4_to_fp8_switch_iter, whose patch would otherwise build a plan over zero layers and silently never switch. Every MXFP6 GEMM needs M, N and K to be multiples of 256, K included because the backward GEMMs use it as an output dimension. Flux 12B satisfies this at every replaced linear. Two test-quality issues the MXFP4 suite this was cloned from also has, hence the different assertions here: - Flux zero-inits proj_out and every adaLN_modulation tail, so a fresh model outputs exactly zero and gives all 12 MXFP6 linears an exactly zero gradient. Asserting "grad is not None" passes in that state whether or not the quantized GEMM ran. The integration test breaks the zero-init first and asserts non-zero gradients, keeping a test for the inert-at-init premise itself. - A BF16 SGD loop cannot show convergence: the update is ~7e-6 relative to the weight while BF16 resolves ~8e-3, so steps round away and the loss is bit-identical regardless of gradient correctness. The loop steps an FP32 master copy, as Megatron's optimizer does. Verified on MI355X: 30 unit + 6 Flux integration + 494 config tests pass. Forward and both gradients land at ~28 dB SNR against BF16, and the linear is bit-identical to Primus-Turbo's FP6GemmMXFunction in both directions.
Splitting the MLP into two autograd Functions forces the activation to exist: fc2's forward needs a real tensor and fc1's backward needs one back. Owning fc1 -> epilogue -> fc2 in a single Function lets the MXFP6 packer take that epilogue as a prologue in both directions, so the activation is packed straight out of LDS and the pre-activation gradient is never assembled at all. The bias gradient comes back as a per-column sum from the packer, because it is a reduction over exactly the tensor that no longer exists. Worth 74.1 -> 78.0 images/s/GPU on Flux 12B, 8x MI355X at micro_batch_size 64: 75.3 ms/step of epilogue and reduction kernels leave, 35.5 ms/step of prologue cost joins the packer, and the step is GPU bound at 96.8% busy so that lands as wall clock. Peak allocated memory moves by the column-sum buffer alone. Engages through a new mlp_module() hook that the Flux specs look for and fall back from, so no other backend is affected. Anything the fused path cannot reproduce exactly -- gated linear units, a non-tanh GELU, an FP8 backward -- falls back per module, and PRIMUS_MXFP6_FUSED_MLP forces either side for A/B measurement. Also ignores throwaway _*.yaml run recipes and ut_out/, which were showing up as untracked in every status.
On the v26.5 image the QK-norm hands back its fp32 accumulation dtype under torch.compile instead of the input dtype. Only Q and K go through a norm, so V still carries the intended dtype while Q and K arrive as fp32, and Turbo's dense flash-attention backends all require fp16/bf16. The dispatcher reports that as "No compatible backend found for FlashAttnDenseDispatcher", listing shapes and never mentioning dtype, which sends you looking at the wrong thing. Cast against V in both attention classes: the joint blocks project QKV themselves, and the single blocks get an override since they reuse Megatron's projection. Both casts stay inside the compiled region, where they fuse into the norm's epilogue rather than costing a separate pass over Q and K -- doing this at the attention call instead measured 14.5ms per step, about 1.8% of a Flux 12B MXFP6 step.
…from the shell The two fallback tests assert that an unusable configuration warns and defers to the stock MLP, which is the behaviour of the default "auto" mode. Under "on" the same configuration raises instead, and the submission container exports PRIMUS_MXFP6_FUSED_MLP=on so that a fallback there cannot pass unnoticed -- so both tests failed in the one environment where running them matters most. Pin the mode per test rather than trusting the ambient value. The kill-switch test still sets it explicitly and is unaffected.
… too Same ambient-environment problem as the unit tests, reached by a different route. FluxConfig's default activation is a hand-written tanh GELU that the fused prologue does not reproduce, so every model these tests build takes the fallback path -- which is an error rather than a fallback under the submission container's PRIMUS_MXFP6_FUSED_MLP=on, failing all six at layer construction.
The MXFP6 work had a development recipe and the MLPerf work had an FP8 recipe, and nothing in between. This is the FP8 MLPerf config with the precision block swapped and nothing else moved, so that comparing the two is comparing numerics rather than comparing two independently drifted configurations. Three settings differ from the MXFP6 development recipe and are deliberate. FP32 optimizer states instead of BF16, matching the FP8 MLPerf recipe: time to train is what is being measured, and FP32 states cost nothing in throughput here while removing a source of convergence difference from the comparison. emulate_precision_casts off, since it models FP8 rounding inside compiled regions and there is no FP8 cast on this path. nemo_aligned_lr_warmup written out explicitly rather than left to the default, because the FP8 MLPerf recipe and the MXFP6 development recipe disagree on it and the disagreement is a different LR trajectory through warmup, not a stylistic difference -- it stays on the setting the convergence evidence was collected under until there is a reason to move it. The recipe cannot produce a submittable log yet, and says so: the compliance checker accepts a fixed vocabulary for lowest_numerical_precision_in_linear and mxfp6 is not in it. Runs under this recipe are valid runs whose logs are rejected on that one key until the format is accepted upstream.
jasainio
force-pushed
the
feat/mxfp6-fused-mlp
branch
from
September 1, 2026 12:23
ad7f482 to
ccd4d18
Compare
The empty T5/CLIP encodings are plain trainer attributes rather than registered buffers, so nothing in the training stack ever moves them, and they are loaded from .npy or torch.randn straight onto the host. The CFG dropout branch in forward_step therefore ran empty_t5_encodings.to(device="cuda", ...) on every step, copying the same never-changing tensors host-to-device for the life of the run. That is wasted work in normal operation and fatal under CUDA graph capture, where an unpinned host-to-device copy inside the captured region raises "Cannot copy between CPU and CUDA tensors during CUDA graph capture" -- so full-iteration capture is unreachable for any run with cfg_dropout_prob > 0 while the copy is there. Moving them once at setup, in the run's compute dtype, turns the per-step .to(...) into an identity and leaves the numerics untouched. Both branches that populate the encodings do it: the discovered-encodings path and the mock_data path. The third branch raises, so there is nothing to move.
The MXFP6 backward allocated each weight gradient and handed it to autograd as
param.grad, leaving Megatron's DDP hook to add it into main_grad and free it. With
gemm_fp6_out_impl the A6W6 asm writes main_grad itself and the hook has nothing left to
do. On the MBS=32 GBS=256 Flux 12B arm (8x MI355X) that is 974.2 -> 958.8 ms per 512
images, 15.4 ms or 1.6%, and it narrows the MBS=32 shortfall against MBS=64 from -15.3%
to -13.9% images/s/GPU. Loss is bit-identical at every iteration, which is the expected
outcome rather than a tolerance: the store computes the same GEMM into the same buffer
the add_ was filling from zeros. The trace confirms the mechanism exactly -- aten::add_
-304, gemm_fp6_impl -304, gemm_fp6_out_impl +304, 17.6 ms less elementwise GPU time, and
total CPU op count unchanged.
Gated on a new Primus-owned mxfp6_fused_wgrad_accum rather than Megatron's
gradient_accumulation_fusion, which was the first thing tried and is the wrong knob.
Megatron's flag is read by every plain linear, so switching it on also routes Flux's 76
AdaLN modulation projections through wgrad_gemm_accum_fp16, and at their M=32 shapes that
kernel costs more than the separate add it replaces: 487 -> 542 ms/step, +11%, with a
changed accumulation order that showed as a drifting loss. That is a finding beyond this
change -- Megatron's fused wgrad is a pessimisation at Flux's AdaLN shapes on this stack.
The MXFP6 linears therefore still reject the Megatron flag outright, now with a message
naming the field to use instead.
grad_added_to_main_grad is set from the forward, in _claim_main_grad, which reads oddly
and is load-bearing. Setting it next to the store, in the Function's backward, makes
dynamo refuse the mutation ("Mutating a variable from outside the scope of this HOP is
not supported") -- and rather than failing it breaks the graph around every MXFP6 linear:
240 breaks in a 14-iteration run, all 152 non-MLP linears eager, 3338 extra elementwise
kernels and +42.7 ms per 512 images against the ~16 ms the fusion saves. From the forward
the same assignment is ordinary traced code that dynamo records and replays. Setting it
before the store is safe in this direction only, because the flag is read by the DDP
backward hook, which cannot run before the backward has stored, and zero_grad_buffer
clears it every step; Megatron relies on the same asymmetry under TE CUDA graphs.
Three constraints are enforced per module rather than assumed, so a misconfiguration
cannot read as a performance result: one microbatch per optimizer step, since the beta=0
store overwrites instead of accumulating; a bf16 main_grad, since the asm writes nothing
else; and the pure-MXFP6 backward, since the FP8 path forms its wgrad with gemm_fp8_impl,
which has no out variant. The backward still returns a placeholder gradient because the
DDP hook asserts param.grad is not None whenever overlap_grad_reduce is on, matching what
Megatron's own fused path does; it is never read, and it costs 159.5 -> 165.1 GB peak.
Ships off. The MBS=64 arm is equally eligible -- one microbatch, bf16 main grads -- but
flipping a shipped default deserves its own measurement.
Six tests cover it: that the field and not Megatron's flag is what enables the fusion,
that the gradient lands in the real main_grad for both the plain linear and both MLP
matrices, that a missing or non-bf16 main_grad is refused, that FP8 backward is refused,
and that the compiled region stays whole. That last one only means something because it
traces with a main_grad attached and an input with requires_grad=True; without the input
there is no autograd graph, dynamo never reaches the backward, and the test reports zero
breaks whatever is in there -- which is how its first version passed while production was
breaking 240 times.
The MXFP6 linears are biased, and the caller was adding the bias after the Function returned. That put the bias outside autograd's view of the op, so the bias gradient became its own reduction over grad_output -- a two-stage Triton reduction per linear, on a tensor the backward is already streaming end to end for quantization. Moving the bias into MXFP6LinearFunction lets the backward take it as a side output of the pack instead. The packer gains an Identity prologue and want_col_sum, and grad_bias is the finishing .sum(0) over its partials. Identity rather than a real prologue because there is no activation to undo here, unlike the MLP's fc1, which already did exactly this. Measured -4.4 ms per 512 images at MBS=32 over three repeats with non-overlapping groups, from -7.5 ms of kernel time: -7.2 ms of Triton elementwise and 152 fewer launches, against +0.6 ms in pack/quantize. Four kernels go outright, the two-stage qkv bias-gradient reduction on both streams. The forward is unchanged arithmetic -- the add simply happens inside the Function now -- and bias grad, input grad, weight grad and forward are all bit-identical to the old arrangement at the real Flux shapes. The fp8-backward path runs no packer, so it keeps paying for its own grad_2d.sum(0); it is not the configuration this is for. Note this does not close the MBS=64 gap. The traffic is token-proportional and MBS=64 gains more from the same change (-7.7 ms), so the ratio moves the wrong way. It is worth having for absolute MBS=32 throughput.
With 2c the bias is an input to MXFP6LinearFunction, but the forward still paid a separate elementwise pass to add it: read the whole [M, N] result, add N values, write it back. aiter's A6W6 GEMM can now do that in its store epilogue, where it is free -- the epilogue is bound by its scatter store rather than by VALU, so the add hides in issue slots the stores were already stalling through. Measured -11.3 ms per 512 images at MBS=32, three repeats per arm interleaved, groups non-overlapping. The trace confirms the mechanism rather than inferring it: triton_poi_fused_add_view_1 (10.6 ms, 150 launches) and _2 are gone entirely, and total kernel time falls 12.5 ms against the 10.8 ms directly deleted. Nothing is gated here and this layer knows no aiter version. Whether the bias can actually be folded is a property of the installed aiter, so Primus-Turbo probes gemm_a6w6's signature for the parameter and adds the separate pass itself when it is absent. Passing the bias unconditionally is therefore always correct, and there is no environment variable. Kernel side: ROCm/aiter#5236. The result is also slightly more accurate than what it replaces, which is why the new test compares against a double-rounding bound rather than bitwise. The separate add rounded the GEMM result to bf16 and then added, rounding twice; the epilogue adds into the fp32 accumulator and rounds once. Like 2c this is token-proportional, so MBS=64 gains comparably and it does not close the gap between the two operating points.
…nels Left to discover the fabric itself, RCCL plans 2 channels between nodes and inter-node bandwidth measures 46 GB/s. The Crusoe topology XML declares which rail is local to which GPU; with it the plan widens to 8 and the same benchmark measures 118 GB/s, a 2.56x that shows up directly in two-node step time. Resolved by GPU PCI device id, never by filename. The mi350x and mi355x files differ only there (0x75a0 against 0x75a3) and a mismatched file still parses, so RCCL applies it and silently loses the affinity the file exists to declare -- a wrong filename is worse than no file, because it looks configured. Read from sysfs rather than lspci, which is a package and not present in every image. /etc/crusoe is not mounted in every container, so a staged copy at /opt/rccl_topo_node.xml is accepted as a fallback. NCCL_NET_GDR_LEVEL=LOC is set in the same branch and must stay coupled to it. GPU-Direct RDMA registration fails on this fabric -- every ibv_reg_mr_iova2 returns EINVAL -- and the topology file is the first configuration that makes RCCL attempt it, so setting one without the other kills every rank during connection setup. The 2.56x comes from the wider plan and does not depend on GDR. PRIMUS_RCCL_TOPO_DISABLE=1 restores the previous behaviour exactly, so the A/B stays measurable from outside the launcher, and as an escape hatch if a future node ships a file that plans worse than RCCL's own discovery. NCCL_DMABUF_ENABLE becomes overridable rather than forced to 0, so the dmabuf registration path can be tried without editing the launcher. Refs: docs/mlperf/rccl_gdr_escalation.md Co-authored-by: Cursor <cursoragent@cursor.com>
…sume Megatron's gradient buffers learn, on their first batch, how many times each parameter registers a ready gradient; from the second batch on they issue a bucket's reduce-scatter only when that golden count is reached again. The MLPerf warmup steps are batches like any other, so they consume the calibration: the golden counts end up describing a synthetic step rather than the first real one. When the real steps then register a different number of times -- a different microbatch count is enough -- the bucket either fires early, and the next registration finds a collective already in flight, or never reaches the golden count and the reduce-scatter is never issued at all. Reset every bucket group back to is_first_batch with empty counts, so the first real step performs the calibration it would have performed without a warmup. Any outstanding collective is drained first. #1069 already waits on the warmup's own reduce-scatter, so in practice the drain count is zero on this stack, but the handle belongs to a synthetic step whose gradients are about to be discarded and leaving it in flight would hand the first real step a bucket that is busy for reasons it cannot see. Draining here also keeps the reset correct if it is ever applied on top of code without #1069. Two-node MXFP6 runs reset 43 bucket groups at MBS 32 and GBS 512. Co-authored-by: Cursor <cursoragent@cursor.com>
Megatron's forward_step takes the first element of what a loss function returns as the tensor to backpropagate and rescales it IN PLACE -- `output_tensor *= cp_group_size`, then `output_tensor /= num_microbatches` -- and only then stores the reported dict. The diffusion validation path put `loss_sum.detach()` in that dict, and a detached tensor shares storage with the one it came from, so the reported loss was rescaled along with it. What reached the caller was therefore the true validation loss divided by the number of microbatches. At one microbatch per rank per step the divisor is 1 and nothing shows, which is why every shape measured so far looked right. At two it halves, and under mlperf_mode the halved value is what the convergence gate is compared against -- so a run reports convergence at roughly half the samples it actually needed, and reports it as a pass. Measured on Flux 12B at matched global batch 512: micro batch 32 reported 0.630950 at step 100 where micro batch 64 reported 1.266634, a ratio of 2.007. With this fix the two agree to 0.7%, and to 0.01% by step 200. Context parallelism has the mirror problem, inflating the reported loss by cp_group_size. The training path a few lines below already clones for this reason. Carries the regression test with it: the test rescales the returned tensor exactly as Megatron's forward_step does and asserts the reported pair does not move, so a future detach-only return fails rather than silently halving the number the convergence gate reads. Co-authored-by: Cursor <cursoragent@cursor.com>
Nine seeds converged under this configuration on 2x8 MI355X -- 20131, 21139, 22147, 23155, 24163, 25171, 26179, 27187, 28195 -- every one reaching val_loss <= 0.586, at a mean of 7,718,684 samples and about 2.4 h per seed. Until now it existed only as a generated file in a scratch directory that .gitignore excludes (examples/megatron/configs/**/_*.yaml), so reproducing a two-node run meant obtaining that file out of band and the recipe the results belong to was not in the repository. Differs from the single-node MLPerf recipe in the batch shape and what the batch shape forces: micro_batch_size 32 at global_batch_size 512 is one microbatch per rank at 16 ranks, and lr 2.0e-4 with lr_warmup_iters 1600 are the flux_ref_512 reference values rather than the GBS 1024 pair. Neither field is constrained by the ruleset, so carrying the wrong pair passes every compliance check and quietly forfeits the reference points the run is scored against. The rest of the deltas are corrections that apply at any shape and are documented inline: do_valid must be set explicitly because the flux path never calls build_train_valid_test_data_loaders, eval_samples must be used instead of eval_iters so the coverage assertion runs, and eval_timestep_source must be dataset or evaluation silently falls back to equidistant timesteps. Seed, experiment name, TensorBoard directory and data path are environment parameters, so one file serves a whole campaign. Co-authored-by: Cursor <cursoragent@cursor.com>
Megatron's 256000000-element default is next to worst of the useful range on Flux 12B. Sweeping seven sizes, 1024000000 is worth -17.4 ms per 512 images at the MBS=64/GBS=512 this recipe ships (78.8 -> 80.6 images/s/GPU) and -27.9 ms at MBS=32/GBS=256, where twice as many optimizer steps make twice as many collective calls. Peak memory is unchanged at 250.8 GB. The mechanism is recovered overlap, not faster math: exposed collective time falls from 17.8 to 4.0 ms per 512 images at MBS=64, so 96.5% of collective time is hidden against 85.2%, while compute stays flat in every arm. Both ends of the sweep lose, which is why the comment in the config says not to push this further in one direction. Small buckets pay per-call overhead and get thin per-call bandwidth. Buckets spanning most of the model cannot start their reduce-scatter until far into backward, so the tail is re-exposed -- 2048000000 measures 13.9 ms exposed against 1024000000's 4.0. 1024000000 is where the two costs cross, and the optimum is flat enough that 2048000000 is within run-to-run drift. Measured with the winner and the default each run twice at both batch sizes; replicate pairs differ by 0.5 to 1.8 ms against a 17-28 ms effect. The MLPerf v6.0 reference moved the same knob in the same direction, 256M to 512M, and 512M is also an improvement here -- just not the best one. Only this recipe changes, and deliberately not the MLPerf one, even though it is the same model at the same batch size. The sweep ran with main_grads_dtype bf16, which is what this recipe uses; the MLPerf recipe sets fp32, so an identical element count is twice the bytes per call. Bytes are what set both per-call bandwidth and how early a bucket can flush, so the optimum in elements should not be assumed to carry across that change -- 1024000000 at fp32 sits closer in bytes to the 2048000000-at-bf16 arm, which is already past the peak. That recipe needs its own sweep. The fp8, mxfp4 and te_spec recipes are untouched for the same reason.
…gime The previous commit deliberately left this recipe alone. ddp_bucket_size counts elements, but what governs per-call bandwidth and how early a bucket can flush is bytes, and this recipe sets main_grads_dtype fp32 where the development recipe sets bf16 -- so 1024000000 here is twice the bytes it is there, close to the 2048000000-at-bf16 arm that had already regressed. Rather than assume either way, swept again in this regime. The concern did not materialize. 1024000000 is the optimum in both, and the fp32 plateau is merely wider: 256000000 1110.2 ms / 512 img 267.0 GB (the default, worst measured) 512000000 1092.2 267.1 1024000000 1086.8 266.0 2048000000 1086.7 266.3 Worth -23.4 ms per 512 images with peak memory very slightly down, measured with the two smaller sizes each run twice; replicate pairs differ by 0.9 and 3.6 ms. 512000000, the value the MLPerf v6.0 reference moved to, does help but is not the best here. Also corrects a claim in the optimizer block that this sweep happened to measure. The comment asserted FP32 optimizer states "cost nothing in throughput here", which was never measured and is not true: two arms differing in nothing but those five lines give 1110.2 ms per 512 images for FP32 against 812.4 for BF16, +36.7%, plus about 21 GB of peak memory. The setting stays, because convergence comparability against FP8 is the whole point of this recipe and time to train is not throughput -- but it is the most expensive line in the file and the comment now says so.
The two commits before this one justified sweeping ddp_bucket_size separately for this recipe on the grounds that it reduces fp32 gradients where the development recipe reduces bf16, so an identical element count would be twice the bytes per call. That reasoning is wrong, and the config is what misled me: main_grads_dtype reads fp32 here, but it belongs to the precision-aware optimizer and is inert when use_precision_aware_optimizer is false, which it is. The resolved runtime arguments say so plainly. Both recipes log DistributedDataParallelConfig(grad_reduce_in_fp32=False, ...) and both build 43 buckets with identical element counts at the same bucket size, so the collective path -- the only thing ddp_bucket_size acts on -- is the same in both. Nothing about the bucket optimum could have differed for the reason I gave, and the two recipes agreeing on 1024000000 is the expected outcome rather than the lucky one. None of the measurements change. The sweep in this regime stands on its own numbers, and re-running it was still the right call, just for a weaker reason than the one recorded: this recipe's optimizer costs 36.7% of step time, so the balance between exposed collective time and everything else is different enough to be worth scoring rather than assuming. The +36.7% figure needs its attribution fixed for the same reason. It is the cost of FP32 m/v and the standard optimizer path, not of FP32 gradients, since the gradients DDP reduces are bf16 in both recipes.
The three numbers 8688146 and 4bc9bc6 put in this file are wrong, and wrong for a dull reason: I averaged the per-iteration times with a throwaway script instead of the harness that every other number in this campaign came from. Each arm contains one profiled iteration of about 3700 ms against a ~800 ms step, the filter I thought I had written was a no-op, and the result was ~290 ms of pure artifact added to all six arms uniformly. Scored with summarize_mbs.py, which drops the profiled steps: 256000000 819.3 ms / 512 img 267.0 GB (the default, still worst) 512000000 806.4 267.1 1024000000 801.9 266.0 2048000000 801.1 266.3 The decision does not change -- 1024000000 is still the optimum and still ships -- but the size does: -17.4 ms per 512 images, not the -23.4 claimed. Which is a better result than it looks, because -17.4 ms is exactly what the development recipe measured, to the decimal. Two regimes that share a collective path agreeing to 0.1 ms is the confirmation that the inflated numbers were quietly denying. The bigger correction is to the optimizer block above. 8688146 overturned a comment saying FP32 states "cost nothing in throughput here", calling it unmeasured and wrong at +36.7%. The +36.7% was the artifact, not the comment. Measured properly the block costs +6.9 ms per 512 images, +0.85%, at both bucket sizes -- so the original claim was right and I should not have touched it. It is restored, with the measurement now attached, and with the cost that is real called out instead: +16.2 GB of peak memory, which matters at 267 of 288 GB. Worth recording how this was caught, because it was not caught by review. It surfaced when someone asked where the fp32-gradient claim in 4bc9bc6 came from. Chasing that question to the resolved runtime args turned up the inert main_grads_dtype, and re-deriving the numbers through the standard harness then turned up the parse bug underneath it. One question about a stray detail, two errors.
Brings the two-node MLPerf work onto the branch the MXFP6 stack is pinned to, so one branch carries both the single-node optimisation work and the multinode recipe. No file is touched by both sides, so the merge is textually clean. What arrives: ac21513 RCCL resolves the node topology XML by GPU PCI device id, widening the inter-node plan from 2 channels to 8 and inter-node bandwidth from 46 to 118 GB/s. NCCL_NET_GDR_LEVEL=LOC must stay coupled to it. 4cdb4ac resets Megatron's DDP grad-ready calibration after the MLPerf warmup steps consume it, so the first real step calibrates its own golden counts instead of inheriting a synthetic step's. 6876783 clones the diffusion validation loss out of Megatron's in-place rescale path. This one is a correctness fix with teeth: the reported loss was divided by the microbatch count, so it only showed at more than one microbatch per rank, and under mlperf_mode the halved value is what the convergence gate reads. d7b0048 the two-node recipe itself, at MBS 32 / GBS 512 over 16 ranks, with nine converged seeds behind it. Two interactions the merge does not resolve, deliberately: The two-node recipe carries ddp_bucket_size 256000000, the old default, because it predates the sweep that moved the single-node recipes to 1024000000. That result does not transfer: the sweep measured 8 ranks inside one node, and the two-node shape puts reduce-scatter on the inter-node fabric, where the overlap that made 1024M win is a different quantity. It needs measuring on the target cluster before the value is changed. See mbs32_e2e/README.md for the method. It also sets use_precision_aware_optimizer true with main_grads_dtype bf16, where the single-node MLPerf recipe sets the flag false and its main_grads_dtype is therefore inert. The two recipes are in different optimizer regimes and should not be assumed comparable field by field. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds MXFP6 (E2M3 with E8M0 block-of-32 scales) as a precision for Flux on gfx950,
via Primus-Turbo and AITER's A6W6 kernels, plus the MLPerf recipe that runs a
submission campaign at it.
A new local spec provider rather than an option on the MXFP4 one. MXFP6 has
one backend, a mandatory Hadamard rotation folded into the packer, and no
preshuffle choice, so MXFP4's preshuffle contract and recipe flags have nothing
to guard and are dropped rather than retargeted.
fp6is a separate configfield for the same reason: overloading
fp4: "mxfp6"would let a config selectthe MXFP4 provider while skipping its preshuffle check and produce a
plausible-looking run in the wrong precision. Megatron never sees
fp6, soBaseDiffusionConfigrejects it alongsidefp4,fp8, andmxfp4_to_fp8_switch_iter.The MLP is one autograd Function, so its activation never reaches HBM.
Splitting it in two forces the activation to exist: fc2's forward needs a real
tensor and fc1's backward needs one back. Owning fc1 through fc2 in a single
Function lets the MXFP6 packer take the epilogue as a prologue in both
directions, packing straight out of LDS and never assembling the pre-activation
gradient at all. Worth 74.1 to 78.0 images/s/GPU on Flux 12B at micro_batch_size
64. It engages through an
mlp_module()hook the Flux specs look for and fallback from, so no other backend is affected, and anything it cannot reproduce
exactly falls back per module.
Q and K are realigned onto V's dtype after the QK-norm. On the v26.5 image
the norm hands back its fp32 accumulation dtype under
torch.compile, and onlyQ and K go through it, so Turbo's dense flash-attention backends see a dtype
they reject. The dispatcher reports that as "No compatible backend found",
listing shapes and never mentioning dtype. The casts stay inside the compiled
region where they fuse into the norm's epilogue.
The fused-MLP tests pin the mode instead of reading it from the shell. The
fallback tests assert that an unusable configuration warns and defers, which is
autobehaviour; underonthe same configuration raises. The submissioncontainer exports
PRIMUS_MXFP6_FUSED_MLP=onso a fallback there cannot passunnoticed, which meant these tests failed in the one environment where running
them matters most.
The MLPerf recipe is the FP8 one with the precision block swapped, so
comparing the two compares numerics rather than two independently drifted
configurations. It cannot produce a submittable log yet and says so: the
compliance checker accepts a fixed vocabulary for
lowest_numerical_precision_in_linearandmxfp6is not in it, so runs arevalid runs whose logs are rejected on that one key until the format is accepted
upstream.
Stacking
Based on
fix/flux-eval-correctness(#1055) rather thanmain, so this diff isthe MXFP6 work alone. Needs retargeting to
mainonce that merges.Note the eval RNG re-keying in #1055 changes which noise every evaluation draws,
so the MXFP6 recipe added here inherits new eval numerics and its convergence
evidence needs regenerating alongside the FP8 recipes'.
Test plan
gradients at ~28 dB SNR against BF16, and the linear bit-identical to
Primus-Turbo's
FP6GemmMXFunctionin both directionsfp6rejections againstfp4,fp8, andmxfp4_to_fp8_switch_iterkernels leave, 35.5 ms/step of prologue cost joins the packer, GPU busy at
96.8% so it lands as wall clock
pre-existing failures in synthetic-dataset and FSDP2-optimizer tests unchanged
pre-commit runover the branchpath, and an aiter lacking them fails at startup rather than quietly training
in BF16
change noted above