Coverage for transformer_lens/pretrained/weight_conversions/neo.py: 100%

1import einops 

2import torch 

3 

4from transformer_lens.HookedTransformerConfig import HookedTransformerConfig 

5 

6 

7def convert_neo_weights(neo, cfg: HookedTransformerConfig): 

8 state_dict = {} 

9 

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

11 

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

13 # positions than the model expects (e.g. TinyStories models were trained with 

14 # seq_len=512 but the HuggingFace config reports max_position_embeddings=2048). 

15 pos_embed = neo.transformer.wpe.weight 

16 if pos_embed.shape[0] > cfg.n_ctx: 

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

18 state_dict["pos_embed.W_pos"] = pos_embed 

19 

20 for l in range(cfg.n_layers): 

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

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

23 

24 W_Q = neo.transformer.h[l].attn.attention.q_proj.weight 

25 W_K = neo.transformer.h[l].attn.attention.k_proj.weight 

26 W_V = neo.transformer.h[l].attn.attention.v_proj.weight 

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

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

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

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

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

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

33 

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

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

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

37 

38 W_O = neo.transformer.h[l].attn.attention.out_proj.weight 

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

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

41 state_dict[f"blocks.{l}.attn.b_O"] = neo.transformer.h[l].attn.attention.out_proj.bias 

42 

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

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

45 

46 state_dict[f"blocks.{l}.mlp.W_in"] = neo.transformer.h[l].mlp.c_fc.weight.T 

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

48 

49 state_dict[f"blocks.{l}.mlp.W_out"] = neo.transformer.h[l].mlp.c_proj.weight.T 

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

51 state_dict["ln_final.w"] = neo.transformer.ln_f.weight 

52 state_dict["ln_final.b"] = neo.transformer.ln_f.bias 

53 

54 state_dict["unembed.W_U"] = neo.lm_head.weight.T 

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

56 return state_dict