Add ragged FP8 MoE path and pinned runtime image for GPT-OSS 20B MLPerf - #923
Open
HuangWei-95 wants to merge 3 commits into
Open
Add ragged FP8 MoE path and pinned runtime image for GPT-OSS 20B MLPerf#923HuangWei-95 wants to merge 3 commits into
HuangWei-95 wants to merge 3 commits into
Conversation
HuangWei-95
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
July 27, 2026 13:34
| from transformer_engine.pytorch.attention.dot_product_attention import backends | ||
|
|
||
| backends.fused_attn_fwd = patched_fused_attn_fwd | ||
| except Exception: |
Enable ragged Turbo FP8 MoE execution and pin the MI355X runtime so convergence performance is reproducible without explicit expert padding. Co-authored-by: Cursor <cursoragent@cursor.com>
Megatron rebuilds its dataset index cache whenever the iteration count changes, but the cache lives inside the read-only dataset mount, so any TRAIN_ITERS override aborted during startup. Co-authored-by: Cursor <cursoragent@cursor.com>
The shim logged once per attention call whenever the root logger sat at INFO, so any verbose run silently paid a per-dispatch logging cost instead of honoring FMHA_HD64_ASM_LOG. Co-authored-by: Cursor <cursoragent@cursor.com>
HuangWei-95
force-pushed
the
feat/gpt-oss-mlperf-turbo-runtime
branch
from
July 30, 2026 00:45
71a4cbd to
7730252
Compare
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
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
Recovers the GPT-OSS 20B MLPerf throughput lost to explicit expert padding, and
makes the measured MI355X software stack reproducible.
PrimusGroupedMLPcan now hand the GPUtokens_per_experttensor straight to PrimusTurbo's tensorwise FP8 groupedGEMM, skipping TE's
Fp8Padding/Fp8Unpaddingand the.tolist()devicesync they require. Gated by
use_turbo_ragged_grouped_gemmand by the actualautocast state at forward time, so other recipes and backends keep the
existing zero-padding path.
use_turbo_grouped_gemm_internal_weight_quantization, BF16 expert weights godirectly to
grouped_gemm_fp8instead of being pre-quantized into buffers androuted through the weight-gradient bridge.
bias_act_func_with_masknowruns whether or not activations are recomputed. Previously the fused kernel
was reached only during recompute, so normal training silently fell back to
the weighted SwiGLU path despite
use_turbo_fused_act_with_probs.symmetric-memory API (the old
hipmodule and async-TP entry points aregone), and the previously hardcoded "first two bucket groups stay on RCCL"
rule became
MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS, defaulting to 2 so behavioris unchanged.
Dockerfile.runtimeplusruntime_versions.shpin the base image by digest along with Triton, Primus-Turbo,
TransformerEngine/AITER, and both HD64 causal ASM attention kernels;
verify_runtime.pyasserts the running stack matches.Both new options default to
falseand are validated at startup: they requireenable_primus_turbo,use_turbo_grouped_gemm, and tensorwise FP8, and theragged path additionally requires
moe_router_padding_for_quantization=false.FP4 is rejected with an explicit message.
Validation
samples (target 3.34), 32.36 samples/s overall and 929 ms/step at
steady state — 7.2% faster than the explicit-padding path at an identical
sample count, with no NaN across 18 evaluations.
(22 tests). All pre-commit hooks pass.
verify_runtime.pypasses on 8×MI355X against the pinned stack.groups.
convergence gap traced to routing splits rather than quantization scale.
Known gap versus the official MLPerf 6.0 image
The official image sustains 36.05 samples/s against 33.16 here (+8.0%). A
same-node A/B plus a grouped-GEMM micro-benchmark traces the entire gap to one
kernel: official Primus-Turbo 0.2.0 dispatches an ASM code-object FP8 grouped
GEMM (
GroupedGEMMFP8ASMCOBackend, 1791 TFLOP/s on the fc1 shape), and thatbackend no longer exists in Turbo 0.3.2, whose best option here reaches 1192
TFLOP/s. Megatron argument dumps are otherwise identical, and 24 MoE layers
times the per-layer delta accounts for the measured per-step difference.
Within this stack there is nothing better to select: CK is slower still, and
hipBLASLt wins the isolated forward shapes but loses 7% end to end, so the
Triton backend this configuration selects is the right choice. Closing the gap
requires the ASM backend to return in Turbo 0.3.x. Pinning 0.2.0 is not viable —
it lacks the symmetric-memory APIs this branch depends on, and its BF16 grouped
GEMM is markedly slower.
Notes for reviewers
run_runtime_longrun.shmounts the Megatron dataset index cache as a writablevolume. The cache key changes with the iteration count, so without this any
TRAIN_ITERSoverride aborted during startup against the read-only datasetmount.
FMHA_HD64_ASM_LOGrather than the inherited logger level, which previouslymade any verbose run pay a logging cost per attention call.