Coverage for transformer_lens/pretrained/weight_conversions/phi3.py: 14%

1import einops 

2import torch 

3 

4from transformer_lens.HookedTransformerConfig import HookedTransformerConfig 

5 

6 

7def convert_phi3_weights(phi, cfg: HookedTransformerConfig): 

8 state_dict = {} 

9 

10 state_dict["embed.W_E"] = phi.model.embed_tokens.weight 

11 

12 for l in range(cfg.n_layers): 

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

14 state_dict[f"blocks.{l}.ln1.b"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype) 

15 

16 W = phi.model.layers[l].self_attn.qkv_proj.weight 

17 W_Q, W_K, W_V = torch.tensor_split(W, 3, dim=0) 

18 W_Q = einops.rearrange( 

19 W_Q, "(n_head d_head) d_model -> n_head d_model d_head", n_head=cfg.n_heads 

20 ) 

21 W_K = einops.rearrange( 

22 W_K, "(n_head d_head) d_model -> n_head d_model d_head", n_head=cfg.n_heads 

23 ) 

24 W_V = einops.rearrange( 

25 W_V, "(n_head d_head) d_model -> n_head d_model d_head", n_head=cfg.n_heads 

26 ) 

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

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

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

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

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

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

33 

34 W_O = phi.model.layers[l].self_attn.o_proj.weight 

35 W_O = einops.rearrange( 

36 W_O, "d_model (n_head d_head) -> n_head d_head d_model", n_head=cfg.n_heads 

37 ) 

38 

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

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

41 

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

43 state_dict[f"blocks.{l}.ln2.b"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype) 

44 

45 W = phi.model.layers[l].mlp.gate_up_proj.weight.T 

46 W_gate, W_in = torch.tensor_split(W, 2, dim=1) 

47 state_dict[f"blocks.{l}.mlp.W_in"] = W_in 

48 state_dict[f"blocks.{l}.mlp.W_gate"] = W_gate 

49 state_dict[f"blocks.{l}.mlp.W_out"] = phi.model.layers[l].mlp.down_proj.weight.T 

50 

51 state_dict["ln_final.w"] = phi.model.norm.weight 

52 

53 state_dict["unembed.W_U"] = phi.lm_head.weight.T 

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

55 

56 return state_dict