Coverage for transformer_lens/model_bridge/sources/native/init.py: 97%

1"""Weight init for NativeModel. 

2 

3Supported modes: ``"gpt2"`` (Normal(0, std) with 1/sqrt(2*n_layers) residual 

4scaling on output projections), ``"xavier_uniform"`` / ``"xavier_normal"``, 

5``"kaiming_uniform"`` / ``"kaiming_normal"`` (relu nonlinearity). Norm weights 

6go to 1, all biases to 0; param-free norms (LNPre / RMSPre) have nothing to 

7initialize. 

8 

9Determinism uses a scoped ``torch.Generator``, not ``torch.manual_seed``, so 

10seeded init does not perturb the caller's global RNG. 

11""" 

12 

13from __future__ import annotations 

14 

15import math 

16from typing import Callable, Optional, cast 

17 

18import torch 

19import torch.nn as nn 

20 

21from transformer_lens.config import TransformerBridgeConfig 

22 

23from .model import ( 

24 NativeAttention, 

25 NativeBlock, 

26 NativeGatedMLP, 

27 NativeLayerNormPre, 

28 NativeMLP, 

29 NativeModel, 

30 NativeRMSNorm, 

31 NativeRMSNormPre, 

32) 

33 

34# Residual-scaled output is gpt2-specific; other modes treat every weight the 

35# same. Each entry takes ``(tensor, generator, gain)`` — gain honors 

36# ``cfg.initializer_range`` like the legacy init did (which passed it as the 

37# xavier/kaiming gain); kaiming has no gain kwarg, so scale after. 

38_NonResidualInit = Callable[[torch.Tensor, Optional[torch.Generator], float], torch.Tensor] 

39_NON_RESIDUAL_MODES: dict[str, _NonResidualInit] = { 

40 "xavier_uniform": lambda t, g, gain: nn.init.xavier_uniform_(t, gain=gain, generator=g), 

41 "xavier_normal": lambda t, g, gain: nn.init.xavier_normal_(t, gain=gain, generator=g), 

42 "kaiming_uniform": lambda t, g, gain: nn.init.kaiming_uniform_( 

43 t, nonlinearity="relu", generator=g 

44 ).mul_(gain), 

45 "kaiming_normal": lambda t, g, gain: nn.init.kaiming_normal_( 

46 t, nonlinearity="relu", generator=g 

47 ).mul_(gain), 

48} 

49 

50_SUPPORTED_MODES = frozenset({"gpt2", *_NON_RESIDUAL_MODES}) 

51 

52 

53def _unwrap_component(module: nn.Module) -> nn.Module: 

54 """Return the native module stored behind a bridge wrapper, if present.""" 

55 original = getattr(module, "original_component", None) 

56 return original if isinstance(original, nn.Module) else module 

57 

58 

59def initialize_native_model( 

60 model: NativeModel, cfg: TransformerBridgeConfig, seed: int | None = None 

61) -> None: 

62 """Initialize ``model`` weights in-place. Honors ``cfg.init_mode`` and ``cfg.seed``.""" 

63 effective_seed = seed if seed is not None else cfg.seed 

64 

65 # Always generate on CPU/fp32 and copy into the parameter: boot initializes 

66 # before .to(device)/.to(dtype) while init_weights() runs after, and a 

67 # generator seeded on the live parameter device produces a different stream 

68 # — the same seed must reproduce the same weights either way. 

69 generator: Optional[torch.Generator] 

70 if effective_seed is not None: 

71 g = torch.Generator() 

72 g.manual_seed(effective_seed) 

73 generator = g 

74 else: 

75 generator = None 

76 

77 def _staged( 

78 fn: Callable[[torch.Tensor], torch.Tensor], 

79 ) -> Callable[[torch.Tensor], torch.Tensor]: 

80 def apply(t: torch.Tensor) -> torch.Tensor: 

81 staging = torch.empty(t.shape, dtype=torch.float32) 

82 fn(staging) 

83 with torch.no_grad(): 

84 t.copy_(staging) 

85 return t 

86 

87 return apply 

88 

89 init_mode = (cfg.init_mode or "gpt2").lower() 

90 if init_mode not in _SUPPORTED_MODES: 

91 raise NotImplementedError( 

92 f"init_mode={init_mode!r} is not supported for NativeModel. " 

93 f"Supported modes: {sorted(_SUPPORTED_MODES)}." 

94 ) 

95 

96 weight_init: Callable[[torch.Tensor], torch.Tensor] 

97 output_init: Callable[[torch.Tensor], torch.Tensor] 

98 if init_mode == "gpt2": 

99 # Default matches the legacy TL scheme: N(0, 0.64/d_model), i.e. 

100 # std = 0.8/sqrt(d_model), not GPT-2's paper 0.02 — toy-model training 

101 # dynamics (e.g. the grokking demo) depend on this scale. 

102 std = cfg.initializer_range if cfg.initializer_range > 0 else 0.8 / math.sqrt(cfg.d_model) 

103 

104 # NOTE: this residual output scaling (1/sqrt(2*n_layers), applied only 

105 # to output projections below) was NOT present in the legacy 

106 # HookedTransformer._init_weights_gpt2 (removed in 4.0). 

107 # Intentional delta for NativeModel: kept because it follows the 

108 # original GPT-2 paper's residual-scaling convention and improves 

109 # training stability at init for deeper models. Flagged in issue #1568 

110 # as a maintainer call; kept + documented rather than removed. 

111 residual_scale = 1.0 / math.sqrt(2 * cfg.n_layers) 

