Skip to content

Add ragged FP8 MoE path and pinned runtime image for GPT-OSS 20B MLPerf - #923

Open
HuangWei-95 wants to merge 3 commits into
mainfrom
feat/gpt-oss-mlperf-turbo-runtime
Open

Add ragged FP8 MoE path and pinned runtime image for GPT-OSS 20B MLPerf#923
HuangWei-95 wants to merge 3 commits into
mainfrom
feat/gpt-oss-mlperf-turbo-runtime

Conversation

@HuangWei-95

Copy link
Copy Markdown
Collaborator

Summary

Recovers the GPT-OSS 20B MLPerf throughput lost to explicit expert padding, and
makes the measured MI355X software stack reproducible.

  • Ragged FP8 grouped MLP. PrimusGroupedMLP can now hand the GPU
    tokens_per_expert tensor straight to PrimusTurbo's tensorwise FP8 grouped
    GEMM, skipping TE's Fp8Padding/Fp8Unpadding and the .tolist() device
    sync they require. Gated by use_turbo_ragged_grouped_gemm and by the actual
    autocast state at forward time, so other recipes and backends keep the
    existing zero-padding path.
  • Internal expert-weight quantization. With
    use_turbo_grouped_gemm_internal_weight_quantization, BF16 expert weights go
    directly to grouped_gemm_fp8 instead of being pre-quantized into buffers and
    routed through the weight-gradient bridge.
  • Fused activation independent of recompute. bias_act_func_with_mask now
    runs 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.
  • SDMA parameter gather. Ported to the current Primus-Turbo
    symmetric-memory API (the old hip module and async-TP entry points are
    gone), and the previously hardcoded "first two bucket groups stay on RCCL"
    rule became MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS, defaulting to 2 so behavior
    is unchanged.
  • Reproducible runtime image. Dockerfile.runtime plus runtime_versions.sh
    pin the base image by digest along with Triton, Primus-Turbo,
    TransformerEngine/AITER, and both HD64 causal ASM attention kernels;
    verify_runtime.py asserts the running stack matches.

Both new options default to false and are validated at startup: they require
enable_primus_turbo, use_turbo_grouped_gemm, and tensorwise FP8, and the
ragged path additionally requires moe_router_padding_for_quantization=false.
FP4 is rejected with an explicit message.

Validation

  • Full convergence on this branch head: eval loss 3.3352 at 221,184
    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.
  • The new argument-validation tests and the existing validate-args suite pass
    (22 tests). All pre-commit hooks pass.
  • verify_runtime.py passes on 8×MI355X against the pinned stack.
  • Ragged grouped-GEMM forward/backward smoke tests pass, including empty expert
    groups.
  • Paired EP=1 runs show identical initial activation amax values, so the earlier
    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 that
backend 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.sh mounts the Megatron dataset index cache as a writable
    volume. The cache key changes with the iteration count, so without this any
    TRAIN_ITERS override aborted during startup against the read-only dataset
    mount.
  • The ASM attention shim now gates its per-dispatch log line on
    FMHA_HD64_ASM_LOG rather than the inherited logger level, which previously
    made any verbose run pay a logging cost per attention call.

from transformer_engine.pytorch.attention.dot_product_attention import backends

backends.fused_attn_fwd = patched_fused_attn_fwd
except Exception:
Wei Huang and others added 3 commits July 28, 2026 07:36
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
HuangWei-95 force-pushed the feat/gpt-oss-mlperf-turbo-runtime branch from 71a4cbd to 7730252 Compare July 30, 2026 00:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant