Coverage for transformer_lens/pretrained/weight_conversions/nanogpt.py: 7%

1import einops 

2import torch 

3 

4from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

5 

6 

7def convert_nanogpt_weights(old_state_dict, cfg: HookedTransformerConfig): 

8 """For https://github.com/karpathy/nanoGPT 

9 There are two complications with converting nanogpt models: 

10 The first is that some state dicts have an unwanted prefix on keys that needs to be removed. 

11 The second is that the models can be saved with or without bias. By default, there 

12 is no bias. This function can handle both cases.""" 

13 # Nanogpt models saved after torch.compile() have this unwanted prefix 

14 unwanted_prefix = "_orig_mod." 

15 for k, v in list(old_state_dict.items()): 

16 if k.startswith(unwanted_prefix): 

17 old_state_dict[k[len(unwanted_prefix) :]] = old_state_dict.pop(k) 

18 

19 new_state_dict = {} 

20 new_state_dict["pos_embed.W_pos"] = old_state_dict["transformer.wpe.weight"] 

21 new_state_dict["embed.W_E"] = old_state_dict["transformer.wte.weight"] 

22 

23 new_state_dict["ln_final.w"] = old_state_dict["transformer.ln_f.weight"] 

24 new_state_dict["ln_final.b"] = torch.zeros_like(old_state_dict["transformer.ln_f.weight"]) 

25 new_state_dict["unembed.W_U"] = old_state_dict["lm_head.weight"].T 

26 

27 bias = False 

28 if "transformer.ln_f.bias" in old_state_dict: 

29 bias = True 

30 new_state_dict["ln_final.b"] = old_state_dict["transformer.ln_f.bias"] 

31 else: 

32 new_state_dict["unembed.b_U"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype) 

33 

34 for layer in range(cfg.n_layers): 

35 layer_key = f"transformer.h.{layer}" 

36 

37 new_state_dict[f"blocks.{layer}.ln1.w"] = old_state_dict[f"{layer_key}.ln_1.weight"] 

38 # A bias of zeros is required for folding layer norm 

39 new_state_dict[f"blocks.{layer}.ln1.b"] = torch.zeros_like( 

40 old_state_dict[f"{layer_key}.ln_1.weight"] 

41 ) 

42 new_state_dict[f"blocks.{layer}.ln2.w"] = old_state_dict[f"{layer_key}.ln_2.weight"] 

43 new_state_dict[f"blocks.{layer}.ln2.b"] = torch.zeros_like( 

44 old_state_dict[f"{layer_key}.ln_2.weight"] 

45 ) 

46 

47 new_state_dict[f"blocks.{layer}.attn.IGNORE"] = torch.tensor(-torch.inf) 

48 

49 W = old_state_dict[f"{layer_key}.attn.c_attn.weight"] 

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

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

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

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

54 new_state_dict[f"blocks.{layer}.attn.W_Q"] = W_Q 

55 new_state_dict[f"blocks.{layer}.attn.W_K"] = W_K 

56 new_state_dict[f"blocks.{layer}.attn.W_V"] = W_V 

57 

58 W_O = old_state_dict[f"{layer_key}.attn.c_proj.weight"] 

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

60 new_state_dict[f"blocks.{layer}.attn.W_O"] = W_O 

61 

62 new_state_dict[f"blocks.{layer}.mlp.W_in"] = old_state_dict[ 

63 f"{layer_key}.mlp.c_fc.weight" 

64 ].T 

65 new_state_dict[f"blocks.{layer}.mlp.W_out"] = old_state_dict[ 

66 f"{layer_key}.mlp.c_proj.weight" 

67 ].T 

68 

69 if bias: 

70 new_state_dict[f"blocks.{layer}.ln1.b"] = old_state_dict[f"{layer_key}.ln_1.bias"] 

71 new_state_dict[f"blocks.{layer}.ln2.b"] = old_state_dict[f"{layer_key}.ln_2.bias"] 

72 new_state_dict[f"blocks.{layer}.mlp.b_in"] = old_state_dict[ 

73 f"{layer_key}.mlp.c_fc.bias" 

74 ] 

75 new_state_dict[f"blocks.{layer}.mlp.b_out"] = old_state_dict[ 

76 f"{layer_key}.mlp.c_proj.bias" 

77 ] 

78 

79 B = old_state_dict[f"{layer_key}.attn.c_attn.bias"] 

80 B_Q, B_K, B_V = torch.tensor_split(B, 3, dim=0) 

81 B_Q = einops.rearrange(B_Q, "(i h)->i h", i=cfg.n_heads) 

82 B_K = einops.rearrange(B_K, "(i h)->i h", i=cfg.n_heads) 

83 B_V = einops.rearrange(B_V, "(i h)->i h", i=cfg.n_heads) 

84 new_state_dict[f"blocks.{layer}.attn.b_Q"] = B_Q 

85 new_state_dict[f"blocks.{layer}.attn.b_K"] = B_K 

86 new_state_dict[f"blocks.{layer}.attn.b_V"] = B_V 

87 new_state_dict[f"blocks.{layer}.attn.b_O"] = old_state_dict[ 

88 f"{layer_key}.attn.c_proj.bias" 

89 ] 

90 else: 

91 if cfg.d_mlp is None: 

92 raise ValueError( 

93 "cfg.d_mlp must be set to convert nanoGPT weights for the no-bias case." 

94 ) 

95 new_state_dict[f"blocks.{layer}.mlp.b_out"] = torch.zeros(cfg.d_model, dtype=cfg.dtype) 

96 new_state_dict[f"blocks.{layer}.mlp.b_in"] = torch.zeros(cfg.d_mlp, dtype=cfg.dtype) 

97 new_state_dict[f"blocks.{layer}.attn.b_Q"] = torch.zeros( 

98 (cfg.n_heads, cfg.d_head), dtype=cfg.dtype 

99 ) 

100 new_state_dict[f"blocks.{layer}.attn.b_K"] = torch.zeros( 

101 cfg.n_heads, cfg.d_head, dtype=cfg.dtype 

102 ) 

103 new_state_dict[f"blocks.{layer}.attn.b_V"] = torch.zeros( 

104 cfg.n_heads, cfg.d_head, dtype=cfg.dtype 

105 ) 

106 new_state_dict[f"blocks.{layer}.attn.b_O"] = torch.zeros(cfg.d_model, dtype=cfg.dtype) 

107 

108 return new_state_dict