Bug/mixtral transformers 5.x

[jlarson4]

Description

• convert_mixtral_weights still read the pre-5.x module layout, every HookedTransformer Mixtral load raised AttributeError: 'MixtralDecoderLayer' object has no attribute 'block_sparse_moe'. 5.x renamed the MoE block (block_sparse_moe → mlp) and replaced the per-expert w1/w2/w3 Linears with batched Parameters on a single MixtralExperts (gate_up_proj [n_experts, 2*d_mlp, d_model], down_proj [n_experts, d_model, d_mlp]). Reachable via all four mistralai/Mixtral-8x7B* aliases in the registry. The converter now reads that layout, mirroring the already-migrated olmoe.py.
• The fused projection is split gate-first (gate_up_proj[e, :d_mlp, :] → W_gate, [d_mlp:, :] → W_in). Confirmed three independent ways, because getting it backwards is silent: MixtralExperts.forward does .chunk(2, dim=-1) after the linear; HF's own legacy-checkpoint conversion concatenates w1 (the gate) first; and a round-trip against HF's expert math is bit-exact while the swapped assignment is off by ~51.
• Second, independent fix in the same load path: HF's MixtralTopKRouter renormalizes the top-k routing weights, while TL's MoE gates that behind cfg.norm_topk_prob, which defaults to False and which the Mixtral config branch never set. MixtralConfig has no such field to read it from, so the branch now pins it to HF's behavior.
• Quantized checkpoints are refused rather than mis-read: slicing packed or scale-separated expert weights would drop the scales silently, so non-floating-point expert tensors raise NotImplementedError naming the fix.
• No legacy fallback is included, deliberately: the fused layout predates TL's transformers>=5.9 floor, and on-disk legacy checkpoints are rewritten to the fused parameters during loading, so a hasattr branch would be unreachable.
• Tests build a tiny Mixtral from config in memory and each was verified to fail with its fix reverted: the conversion completing at all; per-expert weights reproducing HF's expert output with a negative control proving a swapped gate/up would differ; router weights; the quantized refusal; the config pinning norm_topk_prob; TL's MoE block matching HF's; and end-to-end logits parity after loading into a HookedTransformer.
• Two fixture details are load-bearing and commented as such. The weights are amplified to std=0.3 because SiLU is near-linear at small magnitudes, so at HF's default init a swapped gate/up is only ~1.6e-05 off (vs ~8e-02 amplified) and reads as noise. And the model uses 4 experts with top-2 rather than 2-of-2: when every expert is selected the softmax already sums to 1 and the renormalization becomes a no-op that no test can observe — which is exactly how an earlier draft of the routing test passed with the renormalization disabled.
• Audited the sibling converters for the same staleness: olmoe.py and openai.py are current and their split directions were re-verified numerically; every other HF-backed converter's attribute paths were walked on 5.13 with no failures (notably phi.py's self_attn.dense and neox.py's query_key_value are not renamed and must be left alone), and coder.py's q_attn/kv_attn is a false positive — it is dispatched only for santacoder's trust_remote_code class.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility