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

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. 

7 

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

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

10""" 

11 

12from __future__ import annotations 

13 

14import math 

15from typing import Callable, Optional, cast 

16 

17import torch 

18import torch.nn as nn 

19 

20from transformer_lens.config import TransformerBridgeConfig 

21 

22from .model import ( 

23 NativeAttention, 

24 NativeBlock, 

25 NativeGatedMLP, 

26 NativeMLP, 

27 NativeModel, 

28 NativeRMSNorm, 

29) 

30 

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

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

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

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

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

36_NON_RESIDUAL_MODES: dict[str, _NonResidualInit] = { 

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

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

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

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

41 ).mul_(gain), 

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

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

44 ).mul_(gain), 

45} 

46 

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

48 

49 

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

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

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

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

54 

55 

56def initialize_native_model( 

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

58) -> None: 

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

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

61 

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

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

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

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

66 generator: Optional[torch.Generator] 

67 if effective_seed is not None: 

68 g = torch.Generator() 

69 g.manual_seed(effective_seed) 

70 generator = g 

71 else: 

72 generator = None 

73 

74 def _staged( 

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

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

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

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

79 fn(staging) 

80 with torch.no_grad(): 

81 t.copy_(staging) 

82 return t 

83 

84 return apply 

85 

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

87 if init_mode not in _SUPPORTED_MODES: 

88 raise NotImplementedError( 

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

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

91 ) 

92 

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

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

95 if init_mode == "gpt2": 

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

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

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

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

100 

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

102 # to output projections below) is NOT present in HookedTransformer's 

103 # _init_weights_gpt2 (see transformer_lens/HookedTransformer.py). 

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

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

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

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

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

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

110 t, mean=0.0, std=std, generator=generator 

111 ) # noqa: E731 

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

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

114 ) 

115 else: 

116 fn = _NON_RESIDUAL_MODES[init_mode] 

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

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

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

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

121 output_init = weight_init 

122 

123 weight_init = _staged(weight_init) 

124 output_init = _staged(output_init) 

125 

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

127 weight_init(tok_embed.weight) 

128 if model.pos is not None: 

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

130 weight_init(pos.weight) 

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

132 

133 for block in model.layers: 

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

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

136 

137 _init_norm(model.ln_out) 

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

139 weight_init(head.weight) 

140 if head.bias is not None: 

141 nn.init.zeros_(head.bias) 

142 

143 

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

145 norm = _unwrap_component(norm) 

146 if isinstance(norm, NativeRMSNorm): 

147 nn.init.ones_(norm.weight) 

148 elif isinstance(norm, nn.LayerNorm): 

149 nn.init.ones_(norm.weight) 

150 nn.init.zeros_(norm.bias) 

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

152 pass 

153 else: 

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

155 

156 

157def _init_block( 

158 block: NativeBlock, 

159 *, 

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

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

162) -> None: 

163 _init_norm(block.ln1) 

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

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

166 if not block.cfg.attn_only: 

167 _init_norm(block.ln2) 

168 mlp = _unwrap_component(block.mlp) 

169 if isinstance(mlp, NativeGatedMLP): 

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

171 else: 

172 _init_mlp( 

173 cast(NativeMLP, mlp), 

174 weight_init=weight_init, 

175 output_init=output_init, 

176 ) 

177 

178 

179def _init_attention( 

180 attn: NativeAttention, 

181 *, 

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

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

184) -> None: 

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

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

187 weight_init(linear.weight) 

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

189 nn.init.zeros_(linear.bias) 

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

191 output_init(output.weight) 

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

193 nn.init.zeros_(output.bias) 

194 

195 

196def _init_mlp( 

197 mlp: NativeMLP, 

198 *, 

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

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

201) -> None: 

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

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

204 weight_init(fc_in.weight) 

205 nn.init.zeros_(fc_in.bias) 

206 output_init(fc_out.weight) 

207 nn.init.zeros_(fc_out.bias) 

208 

209 

210def _init_gated_mlp( 

211 mlp: NativeGatedMLP, 

212 *, 

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

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

215) -> None: 

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

217 weight_init(gate.weight) 

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

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

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

221 weight_init(in_proj.weight) 

222 output_init(out_proj.weight)