Skip to content

feat(gpt-oss): MXFP4 deosc BF16 snapshots and simplified 20B recipe - #1070

Open
HuangWei-95 wants to merge 3 commits into
mainfrom
feat/gptoss-mxfp4-deosc-bf16-recipe
Open

feat(gpt-oss): MXFP4 deosc BF16 snapshots and simplified 20B recipe#1070
HuangWei-95 wants to merge 3 commits into
mainfrom
feat/gptoss-mxfp4-deosc-bf16-recipe

Conversation

@HuangWei-95

Copy link
Copy Markdown
Collaborator

Summary

  • Store MXFP4 weight de-oscillation snapshots in BF16 to cut sidecar memory; DistRatio window sums and snap write-back stay FP32.
  • Keep QKVO projections in BF16 via a TE precision config so grouped GEMM can stay MXFP4.
  • Align the GPT-OSS 20B example with the llama3.1 launcher: add run_with_docker.sh / config_MI355X_1x8x1.sh, and drop the v26.3/v26.5 runtime Dockerfiles, ASM override, and prewarm path.

Test plan

  • tests/unit_tests/backends/megatron/test_weight_deosc.py (includes legacy fp32 snapshot load)
  • MXFP4 short run: training loss decreases with gpt_oss_20B-MXFP4-deosc-mlperf-pretrain.yaml
  • Confirm FP8 Turbo-attn yaml still launches via run_with_docker.sh / config_MI355X_1x8x1.sh

…y the 20B recipe

Keep QKVO in BF16 via a TE precision config while grouped GEMM stays MXFP4, and
align the GPT-OSS MLPerf example with the llama3.1 launcher instead of the
v26.3/v26.5 runtime Dockerfiles.

Co-authored-by: Cursor <cursoragent@cursor.com>
Wei Huang and others added 2 commits September 1, 2026 09:26
Keep the weight-deosc assignment on one line so code-lint black --line-length=110 passes.

Co-authored-by: Cursor <cursoragent@cursor.com>
use_turbo_norm_te_linear wraps QKV so finish_init never reached the inner
TE module; forward it and match linear_qkv.linear. Keep BF16 grouped GEMM
off FlyDSL, which is only registered for MXFP4.
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