Coverage for transformer_lens/pretrained/weight_conversions/mixtral.py: 100%

1import einops 

2import torch 

3 

4from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

5from transformer_lens.utilities.quantization import require_readable_weight 

6 

7 

8def convert_mixtral_weights(mixtral, cfg: HookedTransformerConfig): 

9 # The same as Mistral, but with the MLP replaced with MoE 

10 # As with Mistral, Mixtral has no biases 

11 

12 state_dict = {} 

13 

14 assert cfg.n_key_value_heads is not None # keep mypy happy 

15 assert cfg.d_mlp is not None 

16 assert cfg.num_experts is not None 

17 

18 state_dict["embed.W_E"] = mixtral.model.embed_tokens.weight 

19 

20 for l in range(cfg.n_layers): 

21 state_dict[f"blocks.{l}.ln1.w"] = mixtral.model.layers[l].input_layernorm.weight 

22 

23 W_Q = mixtral.model.layers[l].self_attn.q_proj.weight 

24 W_K = mixtral.model.layers[l].self_attn.k_proj.weight 

25 W_V = mixtral.model.layers[l].self_attn.v_proj.weight 

26 W_Q = einops.rearrange(W_Q, "(n h) m->n m h", n=cfg.n_heads) 

27 W_K = einops.rearrange(W_K, "(n h) m->n m h", n=cfg.n_key_value_heads) 

28 W_V = einops.rearrange(W_V, "(n h) m->n m h", n=cfg.n_key_value_heads) 

29 state_dict[f"blocks.{l}.attn.W_Q"] = W_Q 

30 state_dict[f"blocks.{l}.attn._W_K"] = W_K 

31 state_dict[f"blocks.{l}.attn._W_V"] = W_V 

32 

33 state_dict[f"blocks.{l}.attn.b_Q"] = torch.zeros(cfg.n_heads, cfg.d_head, dtype=cfg.dtype) 

34 state_dict[f"blocks.{l}.attn._b_K"] = torch.zeros( 

35 cfg.n_key_value_heads, cfg.d_head, dtype=cfg.dtype 

36 ) 

37 state_dict[f"blocks.{l}.attn._b_V"] = torch.zeros( 

38 cfg.n_key_value_heads, cfg.d_head, dtype=cfg.dtype 

39 ) 

40 

41 W_O = mixtral.model.layers[l].self_attn.o_proj.weight 

42 W_O = einops.rearrange(W_O, "m (n h)->n h m", n=cfg.n_heads) 

43 state_dict[f"blocks.{l}.attn.W_O"] = W_O 

44 

45 state_dict[f"blocks.{l}.attn.b_O"] = torch.zeros(cfg.d_model, dtype=cfg.dtype) 

46 

47 state_dict[f"blocks.{l}.ln2.w"] = mixtral.model.layers[l].post_attention_layernorm.weight 

48 

49 # transformers 5.x renamed the MoE block (block_sparse_moe -> mlp) and 

50 # replaced the per-expert w1/w2/w3 Linears with batched Parameters on a 

51 # single MixtralExperts module: 

52 # gate_up_proj: [num_experts, 2 * d_mlp, d_model] (gate fused above up) 

53 # down_proj: [num_experts, d_model, d_mlp] 

54 moe = mixtral.model.layers[l].mlp 

55 # Guarded like the experts below: load_state_dict accepts a SAME-SHAPE 

56 # int8/FP8 router and silently casts it to float32, so nothing 

57 # downstream catches it. 

58 state_dict[f"blocks.{l}.mlp.W_gate.weight"] = require_readable_weight( 

59 moe.gate.weight, operation="convert the Mixtral router weight", owner=mixtral 

60 ) 

61 

62 experts = moe.experts 

63 gate_up = require_readable_weight( 

64 experts.gate_up_proj, 

65 operation="convert Mixtral expert weights (gate_up_proj)", 

66 owner=mixtral, 

67 ) 

68 down = require_readable_weight( 

69 experts.down_proj, 

70 operation="convert Mixtral expert weights (down_proj)", 

71 owner=mixtral, 

72 ) 

73 

74 # MixtralExperts.forward does 

75 # gate, up = F.linear(x, gate_up_proj[e]).chunk(2, dim=-1) 

76 # so the FIRST half of dim 1 is the gate projection and the second is up. 

77 for e in range(cfg.num_experts): 

78 state_dict[f"blocks.{l}.mlp.experts.{e}.W_gate.weight"] = gate_up[e, : cfg.d_mlp, :] 

79 state_dict[f"blocks.{l}.mlp.experts.{e}.W_in.weight"] = gate_up[e, cfg.d_mlp :, :] 

80 state_dict[f"blocks.{l}.mlp.experts.{e}.W_out.weight"] = down[e] 

81 

82 state_dict["ln_final.w"] = mixtral.model.norm.weight.data 

83 

84 state_dict["unembed.W_U"] = mixtral.lm_head.weight.T 

85 state_dict["unembed.b_U"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype) 

86 

87 return state_dict