Conversation
The provider reconstructed the sample count from args.micro_batch_size, but a recipe that overrides the evaluation width evaluates at that width instead -- get_eval_num_microbatches and assert_val_worker_divisibility both already read it. So the reconstruction understated the count by the ratio between the two. An eval width wider than the training one therefore under-reported the budget, and at a large enough ratio the understated count was no longer divisible by the divisor the assert builds from the eval width, so a shape that covers the split exactly was rejected before construction. Read get_eval_micro_batch_size instead, and log the width actually used.
The dual packer produces row and column blobs together, but the column blobs are backward's operands and nothing else reads them. A forward with no backward behind it pays for half a packer it will never use. Evaluation is entirely such a region, and it is where packing hurts most, because one pack no longer amortizes across fwd/dgrad/wgrad. Row-only packing is measurably cheaper than dual across the Flux shapes. grad_enabled has to be sampled by the caller. Inside forward, PyTorch has already cleared grad mode, so torch.is_grad_enabled() reads False even in training, while ctx.needs_input_grad reads True even under no_grad; neither distinguishes the two. Getting it wrong fails loudly in backward on a None operand rather than silently returning a wrong gradient.
val_num_workers was pinned to 0 because a nonzero training num_workers leaves part of the set unread: Energon splits the data across dp_size * num_workers but the batch quota across num_workers alone. 1 is the one nonzero value that cannot do that. Energon clamps the data split to max(1, num_workers), so 1 shards exactly as 0 does and covers the split identically, while 0 additionally forgoes prefetch by reading in the main process. Measured on one node over three replicates each, 1 evaluates the whole split faster than 0 with the val_loss bit-identical. 2 and 4 were no faster than 1 and did shift the loss, by resharding the data across more workers than the split has shards.
Two measurement affordances, both off by default. MXFP6_EVAL_PROFILE wraps the evaluation loop in a torch profiler, writing to MXFP6_EVAL_PROFILE_DIR. Evaluation runs a different kernel mix from training -- forward only, no packer amortization -- so it needs its own trace rather than a window inside a training profile. MXFP6_FORCE_ROPE_FUSION exists because the config key alone cannot turn that path on. Megatron's validate_args clears args.apply_rope_fusion whenever position_embedding_type != "rope" (arguments.py L1232-1233), and Flux never sets that type, so params always arrives as False no matter what the YAML said.
Two fusions on the Flux 12B MXFP6 step, both gated off by default. fused_norm_rope.py folds the QK RMSNorm and the interleaved RoPE rotation into one Triton kernel per direction, for the single blocks and for both streams of the joint blocks. Q and K are read as strided slices of mixed_qkv at stride 3D rather than forced contiguous first, and the whole-QKV entry point moves the split inside the op so backward writes d(mixed_qkv) directly instead of concatenating three tensors. V is bit-exact in both directions; it is copied, never rotated. The QKV-level op is deliberately @custom_op rather than @triton_op. Under @triton_op, Inductor traces in and functionalizes two Triton kernels writing disjoint columns of one output buffer, producing a kernel that read 786304 elements from a 262144-element clone and faulted. Opacity costs nothing there: the input is a GEMM output and the outputs go straight to FMHA, so there are no neighbours to fuse with. normalization.py adds a single-pass LN-modulate backward, one kernel where there were two, which removes a redundant HBM pass over grad_output and x. It scores identically to the two-kernel path against an fp64 reference. Both fusions reduce end-to-end step time, reproduced across two independent 1000-step runs with bit-identical final loss. Measurements are recorded internally rather than here. The two gates are environment variables for now and become FluxConfig keys before this leaves draft.
jasainio
force-pushed
the
feat/mxfp6-flux-fusions
branch
from
September 6, 2026 17:53
b9627db to
f6571be
Compare
MXFP6_ABLATE_ROPE skipped the rotation to put a ceiling on what the RoPE family was worth. It answered that question, and what it leaves behind is a switch that silently produces wrong numerics sitting in a production file, guarding four call sites in two classes. The probe moves to a patch in the internal scratch tree, where turning it on is a deliberate act with a visible diff rather than an environment variable someone can inherit from a shell. Nothing here was reachable without setting it, so this is a no-op for every run that was not already a timing probe.
…mxfp6-stageb-checkpoint-stack
## Summary The Flux distributed-checkpoint round-trip test was added to a `unittest.TestCase` subclass but requested pytest's `tmp_path` fixture. Pytest therefore called the method without that argument and failed before exercising checkpoint I/O. - Replace the fixture argument with a scoped `TemporaryDirectory`. - Keep the checkpoint directory alive through save, load, and tensor comparison. ## Test plan - [x] `isort --check-only` - [x] `black --check` - [x] `py_compile` - [x] Run the corrected distributed-checkpoint tests in the v26.5 development container (12 passed).
wrap_triton costs 0.359 ms of host time per launch. The same kernel with the same arguments, launched directly, costs 0.031 ms. The whole-QKV path was paying that 12.7x dispatch overhead four times per layer, which left the backward host-bound rather than GPU-bound: 0.8555 ms of host enqueue against 0.8593 ms of batched wall time, so the GPU sat idle waiting for Python for essentially the whole call. Repeated across every QKV call a step makes in each direction, that accumulates into a large share of the step's host time. wrap_triton exists so Inductor can trace through a @triton_op into the kernel it launches. `_launch_fwd`/`_launch_bwd` are shared by both entry points, and the whole-QKV ops are @custom_op -- opaque by construction, nothing to trace through -- so there the wrapper buys nothing and still routes every launch through the higher-order-op dispatch path. The helpers now take `traceable`: the two @triton_op call sites keep the wrapper they need, the two @custom_op ones drop it. Same kernel and same arguments, so this is checked as an identity rather than a tolerance: `_launch_bwd`'s partials and `_launch_fwd`'s out and rstd are bit-identical either way. All three scratch gates pass unchanged including every compiled arm, which is what confirms the @triton_op path is still traceable and the @custom_op path still compiles while opaque. The 84-config @triton.autotune was the obvious suspect and is innocent; the direct launch goes through the same autotuner. Found while benchmarking something else: the backward measured 0.8637 ms at seq 512 and 0.8454 at seq 256, a 2% difference for half the work, while the packer arm beside it halved correctly. A cost that does not scale with the problem is not doing the work.
…r stream _slice_rope sliced q and k separately even when they are the same tensor, which is what self-attention passes. Slicing produces two objects, so the `k_pos_emb is q_pos_emb` test that the whole-QKV norm+RoPE op uses to decide whether one frequency table serves both tensors always failed for a sliced table. Only the joint blocks slice, so MXFP6_FUSED_QK_ROPE's whole-QKV op has never engaged in one: it declined and fell back to the per-tensor kernel, silently and without changing any result. That flag's published numbers were all measured with the joint blocks on the fallback path and need re-checking against this.
…into one prologue MXFP6QKVNormRopeFunction spans linear_qkv -> norm -> rope, so the packer computes d(mixed_qkv) during its staging read instead of the norm+RoPE backward writing it to HBM for the quantizer to read straight back. Only a single autograd Function can feed a prologue, which is why the Function has to own the whole chain. The forward is unchanged -- the same projection and the same Triton norm+RoPE -- so its outputs stay bit-identical and only the backward differs. Behind MXFP6_FUSED_QKV (auto/on/off, default off). Both block types are wired, and the forward skips the getters on the fused path so every existing fallback keeps its code. Two checks rather than one: the module's configuration at build time, and this step's tensors at runtime. The norm-weight dtype has to be the runtime one -- torch.nn.RMSNorm is constructed without a dtype and so starts fp32, and Float16Module casts it later, so a check in __init__ would decline in exactly the configuration this is for. Also here, both for measurement rather than for the fusion: PRIMUS_NORM_ROPE_PIN collapses the norm+RoPE autotune spaces to one config each. The key is ["M", "D"], so each distinct M pays a full sweep per block, which cost ~25 minutes of warmup once the whole-QKV path reached all 57 blocks. The sweep earns none of it: the best eight configs are within 2.8% and the winner is not stable between runs, so it also moves the number an A/B is trying to measure. MXFP6_ROPE_SLICE_LEGACY reinstates the pre-fix _slice_rope, so the published baseline can be reproduced from this tree and the identity fix priced separately from the prologue. Scaffolding around a known bug; delete it once that is done.
… branch The MXFP6 container pins one Primus ref, and until now the fusion work and the Stage B checkpoint stack were two branches diverging from 63b9c33, so pinning either dropped the other. This merge makes one ref carry both: the QKV prologue, the _slice_rope identity fix and the wrap_triton removal from this side, and the heterogeneous Flux checkpoint schema of #1116 with its test fix in #1117 from Stage B. Automatic merge with no conflicts -- the two touch disjoint files, fusion in the diffusion attention and MXFP6 extension paths, Stage B in checkpointing. The fusion gates stay off by default, so this changes no behaviour for anyone pinning the merge who does not set them.
…costs Evaluation carried two one-off costs inside the timed region, both repeating rather than amortising, and neither visible to a within-iteration profiler. Compile. Evaluation runs forward-only under no_grad with model.eval(), a different Dynamo guard state from anything the training warmup traces, so all 57 per-block graphs were retraced at the first evaluation -- after run_start. warmup_validation_steps compiles and first-touches the evaluation graph before the clock starts, at get_eval_micro_batch_size rather than the training width, since warming the wrong shape compiles graphs the evaluation will not use. It runs last, after every restore, because forward-only work cannot perturb the state those restores rebuilt. RNG state is saved and restored around it: evaluation noise is drawn from the RNG rather than carried by the val shards, so a warmup that draws from it shifts val_loss, and val_loss is what run_stop is gated on. First fetch. The first fetch of every evaluation is far slower than the rest, which the loader hides. The loop consumes the split exactly, so every evaluation ends on an epoch boundary and the next one's first fetch pays the restart with nothing queued across it. Pulling that batch when an evaluation ends would only move the cost inside the same timed region, so it is issued on a background thread and the training steps that follow absorb it. The consuming side joins before the loop touches the iterator, since a DataLoader iterator is not safe to advance from two threads. The batch it pulls is the one the next evaluation would have read first, and the eval noise is keyed to the eval step index rather than to the fetch, so val_loss is unchanged -- verified bit-identical on every evaluation of a multi-eval run. MXFP6_VAL_PREFETCH=0 disables it. val_num_workers goes to 1. Anything above 1 can under-read the split; 1 is the one nonzero value that cannot, because Energon clamps the data split to max(1, num_workers), so it shards exactly as 0 does while 0 additionally forgoes prefetch by reading in the main process. Two probes, both off by default, because a window inside the iteration reports a faster iteration than the end-to-end rate over the split implies, and cannot see the difference by construction. MXFP6_EVAL_ITER_TIMING=1 measures the iteration period rather than its activity, sync-free in the hot path so it cannot create the serialization it looks for. MXFP6_BATCH_FETCH_TIMING=1 times next(data_iterator) alone, which is what attributed the slow first iteration to the loader rather than to compile or the allocator.
…efer DDP grad-buffer zero Skip redundant contiguous copies on the Flux QKV and AdaLN paths, and overwrite the first gradient into each main-grad slice instead of filling the full DDP buffer. The DDP change lives in a Primus patch so upstream Megatron-LM stays unmodified.
![pr-pundit-code-review[bot]](https://avatars.githubusercontent.com/u/16860860?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 15 files reviewed.
Rules checked: 14 (14 passed, 0 failed).
✅ Passed
All 14 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ⚠️
Verdict: Partial | Categories: config, training, quantization
Test evidence found in this PR:
- In-diff comments claim empirical verification: 'Bit-identical at per_step_rng_reseed: true ... warmed and unwarmed agree on every per-iteration loss'
- In-diff comment claims 'val_loss is unchanged between the two' for val_num_workers 0 vs 1
- In-diff comment cites measurement methodology: 'Measured on one node over three replicates each'
- Reference to an existing runtime check assert_val_worker_divisibility accepting the new value (29696 % (32 * 1 * 32) == 0)
✅ 7 requirement(s) satisfied
- Statement of whether outputs are numerically equivalent or intentionally changed — the diff explicitly states the change is loss-neutral ('bit-identical at per_step_rng_reseed: true', 'warmed and unwarmed agree on every per-iteration loss', 'val_loss is unchanged between the two')
- Measurement methodology is disclosed for the performance claim ('Measured on one node over three replicates each')
- MoE token-dispatcher create_args()/ROCm flag rule — not applicable; this diff contains no MoE token-dispatcher test or dispatcher init changes
- MLflow artifact/trace upload mock side_effect exception-type rule — not applicable; no Primus-LM MLflow error-handling tests are added or modified
- Unique
tagargument for run_script rule — not applicable; no trainer tests are added or modified in this diff - qkv_format 'sbhd'/'bshd' layout test rule — not applicable; no qkv_format-dependent layout/permutation/reshape logic is touched
- mlflow_artifacts @patch target rule (log_rank_0 / warning_rank_0) — not applicable; no patches of mlflow_artifacts helpers appear in this diff
Strong suggestions
🟡 Loss curve / convergence comparison before vs. after the change, backed by reproducible artifacts (log excerpts, per-iteration loss table, or a checked-in comparison script) rather than only prose claims in config comments
The PR changes validation-time behavior (adds validation warmup steps and changes val_num_workers 0 -> 1) in an MXFP6 quantized MLPerf training config, which falls under training_correctness. The diff asserts 'bit-identical at per_step_rng_reseed: true' and 'val_loss is unchanged', but these are unverifiable in-code comments; no loss numbers, run logs, or reproducible comparison harness are included, so a reviewer cannot confirm the convergence claim or re-run it.
Rule: Loss curve or convergence check before/after
Source: repo evidence requirement
🟡 An automated correctness test on a small model run (e.g., a short trainer smoke test that exercises the warmup-validation path and val_num_workers=1 with the MXFP6 extension and compares losses against the unwarmed/0-worker baseline)
Only a full MLPerf-scale, multi-replicate manual measurement is cited. There is no small-scale, repo-runnable test that exercises the new validation warmup path or the non-zero dataloader worker path, so regressions in the warmup/dataloader interaction (e.g., Dynamo re-compilation, worker sharding, RNG reseeding) would not be caught by CI.
Rule: Correctness test on a small model run
Source: repo evidence requirement
Advisory
🔵 Test/assertion coverage for the validation worker divisibility constraint with the new value (e.g., a unit test asserting assert_val_worker_divisibility passes for 29696 % (32 * 1 * 32) == 0 and fails for an invalid combination)
The diff only references an existing runtime check as accepting the new val_num_workers value. Because the change makes the config newly dependent on that divisibility invariant, a direct test (or at least a cited run demonstrating the check firing/passing) would guard against future config drift silently breaking evaluation sharding.
Rule: Correctness test on a small model run
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
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.
Stacked on
feat/mxfp6-fused-mlp(#1064), which it targets. Two fusions on the Flux 12B MXFP6 step, plus the evaluation-cost work that came out of the same campaign.Both fusions reduce end-to-end step time, reproduced across two independent runs with bit-identical final loss, no NaN and no skipped iterations. Measurements are recorded internally rather than here.
What is here
perf(mxfp6): fuse Flux QK norm with RoPE and the LN-modulate backward— the headline change. QK RMSNorm and the interleaved rotation collapse into one Triton kernel per direction, for the single blocks and both streams of the joint blocks. Q and K are read as strided slices ofmixed_qkvinstead of being forced contiguous, and the whole-QKV entry point moves the split inside the op so backward writesd(mixed_qkv)directly. Separately, the LN-modulate backward becomes one kernel where it was two.feat(mxfp6): add an evaluation profiler and a rope-fusion override— two measurement affordances, both off by default.perf(eval): give the validation loader one prefetching worker— evaluates the whole split faster with val_loss bit-identical.perf(mxfp6): pack rows only when no backward will read the columns— evaluation is a forward-only region and pays for half a packer it never uses.fix(eval): size the validation budget from the eval microbatch— the budget was reconstructed from the training width, which understated it and rejected shapes that cover the split exactly.Before this leaves draft
FluxConfigkeys, enabled in the MXFP6 recipes. Note thatapply_rope_fusionalready demonstrates the failure mode to avoid: Megatron'svalidate_argsclears it for Flux, which is whyMXFP6_FORCE_ROPE_FUSIONexists at all. The new keys need a positive assertion that they arrived, not just a default.tests/unit_tests/backends/megatron/, covering both LN-modulate op variants and the joint-block rope slicing, plus a regression test pinning the QKV op as opaque.MXFP6_ABLATE_ROPEinattention.pyis a deliberate wrong-numerics timing probe and comes out before review.