Coverage for transformer_lens/pretrained/weight_conversions/olmoe.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_olmoe_weights(olmoe, cfg: HookedTransformerConfig): 

9 state_dict = {} 

10 

11 assert cfg.n_key_value_heads is not None 

12 assert cfg.d_mlp is not None 

13 assert cfg.num_experts is not None 

14 

15 state_dict["embed.W_E"] = olmoe.model.embed_tokens.weight 

16 

17 for l in range(cfg.n_layers): 

18 olmoe_layer = olmoe.model.layers[l] 

19 state_dict[f"blocks.{l}.ln1.w"] = olmoe_layer.input_layernorm.weight 

20 

21 W_Q = olmoe_layer.self_attn.q_proj.weight 

22 W_K = olmoe_layer.self_attn.k_proj.weight 

23 W_V = olmoe_layer.self_attn.v_proj.weight 

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

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

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

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

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

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

30 state_dict[f"blocks.{l}.attn.q_norm.w"] = olmoe_layer.self_attn.q_norm.weight 

31 state_dict[f"blocks.{l}.attn.k_norm.w"] = olmoe_layer.self_attn.k_norm.weight 

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 = olmoe_layer.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"] = olmoe_layer.post_attention_layernorm.weight 

48 

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

50 olmoe_layer.mlp.gate.weight, 

51 operation="convert the OLMoE router weight", 

52 owner=olmoe, 

53 ) 

54 

55 # HF OLMoE uses batched expert weights: 

56 # gate_up_proj: [num_experts, 2 * intermediate_size, hidden_size] 

57 # down_proj: [num_experts, hidden_size, intermediate_size] 

58 # The gate_up_proj fuses gate and up projections along dim 1. 

59 experts = olmoe_layer.mlp.experts 

60 gate_up = require_readable_weight( 

61 experts.gate_up_proj, 

62 operation="convert OLMoE expert weights (gate_up_proj)", 

63 owner=olmoe, 

64 ) # [num_experts, 2*d_mlp, d_model] 

65 down = require_readable_weight( 

66 experts.down_proj, 

67 operation="convert OLMoE expert weights (down_proj)", 

68 owner=olmoe, 

69 ) # [num_experts, d_model, d_mlp] 

70 

71 for e in range(cfg.num_experts): 

72 # Split fused gate_up into gate and up projections 

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

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

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

76 

77 state_dict["ln_final.w"] = olmoe.model.norm.weight 

78 

79 state_dict["unembed.W_U"] = olmoe.lm_head.weight.T 

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

81 

82 return state_dict