Add MiniMax-M3

[qingzwang]

Note: The below template includes items meant for model contributions only. For other contributions such as bug fixes, features, etc., only fill out the relevant portions of the form.

Description

Add MiniMax-M3 (text backbone) contrib model — NeuronX Distributed Inference port of MiniMaxAI/MiniMax-M3, a ~428B total / ~23B active parameter vision-language MoE model. This contribution ports the text-only causal LM portion — the vision tower, multimodal projector, and Multi-Token Prediction (MTP) modules are not included.

Key implementation highlights:

• GQA attention with per-head Gemma RMSNorm on Q/K and partial RoPE (first half of each head, theta=5M)
• SwiGLU-OAI activation everywhere including MoE experts — gate * sigmoid(alpha*gate) * (up + 1.0) with α=1.702, clamp ±7.0
• Sigmoid routing with e_score_correction_bias as fp32 nn.Parameter + routed_scaling_factor=2.0
• MSA prefill via Lightning Indexer (block_size=128, topk_blocks=16, local_blocks=1) on 57 sparse layers
• Fused gate_up_proj with stride=2 ColumnParallel sharding for TP=64
• HF state-dict converter: strips language_model.model. prefix, fuses per-expert w1/w3, dequantizes MXFP8 to bf16
• Right-padded batching (left-pad causes silent corruption with NxDI's argmax(position_ids) last-token selection)
• End-to-end on trn2.48xlarge: TP=64, moe_ep=64, moe_tp=1, batch=32, seq_len=512

Model Information

Model Name: MiniMax-M3 (text backbone)

Model Architecture: Decoder-only MoE transformer with GQA (64 Q / 4 KV heads), 128 sigmoid-routed experts (top-4), shared expert, Lightning Indexer MSA

Purpose: Text generation (autoregressive LLM)

Checklist

Required Components

• Accuracy Test (test/integration/test_model.py)

  • Integration test validates model accuracy with logit comparison and greedy generation
  • Additional tests: test_full_model.py, test_generate.py, test_generate_right_pad.py, smoke_test_synthetic.py (no checkpoint needed), test_mxfp8_real.py, test_partial_real.py, test_rmsnorm_preshift.py
  • Tested on Neuron (trn2.48xlarge)

• README.md with the following sections:

  • Usage Example: Full Python code showing config setup, compile, load, and generate
  • Compatibility Matrix: Tested instance types and SDK versions
  • Example Checkpoints: Link to MiniMaxAI/MiniMax-M3 on HuggingFace (854 GB)
  • Testing Instructions: Commands for both synthetic smoke test and full integration test

• Source Code (src/)

  • modeling_minimax_m3.py — full model implementation (~61.5 KB)
  • __init__.py

Optional Components

• Unit Tests (CPU or Neuron-based)
  • N/A — tests are integration-level on Neuron hardware

Folder Structure

/contrib/models/MiniMax-M3/
  README.md
  HISTORY.md (15-iteration debug narrative)
  /src
    __init__.py
    modeling_minimax_m3.py
  /test
    __init__.py
    /integration
      __init__.py
      test_model.py
      test_full_model.py
      test_generate.py
      test_generate_right_pad.py
      smoke_test_synthetic.py
      test_mxfp8_real.py
      test_partial_real.py
      test_rmsnorm_preshift.py

Testing

How did you test this change?

Tested on trn2.48xlarge (TP=64, lnc=2) with:

• Synthetic smoke test (random weights, M3-shaped config — validates modeling code end-to-end without downloading 854 GB)
• Full integration test with real checkpoint (compile + shard + load + TTFT + ITL)
• Multiple sanity prompts (batch=32, seq_len=512, greedy decoding)
• LongBench real long-context QA at 2K and 4K sequence lengths

Test Results:

Sanity checks (batch=32, seq_len=512, greedy):

Prompt	Neuron top-1	Continuation
"1+1="	2 ✅	2\) and \(2+1=3\). So the answer is 3.
"The capital of France is"	Paris ✅	Paris.\nThe capital of France is Paris...
"Berlin is the capital of"	Germany ✅	Germany. Madrid is the capital of Spain. Rome...
"The largest planet in our solar system is"	Jupiter ✅	Jupiter. It is the fifth planet from the Sun and is a gas giant.

LongBench (real long-context QA):

Seq	Batch	TTFT	ITL	Answer quality
2K	8	7.3 s	54 ms/tok	✅ Correct
4K	4	17–19 s	43–49 ms/tok	Mixed (edge cases)

Compatibility

Tested with:

• Neuron SDK Version(s): aws_neuronx_venv_pytorch_2_9_nxd_inference (transformers 4.57+)
• Instance Type(s): trn2.48xlarge (TP=64, LNC=2, moe_ep=64, moe_tp=1)
• PyTorch Version: 2.9
• Python Version: 3.10+

Additional Information

• Text-only port — vision tower, multimodal projector, and MTP modules are filtered by the state-dict converter. This is a causal LM contribution.
• Right-padding required — the compiled NEFF uses argmax(position_ids) to locate the last real token; left-padding silently produces garbage.
• Long context (≥ 8K) limited — per-rank 24 GB HBM is insufficient for KV cache at 8K × batch 8. Enabling attention_dp_degree > 1 or cp_degree > 1 would unblock this (future work).
• Decode-side MSA not implemented — prefill runs the Lightning Indexer, but decode falls back to dense causal attention. Not a correctness issue at moderate context; a performance optimization for the 1M-token max_position range.
• Batch=8, seq_len=4096 hangs — the NEFF compiles but deadlocks at collectives OP:0 (Neuron compiler bug at that shape). Workaround: batch=4 for 4K compiles.
• MXFP8 dequantization on host during state-dict conversion.
• 15-iteration debug history preserved in HISTORY.md for reproducibility.

Related Issues

N/A — first contribution of this model architecture.

vLLM Integration

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

vLLM integration is future work — the current contribution focuses on the core NxDI modeling and compilation path.

---

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