Repository navigation
feat(zaya1): add FP8 recipes that train cold on mock data - #1269
Merged
Merged
Conversation
Dispatch model_type zaya1 through the existing GPT training loop and ship MI300X and MI355X mock-data configs for a 30-iteration run. Co-authored-by: Cursor <cursoragent@cursor.com>
The forward math stays the same; comments and the mixer test no longer attribute it to SGLang. Co-authored-by: Cursor <cursoragent@cursor.com>
The new model needs a covered forward path and a trainer launch with a unique log tag, plus the hardware and GSM8K loss evidence from the 8x MI300X run. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
# Conflicts: # primus/backends/megatron/megatron_pretrain_trainer.py # primus/core/utils/import_utils.py
The test builds Zaya1Model directly and hits stock LanguageModule, which rejects a leftover NVTE_UNFUSED_ATTN=1 when the backend is fused.
Drop the GSM8K blend and Gemma tokenizer so the checked-in configs run without a dataset mount. Co-authored-by: Cursor <cursoragent@cursor.com>
…irlee/zaya1. Co-authored-by: Cursor <cursoragent@cursor.com>
CCA projections and expert SwiGLU GEMMs run as Transformer Engine FP8 linears under one autocast; the router, convolutions, norms, attention, embeddings and tied logits stay BF16. The recipes use tensorwise scaling, so FP8 is accurate from the first step without an amax history, and train on mock data like the BF16 recipes. Co-authored-by: Cursor <cursoragent@cursor.com>
clairesonglee
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
October 19, 2026 17:39
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
…rt' and 'import from'' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
6 tasks done
Xiaoming-AMD
previously approved these changes
Oct 20, 2026
) Stacked on #1269 (base `dev/clairlee/zaya1`). Retarget to `main` once #1269 merges. ## Summary - Add MI325X `zaya1_8B-BF16-pretrain.yaml` and `zaya1_8B-FP8-pretrain.yaml`. All six ZAYA1-8B recipes (MI300X, MI325X, MI355X × BF16, FP8) train on mock data. - Size `micro_batch_size` on MI300X and MI325X to the largest multiple of 8 that completes on one 8-GPU node with `rocm/primus:v26.8`, end-of-training evaluation included. `global_batch_size` is micro batch × 8, so each GPU runs one micro batch per step. - FP8 recipes set `NVTE_USE_CAST_TRANSPOSE_TRITON: "0"` in their `env:` block. The runner enables Transformer Engine's Triton cast-transpose, which autotunes for every new tensor shape, and ZAYA1's expert token counts change every step. With the Triton kernel, MI325X FP8 at micro batch 64 spent 307 s on step 1 and still took 2 to 10 s per step at steps 11 to 15. With the HIP cast kernel, steps take 1.8 s from step 3 on. - Docs: MI325X added to the support matrix (config counts updated), and a throughput section in the ZAYA1 validation note. ## Results Recipes run as checked in: 30 steps, mock data, sequence length 128, data parallel 8, `rocm/primus:v26.8` (PyTorch 2.14.0, TE 2.17.1). Step time is the median over steps 6 to 30, tokens/s/GPU the harmonic mean, and TFLOP/s/GPU the mean. Every run finished with 0 skipped and 0 NaN iterations. | GPU | Precision | Micro / global batch | Step time | Tokens/s/GPU | TFLOP/s/GPU | Peak memory per GPU | Tokens/s/GPU before (micro 1, global 8) | |---|---|---|---|---|---|---|---| | MI300X | BF16 | 128 / 1024 | 1635 ms | 10016 | 118.6 | 181 of 192 GB | 171 | | MI300X | FP8 | 120 / 960 | 1814 ms | 8478 | 100.4 | 187 of 192 GB | 115 | | MI325X | BF16 | 216 / 1728 | 2049 ms | 13426 | 159.0 | 249 of 256 GB | 175 | | MI325X | FP8 | 184 / 1472 | 2118 ms | 11154 | 132.1 | 245 of 256 GB | 119 | Next size up: MI300X BF16 136 trains but stalls in the end-of-training evaluation at 98% memory. MI300X FP8 128, MI325X BF16 224, and MI325X FP8 192 run out of memory on the fp32 logits buffer. TFLOP/s/GPU is Megatron's analytic GPT estimate. It models neither CCA nor the alternating attention and MoE stages, so compare it only between ZAYA1 runs. FP8 is slower than BF16 at this size. With hidden size 2048 and top-1 routing over 16 experts, the expert GEMMs are small, and the FP8 casts cost more than the faster GEMMs save. ## Not covered - MI355X was not run. Its batch sizes stay at micro 1, global 8. Its FP8 recipe gets the same `env:` override because the Triton autotuning comes from the TE build, not the GPU. - `PYTORCH_HIP_ALLOC_CONF=expandable_segments:True` let MI325X FP8 fit micro batch 192 (+3% tokens/s), but MI325X BF16 224 still ran out of memory. Not adopted. ## Test plan - [x] Each MI300X (tw011) and MI325X (fremont) recipe run as checked in, without command-line overrides: 30 steps plus validation and test evaluation - [x] `pytest tests/unit_tests/configs/test_example_configs.py -k "zaya1 or every_backend"`: 7 passed - [x] `pytest tests/unit_tests/megatron/transformer/zaya1`: 20 passed - [x] `pytest tests/trainer/test_megatron_trainer.py -k zaya1`: 2 passed (BF16 and FP8 smoke tests on MI300X) - [x] `pre-commit run` on the changed files - [ ] MI355X batch sizing --------- Co-authored-by: Cursor <cursoragent@cursor.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
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
fp8: hybrid). CCA Q/K/V/O projections and expert SwiGLU GEMMs run as Transformer Engine linears (parallel_mode="duplicated") under one FP8 autocast for the whole stack. The router, depthwise convolutions, RMSNorm, attention, embeddings, and the tied logits projection stay in BF16/FP32.zaya1_8B-FP8-pretrain.yamlfor MI300X and MI355X withfp8_recipe: tensorwise. Current scaling takes each scale from the tensor being cast, so FP8 is accurate from step 1 with no amax history. Delayed scaling starts with an empty history; on the released checkpoint a cold delayed-scaling forward scored SGLang rollouts at Pearson r 0.72.vocab_sizestays 262272 fromzaya1.yaml.docs/04-technical-guides/zaya1/validation.mdanddocs/06-developer-guide/model-support-matrix.md.Validation
Run on 1x MI300X node (tw042),
rocm/primus:v26.8, TE 2.17.1, Megatron-LMd3528a21.Log-prob parity against SGLang (sgl-project/sglang#26347): the same 4 GSM8K greedy rollouts used for BF16, 1,024 tokens each, scored by Primus with the
Zyphra/ZAYA1-baseweights and no FP8 warm-up.GSM8K pretraining from scratch with the FP8 recipe (seq 128, global batch 8, lr 1e-5). Early losses track the BF16 validation run:
No skipped or NaN iterations. The 30k-step run was not completed: at 3.2 s/step on 4 GPUs it needs about 27 h, beyond the 8 h job limit.
Test plan
pytest tests/unit_tests/megatron/transformer/zaya1(20 passed on MI300X, includingtest_cold_tensorwise_fp8_matches_high_precision)pytest tests/trainer/test_megatron_trainer.py -k zaya1(BF16 and FP8, 4 stages, 3 iterations, mock data)pre-commit runon the changed fileszaya1_8B-FP8-pretrain.yaml(not yet exercised on gfx950)