Coverage for transformer_lens/pretrained/weight_conversions/coder.py: 11%

1import einops 

2import torch 

3 

4from transformer_lens.HookedTransformerConfig import HookedTransformerConfig 

5 

6 

7def convert_coder_weights(model, cfg: HookedTransformerConfig): 

8 state_dict = {} 

9 

10 state_dict["embed.W_E"] = model.transformer.wte.weight 

11 state_dict["pos_embed.W_pos"] = model.transformer.wpe.weight 

12 

13 for l in range(cfg.n_layers): 

14 state_dict[f"blocks.{l}.ln1.w"] = model.transformer.h[l].ln_1.weight 

15 state_dict[f"blocks.{l}.ln1.b"] = model.transformer.h[l].ln_1.bias 

16 

17 # In GPT-2, q,k,v are produced by one big linear map, whose output is 

18 # concat([q, k, v]) 

19 W_KV = model.transformer.h[l].attn.kv_attn.weight # [d_model, 2 * d_head] 

20 W_K, W_V = torch.tensor_split(W_KV, 2, dim=1) 

21 W_Q = model.transformer.h[l].attn.q_attn.weight # [d_model, d_model] 

22 W_Q = einops.rearrange(W_Q, "m (i h)->i m h", i=cfg.n_heads) 

23 W_K = einops.repeat(W_K, "m h -> i m h", i=cfg.n_heads) 

24 W_V = einops.repeat(W_V, "m h -> i m h", i=cfg.n_heads) 

25 

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

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

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

29 

30 b_Q = einops.rearrange( 

31 model.transformer.h[l].attn.q_attn.bias, 

32 "(index head)-> index head", 

33 index=cfg.n_heads, 

34 head=cfg.d_head, 

35 ) 

36 b_KV = model.transformer.h[l].attn.kv_attn.bias # [2 * d_head] 

37 b_K, b_V = torch.tensor_split(b_KV, 2, dim=0) 

38 b_K = einops.repeat(b_K, "head -> index head", index=cfg.n_heads) 

39 b_V = einops.repeat(b_V, "head -> index head", index=cfg.n_heads) 

40 state_dict[f"blocks.{l}.attn.b_Q"] = b_Q 

41 state_dict[f"blocks.{l}.attn.b_K"] = b_K 

42 state_dict[f"blocks.{l}.attn.b_V"] = b_V 

43 

44 W_O = model.transformer.h[l].attn.c_proj.weight 

45 W_O = einops.rearrange(W_O, "(i h) m->i h m", i=cfg.n_heads) 

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

47 state_dict[f"blocks.{l}.attn.b_O"] = model.transformer.h[l].attn.c_proj.bias 

48 

49 state_dict[f"blocks.{l}.ln2.w"] = model.transformer.h[l].ln_2.weight 

50 state_dict[f"blocks.{l}.ln2.b"] = model.transformer.h[l].ln_2.bias 

51 

52 W_in = model.transformer.h[l].mlp.c_fc.weight 

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

54 state_dict[f"blocks.{l}.mlp.b_in"] = model.transformer.h[l].mlp.c_fc.bias 

55 

56 W_out = model.transformer.h[l].mlp.c_proj.weight 

57 state_dict[f"blocks.{l}.mlp.W_out"] = W_out 

58 state_dict[f"blocks.{l}.mlp.b_out"] = model.transformer.h[l].mlp.c_proj.bias 

59 state_dict["unembed.W_U"] = model.lm_head.weight.T 

60 

61 state_dict["ln_final.w"] = model.transformer.ln_f.weight 

62 state_dict["ln_final.b"] = model.transformer.ln_f.bias 

63 return state_dict