examples/llama: default-on fusion & DP-comm perf toggles - #141
Merged
Conversation
… precision-tagged log dir Add env-var toggles (default on, set 0 to disable) for the fusion and data-parallel-comm optimizations, routed to the arg group matching training/arguments.py (GPT_ARGS for transformer/mixed-precision, EXTRA_ARGS for distributed): - GRADIENT_ACCUMULATION_FUSION - GRAD_REDUCE_IN_BF16 - DDP_AVERAGE_IN_COLLECTIVE - CROSS_ENTROPY_LOSS_FUSION (+ CROSS_ENTROPY_FUSION_IMPL=te) - FUSED_SINGLE_QKV_ROPE Measured ~+3.5-7% throughput across bf16/fp8/fp4 on MI355X (llama3-8B and llama2-7B), no errors; mxfp4 path converges to the MLPerf eval_loss 3.3 target. GRAD_REDUCE_IN_BF16 lowers gradient all-reduce precision (validated to converge); set 0 if convergence-sensitive. train_llama3.sh: mxfp4 now defaults to all layers in FP4 (FP4_SELECTIVE_BF16=0, MLPerf recipe); nvfp4 keeps the first/last-bf16 paper recipe. Log directory is now tagged by precision (bf16 | fp8 | fp8_tensorwise | mxfp8 | mxfp4 | nvfp4) instead of the misleading TE_FP8_<0/1>. arguments.py: expose --fused-single-qkv-rope (drop it from the ArgumentGroupFactory exclude list). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
wenchenvincent
changed the title
…ith fusion toggles) BF16 gradient all-reduce is unsafe in combination with the fusion toggles at the MLPerf reference LR. Bisection on C4 (LLaMA3-8B, BS=32, lr 8e-4) showed the full default-on toggle set plateaus (val ~5.6, grad norm climbing to ~8); a leave-one-out bisect found that removing ONLY grad_reduce_in_bf16 restores convergence (val 3.86 @ 1536 iters, grad norm 0.32), matching the all-off baseline. The other four fusion toggles are convergence-neutral and stay on. Flip the default to 0 in train_llama3.sh and train_llama2.sh and document the finding inline. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ts BF16 only) The previous commit defaulted GRAD_REDUCE_IN_BF16 off (FP32 grad buffer) for convergence. That crashes mxfp4: the aiter a4w4 asm wgrad GEMM only emits BF16 output, so with gradient_accumulation_fusion=1 the fused wgrad-into-FP32-buffer path fails with "gemm_a4w4_asm only support BFloat16 output now!". Guard it: when mxfp4 + gradient_accumulation_fusion are both on, force the grad buffer back to BF16. Validated to converge (LLaMA3-8B C4, BS=32, lr 8e-4) at 1691 TFLOP/s/GPU, grad norm bounded. bf16/fp8 are unaffected (FP32 buffer). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…he cause, not this toggle My earlier commit (f873b67) defaulted GRAD_REDUCE_IN_BF16 off, claiming it caused the lr-8e-4 plateau. That bisection was confounded: the "all-on plateaus" baseline used adam-eps=1e-8 while the leave-one-out runs used adam-eps=1e-5, so the change that actually mattered was the optimizer epsilon, not this toggle. A controlled 2x2 on C4 (LLaMA3-8B, BS=32, lr 8e-4, bf16) settles it: - all-on, eps=1e-8 -> PLATEAU (val ~5.9, grad norm -> 8.8) - all-on, eps=1e-5 -> CONVERGE (val 3.84) - all-but-grad_reduce eps=1e-8 -> PLATEAU (val 5.95, grad norm -> 8.8) - all-but-grad_reduce eps=1e-5 -> CONVERGE (val 3.84) Convergence tracks adam-eps exactly; grad_reduce_in_bf16 has no effect. Restore the original default ON (convergence-neutral perf win; mxfp4 needs it on anyway for the a4w4 BF16-only wgrad GEMM) and fix the comments. The mxfp4 guard from 972faa6 stays (now phrased around an explicit user override to 0). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…note The MLPerf MI355X llama3.1-8B config explicitly sets adam_eps: 1.0e-5 at lr 8e-4 (overriding the Megatron/Primus default 1e-8). Cite it so the convergence note points at the authoritative value rather than just "use 1e-5". Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.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.
Opt-in (default-on) env toggles for fusion + DP-comm optimizations in the LLaMA example scripts, plus mxfp4 fixes and two nits. ~+3.5–7% throughput across bf16/fp8/fp4 on MI355X, convergence-neutral.
Changes
train_llama2.sh,train_llama3.sh) — default on, set0to disable:GRADIENT_ACCUMULATION_FUSION,CROSS_ENTROPY_LOSS_FUSION(+CROSS_ENTROPY_FUSION_IMPL=te),FUSED_SINGLE_QKV_ROPE,GRAD_REDUCE_IN_BF16,DDP_AVERAGE_IN_COLLECTIVE. Each routed to the arg bucket matching itsarguments.pygroup.train_llama3.sh):FP4_SELECTIVE_BF16is recipe-aware — mxfp4 uses all FP4 layers (MLPerf), nvfp4 keeps first/last in bf16 (paper). Added a guard that forcesGRAD_REDUCE_IN_BF16=1for mxfp4: the aiter a4w4 asm wgrad GEMM only emits BF16, so the fused-wgrad grad buffer must be BF16 — otherwise the run crashes withgemm_a4w4_asm only support BFloat16 output now!.bf16|fp8|fp8_tensorwise|mxfp8|mxfp4|nvfp4) instead of the misleadingTE_FP8_<0/1>.--fused-single-qkv-rope.Convergence
All toggles are convergence-neutral. A controlled 2×2 on C4 (LLaMA3-8B, BS=32, lr 8e-4, bf16) showed convergence is gated by
adam_eps, not these toggles:Convergence tracks
adam_epsexactly;grad_reduce_in_bf16has no effect. This matches the MLPerf MI355X llama3.1-8B config (primus/.../MI355X/llama3.1_8B-pretrain-FP8.yaml), which setsadam_eps: 1.0e-5at lr 8e-4. Use--adam-eps 1e-5at this LR.Validation (8×MI355X)
Toggle throughput (mean TFLOP/s, off→on):
End-to-end convergence to val loss ≈3.3 (C4, lr 8e-4,
adam_eps 1e-5, default toggles):(mxfp4 follows the same loss curve but slower per-iter; it reaches 3.33 by iter 6144 and would cross 3.3 with a few hundred more — partly offset by its ~1.8× throughput.)
🤖 Generated with Claude Code