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

1"""Unembedding bridge component. 

2 

3This module contains the bridge component for unembedding layers. 

4""" 

5from typing import Any, Dict, Optional 

6 

7import torch 

8 

9from transformer_lens.model_bridge.generalized_components.base import ( 

10 GeneralizedComponent, 

11) 

12 

13 

14class UnembeddingBridge(GeneralizedComponent): 

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

16 

17 This component provides standardized input/output hooks. 

18 """ 

19 

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

21 

22 def __init__( 

23 self, 

24 name: str, 

25 config: Optional[Any] = None, 

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

27 ): 

28 """Initialize the unembedding bridge. 

29 

30 Args: 

31 name: The name of this component 

32 config: Optional configuration (unused for UnembeddingBridge) 

33 submodules: Dictionary of GeneralizedComponent submodules to register 

34 """ 

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

36 

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

38 """Set the original component and ensure it has bias enabled. 

39 

40 Args: 

41 original_component: The original transformer component to wrap 

42 """ 

43 # If this is a Linear layer without bias, enable it 

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

45 # Get the output features (vocab size) 

46 vocab_size = original_component.weight.shape[0] 

47 device = original_component.weight.device 

48 dtype = original_component.weight.dtype 

49 

50 original_component.bias = torch.nn.Parameter( 

51 torch.zeros(vocab_size, device=device, dtype=dtype) 

52 ) 

53 

54 super().set_original_component(original_component) 

55 

56 @property 

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

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

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

60 processed_W_U = self._parameters["_processed_W_U"] 

61 if processed_W_U is not None: 

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

63 # Transpose to TL format [d_model, d_vocab] 

64 return processed_W_U.T 

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

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

67 assert hasattr( 

68 self.original_component, "weight" 

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

70 weight = self.original_component.weight 

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

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

73 return weight.T 

74 

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

76 """Forward pass through the unembedding bridge. 

77 

78 Args: 

79 hidden_states: Input hidden states 

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

81 

82 Returns: 

83 Unembedded output (logits) 

84 """ 

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

86 raise RuntimeError( 

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

88 ) 

89 hidden_states = self.hook_in(hidden_states) 

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

91 

92 output = self.hook_out(output) 

93 return output 

94 

95 @property 

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

97 """Access the unembedding bias vector.""" 

98 if "_b_U" in self._parameters: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true

99 param = self._parameters["_b_U"] 

100 if param is not None: 

101 return param 

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

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

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

105 bias = self.original_component.bias 

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

107 return bias 

108 else: 

109 assert hasattr( 

110 self.original_component, "weight" 

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

112 weight = self.original_component.weight 

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

114 device = weight.device 

115 dtype = weight.dtype 

116 vocab_size: int = int(weight.shape[0]) 

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