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

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 """String representation of the LinearBridge.""" 

51 if self.original_component is not None: 

52 try: 

53 in_features = self.original_component.in_features 

54 out_features = self.original_component.out_features 

55 bias = self.original_component.bias is not None 

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

57 except AttributeError: 

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

59 else: 

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

61 

62 def set_processed_weights( 

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

64 ) -> None: 

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

66 

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

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

69 

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

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

72 

73 Args: 

74 weights: Dictionary containing: 

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

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

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

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

79 - 1D [out] format 

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

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

82 """ 

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

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

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

86 

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

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

89 weight = weights.get("weight") 

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

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

92 bias = weights.get("bias") 

93 

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

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

96 if bias is not None: 

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

98 

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

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

101 n_heads, dim1, dim2 = weight.shape 

102 if dim1 > dim2: 

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

104 weight = einops.rearrange( 

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

106 ).contiguous() 

107 else: 

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

109 weight = einops.rearrange( 

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

111 ).contiguous() 

112 

113 # Handle 2D bias by flattening to 1D 

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

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

116 

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

118 "weight": weight, 

119 } 

120 

121 if bias is not None: 

122 processed_weights["bias"] = bias 

123 

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