Repository navigation
feat(models): add MiMo-V2.6-Flash-RL support - #3648
Conversation
|
cc cmt E2E RL metrics on the 2x8 H200 devbox, by engine setting All runs: two rollout/train steps with R3 (routing replay) and a step-0 weight check, temperature 1.0, 2048-token responses. "log-prob gap" is the mean |trainer - rollout| log-prob over response tokens (Gate B), with p99 and max. 4-layer partial, one node (8 H200), two TP4 engines, 4 prompts x 2 samples, top-p 1.0 (32,768 tokens per run):
Full 309B model, both nodes (16 H200), trainer TP2 SP PP2 EP8 with NVMe-streamed Adam, two colocated BF16 TP8 engines, dapo-math-17k 4 prompts x 4 samples, top-p 0.95 (62,953 tokens):
|
287c242 to
b265e32
Compare
2b2b1ab to
6366a6e
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
b422e01 to
5acb497
Compare
|
/rerun-test tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py |
|
❌ The execution job failed; inspect the workflow run for the failing step. |
|
@claude review always |
|
/rerun-test tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py |
|
❌ The execution job failed; inspect the workflow run for the failing step. |
There was a problem hiding this comment.
Beyond the inline findings, I also checked two other candidates and ruled them out for this PR: a duplicate index in --layers for convert_mimo_v2_to_bf16.py silently overwrites via the dict comprehension rather than erroring, but that's an operator-misuse edge case on manual re-layering, not a correctness issue in the normal conversion path; and the missing moe_router_num_groups/group_topk branch for HF n_group > 1 in mimo_v2.py doesn't affect the current MiMo-V2.6-Flash checkpoint, which has no group-limited routing.
Extended reasoning...
This PR adds a new model bridge, launch script, and BF16 checkpoint converter for MiMo-V2.6-Flash; the two confirmed findings (a prepare() race on concurrent checkpoint conversion, and an interrupted-conversion leaving a checkpoint that looks complete) are already posted as inline comments. I additionally reviewed the --layers duplicate-index dict comprehension and the n_group>1 MoE routing gap mentioned in the ruled-out list, confirming both are real but out of scope for the shipped checkpoint/usage pattern rather than active bugs. No security-sensitive surface (auth/crypto) is touched; the concerns are data-integrity/concurrency issues in file conversion, already flagged inline.
Findings marked 🟡 are optional suggestions and need no follow-up push.
|
/rerun-test tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py |
|
✅ |
There was a problem hiding this comment.
I reviewed this PR; the automated bug hunt found no new issues this round, and the two issues I flagged on earlier revisions are now fixed: prepare() wraps the whole check-then-download-then-convert step in _exclusive_path_lock(str(target)) (scripts/run_mimo_v2_6_flash.py:122-134), and convert_mimo_v2_to_bf16.py now copies the tokenizer/aux files before atomically publishing model.safetensors.index.json via a .tmp file + replace() (tools/convert_mimo_v2_to_bf16.py). Given this is still a large, complex change (FP8/MXFP4 dequantization, per-layer SWA/global attention, and a new Megatron bridge), a human look is still worthwhile, particularly on the quantization math and per-layer head-dim handling.
Extended reasoning...
The diff adds a new Megatron-Bridge model plugin, a checkpoint dequantization tool (FP8 blockwise + MXFP4), and a launch script/doc set for MiMo-V2.6-Flash; no auth/crypto surface, but correctness of numeric dequantization and checkpoint-completeness handling matters. This run found no new bugs, and the two race/completion-marker issues from my prior reviews are now fixed in commit d033d5b, which tips this toward defer-with-acknowledgment rather than silence; the size and domain complexity (quantization, per-layer attention config) still argue against full approval.
|
|
||
| ### 3.2 SGLang | ||
|
|
||
| The BF16 engine needs the MiMo-V2 fixes that are not in the image yet: |
There was a problem hiding this comment.
This looks stale now: the sglang in the current radixark/miles:dev image (8035002, "[31/32] [sglang-miles] MiMo colocated RL loaders, pausable caches and graph registration (#40648, #40978, #40979, #40980, #40977, #41972)") already carries these fixes. The 4-layer e2e ran on that image as-is. Could this section say which image/sglang version has them instead of "not in the image yet"?
There was a problem hiding this comment.
cc cmt
[P3] Section 3.2 told readers the MiMo-V2 SGLang fixes were missing from the image, which would send them to build a custom SGLang they do not need. Confirmed: sgl-project/sglang#40977 merged into sglang-miles on 09-29, and its picks (#40978 split layout, #40979 SWA KV pools in the memory-saver region, #40980 TopK layer_id, plus #40448 from the v0.5.21 base) are in 8035002, which current sglang-miles and the radixark/miles:dev image (built from sglang-miles HEAD) contain. Fixed in 750acf9f1d: the section now names sglang-miles 8035002 and the dev image instead of "not in the image yet". I also dropped the stale ci-sglang-pr: #40977 pin from the PR body so CI runs the image's sglang-miles.
| import typer | ||
|
|
||
| import miles.utils.external_utils.command_utils as U | ||
| from miles.utils.external_utils.command_utils.base_backend import _exclusive_path_lock |
There was a problem hiding this comment.
_exclusive_path_lock is private to base_backend.py (its only other callers are the backend's own download/convert helpers). Importing it from a launcher couples the script to an internal detail. Could it be made public (e.g. exclusive_path_lock, re-exported from command_utils), or could preparation go through a backend helper that already takes the lock?
There was a problem hiding this comment.
cc cmt
Agreed. Done in 750acf9f1d: _exclusive_path_lock is now public as exclusive_path_lock and re-exported from miles.utils.external_utils.command_utils; the launcher calls U.exclusive_path_lock, the regression test imports it from the package, and the backend's own convert_checkpoint / fp8_cast_bf16 use the new name. No existing backend helper fits this step (they run convert_hf_to_torch_dist.py / fp8_cast_bf16.py), so preparation stays in the launcher.
| # The BF16 engine holds the whole model: 620 GB needs TP8 on 141 GB GPUs. | ||
| rollout_num_gpus_per_engine: int = 4 |
There was a problem hiding this comment.
The comment describes the full model (620 GB, TP8), but it sits on the default, which is 4 and is only used by the 4-layer partial; the full-model recipe overrides it to 8. Maybe move the TP8 note to the MiMo-V2.6-Flash-RL-bf16 entry in _RECIPES and keep this comment about the default.
There was a problem hiding this comment.
cc cmt
Done in 750acf9f1d: the default field comment now only says it is the GPU count per colocated BF16 engine, and the TP8 / 620 GB note sits on rollout_num_gpus_per_engine=8 in the MiMo-V2.6-Flash-RL-bf16 recipe.
Zhichenzzz
left a comment
Zhichenzzz
left a comment
There was a problem hiding this comment.
LGTM after those comments are fixed
|
/rerun-test tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py |
|
✅ |
There was a problem hiding this comment.
Checked back on the two findings from my earlier reviews: the prepare() race (no lock around the checkpoint-missing check and download/convert) is now fixed by wrapping it in U.exclusive_path_lock(str(target)) at scripts/run_mimo_v2_6_flash.py:123, and the partial-conversion marker issue is fixed by writing model.safetensors.index.json to a .tmp path and atomically renaming it only after the auxiliary-file copy loop completes (tools/convert_mimo_v2_to_bf16.py:244-247). Both changes match the PR description and are covered by new/updated tests.
Extended reasoning...
Diff since my last review (commits d033d5b, 323cb45, 750acf9) addresses both inline findings I previously raised: the missing lock around prepare()'s check-then-convert step, and the index.json being written before the auxiliary file copies finished. The lock helper was also renamed from private _exclusive_path_lock to public exclusive_path_lock and reused by existing callers in base_backend.py, a mechanical rename. A human reviewer (Zhichenzzz) already approved after inline comments were addressed. No new concerns found beyond confirming these fixes land as described.
MiMo-V2 attention differs from the stock Megatron GPT layer in ways that TransformerConfig cannot express: per-layer KV heads (8 on sliding-window layers, 4 on global layers), 192-dim Q/K with 128-dim V, a learnable sink on the sliding-window layers only, and V scaled by 0.707 before attention. - miles_plugins/megatron_bridge/mimo_v2.py: Megatron-Bridge plugin for MiMoV2ForCausalLM. Each attention module derives a per-layer config copy and rebuilds linear_qkv/linear_proj with the V width; the core attention passes (k_channels, v_channels) and the sink setting to TE. Per-layer RoPE base and the window reuse rotary_base_per_layer and a list-form window_attn_skip_freq. The provider sets hidden_dropout = 0 because bridge mode does not copy --hidden-dropout onto the provider. - tools/convert_mimo_v2_to_bf16.py: converts the official checkpoint (FP8 block-quantized, fused kv-head-interleaved qkv_proj, MXFP4 experts) into the BF16 split-q/k/v layout the bridge reads; --layers writes partials. - scripts/run_mimo_v2_6_flash.py + model args: a 4-layer partial on one node and the full model on two nodes with NVMe-streamed Adam. - Docs page and a 4-layer CI test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… SGLang does SGLang's RadixAttention window excludes the query token, so a MiMo-V2 SWA layer in the rollout engine attends the query plus sliding_window_size (128) previous keys, 129 in total. The bridge set window_size to (sliding_window_size - 1, 0), 128 keys in total, following the HF modeling code and Megatron-Bridge's mimo_v2_flash, so training and rollout disagreed at every position past the window and showed rare large log-prob gaps. On an 8149-token sample, HF with a 129-key window matched SGLang (|dlogp| mean 0.038, none above 2 nat) while the configured window matched the trainer. SGLang's behavior is the reference (sgl-project/sglang#40981, which moved SGLang to 128 keys, is closed), so the bridge now sets window_size = (sliding_window_size, 0). The docs drop the SGLang fix from section 3.2 and describe the window in section 5.4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MiMoV2SelfAttention let SelfAttention build linear_qkv and linear_proj at the Q/K head size, then built both again at the MiMo sizes with a copy of the base constructor arguments. Both attention classes also re-derived the layer's KV heads and attention sink from a private config copy, while the rest of the layer kept the global config. MiMoV2ModelProvider now sets heterogeneous_block_specs and implements get_config_for_layer, so Megatron's TransformerBlock builds each layer from a config carrying its own num_query_groups and softmax_type. MiMoV2SelfAttention hands SelfAttention linear builders sized for the 128-dim V heads, so each projection is built once, and rejects qk_clip, which reads linear_qkv as q + 2 * kv rows of kv_channels. MiMoV2TEDotProductAttention only passes the Q/K and V channels and scales V. The q/k/v weight mapping is Megatron-Bridge's MiMoV2FlashQKVMapping, whose layout equals the removed copy. RoPE, the SWA window, the sink flags and MTP keep the Miles settings rather than Megatron-Bridge's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The MiMo-V2.6 page presented the 4-layer partial conversion (mimo26-p4-bf16) as a supported variant with its own quick start, though it only serves development checks and the CI smoke test; users train the 48-layer model on two nodes. The page now documents only that model: the variants and parallelism tables lose the partial row, the converter example writes the whole model, the SFT note moves under the two-node launch, and wording that contrasted the full model with the partial is plain. The launcher recipe and the 4-layer CI test keep the partial. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…kend The command_utils package no longer exports execution helpers. Build the backend from ScriptArgs for model preparation, dataset download and training submission. Cover cold and cached checkpoints in RL and SFT, and refresh the MiMo execution snapshot for the backend-owned deploy-component argument.
The MiMo launcher imported the private _exclusive_path_lock from base_backend. Make it public as exclusive_path_lock, re-export it from command_utils, and use U.exclusive_path_lock in the launcher and its regression test. Docs section 3.2 still said the MiMo-V2 SGLang fixes were not in the image. sgl-project/sglang#40977 merged into sglang-miles, whose 8035002 carries them, and radixark/miles:dev installs sglang-miles, so the section now names that commit and the image. Move the TP8 engine note from the recipe default (4, used by the 4-layer partial) to the full-model recipe that sets 8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
750acf9 to
76d5985
Compare
The MiMo-V2.6 4-layer partial e2e test takes about 40 minutes of an 8-GPU H200 runner. Mark its registration nightly=True so regular PR runs of stage-c-8-gpu-h200 skip it, while nightly and weekly cadences keep it and /rerun-test can still run it on demand (file_run.py ignores the nightly gate). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Two passing /rerun-test runs of the MiMo-V2.6 4-layer e2e test took 339 s and 361 s in the test step (runs 37450530812 and 37531719740, warm checkpoint cache). 361 s x 1.25 rounds up to 500 s; 2400 s overstated the test by about 5x in nightly partition balancing. Timeouts are unaffected: /rerun-test uses max(1800, est_time x 1.25) and run_suite.py uses its own per-file timeout. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
3 participants
Summary
Train MiMo-V2.6-Flash-RL (309B MoE) in bridge mode with a BF16 rollout engine.
Motivation
MiMo-V2 needs per-layer KV heads, asymmetric Q/K/V dimensions, sliding-window sinks, and scaled V that stock TransformerConfig cannot express. Its official checkpoint combines fused, kv-head-interleaved FP8
qkv_projwith MXFP4 experts that Miles cannot load directly.Usage
preparedownloads the official checkpoint and converts it to BF16.Design Notes
AutoBridgeselectsMiMoV2Bridge, whose provider supplies per-layer KV heads/sinks throughheterogeneous_block_specs, while its mapping registry loads training weights and synchronizes the BF16 engine.MiMoV2FlashQKVMappingfrom Megatron-Bridge8cd3466d14, with MiMo-specific RoPE, sink, and SWA-window handling.ScriptArgs.create_backend().command_utils.exclusive_path_lock; the backend's own callers use the new name.top-p=1.0; appended--extra-argsretains override precedence.sglang-milescarries since8035002andradixark/miles:devinstalls; full-model training streams Adam state through NVMe (fix(optimizer): convert NVMe-streamed gradients to fp32 one bucket at a time #3743, merged).nightly=True;/rerun-teststill runs it on demand.Verification
run_suite.py --list-onlyforstage-c-8-gpu-h200: regular cadence drops the 4-layer test (11 to 10 files), nightly keeps it.750acf9f1d(same patches as the rebased head), macOS CPU: 22 MiMo launcher, converter and snapshot tests passed.pre-commitpassed for each change.750acf9f1don the image'ssglang-miles.750acf9f1d, 2x8 H200, launcher defaults plus routing replay: three RL steps with log-prob gap 0.0112/0.0105/0.0103 and KL 3.9e-4 to 3.6e-4, peak GPU 132.5 GB.Review Focus
tools/convert_mimo_v2_to_bf16.py: fused-qkv shard dequantization and MXFP4 decoding.