Coverage for transformer_lens/pretrained/weight_conversions/olmo.py: 13%

1import einops 

2import torch 

3 

4from transformer_lens.config.HookedTransformerConfig import HookedTransformerConfig 

5 

6 

7def convert_olmo_weights(olmo, cfg: HookedTransformerConfig): 

8 state_dict = {} 

9 

10 assert cfg.d_mlp is not None 

11 

12 state_dict["embed.W_E"] = olmo.model.embed_tokens.weight 

13 for l in range(cfg.n_layers): 

14 olmo_layer = olmo.model.layers[l] 

15 

16 W_Q = olmo_layer.self_attn.q_proj.weight 

17 W_K = olmo_layer.self_attn.k_proj.weight 

18 W_V = olmo_layer.self_attn.v_proj.weight 

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

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

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

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

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

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

25 

26 W_O = olmo_layer.self_attn.o_proj.weight 

27 W_O = einops.rearrange(W_O, "m (n h)->n h m", n=cfg.n_heads) 

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

29 

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

31 

32 state_dict[f"blocks.{l}.mlp.W_in"] = olmo_layer.mlp.up_proj.weight.T 

33 state_dict[f"blocks.{l}.mlp.W_gate"] = olmo_layer.mlp.gate_proj.weight.T 

34 state_dict[f"blocks.{l}.mlp.b_in"] = torch.zeros(cfg.d_mlp, dtype=cfg.dtype) 

35 

36 state_dict[f"blocks.{l}.mlp.W_out"] = olmo_layer.mlp.down_proj.weight.T 

37 state_dict[f"blocks.{l}.mlp.b_out"] = torch.zeros(cfg.d_model, dtype=cfg.dtype) 

38 

39 state_dict[f"blocks.{l}.ln1.w"] = torch.ones(cfg.d_model, dtype=cfg.dtype) 

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

41 state_dict[f"blocks.{l}.ln2.w"] = torch.ones(cfg.d_model, dtype=cfg.dtype) 

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

43 

44 state_dict["ln_final.w"] = torch.ones(cfg.d_model, dtype=cfg.dtype) 

45 state_dict["ln_final.b"] = torch.zeros(cfg.d_model, dtype=cfg.dtype) 

46 

47 state_dict["unembed.W_U"] = olmo.lm_head.weight.T 

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

49 

50 return state_dict