diff --git a/docs/docs.json b/docs/docs.json
index 7cacbf8a7cb..8b872dd095d 100644
--- a/docs/docs.json
+++ b/docs/docs.json
@@ -205,7 +205,8 @@
],
"expanded": false
},
- "models/gpt-oss/gpt-oss"
+ "models/gpt-oss/gpt-oss",
+ "models/mimo/mimo-v2-6-flash"
]
}
]
diff --git a/docs/models/index.md b/docs/models/index.md
index e799d1d3942..93165e56df1 100644
--- a/docs/models/index.md
+++ b/docs/models/index.md
@@ -22,6 +22,7 @@ Each model name links to its recipe page.
| **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) |
+| **MiMo** | [MiMo-V2.6-Flash](/models/mimo/mimo-v2-6-flash) |
## Diffusion
diff --git a/docs/models/mimo/mimo-v2-6-flash.md b/docs/models/mimo/mimo-v2-6-flash.md
new file mode 100644
index 00000000000..39ec81e9e61
--- /dev/null
+++ b/docs/models/mimo/mimo-v2-6-flash.md
@@ -0,0 +1,98 @@
+---
+title: MiMo-V2.6-Flash
+sidebarTitle: MiMo-V2.6
+description: Launch recipe for Xiaomi's MiMo-V2.6-Flash-RL (309B MoE, hybrid sliding-window attention) — Megatron bridge mode on a BF16 conversion, on two 8-GPU nodes.
+---
+## 1. Model Introduction
+
+[MiMo-V2.6-Flash-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) is Xiaomi's 309B-parameter MoE model (15B active). miles trains its text decoder.
+
+**Key highlights:**
+
+- **Hybrid attention**: 39 sliding-window layers (window 128, learnable attention sink) and 9 global layers, with different KV head counts (8 vs 4) and RoPE bases (1e4 vs 1e7).
+- **Asymmetric heads**: 192-dim Q/K and 128-dim V, V scaled by 0.707, partial RoPE (64 of 192 dims).
+- **MoE**: 256 routed experts, top-8 sigmoid routing with a fixed score-correction bias; layer 0 is dense.
+- **Compressed checkpoint**: FP8 block-quantized dense weights with a fused, kv-head-interleaved `qkv_proj`, and MXFP4 experts. The Megatron bridge reads a BF16 split-q/k/v conversion instead.
+
+## 2. Supported Variants
+
+| `--model-name` | What it is | HF source |
+|---|---|---|
+| `MiMo-V2.6-Flash-RL-bf16` | text decoder, 48 layers | [XiaomiMiMo/MiMo-V2.6-Flash-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) |
+
+The vision, audio and MTP modules stay in the checkpoint but are not trained (text-only RL).
+
+## 3. Environment Setup
+
+### 3.1 Download + BF16 conversion
+
+`prepare` (run by the launcher before training) downloads the official checkpoint to `/MiMo-V2.6-Flash-RL` and converts it unless `/` already holds a finished conversion:
+
+```bash
+hf download XiaomiMiMo/MiMo-V2.6-Flash-RL --local-dir /root/models/MiMo-V2.6-Flash-RL
+python tools/convert_mimo_v2_to_bf16.py --model-dir /root/models/MiMo-V2.6-Flash-RL \
+ --save-dir /root/models/MiMo-V2.6-Flash-RL-bf16 --device cuda
+```
+
+The converter dequantizes each kv-head shard of the fused `qkv_proj` with its own FP8 block scales, splits it into q/k/v in head order, and decodes the MXFP4 experts. The BF16 model is 622 GB.
+
+### 3.2 SGLang
+
+The BF16 engine needs the following MiMo-V2 fixes. The `sglang-miles` branch has them since [`8035002`](https://github.com/sgl-project/sglang/commit/8035002c4ed0360579004791378e5b7d1740503b), and the `radixark/miles:dev` image, built from that branch, includes them:
+
+- upstream [sgl-project/sglang#40448](https://github.com/sgl-project/sglang/pull/40448) (MXFP4 MoE and BF16 router);
+- accepting the `split` attention layout of the BF16 conversion;
+- allocating the sliding-window KV pools inside the memory-saver region, otherwise a colocated engine cannot release its KV cache;
+- passing `layer_id` to the MoE top-k, otherwise `--use-rollout-routing-replay` crashes the CUDA-graph capture.
+
+## 4. Launch
+
+### 4.1 Two nodes
+
+Join the second node to a ray head on the first, then launch from the head:
+
+```bash
+MILES_SCRIPT_EXTERNAL_RAY=1 MASTER_ADDR= python scripts/run_mimo_v2_6_flash.py \
+ --mode rl --model-name MiMo-V2.6-Flash-RL-bf16 --num-nodes 2 --train-offload-disk-dir /scratch/offload
+```
+
+On two 8×H200 nodes with node-local NVMe (4-drive RAID0), a step with 16 samples of up to 2048 response tokens took about 4 minutes of training, of which about 3 minutes is streaming the optimizer state (72 GB read and 144 GB written per GPU), plus 30 s of weight update. Peak GPU memory was 125 GB and peak host memory 1.05 TB per node.
+
+`--mode sft --prompt-data ` trains SFT instead. It reads a `messages` column; put a reasoning trace in `reasoning_content`, since the MiMo chat template renders `{reasoning_content}`.
+
+## 5. Recipe Configuration
+
+### 5.1 Parallelism
+
+| `--model-name` | TP | PP | CP | EP | expert-TP | GPUs | rollout engine |
+|---|---|---|---|---|---|---|---|
+| `MiMo-V2.6-Flash-RL-bf16` | 2 | 2 | 1 | 8 | 1 | 16 (2 × 8) | TP8, `--sglang-mem-fraction-static 0.8` |
+
+TP is capped at 4 by the four global-attention KV heads. Context parallelism is not supported. `--sequence-parallel` is on whenever TP > 1. THD packing with `--use-dynamic-batch-size` is the default; `--qkv-format bshd` runs one sample per micro-batch.
+
+### 5.2 Algorithm
+
+GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --entropy-coef 0.00` and `--rm-type math` on dapo-math-17k. No KL loss: bridge mode has no Megatron-format reference checkpoint for `--ref-load`.
+
+### 5.3 Optimizer
+
+Adam at `--lr 1e-6`. The model's Adam state (3.7 TB) fits neither the GPUs nor two hosts' memory, so it streams through node-local NVMe:
+
+```bash
+--stream-optimizer-state-to-disk
+--stream-optimizer-state-moment-dtype bf16
+--grad-reduce-in-bf16
+--offload-train-target disk
+--offload-train-disk-dir
+```
+
+### 5.4 Notable quirks
+
+- **Fused attention only**: `--attention-backend fused`. The learnable sink on the sliding-window layers is implemented by TE's FusedAttention (THD needs cuDNN ≥ 9.18) and the unfused path, not by FlashAttention.
+- **Sliding window**: a SWA layer attends the query plus `sliding_window_size` (128) previous keys, 129 in total, as SGLang does; the bridge sets `window_size = (sliding_window_size, 0)`. The HF modeling code and Megatron-Bridge's `mimo_v2_flash` attend 128 keys in total (`sliding_window_size - 1` previous), so a log-prob comparison against them differs at positions past the window.
+- **Dropout**: the bridge sets `hidden_dropout = 0`. Bridge mode does not copy `--hidden-dropout` onto the provider, so a provider left at Megatron's default (0.1) trains with dropout silently.
+- **Frozen towers and `--check-weight-update-equal`**: the trainer never sends the vision/audio weights, so exclude them from the post-update check with `--check-weight-update-skip-list visual. audio_tokenizer. input_local_transformer. speech_embeddings. projection.mlp.`.
+
+## 6. Pairs Well With
+
+- [Low Precision RL](/advanced/low-precision)
diff --git a/miles/utils/external_utils/command_utils/__init__.py b/miles/utils/external_utils/command_utils/__init__.py
index 56cadbd5f63..43dee99e231 100644
--- a/miles/utils/external_utils/command_utils/__init__.py
+++ b/miles/utils/external_utils/command_utils/__init__.py
@@ -2,6 +2,7 @@
CommandUtilConfig,
ExecuteTrainConfig,
default_config,
+ exclusive_path_lock,
resolve_extra_env_vars,
resolve_hardware,
)
@@ -30,6 +31,7 @@
"default_config",
"detect_hardware",
"encode_pseudo_file",
+ "exclusive_path_lock",
"get_bool_env_var",
"get_default_wandb_args",
"get_env_enable_infinite_run",
diff --git a/miles/utils/external_utils/command_utils/base_backend.py b/miles/utils/external_utils/command_utils/base_backend.py
index 95f64060700..576c938daec 100644
--- a/miles/utils/external_utils/command_utils/base_backend.py
+++ b/miles/utils/external_utils/command_utils/base_backend.py
@@ -213,7 +213,7 @@ def convert_checkpoint(
# TODO shall we make it in host-mapped folder and thus can cache it to speedup CI
path_dst = f"{dir_dst}/{model_name}_torch_dist"
- with _exclusive_path_lock(path_dst):
+ with exclusive_path_lock(path_dst):
tracker = Path(path_dst) / "latest_checkpointed_iteration.txt"
if tracker.exists() and tracker.read_text().strip() == "release":
logger.info(f"convert_checkpoint skip {path_dst} since tracker is 'release'")
@@ -275,7 +275,7 @@ def hf_download_dataset(self, full_name: str, data_dir: str = "/root/datasets"):
self.exec_command_cpu(f"hf download --repo-type dataset {full_name} --local-dir {data_dir}/{partial_name}")
def fp8_cast_bf16(self, path_src, path_dst):
- with _exclusive_path_lock(path_dst):
+ with exclusive_path_lock(path_dst):
sentinel = Path(path_dst) / "model.safetensors.index.json"
if sentinel.exists():
logger.info(f"fp8_cast_bf16 skip {path_dst} since {sentinel} exists")
@@ -338,7 +338,7 @@ def resolve_extra_env_vars(extra_env_vars: dict[str, str], config: ExecuteTrainC
@contextmanager
-def _exclusive_path_lock(path_dst: str):
+def exclusive_path_lock(path_dst: str):
"""Serialize prepare steps racing on a cache dir shared between runners on one host."""
path = Path(path_dst)
path.parent.mkdir(parents=True, exist_ok=True)
diff --git a/miles_plugins/megatron_bridge/__init__.py b/miles_plugins/megatron_bridge/__init__.py
index c71173d6e32..bed2ee7937d 100644
--- a/miles_plugins/megatron_bridge/__init__.py
+++ b/miles_plugins/megatron_bridge/__init__.py
@@ -76,3 +76,8 @@ def broadcast_obj_from_pp_rank(self, obj, *args, **kwargs):
from . import nemotron_h # noqa: F401
except Exception as _e: # pragma: no cover - defensive
logger.warning("miles nemotron_h plugin failed to load: %s", _e)
+
+try:
+ from . import mimo_v2 # noqa: F401
+except Exception as _e: # pragma: no cover - defensive
+ logger.warning("miles mimo_v2 plugin failed to load: %s", _e)
diff --git a/miles_plugins/megatron_bridge/mimo_v2.py b/miles_plugins/megatron_bridge/mimo_v2.py
new file mode 100644
index 00000000000..0ac3c7c124a
--- /dev/null
+++ b/miles_plugins/megatron_bridge/mimo_v2.py
@@ -0,0 +1,264 @@
+"""Megatron-Bridge support for XiaomiMiMo MiMo-V2 (``MiMoV2ForCausalLM``, e.g. MiMo-V2.6-Flash-RL).
+
+Builds on NVIDIA Megatron-Bridge ``megatron/bridge/models/mimo_v2_flash``, which covers the earlier
+MiMo-V2-Flash checkpoint, and reuses its q/k/v weight mapping. MiMo-V2 attention differs from the
+stock GPT layer in four ways:
+
+* per-layer KV heads: SWA layers use ``swa_num_key_value_heads``, global layers ``num_key_value_heads``;
+* a learnable attention sink on SWA layers only, while ``softmax_type`` is one global setting;
+* asymmetric head dims: Q/K use ``head_dim`` (192) and V uses ``v_head_dim`` (128);
+* V is multiplied by ``attention_value_scale`` before attention.
+
+The first two are per-layer config values: the provider sets ``heterogeneous_block_specs``, so
+``TransformerBlock`` builds each layer from ``MiMoV2ModelProvider.get_config_for_layer``. The other
+two need ``MiMoV2SelfAttention`` and ``MiMoV2TEDotProductAttention``. The per-layer RoPE base and the
+SWA window reuse existing config fields (``rotary_base_per_layer``, ``window_size`` with a per-layer
+``window_attn_skip_freq``) instead of Megatron-Bridge's dual-base RoPE and window rule.
+
+The bridge reads the BF16 split-q/k/v layout written by ``tools/convert_mimo_v2_to_bf16.py``; the
+official checkpoint (FP8/MXFP4, fused kv-head-interleaved ``qkv_proj``) must be converted first.
+Only the text decoder is built: vision/audio towers and the MTP layers are not.
+"""
+
+from __future__ import annotations
+
+import copy
+import dataclasses
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
+from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
+from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping
+from megatron.bridge.models.gpt_provider import GPTModelProvider
+from megatron.bridge.models.mimo_v2_flash.mimo_v2_flash_bridge import MiMoV2FlashQKVMapping
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.transformer import ModuleSpec
+from megatron.core.transformer.attention import SelfAttention
+
+logger = logging.getLogger(__name__)
+
+
+def _sized_linear(build, *, input_size: int | None = None, output_size: int | None = None):
+ """``build`` with the given input/output size in place of the one SelfAttention passes."""
+
+ def build_sized(base_input_size, base_output_size, **kwargs):
+ return build(
+ base_input_size if input_size is None else input_size,
+ base_output_size if output_size is None else output_size,
+ **kwargs,
+ )
+
+ return build_sized
+
+
+class MiMoV2SelfAttention(SelfAttention):
+ """SelfAttention whose V heads (``v_head_dim``) are narrower than its Q/K heads (``kv_channels``).
+
+ The stock layer sizes ``linear_qkv``, ``linear_proj`` and the q/k/v split with ``kv_channels`` for V;
+ this builds the two projections once at the MiMo sizes and splits accordingly. The layer config
+ already carries this layer's KV heads (``MiMoV2ModelProvider.get_config_for_layer``).
+ """
+
+ def __init__(self, config, submodules, layer_number, *args, **kwargs):
+ assert not (config.attention_output_gate or config.head_wise_attn_gate), "MiMo-V2 has no attention gate"
+ # qk_clip reads linear_qkv as q + 2 * kv rows of kv_channels each, which does not hold for a narrower V.
+ assert not config.qk_clip, "qk_clip does not support a V head narrower than Q/K"
+ heads, groups = config.num_attention_heads, config.num_query_groups
+ qkv_out_dim = heads * config.kv_channels + groups * (config.kv_channels + config.v_head_dim)
+ submodules = dataclasses.replace(
+ submodules,
+ linear_qkv=_sized_linear(submodules.linear_qkv, output_size=qkv_out_dim),
+ linear_proj=_sized_linear(submodules.linear_proj, input_size=heads * config.v_head_dim),
+ )
+ super().__init__(config, submodules, layer_number, *args, **kwargs)
+ assert groups >= self.world_size, "TP must not exceed the layer's KV heads"
+ self.linear_qkv_out_dim = qkv_out_dim
+ self.val_hidden_size = config.v_head_dim
+
+ def get_query_key_value_tensors(
+ self, hidden_states, key_value_states=None, output_gate=False, head_wise_gate=False, split_qkv=True
+ ):
+ assert not (output_gate or head_wise_gate) and split_qkv, "MiMo-V2 attention needs the split q/k/v path"
+ mixed_qkv, _ = self.linear_qkv(hidden_states)
+ groups = self.num_query_groups_per_partition
+ heads_per_group = self.num_attention_heads_per_partition // groups
+ qk_dim, v_dim = self.hidden_size_per_attention_head, self.config.v_head_dim
+ # [sq, b, ng * (np/ng * hn + hn + hv)] -> [sq, b, ng, np/ng * hn + hn + hv]
+ mixed_qkv = mixed_qkv.view(*mixed_qkv.shape[:-1], groups, heads_per_group * qk_dim + qk_dim + v_dim)
+ query, key, value = torch.split(mixed_qkv, [heads_per_group * qk_dim, qk_dim, v_dim], dim=3)
+ query = query.reshape(query.size(0), query.size(1), -1, qk_dim)
+ return query, key, value
+
+
+class MiMoV2TEDotProductAttention(TEDotProductAttention):
+ """TE core attention with separate Q/K and V channels and V scaled by ``attention_value_scale``.
+
+ The learnable sink comes from the layer config's ``softmax_type``.
+ """
+
+ def __init__(self, config, layer_number, attn_mask_type, attention_type, attention_dropout=None, **kwargs):
+ super().__init__(
+ config=config,
+ layer_number=layer_number,
+ attn_mask_type=attn_mask_type,
+ attention_type=attention_type,
+ attention_dropout=attention_dropout,
+ k_channels=config.kv_channels,
+ v_channels=config.v_head_dim,
+ **kwargs,
+ )
+ self.attention_value_scale = config.attention_value_scale
+
+ def forward(self, query, key, value, attention_mask, attn_mask_type, **kwargs):
+ if self.attention_value_scale is not None:
+ value = value * self.attention_value_scale
+ return super().forward(query, key, value, attention_mask, attn_mask_type, **kwargs)
+
+
+AutoMapping.register_module_type("MiMoV2TEDotProductAttention", "column")
+
+
+def mimo_v2_layer_spec(config, vp_stage: int | None = None) -> ModuleSpec:
+ spec = get_gpt_decoder_block_spec(config, use_transformer_engine=True, vp_stage=vp_stage)
+ for layer_spec in spec.layer_specs:
+ layer_spec.submodules.self_attention.module = MiMoV2SelfAttention
+ layer_spec.submodules.self_attention.submodules.core_attention = MiMoV2TEDotProductAttention
+ return spec
+
+
+@dataclass
+class MiMoV2ModelProvider(GPTModelProvider):
+ """GPT provider plus the MiMo-V2 attention fields; each layer is built from ``get_config_for_layer``."""
+
+ transformer_layer_spec: ModuleSpec | Callable = field(default_factory=lambda: mimo_v2_layer_spec)
+ # Makes TransformerBlock build each layer from get_config_for_layer.
+ heterogeneous_block_specs: bool = True
+ hybrid_attention_pattern: list[int] | None = None
+ full_attn_num_query_groups: int = 4
+ swa_num_query_groups: int = 8
+ v_head_dim: int = 128
+ attention_value_scale: float | None = None
+ full_attention_sink: bool = False
+ swa_attention_sink: bool = True
+
+ def get_config_for_layer(self, layer_number: int) -> MiMoV2ModelProvider:
+ """Shallow copy with the KV heads and attention sink of global layer ``layer_number`` (1-indexed)."""
+ assert layer_number <= len(self.hybrid_attention_pattern), f"no attention pattern for layer {layer_number}"
+ swa = self.hybrid_attention_pattern[layer_number - 1]
+ config = copy.copy(self)
+ config.num_query_groups = self.swa_num_query_groups if swa else self.full_attn_num_query_groups
+ sink = self.swa_attention_sink if swa else self.full_attention_sink
+ config.softmax_type = "learnable" if sink else "vanilla"
+ return config
+
+ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> GPTModel:
+ assert self.tensor_model_parallel_size <= min(
+ self.full_attn_num_query_groups, self.swa_num_query_groups
+ ), "MiMo-V2 needs TP <= the smallest per-layer KV head count"
+ assert self.context_parallel_size == 1, "MiMo-V2 attention does not support context parallelism yet"
+ return super().provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
+
+
+@MegatronModelBridge.register_bridge(
+ source="MiMoV2ForCausalLM", target=GPTModel, provider=MiMoV2ModelProvider, model_type="mimo_v2"
+)
+class MiMoV2Bridge(MegatronModelBridge):
+ """HF ``MiMoV2ForCausalLM`` (BF16, split q/k/v) <-> Megatron GPTModel with MiMo-V2 attention."""
+
+ def provider_bridge(self, hf_pretrained) -> MiMoV2ModelProvider:
+ hf = hf_pretrained.config
+ if getattr(hf, "attention_projection_layout", "split") != "split" or getattr(hf, "quantization_config", None):
+ raise ValueError(
+ "MiMo-V2 bridge reads BF16 split q/k/v checkpoints; convert the official checkpoint "
+ "with tools/convert_mimo_v2_to_bf16.py first"
+ )
+ for swa_key, full_key in (
+ ("swa_num_attention_heads", "num_attention_heads"),
+ ("swa_head_dim", "head_dim"),
+ ("swa_v_head_dim", "v_head_dim"),
+ ):
+ assert getattr(hf, swa_key) == getattr(hf, full_key), f"{swa_key} must equal {full_key}"
+
+ provider = super().provider_bridge(hf_pretrained)
+ pattern = [int(p) for p in hf.hybrid_layer_pattern]
+ provider.hybrid_attention_pattern = pattern
+ # A SWA layer attends the query plus sliding_window_size previous keys, like SGLang's
+ # RadixAttention (whose window excludes the query token), so rollout and training agree.
+ provider.window_size = (hf.sliding_window_size, 0)
+ provider.window_attn_skip_freq = pattern
+ provider.rotary_base = hf.rope_theta
+ provider.rotary_base_per_layer = [hf.swa_rope_theta if p else hf.rope_theta for p in pattern]
+ provider.apply_rope_fusion = False
+ provider.full_attn_num_query_groups = hf.num_key_value_heads
+ provider.swa_num_query_groups = hf.swa_num_key_value_heads
+ provider.num_query_groups = hf.num_key_value_heads
+ provider.v_head_dim = hf.v_head_dim
+ provider.attention_value_scale = hf.attention_value_scale
+ provider.full_attention_sink = bool(hf.add_full_attention_sink_bias)
+ provider.swa_attention_sink = bool(hf.add_swa_attention_sink_bias)
+
+ provider.normalization = "RMSNorm"
+ provider.layernorm_epsilon = hf.layernorm_epsilon
+ provider.gated_linear_unit = True
+ provider.add_bias_linear = False
+ provider.position_embedding_type = "rope"
+ provider.share_embeddings_and_output_weights = False
+ # The HF model has no hidden dropout; the TransformerConfig default (0.1) would stay on, since
+ # Miles' bridge mode does not copy --hidden-dropout onto the provider (the bundled bridges
+ # set it the same way). attention_dropout comes from the HF config (0.0).
+ provider.hidden_dropout = 0.0
+
+ provider.moe_layer_freq = [int(f) for f in hf.moe_layer_freq]
+ provider.moe_grouped_gemm = True
+ provider.moe_token_dispatcher_type = "alltoall"
+ provider.moe_router_enable_expert_bias = True
+ provider.moe_router_bias_update_rate = 0.0
+ provider.moe_router_load_balancing_type = "none"
+ provider.moe_aux_loss_coeff = 0.0
+ provider.moe_router_dtype = "fp32"
+ if getattr(hf, "n_group", None) in (None, 1):
+ provider.moe_router_num_groups = None
+ provider.moe_router_group_topk = None
+ provider.mtp_num_layers = None
+ return provider
+
+ def mapping_registry(self) -> MegatronMappingRegistry:
+ direct = {
+ "embedding.word_embeddings.weight": "model.embed_tokens.weight",
+ "output_layer.weight": "lm_head.weight",
+ "decoder.final_layernorm.weight": "model.norm.weight",
+ "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight",
+ "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight",
+ "decoder.layers.*.self_attention.core_attention.softmax_offset": "model.layers.*.self_attn.attention_sink_bias",
+ "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight",
+ "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight",
+ "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight",
+ "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight",
+ "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.e_score_correction_bias",
+ "decoder.layers.*.mlp.experts.linear_fc2.weight*": "model.layers.*.mlp.experts.*.down_proj.weight",
+ "decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight": "model.layers.*.mlp.experts.*.down_proj.weight",
+ }
+ gated = {
+ "decoder.layers.*.mlp.linear_fc1.weight": "model.layers.*.mlp",
+ "decoder.layers.*.mlp.experts.linear_fc1.weight*": "model.layers.*.mlp.experts.*",
+ "decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight": "model.layers.*.mlp.experts.*",
+ }
+ mappings = [AutoMapping(megatron_param=m, hf_param=h) for m, h in direct.items()]
+ mappings += [
+ GatedMLPMapping(megatron_param=m, gate=f"{h}.gate_proj.weight", up=f"{h}.up_proj.weight")
+ for m, h in gated.items()
+ ]
+ mappings.append(
+ MiMoV2FlashQKVMapping(
+ megatron_param="decoder.layers.*.self_attention.linear_qkv.weight",
+ q="model.layers.*.self_attn.q_proj.weight",
+ k="model.layers.*.self_attn.k_proj.weight",
+ v="model.layers.*.self_attn.v_proj.weight",
+ )
+ )
+ return MegatronMappingRegistry(*mappings)
diff --git a/scripts/models/mimo-v2.6-flash-4layer.py b/scripts/models/mimo-v2.6-flash-4layer.py
new file mode 100644
index 00000000000..165e5f3cdb0
--- /dev/null
+++ b/scripts/models/mimo-v2.6-flash-4layer.py
@@ -0,0 +1,7 @@
+from model_args_utils import load_sibling_model_args
+
+
+def model_args() -> str:
+ # 4-layer partial from `tools/convert_mimo_v2_to_bf16.py --layers 0,1,5,6`:
+ # global+dense, SWA+MoE, global+MoE, SWA+MoE (every decoder variant, PP=2 splits it in half).
+ return load_sibling_model_args(__file__, "mimo-v2.6-flash", nlayers=4)
diff --git a/scripts/models/mimo-v2.6-flash.py b/scripts/models/mimo-v2.6-flash.py
new file mode 100644
index 00000000000..7aa2f878d2c
--- /dev/null
+++ b/scripts/models/mimo-v2.6-flash.py
@@ -0,0 +1,52 @@
+# XiaomiMiMo MiMo-V2.6-Flash-RL text decoder. Train through Megatron-Bridge
+# (`--megatron-to-hf-mode bridge`) on the BF16 checkpoint written by
+# tools/convert_mimo_v2_to_bf16.py: the bridge (miles_plugins/megatron_bridge/mimo_v2.py) supplies the
+# per-layer SWA/global attention, sink, V head dim and RoPE bases, which these flags cannot express.
+# The flags mirror the global-attention shape and the MoE layout so Miles and Megatron size the model
+# and its pipeline split the same way the bridge builds it.
+from model_args_utils import moe_layer_freq
+
+NHIDDEN = 4096
+FFN_HIDDEN = 16384
+NHEADS = 64
+GLOBAL_KV_HEADS = 4
+QK_HEAD_DIM = 192
+MOE_ROUTED_EXPERTS = 256
+MOE_ACTIVE_ROUTED_EXPERTS = 8
+MOE_FFN_HIDDEN = 2048
+N_LAYERS = 48
+FIRST_K_DENSE_REPLACE = 1
+
+
+def model_args(nlayers: int = N_LAYERS) -> str:
+ return (
+ f"--num-layers {nlayers} "
+ f"--hidden-size {NHIDDEN} "
+ f"--ffn-hidden-size {FFN_HIDDEN} "
+ f"--num-attention-heads {NHEADS} "
+ "--group-query-attention "
+ f"--num-query-groups {GLOBAL_KV_HEADS} "
+ f"--kv-channels {QK_HEAD_DIM} "
+ "--position-embedding-type rope "
+ "--rotary-percent 0.334 "
+ "--rotary-base 10000000 "
+ "--no-rope-fusion "
+ "--max-position-embeddings 1048576 "
+ "--normalization RMSNorm "
+ "--norm-epsilon 1e-6 "
+ "--swiglu "
+ "--disable-bias-linear "
+ "--untie-embeddings-and-output-weights "
+ "--vocab-size 152576 "
+ f"--moe-layer-freq {moe_layer_freq(nlayers=nlayers, first_k_dense_replace=FIRST_K_DENSE_REPLACE)} "
+ f"--num-experts {MOE_ROUTED_EXPERTS} "
+ f"--moe-router-topk {MOE_ACTIVE_ROUTED_EXPERTS} "
+ f"--moe-ffn-hidden-size {MOE_FFN_HIDDEN} "
+ "--moe-router-score-function sigmoid "
+ "--moe-router-enable-expert-bias "
+ "--moe-router-bias-update-rate 0 "
+ "--moe-router-load-balancing-type none "
+ "--moe-aux-loss-coeff 0 "
+ "--moe-router-dtype fp32 "
+ "--moe-grouped-gemm "
+ )
diff --git a/scripts/run_mimo_v2_6_flash.py b/scripts/run_mimo_v2_6_flash.py
new file mode 100644
index 00000000000..14d91582015
--- /dev/null
+++ b/scripts/run_mimo_v2_6_flash.py
@@ -0,0 +1,268 @@
+"""MiMo-V2.6-Flash-RL training (RL with GRPO, or SFT), adapted from run_gpt_oss_20b.py.
+
+Megatron bridge mode on a BF16 checkpoint converted with tools/convert_mimo_v2_to_bf16.py; the
+official checkpoint (fused kv-interleaved FP8 qkv, MXFP4 experts) cannot be read directly. The
+bridge (miles_plugins/megatron_bridge/mimo_v2.py) builds the per-layer SWA/global attention, the SWA
+sink and the 192/128 Q-K/V head dims, so there is no `torch_dist` conversion and no `--ref-load`
+(bridge mode has no Megatron reference checkpoint, hence no KL loss, as for gpt-oss).
+
+`prepare` writes `/` unless it already exists: it downloads the official
+XiaomiMiMo/MiMo-V2.6-Flash-RL (173 GB) to `/MiMo-V2.6-Flash-RL` and converts it on the GPU:
+ MiMo-V2.6-Flash-RL-bf16 the full model (622 GB)
+ mimo26-p4-bf16 `--layers 0,1,5,6` (46 GB), a 4-layer partial with one instance of every
+ decoder variant: global+dense, SWA+MoE, global+MoE, SWA+MoE
+
+Args:
+ --mode: `rl` runs GRPO with full-policy sampling (top-p 1) and a colocated BF16 SGLang engine
+ on the same checkpoint (dapo-math-17k from `--data-dir` unless `--prompt-data` is given);
+ `sft` runs `--debug-train-only` SFT on chat data with a `messages` column (no SGLang),
+ following run_qwen3_sft.py.
+ --model-name: checkpoint directory under `--model-dir`; selects the matching model args.
+ --tensor/pipeline/expert-model-parallel-size: TP (with sequence parallel when > 1), PP and EP;
+ default 2/2/2 for the partial and 2/2/8 for the full model (16 GPUs). TP must not exceed 4,
+ the global-attention KV head count. Context parallel is not supported.
+ --recompute / --no-recompute: full uniform recompute, one layer per checkpoint.
+ --async-train / --no-async-train: SFT only; `train_async.py` prefetches the next batch, so a run
+ resumed from its checkpoint skips one batch, while `train.py` resumes at the next batch.
+ --qkv-format: `thd` packs samples with dynamic batching (`--max-tokens-per-gpu`); `bshd` runs one
+ unpacked sample per micro-batch.
+ --num-rollout, --rollout-batch-size: steps and samples (prompts in RL) per step.
+ --save / --save-interval / --load: Megatron checkpoints under `/checkpoints`.
+ --train-offload-disk-dir: NVMe directory for the full model's streamed Adam state (bf16 moments,
+ bf16 gradient reduction) and, in RL, the actor offloaded while the engines generate.
+ --extra-args: appended verbatim to the train argv.
+
+The full model needs two nodes: start the ray head on the first, join the second, and run with
+`MILES_SCRIPT_EXTERNAL_RAY=1 MASTER_ADDR= --num-nodes 2`.
+
+Examples:
+ python scripts/run_mimo_v2_6_flash.py --mode sft --model-name mimo26-p4-bf16 --prompt-data
+ python scripts/run_mimo_v2_6_flash.py --mode rl --model-name mimo26-p4-bf16
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import typer
+
+import miles.utils.external_utils.command_utils as U
+
+_HF_REPO = "XiaomiMiMo/MiMo-V2.6-Flash-RL"
+
+
+@dataclass(frozen=True)
+class _Recipe:
+ megatron_model_type: str
+ # Source decoder layers kept by tools/convert_mimo_v2_to_bf16.py; None keeps all.
+ layers: str | None = None
+ # Default TP / PP / EP.
+ parallel: tuple[int, int, int] = (2, 2, 2)
+ # GPUs per colocated BF16 engine.
+ rollout_num_gpus_per_engine: int = 4
+ sglang_mem_fraction_static: float = 0.6
+ # Full-parameter Adam state of the full model (3.7 TB) fits neither 16 GPUs nor two hosts'
+ # memory, so it streams through node-local NVMe; that needs bf16 gradient reduction.
+ stream_optimizer_state: bool = False
+
+
+_RECIPES = {
+ "MiMo-V2.6-Flash-RL-bf16": _Recipe(
+ megatron_model_type="mimo-v2.6-flash",
+ parallel=(2, 2, 8),
+ # The BF16 engine holds the whole model: 620 GB needs TP8 on 141 GB GPUs.
+ rollout_num_gpus_per_engine=8,
+ sglang_mem_fraction_static=0.8,
+ stream_optimizer_state=True,
+ ),
+ "mimo26-p4-bf16": _Recipe(megatron_model_type="mimo-v2.6-flash-4layer", layers="0,1,5,6"),
+}
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+ run_id: str = U.create_run_id()
+ mode: Literal["rl", "sft"] = "rl"
+ model_name: Literal["MiMo-V2.6-Flash-RL-bf16", "mimo26-p4-bf16"] = "mimo26-p4-bf16"
+ num_gpus_per_node: int = 8
+ # None takes the recipe default of --model-name.
+ tensor_model_parallel_size: int | None = None
+ pipeline_model_parallel_size: int | None = None
+ expert_model_parallel_size: int | None = None
+ recompute: bool = True
+ # SFT uses train_async.py like run_qwen3_sft.py; --no-async-train runs train.py, whose
+ # checkpoints resume exactly at the next unconsumed batch.
+ async_train: bool = True
+ # thd packs samples with dynamic batching; bshd runs one unpacked sample per micro-batch.
+ qkv_format: Literal["thd", "bshd"] = "thd"
+ max_tokens_per_gpu: int = 9216
+ num_rollout: int = 4
+ rollout_batch_size: int = 16
+ prompt_data: str = ""
+ save: bool = False
+ save_interval: int = 2
+ load: str = ""
+ extra_args: str = ""
+ data_dir: str = "/root/datasets"
+ model_dir: str = "/root/models"
+ megatron_path: str = "/root/Megatron-LM"
+ # NVMe directory for the streamed optimizer state and the offloaded actor (full model).
+ train_offload_disk_dir: str = "/root/shared_data/train_offload"
+
+ def __post_init__(self):
+ recipe = _RECIPES[self.model_name]
+ tp, pp, ep = recipe.parallel
+ self.tensor_model_parallel_size = self.tensor_model_parallel_size or tp
+ self.pipeline_model_parallel_size = self.pipeline_model_parallel_size or pp
+ self.expert_model_parallel_size = self.expert_model_parallel_size or ep
+
+
+def prepare(args: ScriptArgs):
+ backend = args.create_backend()
+ target = Path(args.model_dir) / args.model_name
+ with U.exclusive_path_lock(str(target)):
+ # The converter writes the index last, so it marks a finished conversion.
+ if not (target / "model.safetensors.index.json").exists():
+ source = f"{args.model_dir}/{_HF_REPO.split('/')[1]}"
+ layers = _RECIPES[args.model_name].layers
+ layer_args = f"--layers {layers}" if layers else ""
+ backend.exec_command_cpu(f"mkdir -p {args.model_dir}")
+ backend.exec_command_cpu(f"hf download {_HF_REPO} --local-dir {source}")
+ backend.exec_command_gpu(
+ f"python {U.repo_base_dir}/tools/convert_mimo_v2_to_bf16.py "
+ f"--model-dir {source} --save-dir {target} --device cuda {layer_args}"
+ )
+ if args.mode == "rl" and not args.prompt_data:
+ backend.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir)
+
+
+def execute(args: ScriptArgs):
+ backend = args.create_backend()
+ ckpt_args = f"--hf-checkpoint {args.model_dir}/{args.model_name} " "--megatron-to-hf-mode bridge "
+ if args.save:
+ ckpt_args += f"--save {args.output_dir}/checkpoints " f"--save-interval {args.save_interval} "
+ if args.load:
+ ckpt_args += f"--load {args.load} "
+
+ if args.mode == "sft":
+ data_args = (
+ "--rollout-function-path miles.rollout.sft_rollout.generate_rollout "
+ f"--prompt-data {args.prompt_data} "
+ "--input-key messages "
+ "--rollout-shuffle "
+ f"--num-rollout {args.num_rollout} "
+ f"--rollout-batch-size {args.rollout_batch_size} "
+ f"--global-batch-size {args.rollout_batch_size} "
+ "--loss-type sft_loss "
+ "--calculate-per-token-loss "
+ "--disable-compute-advantages-and-returns "
+ "--debug-train-only "
+ )
+ else:
+ data_args = (
+ f"--prompt-data {args.prompt_data or f'{args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl'} "
+ "--input-key prompt "
+ "--label-key label "
+ "--apply-chat-template "
+ "--rollout-shuffle "
+ "--rm-type math "
+ f"--num-rollout {args.num_rollout} "
+ f"--rollout-batch-size {args.rollout_batch_size} "
+ "--n-samples-per-prompt 8 "
+ "--rollout-max-response-len 8192 "
+ "--rollout-temperature 1.0 "
+ "--rollout-top-p 1.0 "
+ "--num-steps-per-rollout 1 "
+ "--advantage-estimator grpo "
+ "--entropy-coef 0.00 "
+ "--eps-clip 0.2 "
+ "--eps-clip-high 0.28 "
+ )
+
+ tp = args.tensor_model_parallel_size
+ perf_args = (
+ f"--tensor-model-parallel-size {tp} "
+ f"{'--sequence-parallel ' if tp > 1 else ''}"
+ f"--pipeline-model-parallel-size {args.pipeline_model_parallel_size} "
+ "--context-parallel-size 1 "
+ f"--expert-model-parallel-size {args.expert_model_parallel_size} "
+ "--expert-tensor-parallel-size 1 "
+ )
+ if args.qkv_format == "thd":
+ perf_args += f"--use-dynamic-batch-size --max-tokens-per-gpu {args.max_tokens_per_gpu} "
+ else:
+ perf_args += "--qkv-format bshd --micro-batch-size 1 "
+ if args.recompute:
+ perf_args += "--recompute-granularity full --recompute-method uniform --recompute-num-layers 1 "
+
+ recipe = _RECIPES[args.model_name]
+ optimizer_args = (
+ "--optimizer adam "
+ "--lr 1e-6 "
+ "--lr-decay-style constant "
+ "--weight-decay 0.1 "
+ "--adam-beta1 0.9 "
+ "--adam-beta2 0.98 "
+ )
+ if recipe.stream_optimizer_state:
+ optimizer_args += (
+ "--stream-optimizer-state-to-disk "
+ "--stream-optimizer-state-moment-dtype bf16 "
+ "--grad-reduce-in-bf16 "
+ f"--offload-train-disk-dir {args.train_offload_disk_dir} "
+ )
+ if args.mode == "rl":
+ optimizer_args += "--offload-train-target disk "
+ else:
+ optimizer_args += "--accumulate-allreduce-grads-in-fp32 "
+
+ sglang_args = (
+ f"--rollout-num-gpus-per-engine {recipe.rollout_num_gpus_per_engine} "
+ "--sglang-dtype bfloat16 "
+ f"--sglang-mem-fraction-static {recipe.sglang_mem_fraction_static} "
+ "--sglang-decode-log-interval 1000 "
+ )
+
+ misc_args = (
+ "--attention-dropout 0.0 "
+ "--hidden-dropout 0.0 "
+ "--attention-softmax-in-fp32 "
+ "--attention-backend fused "
+ f"--actor-num-nodes {args.num_nodes} "
+ f"--actor-num-gpus-per-node {args.num_gpus_per_node} "
+ f"--num-gpus-per-node {args.num_gpus_per_node} "
+ )
+ if args.mode == "rl":
+ misc_args += "--colocate "
+
+ train_args = (
+ f"{ckpt_args} "
+ f"{data_args} "
+ f"{optimizer_args} "
+ f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} "
+ f"{perf_args} "
+ f"{sglang_args if args.mode == 'rl' else ''} "
+ f"{misc_args} "
+ f"{args.extra_args} "
+ )
+
+ backend.execute_train(
+ train_args=train_args,
+ config=args,
+ num_gpus_per_node=args.num_gpus_per_node,
+ megatron_model_type=recipe.megatron_model_type,
+ megatron_path=args.megatron_path,
+ train_script="train_async.py" if args.mode == "sft" and args.async_train else "train.py",
+ # Colocated RL offloads through torch_memory_saver, which rejects expandable segments.
+ extra_env_vars={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"} if args.mode == "sft" else {},
+ )
+
+
+@U.dataclass_cli
+def main(args: ScriptArgs):
+ prepare(args)
+ execute(args)
+
+
+if __name__ == "__main__":
+ typer.run(main)
diff --git a/tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py b/tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py
new file mode 100644
index 00000000000..fc6be4e0feb
--- /dev/null
+++ b/tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py
@@ -0,0 +1,57 @@
+import os
+
+from scripts.run_mimo_v2_6_flash import ScriptArgs, execute, prepare
+from tests.ci.ci_register import register_cuda_ci
+from tests.ci.metric_history import register_ci_gate
+
+# Smoke test of the MiMo-V2.6 bridge path on the 4-layer partial (source layers 0, 1, 5, 6: global+dense,
+# SWA+MoE, global+MoE, SWA+MoE, all 256 experts): two RL rollout/train/sync cycles with routing replay and
+# the post-update weight check. The partial is not a language model (its NLL is ~16), so the log-prob and
+# entropy checkers of --ci-test do not apply; the metric-history gates below track the train/rollout gap.
+
+register_cuda_ci(
+ est_time=500,
+ suite="stage-c-8-gpu-h200",
+ labels=["megatron", "model-scripts"],
+ hardware=["hopper"],
+ nightly=True,
+)
+
+register_ci_gate(metric_key="train/grad_norm")
+register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")
+register_ci_gate(metric_key="train/train_rollout_kl")
+
+# The trainer is text-only and never sends the vision / audio / speech towers the engine loads.
+_FROZEN_TOWERS = "visual. audio_tokenizer. input_local_transformer. speech_embeddings. projection.mlp."
+
+
+def _args() -> ScriptArgs:
+ return ScriptArgs(
+ mode="rl",
+ model_name="mimo26-p4-bf16",
+ num_rollout=2,
+ rollout_batch_size=4,
+ extra_args=(
+ "--ci-test "
+ "--ci-disable-logprobs-checker "
+ # The eval-mode log-prob pass and the train-mode pass differ at BF16 level on this model (same
+ # input and weights: NLL 16.08760 vs 16.08832), so step-0 ppo_kl is ~1e-5, not bitwise zero.
+ "--ci-disable-kl-checker "
+ "--use-rollout-routing-replay "
+ "--check-weight-update-equal "
+ "--check-weight-update-selector target "
+ f"--check-weight-update-skip-list {_FROZEN_TOWERS} "
+ "--n-samples-per-prompt 2 "
+ "--global-batch-size 8 "
+ "--rollout-max-response-len 2048 "
+ "--sglang-max-running-requests 8 "
+ ),
+ )
+
+
+if __name__ == "__main__":
+ args = _args()
+ prepare(args)
+ for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+ os.environ.pop(proxy_var, None)
+ execute(args)
diff --git a/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py b/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py
new file mode 100644
index 00000000000..1570f82be0c
--- /dev/null
+++ b/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py
@@ -0,0 +1,63 @@
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+
+import pytest
+
+from tests.fast.launch_scripts.py_harness import freeze_environment, import_launch_script, install_command_recorder
+from tests.fast.launch_scripts.sh_harness import REPO_ROOT
+
+from miles.utils.external_utils.command_utils import exclusive_path_lock
+from miles.utils.external_utils.command_utils.base_backend import BaseCommandBackend
+
+
+@pytest.mark.parametrize("mode", ["rl", "sft"])
+@pytest.mark.parametrize("checkpoint_exists", [False, True])
+def test_prepare_and_execute_with_the_configured_backend(monkeypatch, tmp_path, mode, checkpoint_exists):
+ freeze_environment(monkeypatch)
+ recording = install_command_recorder(monkeypatch)
+ launcher = import_launch_script(REPO_ROOT / "scripts/run_mimo_v2_6_flash.py")
+ args = launcher.ScriptArgs(mode=mode, model_dir=str(tmp_path / "models"), data_dir=str(tmp_path / "data"))
+ if checkpoint_exists:
+ checkpoint = tmp_path / "models" / args.model_name
+ checkpoint.mkdir(parents=True)
+ (checkpoint / "model.safetensors.index.json").write_text("{}")
+
+ launcher.prepare(args)
+ launcher.execute(args)
+
+ commands = recording.commands
+ assert any("hf download XiaomiMiMo/MiMo-V2.6-Flash-RL" in cmd for cmd in commands) is not checkpoint_exists
+ assert any("convert_mimo_v2_to_bf16.py" in cmd for cmd in commands) is not checkpoint_exists
+ assert any("hf download --repo-type dataset zhuzilin/dapo-math-17k" in cmd for cmd in commands) is (mode == "rl")
+ train_commands = [cmd for cmd in commands if "ray job submit" in cmd]
+ assert len(train_commands) == 1
+ train_script = "train.py" if mode == "rl" else "train_async.py"
+ assert f"/{train_script} " in train_commands[0]
+
+
+def test_prepare_waits_for_shared_checkpoint_conversion(monkeypatch, tmp_path):
+ freeze_environment(monkeypatch)
+ install_command_recorder(monkeypatch)
+ launcher = import_launch_script(REPO_ROOT / "scripts/run_mimo_v2_6_flash.py")
+ args = launcher.ScriptArgs(mode="sft", model_dir=str(tmp_path / "models"))
+ target = tmp_path / "models" / args.model_name
+ started, converted = Event(), Event()
+ monkeypatch.setattr(BaseCommandBackend, "exec_command_gpu", lambda *args, **kwargs: converted.set())
+
+ def prepare():
+ started.set()
+ launcher.prepare(args)
+
+ # Preparation takes a real filesystem lock even when its shell commands are
+ # recorded, so the launcher snapshots also need a writable model_dir.
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ with exclusive_path_lock(str(target)):
+ future = pool.submit(prepare)
+ assert started.wait(timeout=5)
+ attempted_conversion = converted.wait(timeout=0.2)
+ target.mkdir(parents=True)
+ (target / "model.safetensors.index.json").write_text("{}")
+ future.result(timeout=5)
+
+ assert not attempted_conversion
+ assert not converted.is_set()
diff --git a/tests/fast/tools/test_convert_mimo_v2_to_bf16.py b/tests/fast/tools/test_convert_mimo_v2_to_bf16.py
new file mode 100644
index 00000000000..70f867424ea
--- /dev/null
+++ b/tests/fast/tools/test_convert_mimo_v2_to_bf16.py
@@ -0,0 +1,195 @@
+import json
+
+import pytest
+import tools.convert_mimo_v2_to_bf16 as converter
+import torch
+from safetensors.torch import load_file, save_file
+from tools.convert_mimo_v2_to_bf16 import FP4_TABLE, dequant_fused_qkv, dequant_mxfp4, main, split_fused_qkv
+
+E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+def _pack_mxfp4(codes: torch.Tensor) -> torch.Tensor:
+ """codes: [out, in] 4-bit indices into FP4_TABLE -> [out, in / 2] uint8, element 2i in the low nibble."""
+ return (codes[:, 0::2] | (codes[:, 1::2] << 4)).to(torch.uint8)
+
+
+def _quant_fp8_block(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ """(128, 128) block quantization of one matrix, rows padded to whole blocks."""
+ rows, cols = weight.shape
+ padded = torch.zeros(-(-rows // 128) * 128, -(-cols // 128) * 128)
+ padded[:rows, :cols] = weight
+ blocks = padded.view(padded.shape[0] // 128, 128, padded.shape[1] // 128, 128)
+ scale = blocks.abs().amax(dim=(1, 3)).clamp(min=1e-12) / E4M3_MAX
+ quant = (blocks / scale[:, None, :, None]).view_as(padded)[:rows, :cols].to(torch.float8_e4m3fn)
+ return quant, scale
+
+
+def _interleave_qkv(q, k, v, shards):
+ return torch.cat(
+ [torch.cat(parts) for parts in zip(q.chunk(shards), k.chunk(shards), v.chunk(shards), strict=True)]
+ )
+
+
+def test_dequant_mxfp4_is_exact():
+ codes = torch.randint(0, 16, (8, 64))
+ exponents = torch.randint(120, 134, (8, 2), dtype=torch.uint8)
+ expected = FP4_TABLE[codes] * torch.exp2(exponents.float() - 127).repeat_interleave(32, dim=1)
+ assert torch.equal(dequant_mxfp4(_pack_mxfp4(codes), exponents).float(), expected.to(torch.bfloat16).float())
+
+
+def test_fused_qkv_scales_are_per_shard():
+ # Shards of 224 rows are not 128-aligned: each one carries its own 2 scale rows.
+ shards, rows = 2, 224
+ weight = torch.randn(shards * rows, 256)
+ quant_scale = [_quant_fp8_block(part) for part in weight.chunk(shards)]
+ fp8 = torch.cat([q for q, _ in quant_scale])
+ scale = torch.cat([s for _, s in quant_scale])
+ assert scale.shape == (shards * 2, 2)
+ expected = torch.cat(
+ [q.float() * s.repeat_interleave(128, 0)[:rows].repeat_interleave(128, 1) for q, s in quant_scale]
+ )
+ out = dequant_fused_qkv(fp8, scale, shards)
+ assert torch.equal(out, expected.to(torch.bfloat16))
+ assert (out.float() - weight).abs().max() < 0.07 * weight.abs().max()
+
+
+def test_split_fused_qkv_restores_head_order():
+ config = {
+ "num_key_value_heads": 2,
+ "swa_num_attention_heads": 8,
+ "swa_num_key_value_heads": 4,
+ "swa_head_dim": 3,
+ "swa_v_head_dim": 2,
+ }
+ q, k, v = torch.randn(8 * 3, 5), torch.randn(4 * 3, 5), torch.randn(4 * 2, 5)
+ out = split_fused_qkv(_interleave_qkv(q, k, v, shards=2), config, is_swa=True)
+ for got, want in zip(out, (q, k, v), strict=True):
+ assert torch.equal(got, want)
+
+
+def _tiny_checkpoint(path):
+ """Two layers in the official MiMo-V2 format: global+dense (0) and SWA+MoE (1)."""
+ config = {
+ "architectures": ["MiMoV2ForCausalLM"],
+ "attention_projection_layout": "fused_qkv",
+ "num_hidden_layers": 2,
+ "hybrid_layer_pattern": [0, 1],
+ "moe_layer_freq": [0, 1],
+ "hybrid_block_size": None,
+ "num_attention_heads": 4,
+ "num_key_value_heads": 2,
+ "head_dim": 64,
+ "v_head_dim": 32,
+ "swa_num_attention_heads": 4,
+ "swa_num_key_value_heads": 4,
+ "swa_head_dim": 64,
+ "swa_v_head_dim": 32,
+ "add_swa_attention_sink_bias": True,
+ "n_routed_experts": 4,
+ "num_experts_per_tok": 2,
+ "quantization_config": {
+ "quant_method": "fp8",
+ "store_dtype": "mxfp4",
+ "ignored_layers": ["model.layers.1.self_attn.o_proj"],
+ },
+ }
+ expected, tensors = {}, {"model.embed_tokens.weight": torch.randn(16, 128).bfloat16()}
+ for layer, (kv_heads, is_swa) in enumerate(((2, False), (4, True))):
+ p = f"model.layers.{layer}."
+ q, k, v = torch.randn(4 * 64, 128), torch.randn(kv_heads * 64, 128), torch.randn(kv_heads * 32, 128)
+ fused = _interleave_qkv(q, k, v, shards=2)
+ quant_scale = [_quant_fp8_block(part) for part in fused.chunk(2)]
+ tensors[p + "self_attn.qkv_proj.weight"] = torch.cat([qq for qq, _ in quant_scale])
+ tensors[p + "self_attn.qkv_proj.weight_scale_inv"] = torch.cat([s for _, s in quant_scale])
+ deq = dequant_fused_qkv(
+ tensors[p + "self_attn.qkv_proj.weight"], tensors[p + "self_attn.qkv_proj.weight_scale_inv"], 2
+ )
+ for proj, part in zip("qkv", split_fused_qkv(deq, config, is_swa), strict=True):
+ expected[p + f"self_attn.{proj}_proj.weight"] = part
+ for name in ("input_layernorm.weight", "post_attention_layernorm.weight"):
+ tensors[p + name] = torch.ones(128).bfloat16()
+ tensors[p + "self_attn.o_proj.weight"] = torch.randn(128, 4 * 32).bfloat16()
+ tensors["model.layers.1.self_attn.attention_sink_bias"] = torch.zeros(4).bfloat16()
+ for proj, shape in (("gate_proj", (256, 128)), ("up_proj", (256, 128)), ("down_proj", (128, 256))):
+ quant, scale = _quant_fp8_block(torch.randn(shape))
+ tensors[f"model.layers.0.mlp.{proj}.weight"] = quant
+ tensors[f"model.layers.0.mlp.{proj}.weight_scale_inv"] = scale
+ tensors["model.layers.1.mlp.gate.weight"] = torch.randn(4, 128).bfloat16()
+ tensors["model.layers.1.mlp.gate.e_score_correction_bias"] = torch.randn(4)
+ for expert in range(4):
+ for proj, (rows, cols) in (("gate_proj", (64, 128)), ("up_proj", (64, 128)), ("down_proj", (128, 64))):
+ name = f"model.layers.1.mlp.experts.{expert}.{proj}."
+ codes = torch.randint(0, 16, (rows, cols))
+ exps = torch.randint(124, 130, (rows, cols // 32), dtype=torch.uint8)
+ tensors[name + "weight"], tensors[name + "weight_scale"] = _pack_mxfp4(codes), exps
+ expected[name + "weight"] = dequant_mxfp4(tensors[name + "weight"], exps)
+ save_file(tensors, str(path / "model-00001.safetensors"))
+ (path / "model.safetensors.index.json").write_text(
+ json.dumps({"weight_map": dict.fromkeys(tensors, "model-00001.safetensors")})
+ )
+ (path / "config.json").write_text(json.dumps(config))
+ (path / "tokenizer_config.json").write_text("{}")
+ return expected
+
+
+@pytest.mark.parametrize("layers, num_experts", [(None, None), ([1], 2)])
+def test_convert_end_to_end(tmp_path, layers, num_experts):
+ src, dst = tmp_path / "src", tmp_path / "dst"
+ src.mkdir()
+ expected = _tiny_checkpoint(src)
+ main(str(src), str(dst), layers, num_experts, keep_quant=False, device="cpu")
+
+ config = json.loads((dst / "config.json").read_text())
+ index = json.loads((dst / "model.safetensors.index.json").read_text())["weight_map"]
+ out = {name: t for f in set(index.values()) for name, t in load_file(str(dst / f)).items()}
+ assert "quantization_config" not in config and "attention_projection_layout" not in config
+ assert (dst / "tokenizer_config.json").exists()
+ assert not any(name.endswith(("weight_scale", "weight_scale_inv", "qkv_proj.weight")) for name in out)
+ remap = {old: new for new, old in enumerate(layers)} if layers else {0: 0, 1: 1}
+ assert config["num_hidden_layers"] == len(remap)
+ assert config["hybrid_layer_pattern"] == config["moe_layer_freq"] == list(remap)
+ num_kept_experts = num_experts or 4
+ kept = {}
+ for name, want in expected.items():
+ _, _, old, *rest = name.split(".")
+ if int(old) in remap and not (rest[1] == "experts" and int(rest[2]) >= num_kept_experts):
+ kept[f"model.layers.{remap[int(old)]}.{'.'.join(rest)}"] = want
+ assert len(kept) == (3 if 0 in remap else 0) + 3 + 3 * num_kept_experts
+ for name, want in kept.items():
+ assert torch.equal(out[name], want), name
+ experts = {name for name in out if ".experts." in name}
+ assert len(experts) == 3 * num_kept_experts
+ assert config["n_routed_experts"] == num_kept_experts
+ assert out[f"model.layers.{remap[1]}.mlp.gate.weight"].shape[0] == num_kept_experts
+ assert out[f"model.layers.{remap[1]}.mlp.gate.e_score_correction_bias"].dtype == torch.float32
+
+
+def test_keep_quant_reindexes_ignored_layers(tmp_path):
+ src, dst = tmp_path / "src", tmp_path / "dst"
+ src.mkdir()
+ _tiny_checkpoint(src)
+ main(str(src), str(dst), [1], None, keep_quant=True, device="cpu")
+ config = json.loads((dst / "config.json").read_text())
+ assert config["quantization_config"]["ignored_layers"] == ["model.layers.0.self_attn.o_proj"]
+ assert config["attention_projection_layout"] == "fused_qkv"
+
+
+def test_conversion_is_not_complete_until_auxiliary_files_are_copied(tmp_path, monkeypatch):
+ src, dst = tmp_path / "src", tmp_path / "dst"
+ src.mkdir()
+ _tiny_checkpoint(src)
+ copy2 = converter.shutil.copy2
+
+ def interrupted_copy(source, destination):
+ raise OSError("interrupted auxiliary copy")
+
+ monkeypatch.setattr(converter.shutil, "copy2", interrupted_copy)
+ with pytest.raises(OSError, match="interrupted auxiliary copy"):
+ main(str(src), str(dst), [1], None, keep_quant=True, device="cpu")
+ assert not (dst / "model.safetensors.index.json").exists()
+
+ monkeypatch.setattr(converter.shutil, "copy2", copy2)
+ main(str(src), str(dst), [1], None, keep_quant=True, device="cpu")
+ assert json.loads((dst / "model.safetensors.index.json").read_text())["weight_map"]
+ assert (dst / "tokenizer_config.json").read_text() == (src / "tokenizer_config.json").read_text()
diff --git a/tests/manual/launch_scripts/test_py_launch_scripts.py b/tests/manual/launch_scripts/test_py_launch_scripts.py
index 75d1b92a458..5841ee54744 100644
--- a/tests/manual/launch_scripts/test_py_launch_scripts.py
+++ b/tests/manual/launch_scripts/test_py_launch_scripts.py
@@ -61,6 +61,8 @@ def _nemotron_checkpoint(sandbox: Path) -> dict[str, object]:
"scripts/run_inkling.py": lambda sandbox: {"model_name": "Inkling-4layer"},
# convert_checkpoint locks a file under model_dir, which must be writable while recording
"scripts/run_kimi_k3.py": lambda sandbox: {"model_dir": str(sandbox / "models")},
+ # test_prepare_waits_for_shared_checkpoint_conversion covers the real cache lock.
+ "scripts/run_mimo_v2_6_flash.py": lambda sandbox: {"model_dir": str(sandbox / "models")},
"scripts/run_nemotron_3_nano_4b_fsdp.py": _nemotron_checkpoint,
"scripts/run_nemotron_3_ultra_550b_a55b.py": lambda sandbox: {
"model_name": "NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-4layer"
diff --git a/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/execute.txt b/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/execute.txt
new file mode 100644
index 00000000000..196f1e02c1e
--- /dev/null
+++ b/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/execute.txt
@@ -0,0 +1,103 @@
+### 0
+pkill -9 sglang; sleep 3; ray stop
+ --force; pkill -9 ray; pkill -9 miles; sleep 3; pkill -9 ray; pkill -9 miles; pkill -9 redis; true;
+
+### 1
+export PYTHONUNBUFFERED=1 && ray start
+ --head
+ --node-ip-address 127.0.0.1
+ --num-gpus 8
+ --disable-usage-stats
+
+### 2
+nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l
+
+### 3
+export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit
+ --address="http://127.0.0.1:8265"
+ --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "0", "no_proxy": "127.0.0.1,127.0.0.1", "MASTER_ADDR": "127.0.0.1", "PYTHONPATH": ":/root/Megatron-LM:/frozen/pythonpath"}}'
+ -- python3 /train.py
+ --num-layers 4
+ --hidden-size 4096
+ --ffn-hidden-size 16384
+ --num-attention-heads 64
+ --group-query-attention
+ --num-query-groups 4
+ --kv-channels 192
+ --position-embedding-type rope
+ --rotary-percent 0.334
+ --rotary-base 10000000
+ --no-rope-fusion
+ --max-position-embeddings 1048576
+ --normalization RMSNorm
+ --norm-epsilon 1e-6
+ --swiglu
+ --disable-bias-linear
+ --untie-embeddings-and-output-weights
+ --vocab-size 152576
+ --moe-layer-freq '[0,1,1,1]'
+ --num-experts 256
+ --moe-router-topk 8
+ --moe-ffn-hidden-size 2048
+ --moe-router-score-function sigmoid
+ --moe-router-enable-expert-bias
+ --moe-router-bias-update-rate 0
+ --moe-router-load-balancing-type none
+ --moe-aux-loss-coeff 0
+ --moe-router-dtype fp32
+ --moe-grouped-gemm
+ --hf-checkpoint /models/mimo26-p4-bf16
+ --megatron-to-hf-mode bridge
+ --prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
+ --input-key prompt
+ --label-key label
+ --apply-chat-template
+ --rollout-shuffle
+ --rm-type math
+ --num-rollout 4
+ --rollout-batch-size 16
+ --n-samples-per-prompt 8
+ --rollout-max-response-len 8192
+ --rollout-temperature 1.0
+ --rollout-top-p 1.0
+ --num-steps-per-rollout 1
+ --advantage-estimator grpo
+ --entropy-coef 0.00
+ --eps-clip 0.2
+ --eps-clip-high 0.28
+ --optimizer adam
+ --lr 1e-6
+ --lr-decay-style constant
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.98
+ --accumulate-allreduce-grads-in-fp32
+ --use-wandb
+ --wandb-project miles-run_mimo_v2_6_flash
+ --wandb-group 260101-000000-000
+ --wandb-key 'frozen-wandb-api-key'
+ --disable-wandb-random-suffix
+ --tensor-model-parallel-size 2
+ --sequence-parallel
+ --pipeline-model-parallel-size 2
+ --context-parallel-size 1
+ --expert-model-parallel-size 2
+ --expert-tensor-parallel-size 1
+ --use-dynamic-batch-size
+ --max-tokens-per-gpu 9216
+ --recompute-granularity full
+ --recompute-method uniform
+ --recompute-num-layers 1
+ --rollout-num-gpus-per-engine 4
+ --sglang-dtype bfloat16
+ --sglang-mem-fraction-static 0.6
+ --sglang-decode-log-interval 1000
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --attention-softmax-in-fp32
+ --attention-backend fused
+ --actor-num-nodes 1
+ --actor-num-gpus-per-node 8
+ --num-gpus-per-node 8
+ --colocate
+ --deploy-component all
diff --git a/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/prepare.txt b/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/prepare.txt
new file mode 100644
index 00000000000..30af3580eb6
--- /dev/null
+++ b/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/prepare.txt
@@ -0,0 +1,18 @@
+### 0
+mkdir -p /models
+
+### 1
+hf download XiaomiMiMo/MiMo-V2.6-Flash-RL
+ --local-dir /models/MiMo-V2.6-Flash-RL
+
+### 2
+python /tools/convert_mimo_v2_to_bf16.py
+ --model-dir /models/MiMo-V2.6-Flash-RL
+ --save-dir /models/mimo26-p4-bf16
+ --device cuda
+ --layers 0,1,5,6
+
+### 3
+hf download
+ --repo-type dataset zhuzilin/dapo-math-17k
+ --local-dir /root/datasets/dapo-math-17k
diff --git a/tests/snapshots/model_args/mimo-v2.6-flash-4layer.txt b/tests/snapshots/model_args/mimo-v2.6-flash-4layer.txt
new file mode 100644
index 00000000000..1d4719d2a43
--- /dev/null
+++ b/tests/snapshots/model_args/mimo-v2.6-flash-4layer.txt
@@ -0,0 +1,51 @@
+"--num-layers"
+"4"
+"--hidden-size"
+"4096"
+"--ffn-hidden-size"
+"16384"
+"--num-attention-heads"
+"64"
+"--group-query-attention"
+"--num-query-groups"
+"4"
+"--kv-channels"
+"192"
+"--position-embedding-type"
+"rope"
+"--rotary-percent"
+"0.334"
+"--rotary-base"
+"10000000"
+"--no-rope-fusion"
+"--max-position-embeddings"
+"1048576"
+"--normalization"
+"RMSNorm"
+"--norm-epsilon"
+"1e-6"
+"--swiglu"
+"--disable-bias-linear"
+"--untie-embeddings-and-output-weights"
+"--vocab-size"
+"152576"
+"--moe-layer-freq"
+"[0,1,1,1]"
+"--num-experts"
+"256"
+"--moe-router-topk"
+"8"
+"--moe-ffn-hidden-size"
+"2048"
+"--moe-router-score-function"
+"sigmoid"
+"--moe-router-enable-expert-bias"
+"--moe-router-bias-update-rate"
+"0"
+"--moe-router-load-balancing-type"
+"none"
+"--moe-aux-loss-coeff"
+"0"
+"--moe-router-dtype"
+"fp32"
+"--moe-grouped-gemm"
diff --git a/tests/snapshots/model_args/mimo-v2.6-flash.txt b/tests/snapshots/model_args/mimo-v2.6-flash.txt
new file mode 100644
index 00000000000..e71d4ed4f52
--- /dev/null
+++ b/tests/snapshots/model_args/mimo-v2.6-flash.txt
@@ -0,0 +1,51 @@
+"--num-layers"
+"48"
+"--hidden-size"
+"4096"
+"--ffn-hidden-size"
+"16384"
+"--num-attention-heads"
+"64"
+"--group-query-attention"
+"--num-query-groups"
+"4"
+"--kv-channels"
+"192"
+"--position-embedding-type"
+"rope"
+"--rotary-percent"
+"0.334"
+"--rotary-base"
+"10000000"
+"--no-rope-fusion"
+"--max-position-embeddings"
+"1048576"
+"--normalization"
+"RMSNorm"
+"--norm-epsilon"
+"1e-6"
+"--swiglu"
+"--disable-bias-linear"
+"--untie-embeddings-and-output-weights"
+"--vocab-size"
+"152576"
+"--moe-layer-freq"
+"[0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]"
+"--num-experts"
+"256"
+"--moe-router-topk"
+"8"
+"--moe-ffn-hidden-size"
+"2048"
+"--moe-router-score-function"
+"sigmoid"
+"--moe-router-enable-expert-bias"
+"--moe-router-bias-update-rate"
+"0"
+"--moe-router-load-balancing-type"
+"none"
+"--moe-aux-loss-coeff"
+"0"
+"--moe-router-dtype"
+"fp32"
+"--moe-grouped-gemm"
diff --git a/tools/convert_mimo_v2_to_bf16.py b/tools/convert_mimo_v2_to_bf16.py
new file mode 100644
index 00000000000..ff3f2587ccc
--- /dev/null
+++ b/tools/convert_mimo_v2_to_bf16.py
@@ -0,0 +1,267 @@
+"""
+Convert an official XiaomiMiMo MiMo-V2 checkpoint (`MiMoV2ForCausalLM`, e.g. MiMo-V2.6-Flash-RL)
+to plain BF16 with split q/k/v projections, optionally keeping a subset of decoder layers and
+routed experts for partial-model work.
+
+The source stores `self_attn.qkv_proj` fused and interleaved across `num_key_value_heads`
+shards (shard i is `[q_i | k_i | v_i]` for its heads, the layout SGLang's TP loader expects) in
+blockwise FP8 quantized shard by shard (each shard padded to whole 128-row blocks, so global-attention
+layers carry 4 x 27 scale rows for 13568 weight rows), the dense and MTP linears in blockwise FP8, and the routed experts in MXFP4
+(packed e2m1 in uint8 plus one E8M0 scale per 32 inputs). The output dequantizes all of them to
+BF16, splits qkv into q_proj/k_proj/v_proj in head order, and drops `quantization_config` and
+`attention_projection_layout`, so the HF remote code, Megatron-Bridge and a BF16 SGLang engine
+all read it as-is.
+
+--keep-quant only reindexes layers and copies every tensor in its source format.
+
+python tools/convert_mimo_v2_to_bf16.py --model-dir --save-dir [--layers 0,1,5,6] [--num-experts 32]
+"""
+
+import json
+import re
+import shutil
+from argparse import ArgumentParser
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from tqdm import tqdm
+
+FP4_TABLE = torch.tensor(
+ [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
+ dtype=torch.float32,
+)
+FP8_BLOCK = 128
+MXFP4_BLOCK = 32
+SHARD_BYTES = 5 * 1024**3
+
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+_EXPERT_RE = re.compile(r"^mlp\.experts\.(\d+)\.")
+_ROUTER_KEYS = ("mlp.gate.weight", "mlp.gate.e_score_correction_bias")
+
+
+def dequant_fp8_block(weight: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor:
+ rows, cols = weight.shape
+ scale = scale_inv.float().repeat_interleave(FP8_BLOCK, dim=0)[:rows].repeat_interleave(FP8_BLOCK, dim=1)
+ return (weight.float() * scale[:, :cols]).to(torch.bfloat16)
+
+
+def dequant_mxfp4(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ out_dim, in_dim = packed.shape[0], packed.shape[1] * 2
+ assert scale.shape == (out_dim, in_dim // MXFP4_BLOCK), f"{packed.shape=} {scale.shape=}"
+ packed = packed.view(torch.uint8)
+ table = FP4_TABLE.to(packed.device)
+ values = torch.stack([table[(packed & 0x0F).long()], table[(packed >> 4).long()]], dim=-1).reshape(out_dim, in_dim)
+ return (values * torch.exp2(scale.float() - 127.0).repeat_interleave(MXFP4_BLOCK, dim=1)).to(torch.bfloat16)
+
+
+def dequant_fused_qkv(weight: torch.Tensor, scale_inv: torch.Tensor, shards: int) -> torch.Tensor:
+ """Each kv-head shard of a fused qkv_proj is block-quantized on its own, padded to whole blocks."""
+ rows = weight.shape[0] // shards
+ scale_rows = -(-rows // FP8_BLOCK)
+ assert scale_inv.shape[0] == shards * scale_rows, f"{weight.shape=} {scale_inv.shape=} {shards=}"
+ return torch.cat(
+ [dequant_fp8_block(w, s) for w, s in zip(weight.chunk(shards), scale_inv.split(scale_rows), strict=True)]
+ )
+
+
+def split_fused_qkv(qkv: torch.Tensor, config: dict, is_swa: bool) -> list[torch.Tensor]:
+ """Undo the kv-head interleave of a fused qkv_proj and return [q, k, v] in head order."""
+ shards = config["num_key_value_heads"]
+ prefix = "swa_" if is_swa else ""
+ num_heads, kv_heads = config[f"{prefix}num_attention_heads"], config[f"{prefix}num_key_value_heads"]
+ assert num_heads % shards == 0 and kv_heads % shards == 0, f"{num_heads=} {kv_heads=} {shards=}"
+ sizes = [
+ num_heads // shards * config[f"{prefix}head_dim"],
+ kv_heads // shards * config[f"{prefix}head_dim"],
+ kv_heads // shards * config[f"{prefix}v_head_dim"],
+ ]
+ assert qkv.shape[0] == shards * sum(sizes), f"{qkv.shape=} does not match {shards=} x {sizes=}"
+ return [part.reshape(-1, qkv.shape[1]) for part in qkv.view(shards, sum(sizes), -1).split(sizes, dim=1)]
+
+
+def output_name(name: str, layer_map: dict[int, int] | None, num_experts: int | None) -> str | None:
+ """Output name of a source tensor, or None when the tensor is not kept."""
+ match = _LAYER_RE.match(name)
+ if match is None:
+ return name
+ layer, rest = int(match.group(1)), match.group(2)
+ if layer_map is not None and layer not in layer_map:
+ return None
+ expert = _EXPERT_RE.match(rest)
+ if expert is not None and num_experts is not None and int(expert.group(1)) >= num_experts:
+ return None
+ return f"model.layers.{layer_map[layer] if layer_map is not None else layer}.{rest}"
+
+
+def convert_config(config: dict, layer_map: dict[int, int] | None, num_experts: int | None, keep_quant: bool) -> dict:
+ config = json.loads(json.dumps(config))
+ if layer_map is not None:
+ assert config.get("hybrid_block_size") is None, "per-layer schedules must be explicit lists"
+ kept = list(layer_map)
+ config["num_hidden_layers"] = len(kept)
+ for key in ("hybrid_layer_pattern", "moe_layer_freq"):
+ config[key] = [config[key][layer] for layer in kept]
+ if num_experts is not None:
+ assert config["num_experts_per_tok"] <= num_experts <= config["n_routed_experts"], num_experts
+ config["n_routed_experts"] = num_experts
+ if keep_quant:
+ quant = config["quantization_config"]
+ ignored = []
+ for module in quant.get("ignored_layers", []):
+ renamed = output_name(module + ".weight", layer_map, num_experts)
+ if renamed is not None:
+ ignored.append(renamed.removesuffix(".weight"))
+ quant["ignored_layers"] = ignored
+ else:
+ config.pop("quantization_config", None)
+ config.pop("attention_projection_layout", None)
+ return config
+
+
+class ShardWriter:
+ def __init__(self, save_dir: Path):
+ self.save_dir = save_dir
+ self.buffer: dict[str, torch.Tensor] = {}
+ self.buffer_bytes = 0
+ self.num_shards = 0
+ self.total_bytes = 0
+ self.weight_map: dict[str, str] = {}
+
+ def add(self, name: str, tensor: torch.Tensor) -> None:
+ assert name not in self.weight_map and name not in self.buffer, f"duplicate tensor {name}"
+ self.buffer[name] = tensor.contiguous().cpu()
+ self.buffer_bytes += tensor.numel() * tensor.element_size()
+ if self.buffer_bytes >= SHARD_BYTES:
+ self.flush()
+
+ def flush(self) -> None:
+ if not self.buffer:
+ return
+ self.num_shards += 1
+ file_name = f"model-{self.num_shards:05d}.safetensors"
+ save_file(self.buffer, str(self.save_dir / file_name), metadata={"format": "pt"})
+ self.weight_map.update(dict.fromkeys(self.buffer, file_name))
+ self.total_bytes += self.buffer_bytes
+ self.buffer, self.buffer_bytes = {}, 0
+
+
+def check_output(weight_map: dict[str, str], config: dict) -> None:
+ """Every kept decoder layer carries the tensors its attention/MLP variant needs."""
+ for layer in range(config["num_hidden_layers"]):
+ prefix = f"model.layers.{layer}."
+ names = {name[len(prefix) :] for name in weight_map if name.startswith(prefix)}
+ required = {"input_layernorm.weight", "post_attention_layernorm.weight", "self_attn.o_proj.weight"}
+ required |= (
+ {f"self_attn.{p}_proj.weight" for p in "qkv"} if "self_attn.qkv_proj.weight" not in names else set()
+ )
+ if config["hybrid_layer_pattern"][layer] == 1 and config.get("add_swa_attention_sink_bias"):
+ required.add("self_attn.attention_sink_bias")
+ if config["moe_layer_freq"][layer]:
+ required |= set(_ROUTER_KEYS)
+ experts = {int(m.group(1)) for n in names if (m := _EXPERT_RE.match(n))}
+ assert experts == set(
+ range(config["n_routed_experts"])
+ ), f"layer {layer}: experts {sorted(experts)[:4]}..."
+ else:
+ required |= {f"mlp.{p}_proj.weight" for p in ("gate", "up", "down")}
+ missing = required - names
+ assert not missing, f"layer {layer} is missing {sorted(missing)}"
+
+
+def main(
+ model_dir: str, save_dir: str, layers: list[int] | None, num_experts: int | None, keep_quant: bool, device: str
+):
+ src, dst = Path(model_dir), Path(save_dir)
+ dst.mkdir(parents=True, exist_ok=True)
+ config = json.loads((src / "config.json").read_text())
+ assert config["architectures"] == ["MiMoV2ForCausalLM"], config["architectures"]
+ layer_map = {old: new for new, old in enumerate(layers)} if layers is not None else None
+ assert layer_map is None or all(0 <= layer < config["num_hidden_layers"] for layer in layer_map), layers
+ assert not (keep_quant and num_experts is not None), "--num-experts needs the dequantized output"
+ fused_qkv = config.get("attention_projection_layout") == "fused_qkv"
+
+ weight_map = json.loads((src / "model.safetensors.index.json").read_text())["weight_map"]
+ handles = {}
+
+ def load(name: str) -> torch.Tensor:
+ file_name = weight_map[name]
+ if file_name not in handles:
+ handles[file_name] = safe_open(str(src / file_name), framework="pt", device="cpu")
+ return handles[file_name].get_tensor(name).to(device)
+
+ writer = ShardWriter(dst)
+ counts = {"fp8": 0, "mxfp4": 0, "qkv_split": 0, "copied": 0}
+ scale_suffixes = ("weight_scale_inv", "weight_scale")
+ for file_name in tqdm(sorted(set(weight_map.values())), desc="shards"):
+ for name in [n for n, f in weight_map.items() if f == file_name]:
+ new_name = output_name(name, layer_map, num_experts)
+ if new_name is None or (not keep_quant and name.endswith(scale_suffixes)):
+ continue
+ tensor = load(name)
+ if keep_quant:
+ writer.add(new_name, tensor)
+ counts["copied"] += 1
+ continue
+ if fused_qkv and new_name.endswith("self_attn.qkv_proj.weight"):
+ match = _LAYER_RE.match(name)
+ is_swa = match is None or config["hybrid_layer_pattern"][int(match.group(1))] == 1
+ if tensor.dtype == torch.float8_e4m3fn:
+ tensor = dequant_fused_qkv(tensor, load(name + "_scale_inv"), config["num_key_value_heads"])
+ counts["fp8"] += 1
+ for proj, part in zip("qkv", split_fused_qkv(tensor, config, is_swa), strict=True):
+ writer.add(new_name.replace("qkv_proj", f"{proj}_proj"), part)
+ counts["qkv_split"] += 1
+ continue
+ if tensor.dtype == torch.float8_e4m3fn:
+ tensor = dequant_fp8_block(tensor, load(name + "_scale_inv"))
+ counts["fp8"] += 1
+ elif tensor.dtype == torch.uint8:
+ tensor = dequant_mxfp4(tensor, load(name.removesuffix("weight") + "weight_scale"))
+ counts["mxfp4"] += 1
+ if num_experts is not None and new_name.endswith(_ROUTER_KEYS):
+ tensor = tensor[:num_experts]
+ writer.add(new_name, tensor)
+ counts["copied"] += 1
+ writer.flush()
+
+ new_config = convert_config(config, layer_map, num_experts, keep_quant)
+ if not keep_quant:
+ check_output(writer.weight_map, new_config)
+ (dst / "config.json").write_text(json.dumps(new_config, indent=2) + "\n")
+ index = {"metadata": {"total_size": writer.total_bytes}, "weight_map": writer.weight_map}
+ for path in src.iterdir():
+ if path.name.startswith("."): # e.g. the `hf download` cache
+ continue
+ if path.is_dir():
+ # The DFlash drafter reads hidden states of fixed full-model layers, so it only
+ # stays valid when every layer is kept.
+ if path.name != "dflash" or layer_map is None:
+ shutil.copytree(path, dst / path.name, dirs_exist_ok=True)
+ elif path.suffix != ".safetensors" and path.name not in ("config.json", "model.safetensors.index.json"):
+ shutil.copy2(path, dst / path.name)
+ index_path = dst / "model.safetensors.index.json"
+ pending_index = index_path.with_suffix(".json.tmp")
+ pending_index.write_text(json.dumps(index, indent=2) + "\n")
+ pending_index.replace(index_path)
+ print(
+ f"{counts}; {len(writer.weight_map)} tensors, {writer.total_bytes / 1e9:.2f} GB in {writer.num_shards} shards -> {dst}"
+ )
+
+
+if __name__ == "__main__":
+ parser = ArgumentParser()
+ parser.add_argument("--model-dir", type=str, required=True)
+ parser.add_argument("--save-dir", type=str, required=True)
+ parser.add_argument(
+ "--layers", type=str, default=None, help="Comma-separated source layer indices to keep, in order."
+ )
+ parser.add_argument(
+ "--num-experts", type=int, default=None, help="Keep only the first N routed experts per MoE layer."
+ )
+ parser.add_argument("--keep-quant", action="store_true", help="Reindex layers but keep the source tensor formats.")
+ parser.add_argument("--device", type=str, default="cpu")
+ args = parser.parse_args()
+ layers = [int(x) for x in args.layers.split(",")] if args.layers else None
+ main(args.model_dir, args.save_dir, layers, args.num_experts, args.keep_quant, args.device)