From 9114355125b0987b00f39714b33938f71804b672 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 23 Sep 2026 15:15:41 -0700 Subject: [PATCH 1/9] feat(models): add MiMo-V2.6-Flash-RL support 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 --- docs/docs.json | 3 +- docs/models/index.md | 1 + docs/models/mimo/mimo-v2-6-flash.md | 108 +++++++ miles_plugins/megatron_bridge/__init__.py | 5 + miles_plugins/megatron_bridge/mimo_v2.py | 296 ++++++++++++++++++ scripts/models/mimo-v2.6-flash-4layer.py | 7 + scripts/models/mimo-v2.6-flash.py | 52 +++ scripts/run_mimo_v2_6_flash.py | 263 ++++++++++++++++ .../test_mimo_v2_6_flash_4layer_ci.py | 51 +++ .../tools/test_convert_mimo_v2_to_bf16.py | 175 +++++++++++ .../run_mimo_v2_6_flash.py/execute.txt | 102 ++++++ .../run_mimo_v2_6_flash.py/prepare.txt | 18 ++ .../model_args/mimo-v2.6-flash-4layer.txt | 51 +++ .../snapshots/model_args/mimo-v2.6-flash.txt | 51 +++ tools/convert_mimo_v2_to_bf16.py | 264 ++++++++++++++++ 15 files changed, 1446 insertions(+), 1 deletion(-) create mode 100644 docs/models/mimo/mimo-v2-6-flash.md create mode 100644 miles_plugins/megatron_bridge/mimo_v2.py create mode 100644 scripts/models/mimo-v2.6-flash-4layer.py create mode 100644 scripts/models/mimo-v2.6-flash.py create mode 100644 scripts/run_mimo_v2_6_flash.py create mode 100644 tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py create mode 100644 tests/fast/tools/test_convert_mimo_v2_to_bf16.py create mode 100644 tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/execute.txt create mode 100644 tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/prepare.txt create mode 100644 tests/snapshots/model_args/mimo-v2.6-flash-4layer.txt create mode 100644 tests/snapshots/model_args/mimo-v2.6-flash.txt create mode 100644 tools/convert_mimo_v2_to_bf16.py 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..daf922193d4 --- /dev/null +++ b/docs/models/mimo/mimo-v2-6-flash.md @@ -0,0 +1,108 @@ +--- +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, a 4-layer partial on one 8-GPU node, the full model on two. +--- +## 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` | full text decoder, 48 layers | [XiaomiMiMo/MiMo-V2.6-Flash-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) | +| `mimo26-p4-bf16` | 4-layer partial (source layers 0, 1, 5, 6: one of each decoder variant), all 256 experts | same, converted with `--layers 0,1,5,6` | + +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/mimo26-p4-bf16 --layers 0,1,5,6 --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 (`--num-experts N` keeps the first N experts for smaller tests). The full BF16 model is 622 GB, the 4-layer partial 46 GB. + +### 3.2 SGLang + +The BF16 engine needs the MiMo-V2 fixes that are not in the image yet: + +- 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; +- passing `sliding_window_size - 1` to the attention layers: SGLang's window excludes the query token, so the unmodified model attends 129 keys where the HF reference and Megatron attend 128, which shows up as rare large train/rollout log-prob gaps. + +## 4. Launch + +### 4.1 Quick start (one node, 4-layer partial) + +```bash +cd /root/miles +python scripts/run_mimo_v2_6_flash.py --mode rl --model-name mimo26-p4-bf16 +python scripts/run_mimo_v2_6_flash.py --mode sft --model-name mimo26-p4-bf16 --prompt-data +``` + +`--mode sft` reads a `messages` column; put a reasoning trace in `reasoning_content`, since the MiMo chat template renders `{reasoning_content}`. + +### 4.2 Full model (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. + +## 5. Recipe Configuration + +### 5.1 Parallelism + +| `--model-name` | TP | PP | CP | EP | expert-TP | GPUs | rollout engine | +|---|---|---|---|---|---|---|---| +| `mimo26-p4-bf16` | 2 | 2 | 1 | 2 | 1 | 8 (1 × 8) | TP4, `--sglang-mem-fraction-static 0.6` | +| `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 full 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. +- **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_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..4d59465b375 --- /dev/null +++ b/miles_plugins/megatron_bridge/mimo_v2.py @@ -0,0 +1,296 @@ +"""Megatron-Bridge support for XiaomiMiMo MiMo-V2 (``MiMoV2ForCausalLM``, e.g. MiMo-V2.6-Flash-RL). + +Adapted from NVIDIA Megatron-Bridge ``megatron/bridge/models/mimo_v2_flash`` (Apache-2.0), which +covers the earlier MiMo-V2-Flash checkpoint. MiMo-V2 attention differs from the stock GPT layer +in four ways that the TransformerConfig cannot express on its own: + +* per-layer KV heads: SWA layers use ``swa_num_key_value_heads``, global layers ``num_key_value_heads``; +* asymmetric head dims: Q/K use ``head_dim`` (192) and V uses ``v_head_dim`` (128); +* a learnable attention sink on SWA layers only, while ``softmax_type`` is one global setting; +* V is multiplied by ``attention_value_scale`` before attention. + +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``). + +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 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, QKVMapping +from megatron.bridge.models.conversion.utils import remove_non_pickleables +from megatron.bridge.models.gpt_provider import GPTModelProvider +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 _layer_config(config, layer_number: int): + """Shallow config copy carrying the attention shape and softmax of this 1-indexed layer.""" + assert layer_number <= len(config.hybrid_attention_pattern), f"no attention pattern for layer {layer_number}" + config = copy.copy(config) + if config.hybrid_attention_pattern[layer_number - 1]: + config.num_query_groups = config.swa_num_query_groups + config.softmax_type = "learnable" if config.swa_attention_sink else "vanilla" + else: + config.num_query_groups = config.full_attn_num_query_groups + config.softmax_type = "learnable" if config.full_attention_sink else "vanilla" + return config + + +class MiMoV2SelfAttention(SelfAttention): + """SelfAttention with per-layer KV heads and a V head narrower than Q/K.""" + + def __init__(self, config, submodules, layer_number, *args, **kwargs): + config = _layer_config(config, layer_number) + super().__init__(config, submodules, layer_number, *args, **kwargs) + assert not (config.attention_output_gate or config.head_wise_attn_gate), "MiMo-V2 has no attention gate" + assert config.num_query_groups >= self.world_size, "TP must not exceed the layer's KV heads" + name = kwargs.get("name") + self.val_hidden_size = config.v_head_dim + self.linear_qkv_out_dim = self.query_projection_size + config.num_query_groups * ( + config.kv_channels + config.v_head_dim + ) + self.linear_qkv = submodules.linear_qkv( + config.hidden_size, + self.linear_qkv_out_dim, + config=config, + init_method=config.init_method, + gather_output=False, + bias=config.add_bias_linear or config.add_qkv_bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="qkv", + tp_group=self.pg_collection.tp, + name=(name + ".linear_qkv") if name is not None else None, + ) + self.linear_proj = submodules.linear_proj( + config.v_head_dim * config.num_attention_heads, + config.hidden_size, + config=config, + init_method=config.output_layer_init_method, + bias=config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="proj", + tp_group=self.pg_collection.tp, + name=(name + ".linear_proj") if name is not None else None, + ) + + 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 this layer's sink setting, Q/K vs V channels and V scaling.""" + + def __init__(self, config, layer_number, attn_mask_type, attention_type, attention_dropout=None, **kwargs): + config = _layer_config(config, layer_number) + kwargs["k_channels"] = config.kv_channels + kwargs["v_channels"] = config.v_head_dim + super().__init__( + config=config, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type=attention_type, + attention_dropout=attention_dropout, + **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 per-layer attention fields read by the MiMo-V2 modules.""" + + transformer_layer_spec: ModuleSpec | Callable = field(default_factory=lambda: mimo_v2_layer_spec) + 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 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) + + +class MiMoV2QKVMapping(QKVMapping): + """QKV mapping for per-group [q heads, k head, v head] with a V head of ``v_head_dim``.""" + + def hf_to_megatron(self, hf_weights, megatron_module): + merged = None + if self.tp_rank == 0: + config = self._get_config(megatron_module) + q, k, v = hf_weights["q"], hf_weights["k"], hf_weights["v"] + assert q.ndim == 2, "MiMo-V2 has no qkv bias" + groups, qk_dim, v_dim = config.num_query_groups, config.kv_channels, config.v_head_dim + q = q.view(groups, -1, q.shape[-1]) + merged = torch.cat([q, k.view(groups, qk_dim, -1), v.view(groups, v_dim, -1)], dim=1).flatten(0, 1) + return self._tp_mapping.hf_to_megatron(merged, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module): + if megatron_weights is not None: + megatron_weights = self.maybe_dequantize(megatron_weights) + if megatron_module is None: + config = self.broadcast_obj_from_pp_rank(None, "qkv_config") + else: + config = remove_non_pickleables(self._get_config(megatron_module), max_depth=3) + config = self.broadcast_obj_from_pp_rank(config, "qkv_config") + packed_dict = self._tp_mapping.megatron_to_hf(megatron_weights, megatron_module) + if not packed_dict: + return {} + packed = next(iter(packed_dict.values())) + groups, qk_dim, v_dim = config.num_query_groups, config.kv_channels, config.v_head_dim + q_dim = config.num_attention_heads // groups * qk_dim + q, k, v = packed.view(groups, q_dim + qk_dim + v_dim, -1).split([q_dim, qk_dim, v_dim], dim=1) + return { + self.hf_param["q"]: q.reshape(-1, packed.shape[-1]), + self.hf_param["k"]: k.reshape(-1, packed.shape[-1]), + self.hf_param["v"]: v.reshape(-1, packed.shape[-1]), + } + + +@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 + provider.window_size = (hf.sliding_window_size - 1, 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( + MiMoV2QKVMapping( + 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..4808234e1e4 --- /dev/null +++ b/scripts/run_mimo_v2_6_flash.py @@ -0,0 +1,263 @@ +"""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 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) + # The BF16 engine holds the whole model: 620 GB needs TP8 on 141 GB GPUs. + 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), + 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): + target = Path(args.model_dir) / args.model_name + # 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 "" + U.exec_command_cpu(f"mkdir -p {args.model_dir}") + U.exec_command_cpu(f"hf download {_HF_REPO} --local-dir {source}") + U.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: + U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + + +def execute(args: ScriptArgs): + 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 0.95 " + "--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} " + ) + + U.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..868de5375e2 --- /dev/null +++ b/tests/e2e/megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py @@ -0,0 +1,51 @@ +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=2400, suite="stage-c-8-gpu-h200", labels=["megatron", "model-scripts"], hardware=["hopper"]) + +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/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..a704c2cafcc --- /dev/null +++ b/tests/fast/tools/test_convert_mimo_v2_to_bf16.py @@ -0,0 +1,175 @@ +import json + +import pytest +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" 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..a7b0637e398 --- /dev/null +++ b/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/execute.txt @@ -0,0 +1,102 @@ +### 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 /root/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 0.95 + --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 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..c107a55905a --- /dev/null +++ b/tests/snapshots/launch_scripts/py/scripts/run_mimo_v2_6_flash.py/prepare.txt @@ -0,0 +1,18 @@ +### 0 +mkdir -p /root/models + +### 1 +hf download XiaomiMiMo/MiMo-V2.6-Flash-RL + --local-dir /root/models/MiMo-V2.6-Flash-RL + +### 2 +python /tools/convert_mimo_v2_to_bf16.py + --model-dir /root/models/MiMo-V2.6-Flash-RL + --save-dir /root/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..9251abcf9c4 --- /dev/null +++ b/tools/convert_mimo_v2_to_bf16.py @@ -0,0 +1,264 @@ +""" +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} + (dst / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n") + 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) + 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) From 7c0e3aa551161c7d4803732aa31230f698e9c5ef Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 28 Sep 2026 21:57:12 -0700 Subject: [PATCH 2/9] fix(mimo): attend sliding_window_size previous keys in SWA layers, as 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 --- docs/models/mimo/mimo-v2-6-flash.md | 4 ++-- miles_plugins/megatron_bridge/mimo_v2.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/models/mimo/mimo-v2-6-flash.md b/docs/models/mimo/mimo-v2-6-flash.md index daf922193d4..51315e18ca4 100644 --- a/docs/models/mimo/mimo-v2-6-flash.md +++ b/docs/models/mimo/mimo-v2-6-flash.md @@ -44,8 +44,7 @@ The BF16 engine needs the MiMo-V2 fixes that are not in the image yet: - 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; -- passing `sliding_window_size - 1` to the attention layers: SGLang's window excludes the query token, so the unmodified model attends 129 keys where the HF reference and Megatron attend 128, which shows up as rare large train/rollout log-prob gaps. +- passing `layer_id` to the MoE top-k, otherwise `--use-rollout-routing-replay` crashes the CUDA-graph capture. ## 4. Launch @@ -100,6 +99,7 @@ Adam at `--lr 1e-6`. The full model's Adam state (3.7 TB) fits neither the GPUs ### 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.`. diff --git a/miles_plugins/megatron_bridge/mimo_v2.py b/miles_plugins/megatron_bridge/mimo_v2.py index 4d59465b375..391210bbdf7 100644 --- a/miles_plugins/megatron_bridge/mimo_v2.py +++ b/miles_plugins/megatron_bridge/mimo_v2.py @@ -221,7 +221,9 @@ def provider_bridge(self, hf_pretrained) -> MiMoV2ModelProvider: provider = super().provider_bridge(hf_pretrained) pattern = [int(p) for p in hf.hybrid_layer_pattern] provider.hybrid_attention_pattern = pattern - provider.window_size = (hf.sliding_window_size - 1, 0) + # 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] From eee663d826c37742f62e9775dbab0d1a6a03777a Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 5 Oct 2026 21:00:00 -0700 Subject: [PATCH 3/9] refactor(mimo): build MiMo-V2 attention from per-layer configs 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 --- miles_plugins/megatron_bridge/mimo_v2.py | 156 +++++++++-------------- 1 file changed, 61 insertions(+), 95 deletions(-) diff --git a/miles_plugins/megatron_bridge/mimo_v2.py b/miles_plugins/megatron_bridge/mimo_v2.py index 391210bbdf7..0ac3c7c124a 100644 --- a/miles_plugins/megatron_bridge/mimo_v2.py +++ b/miles_plugins/megatron_bridge/mimo_v2.py @@ -1,16 +1,19 @@ """Megatron-Bridge support for XiaomiMiMo MiMo-V2 (``MiMoV2ForCausalLM``, e.g. MiMo-V2.6-Flash-RL). -Adapted from NVIDIA Megatron-Bridge ``megatron/bridge/models/mimo_v2_flash`` (Apache-2.0), which -covers the earlier MiMo-V2-Flash checkpoint. MiMo-V2 attention differs from the stock GPT layer -in four ways that the TransformerConfig cannot express on its own: +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``; -* asymmetric head dims: Q/K use ``head_dim`` (192) and V uses ``v_head_dim`` (128); * 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 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``). +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. @@ -20,6 +23,7 @@ from __future__ import annotations import copy +import dataclasses import logging from collections.abc import Callable from dataclasses import dataclass, field @@ -27,9 +31,9 @@ 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, QKVMapping -from megatron.bridge.models.conversion.utils import remove_non_pickleables +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 @@ -39,58 +43,42 @@ logger = logging.getLogger(__name__) -def _layer_config(config, layer_number: int): - """Shallow config copy carrying the attention shape and softmax of this 1-indexed layer.""" - assert layer_number <= len(config.hybrid_attention_pattern), f"no attention pattern for layer {layer_number}" - config = copy.copy(config) - if config.hybrid_attention_pattern[layer_number - 1]: - config.num_query_groups = config.swa_num_query_groups - config.softmax_type = "learnable" if config.swa_attention_sink else "vanilla" - else: - config.num_query_groups = config.full_attn_num_query_groups - config.softmax_type = "learnable" if config.full_attention_sink else "vanilla" - return config +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 with per-layer KV heads and a V head narrower than Q/K.""" + """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): - config = _layer_config(config, layer_number) - super().__init__(config, submodules, layer_number, *args, **kwargs) assert not (config.attention_output_gate or config.head_wise_attn_gate), "MiMo-V2 has no attention gate" - assert config.num_query_groups >= self.world_size, "TP must not exceed the layer's KV heads" - name = kwargs.get("name") - self.val_hidden_size = config.v_head_dim - self.linear_qkv_out_dim = self.query_projection_size + config.num_query_groups * ( - config.kv_channels + config.v_head_dim - ) - self.linear_qkv = submodules.linear_qkv( - config.hidden_size, - self.linear_qkv_out_dim, - config=config, - init_method=config.init_method, - gather_output=False, - bias=config.add_bias_linear or config.add_qkv_bias, - skip_bias_add=False, - is_expert=False, - tp_comm_buffer_name="qkv", - tp_group=self.pg_collection.tp, - name=(name + ".linear_qkv") if name is not None else None, - ) - self.linear_proj = submodules.linear_proj( - config.v_head_dim * config.num_attention_heads, - config.hidden_size, - config=config, - init_method=config.output_layer_init_method, - bias=config.add_bias_linear, - input_is_parallel=True, - skip_bias_add=True, - is_expert=False, - tp_comm_buffer_name="proj", - tp_group=self.pg_collection.tp, - name=(name + ".linear_proj") if name is not None else None, + # 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 @@ -108,18 +96,20 @@ def get_query_key_value_tensors( class MiMoV2TEDotProductAttention(TEDotProductAttention): - """TE core attention with this layer's sink setting, Q/K vs V channels and V scaling.""" + """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): - config = _layer_config(config, layer_number) - kwargs["k_channels"] = config.kv_channels - kwargs["v_channels"] = config.v_head_dim 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 @@ -143,9 +133,11 @@ def mimo_v2_layer_spec(config, vp_stage: int | None = None) -> ModuleSpec: @dataclass class MiMoV2ModelProvider(GPTModelProvider): - """GPT provider plus the per-layer attention fields read by the MiMo-V2 modules.""" + """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 @@ -154,6 +146,16 @@ class MiMoV2ModelProvider(GPTModelProvider): 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 @@ -162,42 +164,6 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> GPTMode return super().provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) -class MiMoV2QKVMapping(QKVMapping): - """QKV mapping for per-group [q heads, k head, v head] with a V head of ``v_head_dim``.""" - - def hf_to_megatron(self, hf_weights, megatron_module): - merged = None - if self.tp_rank == 0: - config = self._get_config(megatron_module) - q, k, v = hf_weights["q"], hf_weights["k"], hf_weights["v"] - assert q.ndim == 2, "MiMo-V2 has no qkv bias" - groups, qk_dim, v_dim = config.num_query_groups, config.kv_channels, config.v_head_dim - q = q.view(groups, -1, q.shape[-1]) - merged = torch.cat([q, k.view(groups, qk_dim, -1), v.view(groups, v_dim, -1)], dim=1).flatten(0, 1) - return self._tp_mapping.hf_to_megatron(merged, megatron_module) - - def megatron_to_hf(self, megatron_weights, megatron_module): - if megatron_weights is not None: - megatron_weights = self.maybe_dequantize(megatron_weights) - if megatron_module is None: - config = self.broadcast_obj_from_pp_rank(None, "qkv_config") - else: - config = remove_non_pickleables(self._get_config(megatron_module), max_depth=3) - config = self.broadcast_obj_from_pp_rank(config, "qkv_config") - packed_dict = self._tp_mapping.megatron_to_hf(megatron_weights, megatron_module) - if not packed_dict: - return {} - packed = next(iter(packed_dict.values())) - groups, qk_dim, v_dim = config.num_query_groups, config.kv_channels, config.v_head_dim - q_dim = config.num_attention_heads // groups * qk_dim - q, k, v = packed.view(groups, q_dim + qk_dim + v_dim, -1).split([q_dim, qk_dim, v_dim], dim=1) - return { - self.hf_param["q"]: q.reshape(-1, packed.shape[-1]), - self.hf_param["k"]: k.reshape(-1, packed.shape[-1]), - self.hf_param["v"]: v.reshape(-1, packed.shape[-1]), - } - - @MegatronModelBridge.register_bridge( source="MiMoV2ForCausalLM", target=GPTModel, provider=MiMoV2ModelProvider, model_type="mimo_v2" ) @@ -288,7 +254,7 @@ def mapping_registry(self) -> MegatronMappingRegistry: for m, h in gated.items() ] mappings.append( - MiMoV2QKVMapping( + 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", From cd97a423002bc653d63f3762ed66b8f54369008b Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 5 Oct 2026 21:04:18 -0700 Subject: [PATCH 4/9] docs(mimo): drop the 4-layer partial from the MiMo-V2.6 page 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 --- docs/models/mimo/mimo-v2-6-flash.md | 26 ++++++++------------------ 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/docs/models/mimo/mimo-v2-6-flash.md b/docs/models/mimo/mimo-v2-6-flash.md index 51315e18ca4..5a4eeb16588 100644 --- a/docs/models/mimo/mimo-v2-6-flash.md +++ b/docs/models/mimo/mimo-v2-6-flash.md @@ -1,7 +1,7 @@ --- 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, a 4-layer partial on one 8-GPU node, the full model on two. +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 @@ -18,8 +18,7 @@ description: Launch recipe for Xiaomi's MiMo-V2.6-Flash-RL (309B MoE, hybrid sli | `--model-name` | What it is | HF source | |---|---|---| -| `MiMo-V2.6-Flash-RL-bf16` | full text decoder, 48 layers | [XiaomiMiMo/MiMo-V2.6-Flash-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) | -| `mimo26-p4-bf16` | 4-layer partial (source layers 0, 1, 5, 6: one of each decoder variant), all 256 experts | same, converted with `--layers 0,1,5,6` | +| `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). @@ -32,10 +31,10 @@ The vision, audio and MTP modules stay in the checkpoint but are not trained (te ```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/mimo26-p4-bf16 --layers 0,1,5,6 --device cuda + --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 (`--num-experts N` keeps the first N experts for smaller tests). The full BF16 model is 622 GB, the 4-layer partial 46 GB. +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 @@ -48,17 +47,7 @@ The BF16 engine needs the MiMo-V2 fixes that are not in the image yet: ## 4. Launch -### 4.1 Quick start (one node, 4-layer partial) - -```bash -cd /root/miles -python scripts/run_mimo_v2_6_flash.py --mode rl --model-name mimo26-p4-bf16 -python scripts/run_mimo_v2_6_flash.py --mode sft --model-name mimo26-p4-bf16 --prompt-data -``` - -`--mode sft` reads a `messages` column; put a reasoning trace in `reasoning_content`, since the MiMo chat template renders `{reasoning_content}`. - -### 4.2 Full model (two nodes) +### 4.1 Two nodes Join the second node to a ray head on the first, then launch from the head: @@ -69,13 +58,14 @@ MILES_SCRIPT_EXTERNAL_RAY=1 MASTER_ADDR= python scripts/run_mimo_v2_6_f 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 | |---|---|---|---|---|---|---|---| -| `mimo26-p4-bf16` | 2 | 2 | 1 | 2 | 1 | 8 (1 × 8) | TP4, `--sglang-mem-fraction-static 0.6` | | `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. @@ -86,7 +76,7 @@ GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --entropy-coef 0.00` and `--rm-ty ### 5.3 Optimizer -Adam at `--lr 1e-6`. The full model's Adam state (3.7 TB) fits neither the GPUs nor two hosts' memory, so it streams through node-local NVMe: +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 From 50466b1d5cbf1208224ad86af804248aa35298a7 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 6 Oct 2026 02:19:12 -0700 Subject: [PATCH 5/9] fix(mimo): launch preparation and training through the configured backend 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. --- scripts/run_mimo_v2_6_flash.py | 12 ++++---- .../test_run_mimo_v2_6_flash.py | 29 +++++++++++++++++++ .../run_mimo_v2_6_flash.py/execute.txt | 1 + 3 files changed, 37 insertions(+), 5 deletions(-) create mode 100644 tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py diff --git a/scripts/run_mimo_v2_6_flash.py b/scripts/run_mimo_v2_6_flash.py index 4808234e1e4..cf70028227f 100644 --- a/scripts/run_mimo_v2_6_flash.py +++ b/scripts/run_mimo_v2_6_flash.py @@ -116,23 +116,25 @@ def __post_init__(self): def prepare(args: ScriptArgs): + backend = args.create_backend() target = Path(args.model_dir) / args.model_name # 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 "" - U.exec_command_cpu(f"mkdir -p {args.model_dir}") - U.exec_command_cpu(f"hf download {_HF_REPO} --local-dir {source}") - U.exec_command_gpu( + 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: - U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + 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} " @@ -241,7 +243,7 @@ def execute(args: ScriptArgs): f"{args.extra_args} " ) - U.execute_train( + backend.execute_train( train_args=train_args, config=args, num_gpus_per_node=args.num_gpus_per_node, 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..782eef5aa5d --- /dev/null +++ b/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py @@ -0,0 +1,29 @@ +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 + + +@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] 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 index a7b0637e398..57fcec56492 100644 --- 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 @@ -100,3 +100,4 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --actor-num-gpus-per-node 8 --num-gpus-per-node 8 --colocate + --deploy-component all From a230ed1922c6e6e232825038aa3965907e91a01c Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 6 Oct 2026 03:30:01 -0700 Subject: [PATCH 6/9] fix(mimo): serialize preparation and use valid sampling defaults --- scripts/run_mimo_v2_6_flash.py | 33 ++++++++++--------- .../test_run_mimo_v2_6_flash.py | 33 +++++++++++++++++++ .../tools/test_convert_mimo_v2_to_bf16.py | 22 ++++++++++++- .../launch_scripts/test_py_launch_scripts.py | 2 ++ .../run_mimo_v2_6_flash.py/execute.txt | 4 +-- .../run_mimo_v2_6_flash.py/prepare.txt | 8 ++--- tools/convert_mimo_v2_to_bf16.py | 5 ++- 7 files changed, 84 insertions(+), 23 deletions(-) diff --git a/scripts/run_mimo_v2_6_flash.py b/scripts/run_mimo_v2_6_flash.py index cf70028227f..c71619befe8 100644 --- a/scripts/run_mimo_v2_6_flash.py +++ b/scripts/run_mimo_v2_6_flash.py @@ -13,9 +13,10 @@ decoder variant: global+dense, SWA+MoE, global+MoE, SWA+MoE Args: - --mode: `rl` runs GRPO with 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. + --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, @@ -46,6 +47,7 @@ import typer import miles.utils.external_utils.command_utils as U +from miles.utils.external_utils.command_utils.base_backend import _exclusive_path_lock _HF_REPO = "XiaomiMiMo/MiMo-V2.6-Flash-RL" @@ -118,17 +120,18 @@ def __post_init__(self): def prepare(args: ScriptArgs): backend = args.create_backend() target = Path(args.model_dir) / args.model_name - # 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}" - ) + with _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) @@ -168,7 +171,7 @@ def execute(args: ScriptArgs): "--n-samples-per-prompt 8 " "--rollout-max-response-len 8192 " "--rollout-temperature 1.0 " - "--rollout-top-p 0.95 " + "--rollout-top-p 1.0 " "--num-steps-per-rollout 1 " "--advantage-estimator grpo " "--entropy-coef 0.00 " 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 index 782eef5aa5d..d2988663343 100644 --- a/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py +++ b/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py @@ -1,8 +1,13 @@ +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.base_backend import BaseCommandBackend, _exclusive_path_lock + @pytest.mark.parametrize("mode", ["rl", "sft"]) @pytest.mark.parametrize("checkpoint_exists", [False, True]) @@ -27,3 +32,31 @@ def test_prepare_and_execute_with_the_configured_backend(monkeypatch, tmp_path, 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 index a704c2cafcc..70f867424ea 100644 --- a/tests/fast/tools/test_convert_mimo_v2_to_bf16.py +++ b/tests/fast/tools/test_convert_mimo_v2_to_bf16.py @@ -1,9 +1,9 @@ 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 @@ -173,3 +173,23 @@ def test_keep_quant_reindexes_ignored_layers(tmp_path): 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 index 57fcec56492..196f1e02c1e 100644 --- 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 @@ -46,7 +46,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --moe-aux-loss-coeff 0 --moe-router-dtype fp32 --moe-grouped-gemm - --hf-checkpoint /root/models/mimo26-p4-bf16 + --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 @@ -59,7 +59,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --n-samples-per-prompt 8 --rollout-max-response-len 8192 --rollout-temperature 1.0 - --rollout-top-p 0.95 + --rollout-top-p 1.0 --num-steps-per-rollout 1 --advantage-estimator grpo --entropy-coef 0.00 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 index c107a55905a..30af3580eb6 100644 --- 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 @@ -1,14 +1,14 @@ ### 0 -mkdir -p /root/models +mkdir -p /models ### 1 hf download XiaomiMiMo/MiMo-V2.6-Flash-RL - --local-dir /root/models/MiMo-V2.6-Flash-RL + --local-dir /models/MiMo-V2.6-Flash-RL ### 2 python /tools/convert_mimo_v2_to_bf16.py - --model-dir /root/models/MiMo-V2.6-Flash-RL - --save-dir /root/models/mimo26-p4-bf16 + --model-dir /models/MiMo-V2.6-Flash-RL + --save-dir /models/mimo26-p4-bf16 --device cuda --layers 0,1,5,6 diff --git a/tools/convert_mimo_v2_to_bf16.py b/tools/convert_mimo_v2_to_bf16.py index 9251abcf9c4..ff3f2587ccc 100644 --- a/tools/convert_mimo_v2_to_bf16.py +++ b/tools/convert_mimo_v2_to_bf16.py @@ -231,7 +231,6 @@ def load(name: str) -> torch.Tensor: 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} - (dst / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n") for path in src.iterdir(): if path.name.startswith("."): # e.g. the `hf download` cache continue @@ -242,6 +241,10 @@ def load(name: str) -> torch.Tensor: 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}" ) From 76d59851afac1db9082a0bb8994b74da96a2707e Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 6 Oct 2026 14:02:04 -0700 Subject: [PATCH 7/9] refactor(mimo): export the path lock and update the SGLang note 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 --- docs/models/mimo/mimo-v2-6-flash.md | 2 +- miles/utils/external_utils/command_utils/__init__.py | 2 ++ miles/utils/external_utils/command_utils/base_backend.py | 6 +++--- scripts/run_mimo_v2_6_flash.py | 6 +++--- tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py | 5 +++-- 5 files changed, 12 insertions(+), 9 deletions(-) diff --git a/docs/models/mimo/mimo-v2-6-flash.md b/docs/models/mimo/mimo-v2-6-flash.md index 5a4eeb16588..39ec81e9e61 100644 --- a/docs/models/mimo/mimo-v2-6-flash.md +++ b/docs/models/mimo/mimo-v2-6-flash.md @@ -38,7 +38,7 @@ The converter dequantizes each kv-head shard of the fused `qkv_proj` with its ow ### 3.2 SGLang -The BF16 engine needs the MiMo-V2 fixes that are not in the image yet: +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; 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/scripts/run_mimo_v2_6_flash.py b/scripts/run_mimo_v2_6_flash.py index c71619befe8..14d91582015 100644 --- a/scripts/run_mimo_v2_6_flash.py +++ b/scripts/run_mimo_v2_6_flash.py @@ -47,7 +47,6 @@ import typer import miles.utils.external_utils.command_utils as U -from miles.utils.external_utils.command_utils.base_backend import _exclusive_path_lock _HF_REPO = "XiaomiMiMo/MiMo-V2.6-Flash-RL" @@ -59,7 +58,7 @@ class _Recipe: layers: str | None = None # Default TP / PP / EP. parallel: tuple[int, int, int] = (2, 2, 2) - # The BF16 engine holds the whole model: 620 GB needs TP8 on 141 GB GPUs. + # 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' @@ -71,6 +70,7 @@ class _Recipe: "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, @@ -120,7 +120,7 @@ def __post_init__(self): def prepare(args: ScriptArgs): backend = args.create_backend() target = Path(args.model_dir) / args.model_name - with _exclusive_path_lock(str(target)): + 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]}" 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 index d2988663343..1570f82be0c 100644 --- a/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py +++ b/tests/fast/launch_scripts/test_run_mimo_v2_6_flash.py @@ -6,7 +6,8 @@ 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.base_backend import BaseCommandBackend, _exclusive_path_lock +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"]) @@ -50,7 +51,7 @@ def prepare(): # 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)): + with exclusive_path_lock(str(target)): future = pool.submit(prepare) assert started.wait(timeout=5) attempted_conversion = converted.wait(timeout=0.2) From 6346fb4bf201dd284414d7602ace2d32e7ee6006 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 6 Oct 2026 16:46:37 -0700 Subject: [PATCH 8/9] test(mimo): run the 4-layer e2e test only in nightly CI 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 --- .../model_scripts/test_mimo_v2_6_flash_4layer_ci.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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 index 868de5375e2..5b59b2d17b1 100644 --- 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 @@ -9,7 +9,13 @@ # 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=2400, suite="stage-c-8-gpu-h200", labels=["megatron", "model-scripts"], hardware=["hopper"]) +register_cuda_ci( + est_time=2400, + 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") From 2d3587b4fa97d6ea35acab2646644b48bbca66df Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 6 Oct 2026 16:57:12 -0700 Subject: [PATCH 9/9] test(mimo): tune the 4-layer e2e est_time to observed runtime 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 --- .../megatron/model_scripts/test_mimo_v2_6_flash_4layer_ci.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 5b59b2d17b1..fc6be4e0feb 100644 --- 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 @@ -10,7 +10,7 @@ # entropy checkers of --ci-test do not apply; the metric-history gates below track the train/rollout gap. register_cuda_ci( - est_time=2400, + est_time=500, suite="stage-c-8-gpu-h200", labels=["megatron", "model-scripts"], hardware=["hopper"],