Coverage for transformer_lens/model_bridge/generalized_components/linear.py: 65%

1"""Linear bridge component for wrapping linear layers with hook points.""" 

2from typing import Any, Dict, Mapping 

3 

4import einops 

5import torch 

6 

7from transformer_lens.model_bridge.generalized_components.base import ( 

8 GeneralizedComponent, 

9) 

10 

11 

12class LinearBridge(GeneralizedComponent): 

13 """Bridge component for linear layers. 

14 

15 This component wraps a linear layer (nn.Linear) and provides hook points 

16 for intercepting the input and output activations. 

17 

18 Note: For Conv1D layers (used in GPT-2 style models), use Conv1DBridge instead. 

19 """ 

20 

21 def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: 

22 """Forward pass through the linear layer with hooks. 

23 

24 Args: 

25 input: Input tensor 

26 *args: Additional positional arguments 

27 **kwargs: Additional keyword arguments 

28 

29 Returns: 

30 Output tensor after linear transformation 

31 """ 

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

33 raise RuntimeError( 

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

35 ) 

36 # Projection-hook protocol (see GeneralizedComponent): consume the 

37 # container's hook_in suppression, and record that hook_out fired. 

38 # object.__setattr__ skips the nn.Module setattr machinery on this 

39 # per-call hot path. 

40 if self._suppress_next_hook_in: 

41 object.__setattr__(self, "_suppress_next_hook_in", False) 

42 else: 

43 input = self.hook_in(input) 

44 output = self.original_component(input, *args, **kwargs) 

45 output = self.hook_out(output) 

46 object.__setattr__(self, "_fired_hook_out", True) 

47 return output 

48 

49 def __repr__(self) -> str: 

50 if self.original_component is not None: 50 ↛ 59line 50 didn't jump to line 59 because the condition on line 50 was always true

51 try: 

52 in_features = self.original_component.in_features 

53 out_features = self.original_component.out_features 

54 bias = self.original_component.bias is not None 

55 return f"LinearBridge({in_features} -> {out_features}, bias={bias}, original_component={type(self.original_component).__name__})" 

56 except AttributeError: 

57 return f"LinearBridge(name={self.name}, original_component={type(self.original_component).__name__})" 

58 else: 

59 return f"LinearBridge(name={self.name}, original_component=None)" 

60 

61 def set_processed_weights( 

62 self, weights: Mapping[str, torch.Tensor | None], verbose: bool = False 

63 ) -> None: 

64 """Set the processed weights by loading them into the original component. 

65 

66 This loads the processed weights directly into the original_component's parameters, 

67 so when forward() delegates to original_component, it uses the processed weights. 

68 

69 Handles Linear layers (shape [out, in]). 

70 Also handles 3D weights [n_heads, d_model, d_head] by flattening them first. 

71 

72 Args: 

73 weights: Dictionary containing: 

74 - weight: The processed weight tensor. Can be: 

75 - 2D [in, out] format (will be transposed to [out, in] for Linear) 

76 - 3D [n_heads, d_model, d_head] format (will be flattened to 2D) 

77 - bias: The processed bias tensor (optional). Can be: 

78 - 1D [out] format 

79 - 2D [n_heads, d_head] format (will be flattened to 1D) 

80 verbose: If True, print detailed information about weight setting 

81 """ 

82 if verbose: 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true

83 print(f"\n set_processed_weights: LinearBridge (name={self.name})") 

84 print(f" Received {len(weights)} weight keys") 

85 

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

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

88 weight = weights.get("weight") 

89 if weight is None: 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true

90 raise ValueError("Processed weights for LinearBridge must include 'weight'.") 

91 bias = weights.get("bias") 

92 

93 if verbose: 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true

94 print(f" Found weight key with shape: {weight.shape}") 

95 if bias is not None: 

96 print(f" Found bias key with shape: {bias.shape}") 

97 

98 # Flatten 3D→2D; contiguous() needed for correct bfloat16 matmul order 

99 if weight.ndim == 3: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true

100 n_heads, dim1, dim2 = weight.shape 

101 if dim1 > dim2: 

102 # [n_heads, d_model, d_head] -> [n_heads * d_head, d_model] (nn.Linear format) 

103 weight = einops.rearrange( 

104 weight, "n_heads d_model d_head -> (n_heads d_head) d_model" 

105 ).contiguous() 

106 else: 

107 # [n_heads, d_head, d_model] -> [d_model, n_heads * d_head] 

108 weight = einops.rearrange( 

109 weight, "n_heads d_head d_model -> d_model (n_heads d_head)" 

110 ).contiguous() 

111 

112 # Handle 2D bias by flattening to 1D 

113 if bias is not None and bias.ndim == 2: 

114 bias = einops.rearrange(bias, "n_heads d_head -> (n_heads d_head)") 

115 

116 processed_weights: Dict[str, torch.Tensor] = { 

117 "weight": weight, 

118 } 

119 

120 if bias is not None: 

121 processed_weights["bias"] = bias 

122 

123 super().set_processed_weights(processed_weights, verbose=verbose)