Coverage for transformer_lens/model_bridge/generalized_components/unembedding.py: 84%

1"""Unembedding bridge component. 

2 

3This module contains the bridge component for unembedding layers. 

4""" 

5from collections import ChainMap 

6from typing import Any, Dict, Optional, cast 

7 

8import torch 

9from accelerate.hooks import AlignDevicesHook, SequentialHook 

10from accelerate.utils import align_module_device 

11 

12from transformer_lens.model_bridge.generalized_components.base import ( 

13 GeneralizedComponent, 

14) 

15 

16 

17class UnembeddingBridge(GeneralizedComponent): 

18 """Unembedding bridge that wraps transformer unembedding layers. 

19 

20 This component provides standardized input/output hooks. 

21 """ 

22 

23 property_aliases = {"W_U": "u.weight"} 

24 

25 def __init__( 

26 self, 

27 name: str, 

28 config: Optional[Any] = None, 

29 submodules: Optional[Dict[str, GeneralizedComponent]] = {}, 

30 ): 

31 """Initialize the unembedding bridge. 

32 

33 Args: 

34 name: The name of this component 

35 config: Optional configuration (unused for UnembeddingBridge) 

36 submodules: Dictionary of GeneralizedComponent submodules to register 

37 """ 

38 super().__init__(name, config, submodules=submodules) 

39 

40 def set_original_component(self, original_component: torch.nn.Module) -> None: 

41 """Set the original component with a fixed zero bias when the checkpoint has none. 

42 

43 Args: 

44 original_component: The original transformer component to wrap 

45 """ 

46 # Keep b_U and state_dict compatibility without adding an optimizer parameter. 

47 if isinstance(original_component, torch.nn.Linear) and original_component.bias is None: 

48 # Get the output features (vocab size) 

49 vocab_size = original_component.weight.shape[0] # shape is safe on a meta tensor too 

50 dtype = original_component.weight.dtype # dtype is also safe on a meta tensor 

51 

52 device = self._bias_device(original_component) 

53 

54 del original_component.bias 

55 original_component.register_buffer( 

56 "bias", torch.zeros(vocab_size, device=device, dtype=dtype) 

57 ) 

58 self._update_offloaded_bias(original_component) 

59 

60 super().set_original_component(original_component) 

61 

62 @staticmethod 

63 def _offload_hook(component: torch.nn.Module) -> Optional[AlignDevicesHook]: 

64 pending = [getattr(component, "_hf_hook", None)] 

65 while pending: 

66 hook = pending.pop() 

67 if isinstance(hook, SequentialHook): 

68 pending.extend(hook.hooks) 

69 elif isinstance(hook, AlignDevicesHook) and hook.offload: 

70 return hook 

71 return None 

72 

73 @staticmethod 

74 def _bias_device(component: torch.nn.Module) -> torch.device: 

75 with align_module_device(component): 

76 weight = component.weight 

77 assert isinstance(weight, torch.Tensor) 

78 device = weight.device 

79 if device.type == "meta": 

80 # align_module_device does not recognize SequentialHook offloading. 

81 hook = UnembeddingBridge._offload_hook(component) 

82 if hook is not None and hook.execution_device is not None: 

83 device = torch.device(hook.execution_device) 

84 return device 

85 

86 @staticmethod 

87 def _update_offloaded_bias(component: torch.nn.Module) -> None: 

88 hook = UnembeddingBridge._offload_hook(component) 

89 if hook is not None: 

90 bias = component.bias 

91 assert isinstance(bias, torch.Tensor) 

92 assert hook.weights_map is not None 

93 # Only the first ChainMap layer is written; the offload map stays lazy. 

94 hook.weights_map = ChainMap({"bias": bias.detach().cpu()}, cast(Any, hook.weights_map)) 

95 # Hook removal must restore the latest bias alongside its weight. 

96 hook.original_devices.setdefault("bias", hook.original_devices["weight"]) 

97 

98 def set_processed_weights( 

99 self, weights: Dict[str, torch.Tensor], verbose: bool = False 

100 ) -> None: 

101 """Apply folded output biases while keeping synthetic biases non-trainable.""" 

102 component = self.original_component 

103 if component is not None and "bias" in component._buffers and "bias" in weights: 

104 bias = component._buffers["bias"] 

105 new_bias = weights["bias"] 

106 if bias is not None and bias.shape != new_bias.shape: 

107 raise ValueError( 

108 f"Shape mismatch when setting weight 'bias' in {type(component).__name__}: " 

109 f"existing bias shape {bias.shape} != new tensor shape {new_bias.shape}" 

110 ) 

111 component.bias = new_bias.detach() 

112 self._update_offloaded_bias(component) 

113 super().set_processed_weights(weights, verbose=verbose) 

114 

115 @property 

116 def W_U(self) -> torch.Tensor: 

117 """Return the unembedding weight matrix in TL format [d_model, d_vocab].""" 

118 if "_processed_W_U" in self._parameters: 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true

119 processed_W_U = self._parameters["_processed_W_U"] 

120 if processed_W_U is not None: 

121 # Processed weights are in HF format [vocab, d_model] 

122 # Transpose to TL format [d_model, d_vocab] 

123 return processed_W_U.T 

124 if self.original_component is None: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true

125 raise RuntimeError(f"Original component not set for {self.name}") 

126 assert hasattr( 

127 self.original_component, "weight" 

128 ), f"Component {self.name} has no weight attribute" 

129 weight = self.original_component.weight 

130 assert isinstance(weight, torch.Tensor), f"Weight is not a tensor for {self.name}" 

131 # HF format is [d_vocab, d_model], transpose to TL format [d_model, d_vocab] 

132 return weight.T 

133 

134 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor: 

135 """Forward pass through the unembedding bridge. 

136 

137 Args: 

138 hidden_states: Input hidden states 

139 **kwargs: Additional arguments to pass to the original component 

140 

141 Returns: 

142 Unembedded output (logits) 

143 """ 

144 if self.original_component is None: 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true

145 raise RuntimeError( 

146 f"Original component not set for {self.name}. Call set_original_component() first." 

147 ) 

148 hidden_states = self.hook_in(hidden_states) 

149 output = self.original_component(hidden_states, **kwargs) 

150 

151 output = self.hook_out(output) 

152 return output 

153 

154 @property 

155 def b_U(self) -> torch.Tensor: 

156 """Access the unembedding bias vector.""" 

157 if self.original_component is None: 157 ↛ 158line 157 didn't jump to line 158 because the condition on line 157 was never true

158 raise RuntimeError(f"Original component not set for {self.name}") 

159 if hasattr(self.original_component, "bias") and self.original_component.bias is not None: 159 ↛ 164line 159 didn't jump to line 164 because the condition on line 159 was always true

160 bias = self.original_component.bias 

161 assert isinstance(bias, torch.Tensor), f"Bias is not a tensor for {self.name}" 

162 return bias 

163 else: 

164 assert hasattr( 

165 self.original_component, "weight" 

166 ), f"Component {self.name} has no weight attribute" 

167 weight = self.original_component.weight 

168 assert isinstance(weight, torch.Tensor), f"Weight is not a tensor for {self.name}" 

169 dtype = weight.dtype # safe on a meta tensor 

170 vocab_size: int = int(weight.shape[0]) # safe on a meta tensor 

171 device = self._bias_device(self.original_component) 

172 return torch.zeros(vocab_size, device=device, dtype=dtype)