Bug: MiniMax-M2.7 JANG MoE model fails to load — switch_mlp.gate_proj shape mismatch (config says 8-bit, actual is 4-bit)

[Keenni]

Bug: MiniMax-M2.7 JANG MoE model fails to load

Model: MiniMax-M2.7-JANG_2L (JANG format v2.0, MoE, 256 local experts)
oMLX Version: latest (via App / Homebrew)
vMLX: loads the same model successfully

Error

Expected shape (256, 1536, 768) but received shape (256, 1536, 192)
for parameter model.layers.0.block_sparse_moe.switch_mlp.gate_proj.weight

Root Cause Analysis

The config.json declares switch_mlp layer quantization as:

"model.layers.0.block_sparse_moe.switch_mlp.gate_proj": {"bits": 8, "group_size": 32, "mode": "affine"}

But the actual stored tensor in safetensors is:

weight: shape=(256, 1536, 192), dtype=uint32
scales: shape=(256, 1536, 24), dtype=float16
biases: shape=(256, 1536, 24), dtype=float16

The config claims 8-bit quantization with group_size=32, but the actual storage is 4-bit:

Parameter	Claimed (config)	Actual (stored)
bits	8	4
group_size	32	32
expected shape	(256, 1536, 768)	(256, 1536, 192) ✅

The dequantize logic in oMLX reads bits: 8 from config and computes a shape that doesn't match the actual tensor.

Meanwhile the actual tensor was stored with 4-bit quantization (4 values per uint32), giving:

• 3072 elements / 32 groups = 96 groups
• Each group: 32 × 4 bits / 8 = 16 bytes → but packed as uint32: 96 × 2 = 192 uint32

The config.json quantization metadata is wrong for switch_mlp layers — it reports bits: 8 but the safetensors are actually 4-bit packed. oMLX should either:

1. Auto-detect the actual bits from the tensor shape vs. config shape, or
2. Allow the JANG format to store the correct per-tensor bits in a separate field

Reproduction Steps

1. Download MiniMax-M2.7-JANG_2L model
2. Load in oMLX → fails with shape mismatch error
3. Load in vMLX → works fine

Model Config Snippet

{
  "model_type": "minimax_m2",
  "hidden_size": 3072,
  "intermediate_size": 1536,
  "num_hidden_layers": 62,
  "num_local_experts": 256,
  "num_experts_per_tok": 8,
  "head_dim": 128,
  "quantization": {
    "bits": 8,
    "group_size": 128,
    "mode": "affine",
    "model.layers.0.block_sparse_moe.switch_mlp.gate_proj": {"bits": 8, "group_size": 32, "mode": "affine"},
    "model.layers.0.block_sparse_moe.switch_mlp.up_proj": {"bits": 8, "group_size": 32, "mode": "affine"},
    "model.layers.0.block_sparse_moe.switch_mlp.down_proj": {"bits": 8, "group_size": 32, "mode": "affine"},
    ...
  }
}

Suggested Fix

In the MoE expert dequantization path, before applying the config's bits value, verify the actual tensor shape against what the config claims. If there's a mismatch, auto-detect the correct bit width from the tensor dimensions.

Alternatively, store the correct per-tensor bits value in the JANG format metadata so oMLX reads the right quantization parameters.

[Keenni]

Further Investigation: Confirmed JANG format mismatch

I compared loading two JANG models:

• Qwen3.6-27B-JANG_4M-CRACK ✅ loads fine in oMLX
• MiniMax-M2.7-JANG_2L ❌ fails with shape mismatch

Root Cause: config.json quantization metadata is wrong for MiniMax M2

The JANG_2L profile from HuggingFace explicitly states:

Tier	Components	Bits
CRITICAL	Attention (Q/K/V/O), lm_head	8
IMPORTANT	Embeddings	6
COMPRESS	Expert MLP (w1/w2/w3)	2
Passthrough	MoE router/gate, norms	16

But the config.json declares ALL layers uniformly:

"model.layers.0.block_sparse_moe.switch_mlp.gate_proj": {"bits": 8, "group_size": 32}

The actual safetensors tensor for switch_mlp.gate_proj is 4-bit packed (not 2-bit, not 8-bit):

weight: shape=(256, 1536, 192), dtype=uint32
scales: shape=(256, 1536, 24), dtype=float16

Verification: 3072 hidden_size / 128 group_size = 24 groups. For 4-bit: 24 groups × (128/16 elements_per_uint32) = 24 × 8 = 192 uint32. ✅

Why Qwen3.6-27B works

Qwen3.6-27B uses qwen3_5 architecture with standard experts MoE structure. Its config.json quantization metadata matches the actual tensor shapes, so oMLX's oq.py dequantize path works correctly.

Why MiniMax-M2.7 fails

oMLX reads bits: 8 from config for switch_mlp.gate_proj and computes expected shape (256, 1536, 768), but the actual tensor is (256, 1536, 192).

Recommendation

Fix for JANG model authors: The config.json quantization metadata should correctly reflect per-tensor bit widths. MiniMax M2.7 JANG models have incorrect bits values in their config — the switch_mlp layers should report bits: 4 (or 2) not bits: 8.

Fix for oMLX: Auto-detect the actual bit width from tensor shape vs config shape, or add a fallback that reads jang_config.json for JANG-format models instead of trusting config.json quantization metadata.

Workaround for users

Until this is fixed, you can try:

1. Manually edit config.json to change bits: 8 to bits: 4 for all switch_mlp entries
2. Use vMLX/MLX Studio which correctly reads jang_config.json