[contrib] Add MiMo-V2.5-Pro (Xiaomi, 384 experts MoE, FP8 on Trn2)

[whn09]

Description

Adds a contrib port of XiaomiMiMo/MiMo-V2.5-Pro targeting Trn2 with the FP8 blockwise recipe already used by the MiMo-V2-Flash and MiMo-V2.5 contrib ports (moe_tp_degree=1, moe_ep_degree=64). Self-contained under contrib/models/MiMo-V2.5-Pro/; no upstream source is modified.

Model Information

Model Name: MiMo-V2.5-Pro (Xiaomi)

Model Architecture: Decoder-only MoE transformer, 70 layers, 6144 hidden, 128 Q heads / 8 KV heads, asymmetric Q-K head_dim=192 vs V head_dim=128, hybrid attention (10 full + 60 sliding-window layers), attention sink bias on SWA layers, fused qkv_proj, sigmoid router with noaux_tc, 384 routed experts (top-8), no shared expert.

Purpose: Text generation (general-purpose Chinese/English LLM).

Checklist

Required Components

• Accuracy Test (contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py)
  • Compiles + loads the model on Neuron and asserts the generated tokens are deterministic.
• README.md with:
  • Usage Example (direct NxDI + vLLM)
  • Compatibility Matrix (Trn2 with SDK 2.29)
  • Example Checkpoints (HF link)
  • Testing Instructions (pytest ...)
• Source Code (contrib/models/MiMo-V2.5-Pro/src/)
  • modeling_mimo_v2.py (NxDI modeling wrapper)
  • conversion_script/preprocess_mimo_v2_fp8.py (HF OCP FP8 → Neuron FP8 streaming preprocess)

Optional Components

• Unit tests — not included; covered by the integration test.

Folder Structure

contrib/models/MiMo-V2.5-Pro/
├── README.md
├── perf_test/
│   ├── 0_setup.sh
│   ├── bench_mimo_v2.sh
│   ├── run_bench_single.sh
│   ├── sanity_check.sh
│   ├── smoke_compile_mimo_v2.py
│   ├── smoke_generate_mimo_v2.py
│   └── vllm-neuron-patch.patch
├── src/
│   ├── modeling_mimo_v2.py
│   └── conversion_script/
│       ├── preprocess_mimo_v2_fp8.py
│       ├── preprocess_mimo_v2_parallel.py
│       └── run_preprocess_parallel.sh
└── test/
    ├── __init__.py
    ├── unit/__init__.py
    └── integration/
        ├── __init__.py
        └── test_model.py

Testing

How did you test this change?

Smoke + benchmark runs on a trn2.48xlarge (Neuron SDK 2.29, PyTorch 2.9, Python 3.12):

1. Preprocess the HF OCP-FP8 checkpoint to Neuron-FP8 (~20 min).
2. smoke_compile_mimo_v2.py — compile the model (TP=64, moe_tp=1/moe_ep=64, BS=48, seq_len=1024). First compile ~60 min TKG + ~15 min CTE; cached compile ~1 min.
3. smoke_generate_mimo_v2.py — 20-token generation via HuggingFaceGenerationAdapter.
4. bench_mimo_v2.sh — vLLM serving via vllm-neuron 0.5.0 + the patch in perf_test/vllm-neuron-patch.patch. Benchmarked with vllm bench serve --dataset-name random --random-input-len 900 --random-output-len 90.

Test Results:

vLLM serving throughput on trn2.48xlarge, FP8, BS=48, TP=64 / moe_tp=1 / moe_ep=64, continuous batching:

Concurrency	Total tok/s	Output tok/s	TTFT median (ms)	TPOT median (ms)
1	47	4.3	1,392	220
16	391	35.6	2,361	422
48	606	55	7,322	752

Per-stream ITL median holds at ~220 ms across all concurrency levels; growth at higher concurrency is from continuous-batching queue pressure.

Integration test: pytest contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py -v — passes locally on the DLAMI venv (requires the preprocessed checkpoint path; see README).

Compatibility

Tested with:

• Neuron SDK Version(s): 2.29
• Instance Type(s): trn2.48xlarge (128 physical NeuronCores, logical_nc_config=2 → 64 logical cores)
• PyTorch Version: 2.9.1
• Python Version: 3.12

Additional Information

• HBM footprint: Pro's FP8 expert weights are ~500 GB across 64 ranks; per-rank tensors reach ~20 GB on Trn2's 24 GB HBM, so the recipe is sensitive to collective DMA ring headroom. Any recipe change that adds per-rank tensor size (e.g. BF16 attention weights) can OOM at load.
• BS constraint: NxDI's TKG path refuses Expert Parallelism when batch_size < num_experts / top_k. For Pro that is 384 / 8 = 48, so the smallest working BS on the FP8 path is 48.
• FP8 numerical behaviour: under the all-FP8 recipe, Pro's attention weights (abs_mean ≈ 0.00124, 4× smaller than V2.5) cause the NKI blockwise FP8 accumulator to drift over 70 layers, producing prompt-dependent gibberish. A BF16-attn recipe (MoE kept FP8, q/k/v dequanted to BF16 via src/conversion_script/repatch_qkv_bf16.py, compiled at seq_len=256 to fit HBM) restores coherent output — verified end-to-end on smoke_generate_mimo_v2.py. This narrows the root cause to the attention path, not the MoE experts. The vLLM scripts in this PR still use the all-FP8 recipe so the bench numbers are from that configuration; re-benchmarking on BF16-attn is queued. See the README Status section for the full write-up.
• Preprocess: src/conversion_script/preprocess_mimo_v2_fp8.py streams layer-by-layer (~24 GB peak RAM, ~20 min) and writes <save_path>/model_layer{N}.safetensors plus model_extras.safetensors. A parallel variant is provided in preprocess_mimo_v2_parallel.py.

Related Issues

• Related contrib ports in the same family: #137 MiMo-V2-Flash, #148 MiMo-V2.5. Pro reuses the same modeling_mimo_v2.py wrapper and preprocess pipeline; the contrib folder is self-contained to match the per-model layout of the MiMo series.

vLLM Integration

• This model is intended for use with vLLM.
• Documentation includes vLLM registration instructions.

The perf_test/vllm-neuron-patch.patch adds a _register_contrib_models() hook to vllm-neuron 0.5.0's neuronx_distributed_model_loader.py that:

• Registers NeuronMiMoV2ForCausalLM into NxDI's MODEL_TYPES under mimov2flash and mimov2pro.
• Overrides vLLM's built-in GPU stubs for MiMoV2FlashForCausalLM / MiMoV2ProForCausalLM in ModelRegistry (they otherwise block ModelConfig validation).
• Patches AutoConfig.from_pretrained to default trust_remote_code=True so NxDI's hf_adapter.load_config can load the custom MiMoV2Config that ships with the checkpoint.

No upstream vllm-neuron code is modified — the patch lives in the contrib folder and is applied at install time by perf_test/0_setup.sh.

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included

🤖 Generated with Claude Code