From 4689ccfb2b57e1db2efe0ccbaadd18c9aa8020d2 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 11 Aug 2026 11:11:33 -0700 Subject: [PATCH 1/6] docs: bring the supported-models page in line with the launch scripts The page drifted from what scripts/ actually ships, in both directions. Missing entirely, though each has had a launch script for a while, and each is already advertised on the docs home page: Gemma-4 26B-A4B, Gemma-4 31B JoyAI-LLM-Flash Nemotron-3-Ultra-550B-A55B They have no recipe page yet, so they link to the script the way the home page already links them. The intro now says that a bare script link means the recipe page is not written yet, since the page otherwise promises one per model. Inkling-Small has a recipe page and was on the home page, but was never listed here. Removed, because #2356 and the branch under it retire both the launch script and the recipe page: GLM-Z1-9B-0414, Moonlight-16B-A3B and MiMo-7B-RL. MiMo had no other entry, so the family row goes with it. That stack already drops these from the docs.json nav and from the GLM and Kimi family pages; the top-level table is the one place it missed. Also: every launch script is Python after #2356, so the recipe-structure section no longer offers .sh as an alternative. The home page is untouched. --- docs/models/index.md | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/models/index.md b/docs/models/index.md index 415bed20b8a..81a6f091649 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -8,17 +8,19 @@ them. ## By family -Each model name links to its recipe page. +Each model name links to its recipe page. A few models ship a launch script before their +recipe page is written; those link to the script itself. | Family | Models | |---|---| | **DeepSeek** | [DeepSeek-V4 Pro](/models/deepseek/deepseek-v4-pro)
[DeepSeek-V4 Flash](/models/deepseek/deepseek-v4-flash)
[DeepSeek-R1](/models/deepseek/deepseek)
[DeepSeek-V3](/models/deepseek/deepseek) | -| **Thinking Machines** | [Inkling](/models/thinkingmachines/inkling) | +| **Thinking Machines** | [Inkling](/models/thinkingmachines/inkling)
[Inkling-Small](/models/thinkingmachines/inkling-small) | | **Qwen** | [Qwen3.6 MoE](/models/qwen/qwen3-6-moe)
[Qwen3.6](/models/qwen/qwen3-6)
[Qwen3.5-35B-A3B](/models/qwen/qwen3-5-moe)
[Qwen3.5-4B / 9B / 27B](/models/qwen/qwen3-5)
[Qwen3-Next-80B-A3B-Thinking](/models/qwen/qwen3-next)
[Qwen3-30B-A3B / 235B-A22B](/models/qwen/qwen3-moe)
[Qwen3-0.6B / 1.7B / 4B / 8B / 14B / 32B](/models/qwen/qwen3) | -| **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash)
[GLM-4.5](/models/glm/glm4-5)
[GLM-Z1-9B-0414](/models/glm/glm4) | -| **Kimi** | [Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2)
[Moonlight-16B-A3B](/models/kimi/moonlight) | -| **Nemotron** | [Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | -| **MiMo** | [MiMo-7B-RL](/models/mimo/mimo) | +| **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash)
[GLM-4.5](/models/glm/glm4-5) | +| **Kimi** | [Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2) | +| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](https://github.com/radixark/miles/blob/main/scripts/run_nemotron_3_ultra_550b_a55b.py)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | +| **Gemma** | [Gemma-4 26B-A4B](https://github.com/radixark/miles/blob/main/scripts/run_gemma_4_26b_a4b.py)
[Gemma-4 31B](https://github.com/radixark/miles/blob/main/scripts/run_gemma_4_31b.py) | +| **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) | | **GPT-OSS** | [gpt-oss-20b](/models/gpt-oss/gpt-oss) | ## How a recipe is structured @@ -28,7 +30,7 @@ Every recipe page follows the same six sections: 1. **Model Introduction** — what the model is and why miles supports it. 2. **Supported Variants** — model sizes + HF links. 3. **Environment Setup** — env vars, downloads, and HF → Megatron conversion. -4. **Launch** — the `scripts/run-.sh` (or `run_.py`) invocation. +4. **Launch** — the `scripts/run_.py` invocation. 5. **Recipe Configuration** — parallelism, algorithm, rollout/SGLang, optimizer. 6. **Pairs Well With** — links to the advanced features that complement this recipe. From f9905c806bd33355b33e797b5b35ea6b9b9a3e17 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 11 Aug 2026 11:21:51 -0700 Subject: [PATCH 2/6] docs: add recipe pages for Nemotron-3-Ultra and Gemma-4 Three model families ship launch scripts and are advertised on the docs home page, but had no recipe page and no entry in the supported-models table. Nemotron-3-Ultra-550B-A55B and Gemma-4 now have pages, written from their launch scripts and model configs. Both follow the six-section shape the rest of the model docs use. The Ultra page spends most of its length on one constraint, because it explains the whole layout: Mamba n_groups=8 forces attention and Mamba tensor-parallel to divide 8, so scale comes from PP and EP rather than TP, and the rollout side runs 32-GPU engines with DP-attention sized so attn_tp lands back on 8. It also records that routing replay is off for the 108-layer model and what that costs in train/rollout logprob drift. The Gemma page covers both sizes on one page, since 26B-A4B and 31B differ only in expert parallelism, token budget and a few sglang knobs. It notes that miles trains the language stack on the base VLM checkpoint without stripping the vision tower, and that the dense 31B needs the gemma4-dense branch of Megatron-Bridge. JoyAI-LLM-Flash keeps a script link rather than a page. Its script is a single recipe with no variants to compare, and nothing in the repo documents the model itself, so a page would be a paraphrase of the script. Inkling-Small had a page and a home-page entry but was never in the table. Removals are left to #2356, which already drops GLM-Z1-9B-0414, Moonlight-16B, MiMo and DeepSeek-R1 from this table along with their scripts. Keeping this PR purely additive means the two do not conflict. --- docs/docs.json | 11 +- docs/models/gemma/gemma-4.md | 151 +++++++++++++++++++++ docs/models/gemma/index.md | 39 ++++++ docs/models/index.md | 15 ++- docs/models/nemotron/index.md | 2 + docs/models/nemotron/nemotron-3-ultra.md | 164 +++++++++++++++++++++++ 6 files changed, 374 insertions(+), 8 deletions(-) create mode 100644 docs/models/gemma/gemma-4.md create mode 100644 docs/models/gemma/index.md create mode 100644 docs/models/nemotron/nemotron-3-ultra.md diff --git a/docs/docs.json b/docs/docs.json index 3913f0a46d7..c72ff0958c0 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -147,7 +147,16 @@ ], "expanded": false }, - "models/nemotron/nemotron-3-super" + "models/nemotron/nemotron-3-super", + "models/nemotron/nemotron-3-ultra" + ], + "expanded": false + }, + { + "group": "Gemma", + "root": "models/gemma/index", + "pages": [ + "models/gemma/gemma-4" ], "expanded": false }, diff --git a/docs/models/gemma/gemma-4.md b/docs/models/gemma/gemma-4.md new file mode 100644 index 00000000000..05d46658b67 --- /dev/null +++ b/docs/models/gemma/gemma-4.md @@ -0,0 +1,151 @@ +--- +title: Gemma-4 +description: Launch recipes for Gemma-4 26B-A4B-it (MoE) and 31B-it (dense) via the HF to Megatron bridge. +--- +## 1. Model Introduction + +[Gemma-4](https://huggingface.co/google) is Google's multimodal model line. +miles trains both released instruction-tuned sizes as language models, on the +base VLM checkpoint directly. + +Both go through the HF to Megatron bridge (`--megatron-to-hf-mode bridge`), and +on the rollout side sglang runs `Gemma4ForConditionalGeneration`, which loads +Gemma-4's hybrid `head_dim` weights correctly. There is no offline `torch_dist` +conversion and no LLM-view rewrite of the checkpoint. + +**Key highlights:** + +- **Two shapes, one recipe family**: 26B-A4B is MoE (128 experts, top-8), 31B is + dense. They differ mainly in expert parallelism and the token budget. +- **Bridge-mode load** straight from the VLM checkpoint. +- **Tied embeddings**: neither config passes + `--untie-embeddings-and-output-weights`. +- **Single node**: both recipes target 8 × H200. + +## 2. Supported Variants + +| Model | Class | Active / Total | Layers | Hidden | HF ID | +|---|---|---|---|---|---| +| Gemma-4 26B-A4B-it | MoE, 128 experts top-8 | 4 B / 26 B | 30 | 2816 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it) | +| Gemma-4 31B-it | Dense | 31 B | 60 | 5376 | [google/gemma-4-31B-it](https://huggingface.co/google/gemma-4-31B-it) | + +Both use GQA with `kv_channels=256`, RoPE base 1e6, and a 262144-token vocab. + +The 31B recipe requires the `gemma4-dense` branch of `radixark/Megatron-Bridge`. + +## 3. Environment Setup + +### 3.1 Download model + datasets + +```bash +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k +hf download google/gemma-4-26B-A4B-it --local-dir /root/models/gemma-4-26B-A4B-it +``` + +`--model-dir` and `--data-dir` default to `/root/models` and `/root/datasets`. +`prepare` performs these downloads for you. + +### 3.2 No `torch_dist` conversion + +The bridge reads the HF checkpoint directly, so `--hf-checkpoint` and +`--ref-load` both point at the download: + +```bash +--hf-checkpoint / +--ref-load / +--megatron-to-hf-mode bridge +``` + +## 4. Launch + +```bash +cd /root/miles + +# MoE, single node +python scripts/run_gemma_4_26b_a4b.py full-train --num-nodes 1 + +# dense, single node +python scripts/run_gemma_4_31b.py full-train --num-nodes 1 +``` + +Passing `--num-nodes 1` puts the recipe into `debug_minimal` mode, which shortens +`--rollout-max-response-len` to 256 for a quick smoke test. Multi-node runs use +the full 8192. + +## 5. Recipe Configuration + +### 5.1 Parallelism + +| Model | TP | PP | CP | EP | ETP | `max_tokens_per_gpu` | GPUs | +|---|---|---|---|---|---|---|---| +| 26B-A4B (MoE) | 4 | 1 | 1 | 8 | 1 | 1024 | 8 (1 × 8) | +| 31B (dense) | 4 | 1 | 1 | — | — | 512 | 8 (1 × 8) | + +Activation checkpointing is on for the MoE recipe +(`--recompute-granularity full --recompute-method uniform +--recompute-num-layers 1`). The dense 31B runs a smaller token budget because +its 60 dense layers at hidden 5376 cost more activation memory per token than +the MoE's 30 layers at 2816. + +### 5.2 Algorithm + +GRPO. The MoE recipe adds low-variance KL; the dense one runs without it: + +```bash +--advantage-estimator grpo +--entropy-coef 0.00 +--eps-clip 0.2 +--eps-clip-high 0.28 +--rm-type gemma_math +--balance-data + +# 26B-A4B only +--use-kl-loss +--kl-loss-coef 0.00 +--kl-loss-type low_var_kl +``` + +Rollout batch 32 at 8 samples per prompt, global batch 256, `--lr 1e-6`. AIME +evaluation every 20 steps is available behind `--enable-eval` and is off by +default. + +### 5.3 Rollout & SGLang + +```bash +--rollout-num-gpus-per-engine 4 +--sglang-mem-fraction-static 0.55 # 0.5 for the dense 31B +``` + +The MoE recipe pins sglang to conservative kernels: + +```bash +--sglang-attention-backend triton +--sglang-moe-runner-backend triton +--sglang-disable-custom-all-reduce +--sglang-disable-cuda-graph +--sglang-disable-overlap-schedule +--sglang-disable-radix-cache +--use-rollout-routing-replay +``` + +`--use-rollout-routing-replay` replays the rollout's expert routing during the +training forward pass, so train log-probs match rollout log-probs. Every +sigmoid- or softmax-routed MoE recipe in miles needs this; the dense 31B does +not. + +### 5.4 Notable quirks + +- **Trained on the VLM checkpoint.** miles does not strip the vision tower; the + bridge and sglang both handle the multimodal config, and the RL recipe simply + trains the language stack. +- `--attention-backend unfused` on the training side for the MoE recipe. +- Routing is softmax with `seq_aux_loss` balancing and the bias update rate at 0 + (`--moe-router-bias-update-rate 0 --moe-aux-loss-coeff 0`), plus + `--moe-grouped-gemm` and `--moe-router-dtype fp32`. +- The 31B recipe needs the `gemma4-dense` branch of `radixark/Megatron-Bridge`, + because the dense config is driven straight through `Gemma4VLBridge`. + +## 6. Pairs Well With + +- [Backends Beyond Megatron](/advanced/architecture-support) +- [P2P Weight Transfer](/advanced/p2p-weight-transfer) diff --git a/docs/models/gemma/index.md b/docs/models/gemma/index.md new file mode 100644 index 00000000000..7997967fdf5 --- /dev/null +++ b/docs/models/gemma/index.md @@ -0,0 +1,39 @@ +--- +title: Gemma +description: Miles recipes for Google's Gemma-4 line, trained on the base VLM checkpoint through the HF to Megatron bridge. +--- +Miles supports Google's Gemma-4 in both released instruction-tuned sizes. Both train as +language models on the base VLM checkpoint, through the HF to Megatron bridge +(`--megatron-to-hf-mode bridge`), so there is no offline `torch_dist` conversion. + +## Variants + +| Model | Class | Active / Total | HF ID | Recipe | +|---|---|---|---|---| +| Gemma-4 26B-A4B-it | MoE, 128 experts top-8 | 4 B / 26 B | `google/gemma-4-26B-A4B-it` | [gemma-4](/models/gemma/gemma-4) | +| Gemma-4 31B-it | Dense | 31 B | `google/gemma-4-31B-it` | [gemma-4](/models/gemma/gemma-4) | + +## Fastest path to train + +Both recipes run on a single 8-GPU node: + +```bash +cd /root/miles +python scripts/run_gemma_4_26b_a4b.py full-train --num-nodes 1 +``` + +`--num-nodes 1` shortens the response length for a smoke test. See +[Gemma-4](/models/gemma/gemma-4) for the full walkthrough. + +## Which variant do I pick? + +- **Cheaper to train, sparse** → 26B-A4B. Four billion active parameters, and expert + parallelism carries the width. +- **Dense, no routing to reason about** → 31B. It needs the `gemma4-dense` branch of + `radixark/Megatron-Bridge`, and runs a smaller per-GPU token budget because its 60 dense + layers cost more activation memory per token. + +## Pairs well with + +- [Backends Beyond Megatron](/advanced/architecture-support) — the bridge path Gemma rides on. +- [P2P Weight Transfer](/advanced/p2p-weight-transfer) diff --git a/docs/models/index.md b/docs/models/index.md index 81a6f091649..6510d5a908d 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -8,18 +8,19 @@ them. ## By family -Each model name links to its recipe page. A few models ship a launch script before their -recipe page is written; those link to the script itself. +Each model name links to its recipe page. JoyAI-LLM-Flash has a launch script but no +recipe page yet, so it links to the script. | Family | Models | |---|---| | **DeepSeek** | [DeepSeek-V4 Pro](/models/deepseek/deepseek-v4-pro)
[DeepSeek-V4 Flash](/models/deepseek/deepseek-v4-flash)
[DeepSeek-R1](/models/deepseek/deepseek)
[DeepSeek-V3](/models/deepseek/deepseek) | | **Thinking Machines** | [Inkling](/models/thinkingmachines/inkling)
[Inkling-Small](/models/thinkingmachines/inkling-small) | | **Qwen** | [Qwen3.6 MoE](/models/qwen/qwen3-6-moe)
[Qwen3.6](/models/qwen/qwen3-6)
[Qwen3.5-35B-A3B](/models/qwen/qwen3-5-moe)
[Qwen3.5-4B / 9B / 27B](/models/qwen/qwen3-5)
[Qwen3-Next-80B-A3B-Thinking](/models/qwen/qwen3-next)
[Qwen3-30B-A3B / 235B-A22B](/models/qwen/qwen3-moe)
[Qwen3-0.6B / 1.7B / 4B / 8B / 14B / 32B](/models/qwen/qwen3) | -| **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash)
[GLM-4.5](/models/glm/glm4-5) | -| **Kimi** | [Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2) | -| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](https://github.com/radixark/miles/blob/main/scripts/run_nemotron_3_ultra_550b_a55b.py)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | -| **Gemma** | [Gemma-4 26B-A4B](https://github.com/radixark/miles/blob/main/scripts/run_gemma_4_26b_a4b.py)
[Gemma-4 31B](https://github.com/radixark/miles/blob/main/scripts/run_gemma_4_31b.py) | +| **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash)
[GLM-4.5](/models/glm/glm4-5)
[GLM-Z1-9B-0414](/models/glm/glm4) | +| **Kimi** | [Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2)
[Moonlight-16B-A3B](/models/kimi/moonlight) | +| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/models/nemotron/nemotron-3-ultra)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | +| **MiMo** | [MiMo-7B-RL](/models/mimo/mimo) | +| **Gemma** | [Gemma-4 26B-A4B / 31B](/models/gemma/gemma-4) | | **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) | | **GPT-OSS** | [gpt-oss-20b](/models/gpt-oss/gpt-oss) | @@ -30,7 +31,7 @@ Every recipe page follows the same six sections: 1. **Model Introduction** — what the model is and why miles supports it. 2. **Supported Variants** — model sizes + HF links. 3. **Environment Setup** — env vars, downloads, and HF → Megatron conversion. -4. **Launch** — the `scripts/run_.py` invocation. +4. **Launch** — the `scripts/run-.sh` (or `run_.py`) invocation. 5. **Recipe Configuration** — parallelism, algorithm, rollout/SGLang, optimizer. 6. **Pairs Well With** — links to the advanced features that complement this recipe. diff --git a/docs/models/nemotron/index.md b/docs/models/nemotron/index.md index 82e64ff6d80..f3c79d14d3a 100644 --- a/docs/models/nemotron/index.md +++ b/docs/models/nemotron/index.md @@ -11,6 +11,7 @@ Miles supports NVIDIA's Nemotron-3 line: a Mamba + Attention hybrid that, in the | Nemotron-3-Nano | 4 B / 4 B (dense) | `nvidia/Nemotron-3-Nano-4B` | [nemotron-3-nano](/models/nemotron/nemotron-3-nano) | | Nemotron-3-Nano MoE | 3 B / 30 B | `nvidia/Nemotron-3-Nano-30B-A3B` | [nemotron-3-nano-moe](/models/nemotron/nemotron-3-nano-moe) | | Nemotron-3-Super | 12 B / 120 B (FP8) | `nvidia/Nemotron-3-Super-120B-A12B-FP8` | [nemotron-3-super](/models/nemotron/nemotron-3-super) | +| Nemotron-3-Ultra | 55 B / 550 B | `nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16` | [nemotron-3-ultra](/models/nemotron/nemotron-3-ultra) | ## Fastest path to train @@ -28,6 +29,7 @@ See the [Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) page for the dense w - **Smallest, single-node smoke test** → Nemotron-3-Nano ([nemotron-3-nano](/models/nemotron/nemotron-3-nano)). - **Mid-scale hybrid MoE** → Nemotron-3-Nano MoE ([nemotron-3-nano-moe](/models/nemotron/nemotron-3-nano-moe)). - **Frontier-scale FP8-native MoE** → Nemotron-3-Super ([nemotron-3-super](/models/nemotron/nemotron-3-super)). +- **Largest, latent MoE across 16 nodes** → Nemotron-3-Ultra ([nemotron-3-ultra](/models/nemotron/nemotron-3-ultra)). ## Pairs well with diff --git a/docs/models/nemotron/nemotron-3-ultra.md b/docs/models/nemotron/nemotron-3-ultra.md new file mode 100644 index 00000000000..16f272bffd2 --- /dev/null +++ b/docs/models/nemotron/nemotron-3-ultra.md @@ -0,0 +1,164 @@ +--- +title: Nemotron-3-Ultra +description: Launch recipe for NVIDIA Nemotron-3-Ultra-550B-A55B (hybrid Mamba2 + Attention + latent-MoE) via Megatron AutoBridge. +--- +## 1. Model Introduction + +[NVIDIA Nemotron-3-Ultra-550B-A55B](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) +is the Ultra tier of the `nemotron_h` family: the same hybrid block pattern as +Nano and Super, scaled to **550 B total / 55 B active** across 108 layers, with +a **latent MoE** (512 experts, top-22, `moe_latent_size=2048`) and one shared +expert. + +miles loads it through the `megatron.bridge` AutoBridge with the shared +NemotronH MoE shim (`miles_plugins/megatron_bridge/nemotron_h.py`), the same +path Super-120B uses. There is no offline `torch_dist` conversion. + +**Key highlights:** + +- **Hybrid + latent MoE**: Mamba2 and attention blocks with a latent-projection + MoE FFN, 512 experts at top-22. +- **Bridge-mode load**: `--megatron-to-hf-mode bridge`, straight from the HF + checkpoint. +- **Sigmoid routing** with aux-free expert-bias load balancing, plus an MTP head + in the checkpoint. +- **A single-node 4-layer slice** is published alongside the full model, so the + recipe can be smoke-tested without 16 nodes. + +## 2. Supported Variants + +| Model | Active / Total | Layers | GPUs | HF ID | +|---|---|---|---|---| +| Nemotron-3-Ultra-550B-A55B-BF16 | 55 B / 550 B | 108 | 128 (16 × 8) | [nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) | +| Nemotron-3-Ultra-550B-A55B-BF16-4layer | slice | 4 | 8 (1 × 8) | pruned slice of the above, for smoke tests | + +Tested on H200. Use the `radixark/miles:dev` image. + +## 3. Environment Setup + +### 3.1 Download model + datasets + +```bash +hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k +hf download nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 \ + --local-dir /root/models/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 +``` + +`--model-dir` and `--data-dir` default to `/root/models` and `/root/datasets`. +`--model-name` names the checkpoint directory inside `--model-dir`; passing a +name matching `layer` selects the pruned slice and switches the recipe to its +single-node parallelism automatically. + +### 3.2 No `torch_dist` conversion + +AutoBridge plus the NemotronH shim read the HF checkpoint directly, so +`--hf-checkpoint` and `--ref-load` both point at the download: + +```bash +--hf-checkpoint / +--ref-load / +--megatron-to-hf-mode bridge +``` + +## 4. Launch + +### 4.1 Single-node smoke test + +```bash +cd /root/miles +python scripts/run_nemotron_3_ultra_550b_a55b.py full-train \ + --model-name NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-4layer --num-nodes 1 +``` + +### 4.2 Full model + +The full 108-layer model needs **16 nodes × 8 GPU**. Worker pods join the head's +ray cluster; the head submits: + +```bash +# on each worker pod +python scripts/run_nemotron_3_ultra_550b_a55b.py worker --head-ip + +# on the head pod +python scripts/run_nemotron_3_ultra_550b_a55b.py train --head-ip +``` + +The recipe defaults to a 30-rollout run (`--num-rollout`), rollout batch 32 at 8 +samples per prompt, global batch 128. + +## 5. Recipe Configuration + +### 5.1 Parallelism + +| Variant | TP | PP | CP | EP | ETP | GPUs | +|---|---|---|---|---|---|---| +| Full 108-layer | 8 | 4 | 1 | 32 | 1 | 128 (16 × 8) | +| 4-layer slice | 1 | 1 | 1 | 8 | 1 | 8 (1 × 8) | + +**Mamba `n_groups=8` caps attention and Mamba tensor parallelism at 8**, because +Megatron requires `n_groups % tp == 0`. That constraint drives the whole layout: +TP cannot grow past 8, so scale comes from PP and EP instead. + +The 4-layer slice fits on one node, so it gives every rank to expert parallelism +(512 experts over EP=8 is 64 per rank) and keeps attention and Mamba at TP=1. + +`--sequence-parallel` is on when TP > 1. Activation checkpointing is enabled, and +`--log-probs-chunk-size 128` keeps the log-prob pass inside the memory budget. + +### 5.2 Algorithm + +GRPO with low-variance KL: + +```bash +--advantage-estimator grpo +--kl-loss-coef 0.00 +--kl-loss-type low_var_kl +--entropy-coef 0.00 +--eps-clip 0.2 +--eps-clip-high 0.28 +--rm-type deepscaler +``` + +### 5.3 Rollout & SGLang + +| Variant | GPUs per engine | `--sglang-ep-size` | `--sglang-dp-size` | mem fraction | +|---|---|---|---|---| +| Full 108-layer | 32 | 32 | 4 | 0.7 | +| 4-layer slice | 8 | 8 | 2 | 0.6 | + +The 550 B model is roughly 1.1 TB in BF16 and **does not fit one 8-GPU engine**, +so rollout runs 32-GPU engines with EP=32 and DP-attention. The DP size is chosen +so that `attn_tp = gpus_per_engine / dp_size` lands on 8, satisfying the same +Mamba `n_groups` constraint the training side has. The launcher asserts both +divisibility rules rather than letting a bad combination fail deep in SGLang. + +### 5.4 Optimizer + +CPU Adam, with the host transfer overlapped: + +```bash +--optimizer-cpu-offload +--overlap-cpu-optimizer-d2h-h2d +--use-precision-aware-optimizer +--lr 1e-6 +``` + +### 5.5 Notable quirks + +- **Routing replay is not enabled yet.** `--use-rollout-routing-replay` is off + for the 108-layer model: the routing capturer needs a fix for per-layer top-22 + under DP-attention. Train and rollout log-probs differ by about 0.01 without + it. The Super-120B recipe does enable it. +- **No `--spec`**: AutoBridge and the NemotronH shim synthesize the Megatron MoE + spec from the HF config. +- The shim is what wires `routed_scaling_factor`, `n_group` and `topk_group` + onto the Megatron provider. Without it the routed output is silently scaled + 1.0×, the same drift class the Nano-MoE and Super recipes call out. +- The checkpoint carries an MTP head (`num_nextn_predict_layers=1`); the RL + recipe does not train it. + +## 6. Pairs Well With + +- [Backends Beyond Megatron](/advanced/architecture-support) +- [Nemotron-3-Super](/models/nemotron/nemotron-3-super) +- [P2P Weight Transfer](/advanced/p2p-weight-transfer) From ddd9f4d839cb5d943a959a0988de79fb285f435d Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 11 Aug 2026 11:30:27 -0700 Subject: [PATCH 3/6] docs: point the home page at the new recipes, and align the two model tables The home page linked Nemotron-3-Ultra and both Gemma-4 sizes straight at their launch scripts, because that was the only thing to link at the time. Those three now have recipe pages, so the home page links to them instead. The supported-models table now carries Kimi-K3 and matches the home page's family order and entry order, so the two read as the same list at two levels of detail rather than two lists that disagree. Gemma is split into its two sizes the way the home page has it, both pointing at the shared recipe page. Two entries still link outside the docs, and the intro now says why. JoyAI-LLM- Flash links to its launch script: one recipe, no variants, and nothing in the repo describing the model, so a page would restate the script. Kimi-K3 links to the pull request adding its support, which is backend plumbing with no launch script or model config yet, so there is no recipe to write down. The one ordering difference left is the MiMo row, which #2356 removes. Once that lands the two tables agree entry for entry. --- docs/index.md | 4 ++-- docs/models/index.md | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/index.md b/docs/index.md index 7db0aec8a4f..98a87e34c00 100644 --- a/docs/index.md +++ b/docs/index.md @@ -80,8 +80,8 @@ of the box, including older generations of the families below. | **Qwen** | [Qwen3.6 MoE](/models/qwen/qwen3-6-moe)
[Qwen3.6](/models/qwen/qwen3-6)
[Qwen3.5-35B-A3B](/models/qwen/qwen3-5-moe)
[Qwen3.5-4B / 9B / 27B](/models/qwen/qwen3-5) | | **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash) | | **Kimi** | [Kimi-K3](https://github.com/radixark/miles/pull/1825)
[Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5) | -| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](https://github.com/radixark/miles/blob/main/scripts/run_nemotron_3_ultra_550b_a55b.py)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | -| **Gemma** | [Gemma-4 26B-A4B](https://github.com/radixark/miles/blob/main/scripts/run_gemma_4_26b_a4b.py)
[Gemma-4 31B](https://github.com/radixark/miles/blob/main/scripts/run_gemma_4_31b.py) | +| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/models/nemotron/nemotron-3-ultra)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | +| **Gemma** | [Gemma-4 26B-A4B](/models/gemma/gemma-4)
[Gemma-4 31B](/models/gemma/gemma-4) | | **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) | See [Models](/models/index) for exact conversion commands, launch scripts, and diff --git a/docs/models/index.md b/docs/models/index.md index 6510d5a908d..d399a8d8c59 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -8,8 +8,8 @@ them. ## By family -Each model name links to its recipe page. JoyAI-LLM-Flash has a launch script but no -recipe page yet, so it links to the script. +Each model name links to its recipe page. Two entries do not have one yet: JoyAI-LLM-Flash +links to its launch script, and Kimi-K3 links to the pull request adding its support. | Family | Models | |---|---| @@ -17,10 +17,10 @@ recipe page yet, so it links to the script. | **Thinking Machines** | [Inkling](/models/thinkingmachines/inkling)
[Inkling-Small](/models/thinkingmachines/inkling-small) | | **Qwen** | [Qwen3.6 MoE](/models/qwen/qwen3-6-moe)
[Qwen3.6](/models/qwen/qwen3-6)
[Qwen3.5-35B-A3B](/models/qwen/qwen3-5-moe)
[Qwen3.5-4B / 9B / 27B](/models/qwen/qwen3-5)
[Qwen3-Next-80B-A3B-Thinking](/models/qwen/qwen3-next)
[Qwen3-30B-A3B / 235B-A22B](/models/qwen/qwen3-moe)
[Qwen3-0.6B / 1.7B / 4B / 8B / 14B / 32B](/models/qwen/qwen3) | | **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash)
[GLM-4.5](/models/glm/glm4-5)
[GLM-Z1-9B-0414](/models/glm/glm4) | -| **Kimi** | [Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2)
[Moonlight-16B-A3B](/models/kimi/moonlight) | +| **Kimi** | [Kimi-K3](https://github.com/radixark/miles/pull/1825)
[Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2)
[Moonlight-16B-A3B](/models/kimi/moonlight) | | **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/models/nemotron/nemotron-3-ultra)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | | **MiMo** | [MiMo-7B-RL](/models/mimo/mimo) | -| **Gemma** | [Gemma-4 26B-A4B / 31B](/models/gemma/gemma-4) | +| **Gemma** | [Gemma-4 26B-A4B](/models/gemma/gemma-4)
[Gemma-4 31B](/models/gemma/gemma-4) | | **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) | | **GPT-OSS** | [gpt-oss-20b](/models/gpt-oss/gpt-oss) | From 9535d98c56c1cd51f4172522fa4f8d67f3eb1176 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 11 Aug 2026 11:37:00 -0700 Subject: [PATCH 4/6] docs: write the Kimi-K3 recipe, and drop the merged-PR pointer from Inkling Kimi-K3 now has a page. It documents the day-0 LoRA RL recipe from #1825: the MXFP4 to BF16 upcast, the 32-rank torch_dist conversion, and the validated 16-node launch, with the LoRA target modules and what a healthy run looks like on the GB300 validation cluster. A line at the top says the branch it comes from, since none of it is on main yet. The page leans on two things worth calling out. LoRA is not a nicety here, it is what makes the recipe fit: the base stays frozen and only the adapters carry gradients, with one shared A factor and per-expert B factors across the 896 experts. And the memory figure in "what a healthy run looks like" is the one to watch, because an adapter export that leaks shows up as allocated memory not returning to baseline after a weight sync, long before the run dies. Inkling still said its implementation "is open at" #1683. That merged on 2026-08-03 and the code is on main, so the line is gone; the page is the documentation now. Both tables link Kimi-K3 at the page instead of the pull request. MiMo and GPT-OSS move below JoyAI so the families the home page lists appear in the same order in both, with the two it does not carry after them. --- docs/docs.json | 1 + docs/index.md | 2 +- docs/models/gemma/index.md | 2 +- docs/models/index.md | 8 +- docs/models/kimi/kimi-k3.md | 176 ++++++++++++++++++++++++ docs/models/thinkingmachines/inkling.md | 2 - 6 files changed, 183 insertions(+), 8 deletions(-) create mode 100644 docs/models/kimi/kimi-k3.md diff --git a/docs/docs.json b/docs/docs.json index c72ff0958c0..1e1a423120c 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -129,6 +129,7 @@ "group": "Kimi", "root": "models/kimi/index", "pages": [ + "models/kimi/kimi-k3", "models/kimi/kimi-k2.5", "models/kimi/kimi-k2", "models/kimi/moonlight" diff --git a/docs/index.md b/docs/index.md index 98a87e34c00..9985bd6354f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -79,7 +79,7 @@ of the box, including older generations of the families below. | **Thinking Machines** | [Inkling](/models/thinkingmachines/inkling)
[Inkling-Small](/models/thinkingmachines/inkling-small) | | **Qwen** | [Qwen3.6 MoE](/models/qwen/qwen3-6-moe)
[Qwen3.6](/models/qwen/qwen3-6)
[Qwen3.5-35B-A3B](/models/qwen/qwen3-5-moe)
[Qwen3.5-4B / 9B / 27B](/models/qwen/qwen3-5) | | **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash) | -| **Kimi** | [Kimi-K3](https://github.com/radixark/miles/pull/1825)
[Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5) | +| **Kimi** | [Kimi-K3](/models/kimi/kimi-k3)
[Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5) | | **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/models/nemotron/nemotron-3-ultra)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | | **Gemma** | [Gemma-4 26B-A4B](/models/gemma/gemma-4)
[Gemma-4 31B](/models/gemma/gemma-4) | | **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) | diff --git a/docs/models/gemma/index.md b/docs/models/gemma/index.md index 7997967fdf5..a6a615874db 100644 --- a/docs/models/gemma/index.md +++ b/docs/models/gemma/index.md @@ -35,5 +35,5 @@ python scripts/run_gemma_4_26b_a4b.py full-train --num-nodes 1 ## Pairs well with -- [Backends Beyond Megatron](/advanced/architecture-support) — the bridge path Gemma rides on. +- [Backends Beyond Megatron](/advanced/architecture-support), the bridge path Gemma rides on. - [P2P Weight Transfer](/advanced/p2p-weight-transfer) diff --git a/docs/models/index.md b/docs/models/index.md index d399a8d8c59..8e13946fbf0 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -8,8 +8,8 @@ them. ## By family -Each model name links to its recipe page. Two entries do not have one yet: JoyAI-LLM-Flash -links to its launch script, and Kimi-K3 links to the pull request adding its support. +Each model name links to its recipe page. JoyAI-LLM-Flash is the one exception: it has a +launch script but no page yet, so it links to the script. | Family | Models | |---|---| @@ -17,11 +17,11 @@ links to its launch script, and Kimi-K3 links to the pull request adding its sup | **Thinking Machines** | [Inkling](/models/thinkingmachines/inkling)
[Inkling-Small](/models/thinkingmachines/inkling-small) | | **Qwen** | [Qwen3.6 MoE](/models/qwen/qwen3-6-moe)
[Qwen3.6](/models/qwen/qwen3-6)
[Qwen3.5-35B-A3B](/models/qwen/qwen3-5-moe)
[Qwen3.5-4B / 9B / 27B](/models/qwen/qwen3-5)
[Qwen3-Next-80B-A3B-Thinking](/models/qwen/qwen3-next)
[Qwen3-30B-A3B / 235B-A22B](/models/qwen/qwen3-moe)
[Qwen3-0.6B / 1.7B / 4B / 8B / 14B / 32B](/models/qwen/qwen3) | | **GLM** | [GLM-5.2](/models/glm/glm5-2)
[GLM-5.1](/models/glm/glm5)
[GLM-5](/models/glm/glm5)
[GLM-4.7-Flash](/models/glm/glm4-7-flash)
[GLM-4.5](/models/glm/glm4-5)
[GLM-Z1-9B-0414](/models/glm/glm4) | -| **Kimi** | [Kimi-K3](https://github.com/radixark/miles/pull/1825)
[Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2)
[Moonlight-16B-A3B](/models/kimi/moonlight) | +| **Kimi** | [Kimi-K3](/models/kimi/kimi-k3)
[Kimi-K2.6](/models/kimi/kimi-k2.5)
[Kimi-K2.5](/models/kimi/kimi-k2.5)
[Kimi-K2-Instruct / Thinking](/models/kimi/kimi-k2)
[Moonlight-16B-A3B](/models/kimi/moonlight) | | **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/models/nemotron/nemotron-3-ultra)
[Nemotron-3-Super-120B-A12B-FP8](/models/nemotron/nemotron-3-super)
[Nemotron-3-Nano MoE](/models/nemotron/nemotron-3-nano-moe)
[Nemotron-3-Nano](/models/nemotron/nemotron-3-nano) | -| **MiMo** | [MiMo-7B-RL](/models/mimo/mimo) | | **Gemma** | [Gemma-4 26B-A4B](/models/gemma/gemma-4)
[Gemma-4 31B](/models/gemma/gemma-4) | | **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) | +| **MiMo** | [MiMo-7B-RL](/models/mimo/mimo) | | **GPT-OSS** | [gpt-oss-20b](/models/gpt-oss/gpt-oss) | ## How a recipe is structured diff --git a/docs/models/kimi/kimi-k3.md b/docs/models/kimi/kimi-k3.md new file mode 100644 index 00000000000..c86c1ea96be --- /dev/null +++ b/docs/models/kimi/kimi-k3.md @@ -0,0 +1,176 @@ +--- +title: Kimi-K3 +description: LoRA RL recipe for Kimi-K3, a KDA + MLA hybrid with an 896-expert latent MoE, trained colocated with SGLang. +--- + +The complete Kimi-K3 LoRA RL implementation is open at the Miles pull request: +[`radixark/miles#1825`](https://github.com/radixark/miles/pull/1825). The scripts and image +below come from that branch. Results and background are in the +[LMSYS day-0 write-up](https://www.lmsys.org/blog/2026-07-27-kimi-k3-day0-support). + +## 1. Model Introduction + +Kimi-K3 pairs two attention mechanisms in one stack, **KDA and MLA chosen per layer**, +with an **896-expert latent MoE** at top-16. The checkpoint ships in **MXFP4**. + +miles trains it with **native LoRA adapters** rather than full fine-tuning, which is what +makes the recipe fit at all: the base weights stay frozen and only the adapters carry +gradients. The adapters are implemented under TP, EP, PP and CP, with **shared-A and +per-expert-B factors** across the 896 experts, and they are exported to the rollout engines +as HF-named chunks over CUDA IPC. + +**Key highlights:** + +- **Two attention types per layer**: KDA and MLA, with an attention-residual snapshot bank. +- **896-expert latent MoE**, top-16, `moe_latent_size=3584`, plus a shared expert. +- **LoRA RL, not full fine-tuning.** Rank 16 by default, 32 in the validated run. +- **Colocated rollout**: trainer and SGLang share the GPUs, with adapters synced over CUDA IPC. +- **MXFP4 checkpoint** upcast to BF16 once, offline. + +## 2. Supported Variants + +| Variant | Layers | Purpose | GPUs | +|---|---|---|---| +| `full` | full stack | the real model | 64 (16 × 4), validated | +| `4layer` | 4 | smoke test, default | single node | + +`--model-variant` selects between them and sets the matching checkpoint paths and +`megatron_model_type`. + +Architecture, from `scripts/models/kimi-k3.sh`: hidden 7168, FFN 33792, 96 attention heads, +`kv_channels=256`, MLA with `q_lora_rank=1536` / `kv_lora_rank=512` / +`qk_head_dim=128` / `qk_pos_emb_head_dim=64` / `v_head_dim=128`, 896 experts at +`moe_ffn_hidden_size=3072`, shared expert 6144, vocab 163840, no position embedding. + +## 3. Environment Setup + +Use the `docker.io/radixark/miles:kimi-k3` image, which pins miles, SGLang (the +[`sglang-miles-k3`](https://github.com/sgl-project/sglang/tree/sglang-miles-k3) branch) and +flashinfer `0.6.15.post1` at the validated versions. On Hopper set +`SGLANG_K3_ATTN_RES_MODE=jit`. + +The only external asset is the Kimi-K3 MXFP4 HF checkpoint. Everything else derives in-repo. + +### 3.1 Data + +```bash +python scripts/run_kimi_k3_lora.py prepare-data --task dapo-math --data-dir +``` + +### 3.2 MXFP4 to BF16 + +```bash +python tools/convert_mxfp4_to_bf16.py --model-dir --save-dir +``` + +### 3.3 BF16 to `torch_dist` + +Unlike the bridge-mode recipes, K3 needs an offline conversion. Run it on 32 ranks; the +output re-shards at load, so the conversion layout does not have to match the training one: + +```bash +source scripts/models/kimi-k3.sh # defines MODEL_ARGS +torchrun --nnodes=8 --nproc-per-node=4 ... \ + tools/convert_hf_to_torch_dist.py "${MODEL_ARGS[@]}" \ + --hf-checkpoint --save \ + --bf16 --tensor-model-parallel-size 32 --sequence-parallel \ + --pipeline-model-parallel-size 1 --context-parallel-size 1 \ + --expert-model-parallel-size 32 --expert-tensor-parallel-size 1 \ + --megatron-to-hf-mode raw +``` + +Training then takes the **MXFP4** directory as `--hf-checkpoint` and the converted +`torch_dist` as `--ref-load`. + +## 4. Launch + +Validated on **16 nodes × 4 GPUs**. One container per node; bring up a ray cluster across +them, `export MILES_SCRIPT_EXTERNAL_RAY=1`, then: + +```bash +python scripts/run_kimi_k3_lora.py train \ + --mode normal --model-variant full --task dapo-math --reward-model deepscaler \ + --num-nodes 16 --num-gpus-per-node 4 \ + --pipeline-parallel-size 8 --context-parallel-size 2 \ + --rollout-tp-size 16 --rollout-max-concurrency 8 \ + --lora-rank 32 --lora-alpha 64 \ + --num-rollout 1000 --rollout-batch-size 8 --n-samples-per-prompt 8 \ + --rollout-max-response-len 4096 --sglang-max-total-tokens 65536 \ + --global-batch-size 64 --lr 1e-05 --eval-interval 10 \ + --distributed-timeout-minutes 60 \ + --hf-checkpoint --ref-load \ + --sglang-path /sgl-workspace/sglang/python \ + --data-dir --enable-wandb +``` + +`--rollout-max-concurrency 8` is passed explicitly: the field default is 64, and the +validated runs pin 8. + +For a single-node smoke test, drop to the default `--model-variant 4layer`. + +## 5. Recipe Configuration + +### 5.1 Parallelism + +The resolved config at startup should show `expert_model_parallel_size 8`, +`max_tokens_per_gpu 8192`, `colocate_memory_peak_device gpu` and +`lora_base_cpu_backup True`. Checking those four lines is the fastest way to confirm the +run came up in the intended shape. + +| Knob | Validated value | +|---|---| +| Pipeline parallel | 8 | +| Context parallel | 2 | +| Expert parallel | 8 (resolved) | +| Rollout TP | 16 | +| `max_tokens_per_gpu` | 8192 | + +### 5.2 LoRA + +Rank 32 / alpha 64 in the validated run; the script defaults to 16 / 32. Adapters attach to +attention output and both MLA down-projections, the dense MLP, and both expert projections: + +``` +self_attention.o_proj +self_attention.q_a_proj +self_attention.kv_a_proj_with_mqa +mlp.linear_fc1, mlp.linear_fc2 +mlp.experts.linear_fc1, mlp.experts.linear_fc2 +``` + +The 896 experts share one A factor and carry per-expert B factors, which is what keeps the +adapter count tractable at this expert width. + +### 5.3 Rollout + +Rollout is **colocated**: the trainer and SGLang share GPUs, and adapters sync over CUDA +IPC as HF-named chunks. `lora_base_cpu_backup` keeps a host copy of the frozen base so the +GPU copy can be reclaimed during rollout. + +### 5.4 What a healthy run looks like + +From the GB300 validation runs: + +- 11 to 13 minutes per rollout cycle +- trainer allocated memory returns to about 91 GB after every weight sync +- `rollout/raw_reward` between 0.5 and 0.75 from rollout 0 +- `eval/aime` 0.37 to 0.43 at eval@0 (that spread is temperature-0 nondeterminism), rising + by at least 0.06 by eval@9; the measured run went 0.367 to 0.467 + +The memory figure is the one to watch. If allocated does not return to its baseline after a +weight sync, the adapter export is leaking and the run will die later rather than sooner. + +### 5.5 Notable quirks + +- The image preloads a small `shm_unlink` shim through `/etc/ld.so.preload`. It tolerates a + benign PyTorch CUDA-IPC unlink race that otherwise aborts colocated weight sync at scale. +- On Hopper, set `SGLANG_K3_ATTN_RES_MODE=jit`. +- Weight conversion for K3 lives in + `miles/backends/megatron_utils/megatron_to_hf/kimi_k3.py`, and the model itself in + `miles_plugins/models/kimi_k3/`. + +## 6. Pairs Well With + +- [LoRA](/advanced/lora) +- [Backends Beyond Megatron](/advanced/architecture-support) +- [Kimi-K2.5](/models/kimi/kimi-k2.5) diff --git a/docs/models/thinkingmachines/inkling.md b/docs/models/thinkingmachines/inkling.md index 801f9653f00..18dd90f0169 100644 --- a/docs/models/thinkingmachines/inkling.md +++ b/docs/models/thinkingmachines/inkling.md @@ -3,8 +3,6 @@ title: Inkling description: Launch recipe for Inkling (975 B), Thinking Machines' multimodal MoE with short convolution, relative attention, and a shared-expert sink. --- -The complete Inkling RL implementation is open at the Miles pull request: [`radixark/miles#1683`](https://github.com/radixark/miles/pull/1683). - ## 1. Model Introduction [Inkling](https://huggingface.co/thinkingmachines/Inkling) is a mixture-of-experts transformer released by Thinking Machines Lab, with 975 B total parameters and 41 B active, a context window of up to 1 M tokens, and pretraining on 45 trillion tokens of text, images, audio and video. Its architecture introduces short convolution, attention with relative positional embedding, and a novel MoE design with a shared-expert sink. Miles implements Inkling as a native Megatron model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders, and the same backend drives both full-parameter and LoRA RL. From 664cbd74d30062e6e8ad731facb99a0834c98abb Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 11 Aug 2026 11:39:52 -0700 Subject: [PATCH 5/6] docs: keep the supported-models intro as it was --- docs/models/index.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/models/index.md b/docs/models/index.md index 8e13946fbf0..ead05f1cf4c 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -8,8 +8,7 @@ them. ## By family -Each model name links to its recipe page. JoyAI-LLM-Flash is the one exception: it has a -launch script but no page yet, so it links to the script. +Each model name links to its recipe page. | Family | Models | |---|---| From a2147fa7339b0d3a685c719e6e788ee8a88754f3 Mon Sep 17 00:00:00 2001 From: Shi Dong Date: Tue, 11 Aug 2026 16:01:38 -0700 Subject: [PATCH 6/6] docs: fix the Nemotron-3-Ultra launch commands and the Gemma bridge branch name The launcher has no worker subcommand and no --head-ip option; multi-node runs join an external ray cluster and set MILES_SCRIPT_EXTERNAL_RAY, like the other recipe launchers. The Megatron-Bridge branch for the dense 31B recipe is zhichen/gemma4-dense, not gemma4-dense. --- docs/models/gemma/gemma-4.md | 2 +- docs/models/nemotron/nemotron-3-ultra.md | 20 ++++++++++++++------ 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/docs/models/gemma/gemma-4.md b/docs/models/gemma/gemma-4.md index 05d46658b67..bdf0f3312b2 100644 --- a/docs/models/gemma/gemma-4.md +++ b/docs/models/gemma/gemma-4.md @@ -31,7 +31,7 @@ conversion and no LLM-view rewrite of the checkpoint. Both use GQA with `kv_channels=256`, RoPE base 1e6, and a 262144-token vocab. -The 31B recipe requires the `gemma4-dense` branch of `radixark/Megatron-Bridge`. +The 31B recipe requires the `zhichen/gemma4-dense` branch of `radixark/Megatron-Bridge`. ## 3. Environment Setup diff --git a/docs/models/nemotron/nemotron-3-ultra.md b/docs/models/nemotron/nemotron-3-ultra.md index 16f272bffd2..2aea61b3b99 100644 --- a/docs/models/nemotron/nemotron-3-ultra.md +++ b/docs/models/nemotron/nemotron-3-ultra.md @@ -72,17 +72,25 @@ python scripts/run_nemotron_3_ultra_550b_a55b.py full-train \ ### 4.2 Full model -The full 108-layer model needs **16 nodes × 8 GPU**. Worker pods join the head's -ray cluster; the head submits: +The full 108-layer model needs **16 nodes × 8 GPU**. Bring up the ray cluster +yourself, tell the launcher it is external, and submit from the head — the +launcher has no worker-side subcommand: ```bash -# on each worker pod -python scripts/run_nemotron_3_ultra_550b_a55b.py worker --head-ip - # on the head pod -python scripts/run_nemotron_3_ultra_550b_a55b.py train --head-ip +ray start --head --num-gpus 8 --disable-usage-stats +# on every worker pod +ray start --address=${HEAD_IP}:6379 --num-gpus 8 --disable-usage-stats + +# then, on the head pod +export MILES_SCRIPT_EXTERNAL_RAY=1 +export RAY_ADDRESS=http://${HEAD_IP}:8265 +python scripts/run_nemotron_3_ultra_550b_a55b.py train --num-nodes 16 ``` +Without `MILES_SCRIPT_EXTERNAL_RAY=1` the launcher runs `ray stop --force` and +starts a fresh single-node head, tearing down the cluster the workers joined. + The recipe defaults to a 30-rollout run (`--num-rollout`), rollout batch 32 at 8 samples per prompt, global batch 128.