112 weight_init = lambda t: nn.init.normal_( 

113 t, mean=0.0, std=std, generator=generator 

114 ) # noqa: E731 

115 output_init = lambda t: nn.init.normal_( # noqa: E731 

116 t, mean=0.0, std=std * residual_scale, generator=generator 

117 ) 

118 else: 

119 fn = _NON_RESIDUAL_MODES[init_mode] 

120 # Honor an explicitly-set initializer_range as the gain (legacy 

121 # behavior); the sentinel/default keeps plain xavier/kaiming scaling. 

122 gain = cfg.initializer_range if cfg.initializer_range > 0 else 1.0 

123 weight_init = lambda t: fn(t, generator, gain) # noqa: E731 

124 output_init = weight_init 

125 

126 weight_init = _staged(weight_init) 

127 output_init = _staged(output_init) 

128 

129 tok_embed = cast(nn.Embedding, _unwrap_component(model.tok_embed)) 

130 weight_init(tok_embed.weight) 

131 if model.pos is not None: 

132 pos = cast(nn.Embedding, _unwrap_component(model.pos)) 

133 weight_init(pos.weight) 

134 # Rotary has only registered buffers (cos/sin), no parameters to init. 

135 

136 for block in model.layers: 

137 native_block = cast(NativeBlock, _unwrap_component(block)) 

138 _init_block(native_block, weight_init=weight_init, output_init=output_init) 

139 

140 _init_norm(model.ln_out) 

141 head = cast(nn.Linear, _unwrap_component(model.head)) 

142 weight_init(head.weight) 

143 if head.bias is not None: 

144 nn.init.zeros_(head.bias) 

145 

146 

147def _init_norm(norm: nn.Module) -> None: 

148 norm = _unwrap_component(norm) 

149 if isinstance(norm, NativeRMSNorm): 

150 nn.init.ones_(norm.weight) 

151 elif isinstance(norm, (NativeRMSNormPre, NativeLayerNormPre)): 

152 pass 

153 elif isinstance(norm, nn.LayerNorm): 

154 nn.init.ones_(norm.weight) 

155 nn.init.zeros_(norm.bias) 

156 elif isinstance(norm, nn.Identity): 156 ↛ 159line 156 didn't jump to line 159 because the condition on line 156 was always true

157 pass 

158 else: 

159 raise TypeError(f"Unknown normalization type: {type(norm).__name__}") 

160 

161 

162def _init_block( 

163 block: NativeBlock, 

164 *, 

165 weight_init: Callable[[torch.Tensor], torch.Tensor], 

166 output_init: Callable[[torch.Tensor], torch.Tensor], 

167) -> None: 

168 _init_norm(block.ln1) 

169 attn = cast(NativeAttention, _unwrap_component(block.attn)) 

170 _init_attention(attn, weight_init=weight_init, output_init=output_init) 

171 if not block.cfg.attn_only: 

172 _init_norm(block.ln2) 

173 mlp = _unwrap_component(block.mlp) 

174 if isinstance(mlp, NativeGatedMLP): 

175 _init_gated_mlp(mlp, weight_init=weight_init, output_init=output_init) 

176 else: 

177 _init_mlp( 

178 cast(NativeMLP, mlp), 

179 weight_init=weight_init, 

180 output_init=output_init, 

181 ) 

182 

183 

184def _init_attention( 

185 attn: NativeAttention, 

186 *, 

187 weight_init: Callable[[torch.Tensor], torch.Tensor], 

188 output_init: Callable[[torch.Tensor], torch.Tensor], 

189) -> None: 

190 for component in (attn.q, attn.k, attn.v): 

191 linear = cast(nn.Linear, _unwrap_component(component)) 

192 weight_init(linear.weight) 

193 if linear.bias is not None: 193 ↛ 190line 193 didn't jump to line 190 because the condition on line 193 was always true

194 nn.init.zeros_(linear.bias) 

195 output = cast(nn.Linear, _unwrap_component(attn.o)) 

196 output_init(output.weight) 

197 if output.bias is not None: 197 ↛ exitline 197 didn't return from function '_init_attention' because the condition on line 197 was always true

198 nn.init.zeros_(output.bias) 

199 

200 

201def _init_mlp( 

202 mlp: NativeMLP, 

203 *, 

204 weight_init: Callable[[torch.Tensor], torch.Tensor], 

205 output_init: Callable[[torch.Tensor], torch.Tensor], 

206) -> None: 

207 fc_in = cast(nn.Linear, _unwrap_component(mlp.fc_in)) 

208 fc_out = cast(nn.Linear, _unwrap_component(mlp.fc_out)) 

209 weight_init(fc_in.weight) 

210 nn.init.zeros_(fc_in.bias) 

211 output_init(fc_out.weight) 

212 nn.init.zeros_(fc_out.bias) 

213 

214 

215def _init_gated_mlp( 

216 mlp: NativeGatedMLP, 

217 *, 

218 weight_init: Callable[[torch.Tensor], torch.Tensor], 

219 output_init: Callable[[torch.Tensor], torch.Tensor], 

220) -> None: 

221 gate = cast(nn.Linear, _unwrap_component(mlp.gate)) 

222 weight_init(gate.weight) 

223 # ``in`` is registered via add_module; getattr resolves it from _modules. 

224 in_proj = cast(nn.Linear, _unwrap_component(getattr(mlp, "in"))) 

225 out_proj = cast(nn.Linear, _unwrap_component(mlp.out)) 

226 weight_init(in_proj.weight) 

227 output_init(out_proj.weight)