Coverage for transformer_lens/utilities/tl_checkpoint_conversion.py: 100%

1"""One-time converter for legacy TL-property-format checkpoints (#1588). 

2 

3Rotary-model checkpoints are unsupported: their ``rotary_sin``/``rotary_cos`` 

4buffer keys fail loudly as unrecognized keys (drop them first if you need to 

5convert one — the bridge recomputes rotary embeddings from the config). 

6 

7Historical training runs (OthelloGPT, grokking demos, ARENA content) were 

8saved via ``HookedTransformer.state_dict()`` before ``TransformerBridge`` 

9existed, using property-style keys ("blocks.0.attn.W_Q", "embed.W_E", ...) 

10and per-head tensor shapes. ``convert_tl_checkpoint`` maps those onto the 

11key/tensor format ``TransformerBridge.boot_native(cfg).load_state_dict`` 

12accepts natively, so these checkpoints can be converted once and re-saved in 

13bridge format. This is deliberately a standalone converter rather than a 

14second key convention taught to ``load_state_dict`` itself: convert once, 

15``bridge.load_state_dict(converted)``, then re-save with ``bridge.state_dict()``. 

16""" 

17 

18from __future__ import annotations 

19 

20from typing import Callable, Optional 

21 

22import einops 

23import torch 

24 

25from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig 

26 

27# Buffers that live on HookedTransformer's attention blocks but have no 

28# Parameter counterpart on the bridge side (causal mask, IGNORE sentinel). 

29_DROPPED_BUFFER_SUFFIXES = (".mask", ".IGNORE") 

30 

31TensorConvert = Callable[[torch.Tensor, TransformerBridgeConfig, str], torch.Tensor] 

32 

33 

34def _validate_shape(tensor: torch.Tensor, expected: tuple[int, ...], key: str) -> None: 

35 # Merging/splitting per-head dims (unlike a plain transpose) produces a 

36 # validly-shaped result for *any* head count, since d_model == n_heads * 

37 # d_head for any factoring of it — a wrong cfg silently mis-groups heads 

38 # without ever tripping a downstream shape-mismatch error. Check the 

39 # untouched per-head shape explicitly before reshaping. 

40 if tuple(tensor.shape) != expected: 

41 raise ValueError( 

42 f"convert_tl_checkpoint: {key!r} has shape {tuple(tensor.shape)}, " 

43 f"expected {expected} for the given cfg. The checkpoint may not " 

44 "match this cfg (n_heads/n_key_value_heads/d_head/d_model)." 

45 ) 

46 

47 

48def _kv_heads(cfg: TransformerBridgeConfig) -> int: 

49 return cfg.n_key_value_heads or cfg.n_heads 

50 

51 

