Repository navigation
[1/4] Add MiniMax-M3 support - #4581
Conversation
Megatron plugin, mbridge bridge, HF converter and launcher for MiniMax-M3: GQA with MSA block-sparse attention on the sparse layers (a frozen 4-head indexer picks each query's top-16 key blocks) and plain dense attention on the first three, swigluoai MoE with a sigmoid router and one shared expert. - miles_plugins/models/minimax_m3: the attention module and block spec; the sparse layers run the QSA block-sparse kernel, whose tile map now covers blocks wider than a tile. - miles_plugins/mbridge/minimax_m3.py and megatron_to_hf/minimax_m3.py: weight mapping both ways. - scripts/run_minimax_m3.py: the full model (8 x 4 GB300, TP4/PP4/EP8, Adam state streamed through disk) and a 4-layer cut for one node. - tools/convert_minimax_m3_to_bf16.py: cuts a layer subset from a release and dequantizes the MXFP8/NVFP4 MiniMax-M3.1 checkpoints to BF16.
![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.
- block_sparse_attention: launch configs per block width (narrow = qwen3.8-next's 4-token QSA blocks, wide = MiniMax-M3's 128-token MSA blocks), and the dK/dV kernel splits a KV head's query heads over several programs, each writing its own slice that the caller sums, instead of one program walking all of them. - MiniMax-M3 indexer: the per-block max logits come from a Triton kernel (miles/kernels/attention/msa) and the init/local/causal rules and top-k run vectorized, replacing the per-sequence, per-query-chunk torch loop. GB300, one GPU, 16k-token packs (fwd+bwd of the attention kernel; indexer select): qwen3.8-next (12 q / 1 kv head, D 256, block 4): 22.5 -> 16.4 ms MiniMax-M3 (16 q / 1 kv head, D 128, block 128): 12.7 -> 6.7 ms (32k: 38.6 -> 20.2 ms) MiniMax-M3 select: 5.0 -> 0.62 ms (32k: 16.3 -> 1.2 ms), identical block sets on >= 99.997% of rows (the rest are near-ties) Gradients agree with the previous launch configs to bf16 noise (rel <= 8e-3). MiniMax-M3 4-layer e2e: log-prob abs diff unchanged (0.0050), weight check passes.
Look the (forward, backward) configs up through one function instead of indexing two bool-keyed dicts at each call site, and store dV with V's strides rather than relying on K and V sharing a layout. No behavior change.
Runs the MiniMax-M3-4layer recipe on 4 H200s from a pre-cut BF16 checkpoint (CharyZeng/MiniMax-M3-4layer: layers 0,1,3,4, two dense-attention and two sparse MoE layers), with a deterministic pseudo-random reward so every weight sync carries real deltas, and gates the train/rollout log-prob gap, grad norm and KL.
|
/rerun-test tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py |
|
⚪ The run was cancelled. |
|
/rerun-test tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py |
|
✅ |
|
/rerun-test tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py |
|
❌ The execution job failed; inspect the workflow run for the failing step. |
|
/rerun-test tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py |
|
✅ |
The check now also treats quick-GeGLU (MiniMax-M3's swigluoai) as a gated FC1, so the test's backend namespace carries quick_geglu and gains that case.
MiniMax-M3 support: Megatron plugin, mbridge bridge, HF converter, launcher and a BF16 conversion tool.
Stack: #4581 (this) → #4582 → #4583 → #4584.
What
miles_plugins/models/minimax_m3: GQA with MSA block-sparse attention on the sparse layers (a frozen 4-head indexer picks each query's top-16 key blocks) and plain dense attention on the first three layers; swigluoai MoE with a sigmoid router and one shared expert.miles/kernels/attention/qsa/block_sparse_attention.py: the tile map now covers key blocks wider than a tile (MiniMax uses 128-token blocks, the kernel walks 64/32-token tiles), with no host sync.miles/kernels/attention/msa: Triton per-block max-logit scores for the MiniMax indexer, replacing a per-sequence torch loop.miles_plugins/mbridge/minimax_m3.py,megatron_to_hf/minimax_m3.py: weight mapping both ways.scripts/run_minimax_m3.py:MiniMax-M3(8 × 4 GB300, TP4/PP4/EP8, Adam state streamed through disk) andMiniMax-M3-4layer(one node).tools/convert_minimax_m3_to_bf16.py: cuts a layer subset from a release, and dequantizes the MXFP8/NVFP4 MiniMax-M3.1 checkpoints to BF16 (used by [3/4]).Depends on
SGLang must accept online weight updates into MiniMax-M3's fused q/k/v + index GEMM: sgl-project/sglang#47612 (into
sglang-miles). Without it the first weight sync fails withIndexError: start out of range.Validation (GB300, colocated, BF16 rollout, sglang with #47612)
MiniMax-M3-4layer, 3 steps: train/rollout log-prob abs diff 0.0050 / 0.0050 / 0.0050,--check-weight-update-equalpasses after every update.MiniMax-M3, 8 × 4 GB300, 10 steps: abs diff 0.0070–0.0096, dapo-math reward 0.52 → 0.79, 10–12 min per step after the first.tests/fast/test_qsa_tile_index.pygains a blocks-wider-than-a-tile case: 16 of its cases fail on the previous two-tile_tile_maskand pass now.pytest tests/manual/launch_scripts tests/fast/launch_scripts: the minimax cases pass; the remaining 19 failures / 4 errors are identical on main.Kernel tuning (GB300, one GPU, 16k-token packs)
Gradients match the previous configs to bf16 noise (rel ≤ 8e-3); the new select picks identical block sets on ≥ 99.997% of rows (the rest are near-ties). MiniMax-M3 4-layer e2e after tuning: abs diff 0.0050, weight check passes.
tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.pypasses on 4 × H200 (stage-c-4-gpu-h200) on the SGLang v0.5.23radixark/miles:devimage with [MiniMax-M3] Keep the fused q/k/v + index GEMM writable by online weight updates sgl-project/sglang#47612: https://github.com/radixark/miles/actions/runs/41275852481.ci-sglang-pr: #47612