Coverage for transformer_lens/pretrained/weight_conversions/apertus.py: 87%

1"""Apertus weight conversion. 

2 

3Converts Apertus (Swiss AI) weights to HookedTransformer format. Apertus is 

4structurally similar to Llama but uses non-gated MLP with XIeLU activation, 

5and different layer norm names (attention_layernorm / feedforward_layernorm). 

6""" 

7 

8import logging 

9from typing import cast 

10 

11import einops 

12import torch 

13 

14from transformer_lens.HookedTransformerConfig import HookedTransformerConfig 

15 

16logger = logging.getLogger(__name__) 

17 

18 

19def convert_apertus_weights(apertus, cfg: HookedTransformerConfig): 

20 state_dict = {} 

21 

22 state_dict["embed.W_E"] = apertus.model.embed_tokens.weight 

23 

24 using_gqa = cfg.n_key_value_heads is not None 

25 gqa_uscore = "_" if using_gqa else "" 

26 n_kv_heads = cast(int, cfg.n_key_value_heads if using_gqa else cfg.n_heads) 

27 

28 assert cfg.d_mlp is not None # keep mypy happy 

29 

30 for l in range(cfg.n_layers): 

31 state_dict[f"blocks.{l}.ln1.w"] = apertus.model.layers[l].attention_layernorm.weight 

32 

33 W_Q = apertus.model.layers[l].self_attn.q_proj.weight 

34 W_K = apertus.model.layers[l].self_attn.k_proj.weight 

35 W_V = apertus.model.layers[l].self_attn.v_proj.weight 

36 

37 if not cfg.load_in_4bit: 37 ↛ 42line 37 didn't jump to line 42 because the condition on line 37 was always true

38 W_Q = einops.rearrange(W_Q, "(n h) m->n m h", n=cfg.n_heads) 

39 W_K = einops.rearrange(W_K, "(n h) m->n m h", n=n_kv_heads) 

40 W_V = einops.rearrange(W_V, "(n h) m->n m h", n=n_kv_heads) 

41 

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

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

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

45 

46 # QK normalization weights 

47 if cfg.use_qk_norm: 

48 state_dict[f"blocks.{l}.attn.q_norm.w"] = apertus.model.layers[ 

49 l 

50 ].self_attn.q_norm.weight 

51 state_dict[f"blocks.{l}.attn.k_norm.w"] = apertus.model.layers[ 

52 l 

53 ].self_attn.k_norm.weight 

54 

55 state_dict[f"blocks.{l}.attn.b_Q"] = torch.zeros( 

56 cfg.n_heads, cfg.d_head, dtype=cfg.dtype, device=cfg.device 

57 ) 

58 state_dict[f"blocks.{l}.attn.{gqa_uscore}b_K"] = torch.zeros( 

59 n_kv_heads, 

60 cfg.d_head, 

61 dtype=cfg.dtype, 

62 device=cfg.device, 

63 ) 

64 state_dict[f"blocks.{l}.attn.{gqa_uscore}b_V"] = torch.zeros( 

65 n_kv_heads, 

66 cfg.d_head, 

67 dtype=cfg.dtype, 

68 device=cfg.device, 

69 ) 

70 

71 W_O = apertus.model.layers[l].self_attn.o_proj.weight 

72 

73 if not cfg.load_in_4bit: 73 ↛ 76line 73 didn't jump to line 76 because the condition on line 73 was always true

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

75 

76 state_dict[f"blocks.{l}.attn.W_O"] = W_O.to(device=cfg.device) 

77 

78 state_dict[f"blocks.{l}.attn.b_O"] = torch.zeros( 

79 cfg.d_model, dtype=cfg.dtype, device=cfg.device 

80 ) 

81 

82 state_dict[f"blocks.{l}.ln2.w"] = apertus.model.layers[l].feedforward_layernorm.weight 

83 

84 if not cfg.load_in_4bit: 84 ↛ 88line 84 didn't jump to line 88 because the condition on line 84 was always true

85 state_dict[f"blocks.{l}.mlp.W_in"] = apertus.model.layers[l].mlp.up_proj.weight.T 

86 state_dict[f"blocks.{l}.mlp.W_out"] = apertus.model.layers[l].mlp.down_proj.weight.T 

87 else: 

88 state_dict[f"blocks.{l}.mlp.W_in"] = apertus.model.layers[l].mlp.up_proj.weight 

89 state_dict[f"blocks.{l}.mlp.W_out"] = apertus.model.layers[l].mlp.down_proj.weight 

90 

91 state_dict[f"blocks.{l}.mlp.b_in"] = torch.zeros( 

92 cfg.d_mlp, dtype=cfg.dtype, device=cfg.device 

93 ) 

94 state_dict[f"blocks.{l}.mlp.b_out"] = torch.zeros( 

95 cfg.d_model, dtype=cfg.dtype, device=cfg.device 

96 ) 

97 

98 # Extract trainable XIeLU activation parameters 

99 mlp = apertus.model.layers[l].mlp 

100 try: 

101 if hasattr(mlp, "act_fn"): 

102 alpha_p = mlp.act_fn.alpha_p 

103 alpha_n = mlp.act_fn.alpha_n 

104 beta = mlp.act_fn.beta 

105 elif hasattr(mlp, "act"): 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 alpha_p = mlp.act.alpha_p 

107 alpha_n = mlp.act.alpha_n 

108 beta = mlp.act.beta 

109 else: 

110 alpha_p = mlp.alpha_p 

111 alpha_n = mlp.alpha_n 

112 beta = mlp.beta 

113 state_dict[f"blocks.{l}.mlp.act_fn.alpha_p"] = alpha_p 

114 state_dict[f"blocks.{l}.mlp.act_fn.alpha_n"] = alpha_n 

115 state_dict[f"blocks.{l}.mlp.act_fn.beta"] = beta 

116 except AttributeError: 

117 logger.warning("XIeLU activation parameters not found in layer %d, using defaults", l) 

118 state_dict[f"blocks.{l}.mlp.act_fn.alpha_p"] = torch.tensor( 

119 0.8, dtype=cfg.dtype, device=cfg.device 

120 ) 

121 state_dict[f"blocks.{l}.mlp.act_fn.alpha_n"] = torch.tensor( 

122 0.8, dtype=cfg.dtype, device=cfg.device 

123 ) 

124 state_dict[f"blocks.{l}.mlp.act_fn.beta"] = torch.tensor( 

125 0.5, dtype=cfg.dtype, device=cfg.device 

126 ) 

127 

128 state_dict["ln_final.w"] = apertus.model.norm.weight 

129 

130 state_dict["unembed.W_U"] = apertus.lm_head.weight.T 

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

132 

133 return state_dict