Coverage for transformer_lens/pretrained/weight_conversions/gpt2.py: 96%

1import einops 

2import torch 

3 

4from transformer_lens.HookedTransformerConfig import HookedTransformerConfig 

5 

6 

7def convert_gpt2_weights(gpt2, cfg: HookedTransformerConfig): 

8 state_dict = {} 

9 

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

11 

12 # Trim positional embeddings to n_ctx if the pretrained weights have more 

13 # positions than the model expects. 

14 pos_embed = gpt2.transformer.wpe.weight 

15 if pos_embed.shape[0] > cfg.n_ctx: 15 ↛ 16line 15 didn't jump to line 16 because the condition on line 15 was never true

16 pos_embed = pos_embed[: cfg.n_ctx, :] 

17 state_dict["pos_embed.W_pos"] = pos_embed 

18 

19 for l in range(cfg.n_layers): 

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

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

22 

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

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

25 W = gpt2.transformer.h[l].attn.c_attn.weight 

26 W_Q, W_K, W_V = torch.tensor_split(W, 3, dim=1) 

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

28 W_K = einops.rearrange(W_K, "m (i h)->i m h", i=cfg.n_heads) 

29 W_V = einops.rearrange(W_V, "m (i h)->i m h", i=cfg.n_heads) 

30 

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

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

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

34 

35 qkv_bias = gpt2.transformer.h[l].attn.c_attn.bias 

36 qkv_bias = einops.rearrange( 

37 qkv_bias, 

38 "(qkv index head)->qkv index head", 

39 qkv=3, 

40 index=cfg.n_heads, 

41 head=cfg.d_head, 

42 ) 

43 state_dict[f"blocks.{l}.attn.b_Q"] = qkv_bias[0] 

44 state_dict[f"blocks.{l}.attn.b_K"] = qkv_bias[1] 

45 state_dict[f"blocks.{l}.attn.b_V"] = qkv_bias[2] 

46 

47 W_O = gpt2.transformer.h[l].attn.c_proj.weight 

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

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

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

51 

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

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

54 

55 W_in = gpt2.transformer.h[l].mlp.c_fc.weight 

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

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

58 

59 W_out = gpt2.transformer.h[l].mlp.c_proj.weight 

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

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

62 state_dict["unembed.W_U"] = gpt2.lm_head.weight.T 

63 

64 state_dict["ln_final.w"] = gpt2.transformer.ln_f.weight 

65 state_dict["ln_final.b"] = gpt2.transformer.ln_f.bias 

66 return state_dict