52def _convert_w_q(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

53 _validate_shape(t, (cfg.n_heads, cfg.d_model, cfg.d_head), key) 

54 return einops.rearrange(t, "n_heads d_model d_head -> (n_heads d_head) d_model") 

55 

56 

57def _convert_w_kv(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

58 _validate_shape(t, (_kv_heads(cfg), cfg.d_model, cfg.d_head), key) 

59 return einops.rearrange(t, "n_heads d_model d_head -> (n_heads d_head) d_model") 

60 

61 

62def _convert_w_o(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

63 _validate_shape(t, (cfg.n_heads, cfg.d_head, cfg.d_model), key) 

64 return einops.rearrange(t, "n_heads d_head d_model -> d_model (n_heads d_head)") 

65 

66 

67def _convert_b_q(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

68 _validate_shape(t, (cfg.n_heads, cfg.d_head), key) 

69 return einops.rearrange(t, "n_heads d_head -> (n_heads d_head)") 

70 

71 

72def _convert_b_kv(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

73 _validate_shape(t, (_kv_heads(cfg), cfg.d_head), key) 

74 return einops.rearrange(t, "n_heads d_head -> (n_heads d_head)") 

75 

76 

77def _identity(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

78 return t 

79 

80 

81def _transpose(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: 

82 return t.T.contiguous() 

83 

84 

85# Old TL-property suffix -> (new bridge-key suffix, tensor conversion). 

86# Checked in order, longest/most-specific first, so e.g. ".b_Q" is matched 

87# before the generic ".b" LayerNorm-bias fallback. 

88_SUFFIX_CONVERSIONS: list[tuple[str, str, TensorConvert]] = [ 

89 (".W_Q", ".q.weight", _convert_w_q), 

90 # GQA stores K/V under a leading-underscore name (the raw Parameter); 

91 # plain ".W_K"/".W_V" become expanding (non-Parameter) properties instead. 

92 ("._W_K", ".k.weight", _convert_w_kv), 

93 ("._W_V", ".v.weight", _convert_w_kv), 

94 (".W_K", ".k.weight", _convert_w_kv), 

95 (".W_V", ".v.weight", _convert_w_kv), 

96 (".W_O", ".o.weight", _convert_w_o), 

97 (".b_Q", ".q.bias", _convert_b_q), 

98 ("._b_K", ".k.bias", _convert_b_kv), 

99 ("._b_V", ".v.bias", _convert_b_kv), 

100 (".b_K", ".k.bias", _convert_b_kv), 

101 (".b_V", ".v.bias", _convert_b_kv), 

102 (".b_O", ".o.bias", _identity), 

103 (".W_in", ".in.weight", _transpose), 

104 (".b_in", ".in.bias", _identity), 

105 (".W_out", ".out.weight", _transpose), 

106 (".b_out", ".out.bias", _identity), 

107 (".W_gate", ".gate.weight", _transpose), 

108 (".b_gate", ".gate.bias", _identity), 

109 (".W_U", ".weight", _transpose), 

110 (".b_U", ".bias", _identity), 

111 (".W_E", ".weight", _identity), 

112 (".W_pos", ".weight", _identity), 

113 (".w", ".weight", _identity), 

114 (".b", ".bias", _identity), 

115] 

116 

117 

118def _convert_key_and_tensor( 

119 key: str, tensor: torch.Tensor, cfg: TransformerBridgeConfig 

120) -> Optional[tuple[str, torch.Tensor]]: 

121 for old_suffix, new_suffix, convert in _SUFFIX_CONVERSIONS: 

122 if key.endswith(old_suffix): 

123 new_key = key[: -len(old_suffix)] + new_suffix 

124 return new_key, convert(tensor, cfg, key) 

125 return None 

126 

127 

128def convert_tl_checkpoint( 

129 state_dict: dict[str, torch.Tensor], 

130 cfg: TransformerBridgeConfig, 

131) -> dict[str, torch.Tensor]: 

132 """Convert a legacy TL-property-format state dict to the key/tensor 

133 format ``TransformerBridge.boot_native(cfg).load_state_dict`` accepts. 

134 

135 Args: 

136 state_dict: A state dict in the old ``HookedTransformer`` convention 

137 (e.g. from ``HookedTransformer.state_dict()``), with keys like 

138 ``"blocks.0.attn.W_Q"`` and per-head tensor shapes. 

139 cfg: The config the checkpoint was trained/saved under. Used both to 

140 reshape per-head attention weights and to validate that the 

141 checkpoint's per-head shapes actually match this cfg — a 

142 mismatched cfg would otherwise silently mis-group heads without 

143 ever tripping a shape error, since d_model == n_heads * d_head 

144 holds for any wrong factoring too. 

145 

146 Returns: 

147 A state dict with modern bridge keys (e.g. ``"blocks.0.attn.q.weight"``) 

148 and flat ``nn.Linear``-oriented tensor shapes, ready for 

149 ``bridge.load_state_dict(converted, strict=True)``. 

150 """ 

151 converted: dict[str, torch.Tensor] = {} 

152 for key, tensor in state_dict.items(): 

153 if key.endswith(_DROPPED_BUFFER_SUFFIXES): 

154 continue 

155 result = _convert_key_and_tensor(key, tensor, cfg) 

156 if result is None: 

157 raise ValueError( 

158 f"convert_tl_checkpoint: don't know how to convert key {key!r} " 

159 "(not a recognized TL-property parameter or buffer suffix)." 

160 ) 

161 new_key, new_tensor = result 

162 converted[new_key] = new_tensor 

163 return converted