|
| 1 | +"""convert_mixtral_weights against the transformers 5.x Mixtral layout. |
| 2 | +
|
| 3 | +5.x renamed the MoE block (``block_sparse_moe`` -> ``mlp``) and replaced the |
| 4 | +per-expert ``w1``/``w2``/``w3`` Linears with batched Parameters on a single |
| 5 | +``MixtralExperts`` module, so the previous converter raised AttributeError on |
| 6 | +every Mixtral load. |
| 7 | +
|
| 8 | +The model is built from a tiny config in memory — no download, no hub access. |
| 9 | +""" |
| 10 | + |
| 11 | +from __future__ import annotations |
| 12 | + |
| 13 | +import pytest |
| 14 | +import torch |
| 15 | +import torch.nn.functional as F |
| 16 | +from transformers import MixtralConfig, MixtralForCausalLM |
| 17 | + |
| 18 | +from transformer_lens.config import HookedTransformerConfig |
| 19 | +from transformer_lens.pretrained.weight_conversions import convert_mixtral_weights |
| 20 | + |
| 21 | +D_MODEL, D_MLP, N_EXPERTS, N_LAYERS = 8, 16, 4, 1 |
| 22 | +# top-k must be < N_EXPERTS or the softmax already sums to 1 and the |
| 23 | +# top-k renormalization becomes a no-op no test could observe. |
| 24 | +EXPERTS_PER_TOKEN = 2 |
| 25 | + |
| 26 | + |
| 27 | +@pytest.fixture(scope="module") |
| 28 | +def hf_model() -> MixtralForCausalLM: |
| 29 | + """Tiny Mixtral with AMPLIFIED weights. |
| 30 | +
|
| 31 | + The amplification is load-bearing, not cosmetic: SiLU is near-linear at small |
| 32 | + magnitudes, so ``act(gate) * up ~= act(up) * gate`` and a swapped gate/up |
| 33 | + mapping is only ~1.6e-05 off at HF's default init (0.02) — close enough to |
| 34 | + read as noise. At std=0.3 the same swap is ~8e-02, a ~5000x margin. Do not |
| 35 | + lower this without re-measuring the negative control below. |
| 36 | + """ |
| 37 | + with torch.random.fork_rng(devices=[]): |
| 38 | + torch.manual_seed(0) |
| 39 | + model = MixtralForCausalLM( |
| 40 | + MixtralConfig( |
| 41 | + hidden_size=D_MODEL, |
| 42 | + intermediate_size=D_MLP, |
| 43 | + num_hidden_layers=N_LAYERS, |
| 44 | + num_attention_heads=2, |
| 45 | + num_key_value_heads=1, |
| 46 | + vocab_size=32, |
| 47 | + num_local_experts=N_EXPERTS, |
| 48 | + num_experts_per_tok=EXPERTS_PER_TOKEN, |
| 49 | + max_position_embeddings=32, |
| 50 | + ) |
| 51 | + ).eval() |
| 52 | + for param in model.parameters(): |
| 53 | + torch.nn.init.normal_(param, std=0.3) |
| 54 | + return model |
| 55 | + |
| 56 | + |
| 57 | +@pytest.fixture(scope="module") |
| 58 | +def tl_cfg() -> HookedTransformerConfig: |
| 59 | + return HookedTransformerConfig( |
| 60 | + d_model=D_MODEL, |
| 61 | + d_head=4, |
| 62 | + n_heads=2, |
| 63 | + n_key_value_heads=1, |
| 64 | + n_layers=N_LAYERS, |
| 65 | + n_ctx=32, |
| 66 | + d_vocab=32, |
| 67 | + d_mlp=D_MLP, |
| 68 | + num_experts=N_EXPERTS, |
| 69 | + experts_per_token=EXPERTS_PER_TOKEN, |
| 70 | + act_fn="silu", |
| 71 | + normalization_type="RMS", |
| 72 | + positional_embedding_type="rotary", |
| 73 | + # Mirrors what the MixtralForCausalLM config branch produces; that the |
| 74 | + # branch really does set it is pinned separately by |
| 75 | + # test_config_pins_topk_renormalization, so the two together chain the |
| 76 | + # real load path to the behavior below. |
| 77 | + norm_topk_prob=True, |
| 78 | + ) |
| 79 | + |
| 80 | + |
| 81 | +@pytest.fixture(scope="module") |
| 82 | +def state_dict(hf_model, tl_cfg) -> dict: |
| 83 | + return convert_mixtral_weights(hf_model, tl_cfg) |
| 84 | + |
| 85 | + |
| 86 | +def test_converts_the_5x_batched_expert_layout(state_dict) -> None: |
| 87 | + """The whole conversion runs — it previously raised AttributeError looking |
| 88 | + for the removed `block_sparse_moe` attribute.""" |
| 89 | + for expert in range(N_EXPERTS): |
| 90 | + for name, shape in ( |
| 91 | + ("W_gate", (D_MLP, D_MODEL)), |
| 92 | + ("W_in", (D_MLP, D_MODEL)), |
| 93 | + ("W_out", (D_MODEL, D_MLP)), |
| 94 | + ): |
| 95 | + key = f"blocks.0.mlp.experts.{expert}.{name}.weight" |
| 96 | + assert key in state_dict, f"missing {key}" |
| 97 | + assert state_dict[key].shape == shape |
| 98 | + |
| 99 | + |
| 100 | +def test_expert_weights_reproduce_hf_expert_output(hf_model, state_dict) -> None: |
| 101 | + """The decisive check on the fused-projection split: the extracted weights |
| 102 | + must reproduce HF's own expert computation. |
| 103 | +
|
| 104 | + Shapes cannot catch a swapped gate/up — both halves are [d_mlp, d_model] — |
| 105 | + and the swap is silent, so this compares numerically and carries a negative |
| 106 | + control proving the swapped assignment would differ. |
| 107 | + """ |
| 108 | + with torch.random.fork_rng(devices=[]): |
| 109 | + torch.manual_seed(1) |
| 110 | + x = torch.randn(1, D_MODEL) |
| 111 | + |
| 112 | + experts = hf_model.model.layers[0].mlp.experts |
| 113 | + for expert in range(N_EXPERTS): |
| 114 | + # HF's own math (MixtralExperts.forward): the fused projection is split |
| 115 | + # with .chunk(2, dim=-1) AFTER the linear, so the first half is the gate. |
| 116 | + gate_hf, up_hf = F.linear(x, experts.gate_up_proj[expert]).chunk(2, dim=-1) |
| 117 | + expected = F.linear(F.silu(gate_hf) * up_hf, experts.down_proj[expert]) |
| 118 | + |
| 119 | + w_gate = state_dict[f"blocks.0.mlp.experts.{expert}.W_gate.weight"] |
| 120 | + w_in = state_dict[f"blocks.0.mlp.experts.{expert}.W_in.weight"] |
| 121 | + w_out = state_dict[f"blocks.0.mlp.experts.{expert}.W_out.weight"] |
| 122 | + actual = F.linear(F.silu(F.linear(x, w_gate)) * F.linear(x, w_in), w_out) |
| 123 | + torch.testing.assert_close(actual, expected) |
| 124 | + |
| 125 | + swapped = F.linear(F.silu(F.linear(x, w_in)) * F.linear(x, w_gate), w_out) |
| 126 | + assert not torch.allclose(swapped, expected, atol=1e-6), ( |
| 127 | + f"expert {expert}: gate and up are interchangeable in this fixture, " |
| 128 | + "so the assertion above cannot detect a swapped mapping" |
| 129 | + ) |
| 130 | + |
| 131 | + |
| 132 | +def test_router_weights_come_from_the_moe_gate(hf_model, state_dict) -> None: |
| 133 | + torch.testing.assert_close( |
| 134 | + state_dict["blocks.0.mlp.W_gate.weight"], |
| 135 | + hf_model.model.layers[0].mlp.gate.weight, |
| 136 | + ) |
| 137 | + |
| 138 | + |
| 139 | +def test_quantized_expert_weights_are_refused(hf_model, tl_cfg) -> None: |
| 140 | + """Slicing packed/scaled expert weights would silently drop their scales, so |
| 141 | + the converter must refuse rather than emit plausible garbage.""" |
| 142 | + experts = hf_model.model.layers[0].mlp.experts |
| 143 | + original = experts.gate_up_proj |
| 144 | + try: |
| 145 | + experts.gate_up_proj = torch.nn.Parameter( |
| 146 | + original.detach().to(torch.int8), requires_grad=False |
| 147 | + ) |
| 148 | + with pytest.raises(NotImplementedError, match="floating-point"): |
| 149 | + convert_mixtral_weights(hf_model, tl_cfg) |
| 150 | + finally: |
| 151 | + experts.gate_up_proj = original |
| 152 | + |
| 153 | + |
| 154 | +def test_config_pins_topk_renormalization() -> None: |
| 155 | + """HF's MixtralTopKRouter renormalizes top-k weights unconditionally, and |
| 156 | + MixtralConfig has no field to read that from, so TL's config branch must pin |
| 157 | + it — otherwise TL skips the renormalization and routing is silently wrong.""" |
| 158 | + from transformer_lens.loading_from_pretrained import convert_hf_model_config |
| 159 | + |
| 160 | + cfg = convert_hf_model_config("mistralai/Mixtral-8x7B-v0.1") |
| 161 | + assert cfg["norm_topk_prob"] is True |
| 162 | + |
| 163 | + |
| 164 | +def test_converted_model_reproduces_hf_logits(hf_model, tl_cfg, state_dict) -> None: |
| 165 | + """End-to-end: the converted weights must drive a HookedTransformer to the |
| 166 | + same logits as the HF model they came from. |
| 167 | +
|
| 168 | + This is what catches an orientation or routing error that the per-tensor |
| 169 | + assertions above would miss — notably the top-k renormalization, which is a |
| 170 | + config-level behavior no weight assertion can see. |
| 171 | + """ |
| 172 | + from transformer_lens import HookedTransformer |
| 173 | + |
| 174 | + tl_model = HookedTransformer(tl_cfg, tokenizer=None) |
| 175 | + tl_model.load_state_dict(state_dict, strict=False) |
| 176 | + tl_model.eval() |
| 177 | + |
| 178 | + ids = torch.tensor([[1, 5, 9, 2]]) |
| 179 | + with torch.no_grad(): |
| 180 | + tl_logits = tl_model(ids) |
| 181 | + hf_logits = hf_model(ids).logits |
| 182 | + |
| 183 | + max_diff = (tl_logits - hf_logits).abs().max().item() |
| 184 | + scale = max(1.0, hf_logits.abs().max().item()) |
| 185 | + assert max_diff < 1e-4 * scale, ( |
| 186 | + f"converted Mixtral drifts {max_diff:.3e} from HF (scale {scale:.3f}) — " |
| 187 | + "a weight orientation or routing term is wrong" |
| 188 | + ) |
| 189 | + |
| 190 | + |
| 191 | +def test_routing_renormalization_matches_hf(hf_model, tl_cfg, state_dict) -> None: |
| 192 | + """TL's MoE block must reproduce HF's, which renormalizes the top-k routing |
| 193 | + weights unconditionally. |
| 194 | +
|
| 195 | + Compares the blocks directly rather than hooking `hook_expert_weights`: that |
| 196 | + hook fires on the full pre-top-k softmax, whose top-k slice sums to 1 by |
| 197 | + construction, so an assertion on it holds whether or not the renormalization |
| 198 | + ran. |
| 199 | + """ |
| 200 | + from transformer_lens import HookedTransformer |
| 201 | + |
| 202 | + tl_model = HookedTransformer(tl_cfg, tokenizer=None) |
| 203 | + tl_model.load_state_dict(state_dict, strict=False) |
| 204 | + tl_model.eval() |
| 205 | + |
| 206 | + with torch.random.fork_rng(devices=[]): |
| 207 | + torch.manual_seed(3) |
| 208 | + hidden = torch.randn(1, 4, D_MODEL) |
| 209 | + |
| 210 | + with torch.no_grad(): |
| 211 | + tl_out = tl_model.blocks[0].mlp(hidden) |
| 212 | + hf_out = hf_model.model.layers[0].mlp(hidden) |
| 213 | + hf_out = hf_out[0] if isinstance(hf_out, tuple) else hf_out |
| 214 | + |
| 215 | + torch.testing.assert_close(tl_out, hf_out, atol=1e-5, rtol=1e-4) |
0 commit comments