You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Qwen3-30B-A3B FP8 pretraining on 8× MI355X goes from a memory fault at iteration 3 to 40,283 tokens/s/GPU (6507.6 ms/iter), 47% above the BF16 config. This PR switches the FP8 example config to Primus-Turbo tensorwise FP8 and turns on every Turbo fusion that got there. Two of those fusions are new in this PR:
head_dim 128 for the fused QKV split + Q/K RMSNorm + RoPE patch. The patch was GPT-OSS-only before.
turbo_fp8_permute: the DeepEP permute hands the fused grouped MLP FP8 tokens, and takes its gradient back in FP8.
The PR also fixes enable_turbo_attention_float8, which failed on its first call.
Primus-Turbo requirements.
turbo_fp8_permute needs a Primus-Turbo whose moe_permute takes quantize_dtype (AMD-AGI/Primus-Turbo#555). With an older Turbo, the first dispatch fails.
The head_dim 128 fusion needs AMD-AGI/Primus-Turbo#556. Without it, the patch keeps the unfused path.
The other Turbo rows in the table below come from #554, #557 and #558. They need no Primus change.
End-to-end. Qwen3-30B-A3B FP8 pretrain on 8× MI355X (EP8, MBS 8, GBS 512, seq 4096, even routing), mean of iterations 11–20 of 20-iteration runs. Each row adds one change on top of the row above:
From the FP8 tensorwise config to the last row: −26.8% ms/iter, +36.7% tokens/s/GPU. Loss differences are within run-to-run noise (about 4e-4 at iteration 20). This branch, rebased on main 0a68e5c, with the switches on: 6518.7 ms/iter, 40,214 tokens/s/GPU, loss@20 11.33415.
KNOWN_ISSUE
Changes
fix(turbo): make enable_turbo_attention_float8 runnable again
Pass sink= only when a sink tensor exists; flash_attn_fp8_func has no sink parameter.
Make q/k/v bshd-contiguous on the FP8 path. flash_attn_fp8_func permutes sbhd storage as if it were the logical layout (fixed in Turbo by fix/fp8-attn-strided-layout).
perf(qwen3): tune Qwen3-30B-A3B FP8 MI355X config to Turbo tensorwise
examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml: fp8_recipe: tensorwise, use_turbo_gemm / use_turbo_grouped_gemm, and the TE grouped-MLP path instead of the legacy grouped GEMM.
perf(megatron): enable the fused packed-QKV RMSNorm + RoPE patch for head_dim 128
Accept any head_dim in Turbo's QK_RMSNORM_ROPE_HEAD_DIMS, falling back to (64,) on older Turbo.
Accept TE RMSNorm as well as PrimusTurboRMSNorm. The fused path only reads weight / eps; zero_centered_gamma is still rejected.
Fused QKV split + Q/K RMSNorm + RoPE at the Qwen3 attention shape: fwd+bwd 1243.6 → 369.9 µs (3.36×).
perf(megatron): hand the Turbo fused grouped MLP FP8 tokens from the DeepEP permute
New opt-in flag turbo_fp8_permute (default false in primus_turbo.yaml).
Under Turbo FP8 tensorwise current scaling, PrimusTurboDeepEPTokenDispatcher passes quantize_dtype to _post_dispatch and grad_quantize_dtype to _pre_combine. The dtypes follow the FP8 format (HYBRID: e4m3 input, e5m2 gradient), and results are bit-identical to the default path. Any other recipe, or a layer with Turbo FP8 off, keeps the bf16 path.
Argument validation requires enable_primus_turbo, use_turbo_deepep and turbo_fused_grouped_gemm, and rejects selective recompute of moe_act.
Dispatcher kernel time per call at the Qwen3 EP8 shape: forward 645 → 167 µs, backward 687 → 174 µs.
perf(qwen3): turn on the Turbo fusions in the Qwen3-30B-A3B FP8 MI355X config
Top-level env: PRIMUS_FUSED_QK_RMSNORM_ROPE: "1".
turbo_fused_grouped_gemm, use_turbo_fused_act_with_probs and turbo_fp8_permute set to true.
turbo_deepep_num_cu 80 → 160. 192 gave a NaN loss on one rank at iteration 5; 256 is slower (7028 ms).
The force-"even" routing fix that this work also needed is already on main (#1243), so it is not part of this PR.
Tests
New in tests/unit_tests/backends/megatron/test_rocm_arg_validation.py: validate_turbo_fp8_permute (each required flag, selective recompute of moe_act, disabled no-op).
New tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py: _fp8_permute_dtypes returns (e4m3, e5m2) for HYBRID and (e4m3, e4m3) for E4M3 tensorwise, and (None, None) with the flag off, Turbo FP8 off or blockwise scaling.
test_qk_rmsnorm_rope_patches.py, test_rocm_arg_validation.py, test_validate_args_patches.py, test_router_force_even_routing.py, test_turbo_fp8_permute_dtypes.py: all pass on MI355X.
…ripts
Handoff for the 8x MI355X Qwen3-30B-A3B FP8 tensorwise campaign
(29,474 -> 40,283 tokens/s/GPU): cumulative e2e table, the Primus switches and
Turbo/Primus branches the final stack needs, PR order and ready-to-paste PR
bodies, how to rebuild the stack, rejected options and open items, plus the
round log, the microbenchmark / trace / e2e scripts and the task prompt.
Adds an archival handoff for the Qwen3 FP8 optimization campaign, including benchmarks, diagnostics, run orchestration, results, and draft Turbo PR descriptions. It does not include the Primus implementation advertised in the PR description.
Changes:
Documents performance findings and required optimization branches.
Adds profiling, benchmarking, and correctness utilities.
Adds end-to-end experiment scripts and draft upstream PR bodies.
run always returns the successful status of this echo, so the && attribution chain continues and can report completion after a failed training run. Capture the command status before logging and abort the script on failure so later benchmark rows are not attributed to an invalid base.
Unstable optimization is enabled by default despite NaN loss
The final configuration should not be enabled by default while identical runs still produce NaN loss (1/13 at the selected 160-CU setting). This is a training-correctness failure, not benchmark jitter; keep the unstable optimization opt-in or revert to the known-stable setting until the proposed stress test isolates the cause and a regression test covers it.
The parser records "-" when an iteration line has no loss, but then unconditionally converts that sentinel to float, causing the summary command to fail instead of reporting the missing value. Format the loss conditionally, as already done conceptually for absent fields.
Qwen3-30B-A3B FP8 pretraining on 8× MI355X goes from a memory fault at iteration 3 to **40,283 tokens/s/GPU** (6507.6 ms/iter), 47% above the BF16 config. This PR switches the FP8 example config to Primus-Turbo tensorwise FP8 and turns on every Turbo fusion that got there. Two of those fusions are new in this PR:
- **head_dim 128 for the fused QKV split + Q/K RMSNorm + RoPE patch.** The patch was GPT-OSS-only before.
- **`turbo_fp8_permute`**: the DeepEP permute hands the fused grouped MLP FP8 tokens, and takes its gradient back in FP8.
| `perf/qwen3-30b-a3b-fp8-turbo` | bef41e8a | opening | the one Primus PR, on main 0a68e5cd; body: [pr_bodies/primus-perf-qwen3-30b-a3b-fp8-turbo.md](pr_bodies/primus-perf-qwen3-30b-a3b-fp8-turbo.md) |
It holds five commits: the FP8 attention fix, the tensorwise config, the head_dim 128 qk-norm patch, `turbo_fp8_permute` (with new unit tests), and the config turning every switch on. It replaces `perf/qwen3-30b-a3b-tuning`, `perf/qwen3-qk-rmsnorm-rope-hd128` and `perf/megatron/turbo-fp8-permute`. Their force-"even" routing fix is already on main (#1243).
Record the Primus DISABLE_CHEAP_FENCE config fix and the Turbo high-CU
fence default branch for the intermittent forward-loss NaN.
Co-authored-by: Cursor <cursoragent@cursor.com>
The timed callbacks reuse leaf tensors without clearing their gradients, so after warmup every measured backward includes large AccumulateGrad additions into existing buffers. This inflates the reported permute/unpermute backward timings independently of the backend; use autograd.grad to measure only the operation's backward path.
Update GLM5 documentation to match the executed batch size
The documented GLM5 shape says B=4, but the benchmark actually runs B=8 in SHAPES. This makes the advertised workload and any comparison based on it ambiguous; update the docstring to match the executed shape.
🧠 Review effort: Balanced
This branch has not been deployed
No deployments
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
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.
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Description
Qwen3-30B-A3B FP8 pretraining on 8× MI355X goes from a memory fault at iteration 3 to 40,283 tokens/s/GPU (6507.6 ms/iter), 47% above the BF16 config. This PR switches the FP8 example config to Primus-Turbo tensorwise FP8 and turns on every Turbo fusion that got there. Two of those fusions are new in this PR:
turbo_fp8_permute: the DeepEP permute hands the fused grouped MLP FP8 tokens, and takes its gradient back in FP8.The PR also fixes
enable_turbo_attention_float8, which failed on its first call.Primus-Turbo requirements.
turbo_fp8_permuteneeds a Primus-Turbo whosemoe_permutetakesquantize_dtype(AMD-AGI/Primus-Turbo#555). With an older Turbo, the first dispatch fails.End-to-end. Qwen3-30B-A3B FP8 pretrain on 8× MI355X (EP8, MBS 8, GBS 512, seq 4096, even routing), mean of iterations 11–20 of 20-iteration runs. Each row adds one change on top of the row above:
use_turbo_fused_act_with_probsget_dispatch_layoutDeepEPTokenDispatcherturbo_deepep_num_cu80 → 160turbo_fused_grouped_gemm(mean of 2 runs)turbo_fp8_permute(mean of 4 runs)From the FP8 tensorwise config to the last row: −26.8% ms/iter, +36.7% tokens/s/GPU. Loss differences are within run-to-run noise (about 4e-4 at iteration 20). This branch, rebased on main 0a68e5c, with the switches on: 6518.7 ms/iter, 40,214 tokens/s/GPU, loss@20 11.33415.
KNOWN_ISSUE
Changes
fix(turbo): make enable_turbo_attention_float8 runnable againprimus/backends/megatron/core/extensions/primus_turbo.py:sink=only when a sink tensor exists;flash_attn_fp8_funchas nosinkparameter.flash_attn_fp8_funcpermutes sbhd storage as if it were the logical layout (fixed in Turbo byfix/fp8-attn-strided-layout).perf(qwen3): tune Qwen3-30B-A3B FP8 MI355X config to Turbo tensorwiseexamples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml:fp8_recipe: tensorwise,use_turbo_gemm/use_turbo_grouped_gemm, and the TE grouped-MLP path instead of the legacy grouped GEMM.perf(megatron): enable the fused packed-QKV RMSNorm + RoPE patch for head_dim 128primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py:QK_RMSNORM_ROPE_HEAD_DIMS, falling back to(64,)on older Turbo.PrimusTurboRMSNorm. The fused path only reads weight / eps;zero_centered_gammais still rejected.perf(megatron): hand the Turbo fused grouped MLP FP8 tokens from the DeepEP permuteturbo_fp8_permute(default false inprimus_turbo.yaml).PrimusTurboDeepEPTokenDispatcherpassesquantize_dtypeto_post_dispatchandgrad_quantize_dtypeto_pre_combine. The dtypes follow the FP8 format (HYBRID: e4m3 input, e5m2 gradient), and results are bit-identical to the default path. Any other recipe, or a layer with Turbo FP8 off, keeps the bf16 path.enable_primus_turbo,use_turbo_deepepandturbo_fused_grouped_gemm, and rejects selective recompute ofmoe_act.perf(qwen3): turn on the Turbo fusions in the Qwen3-30B-A3B FP8 MI355X configenv: PRIMUS_FUSED_QK_RMSNORM_ROPE: "1".turbo_fused_grouped_gemm,use_turbo_fused_act_with_probsandturbo_fp8_permuteset to true.turbo_deepep_num_cu80 → 160. 192 gave a NaN loss on one rank at iteration 5; 256 is slower (7028 ms).The force-"even" routing fix that this work also needed is already on main (#1243), so it is not part of this PR.
Tests
tests/unit_tests/backends/megatron/test_rocm_arg_validation.py:validate_turbo_fp8_permute(each required flag, selective recompute ofmoe_act, disabled no-op).tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py:_fp8_permute_dtypesreturns (e4m3, e5m2) for HYBRID and (e4m3, e4m3) for E4M3 tensorwise, and (None, None) with the flag off, Turbo FP8 off or blockwise scaling.test_qk_rmsnorm_rope_patches.py,test_rocm_arg_validation.py,test_validate_args_patches.py,test_router_force_even_routing.py,test_turbo_fp8_permute_dtypes.py: all pass on MI355X.