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

1"""Unembedding bridge component. 

2 

3This module contains the bridge component for unembedding layers. 

4""" 

5from typing import Any, Dict, Optional 

6 

7import torch 

8from accelerate.utils import align_module_device 

9 

10from transformer_lens.model_bridge.generalized_components.base import ( 

11 GeneralizedComponent, 

12) 

13 

14 

15class UnembeddingBridge(GeneralizedComponent): 

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

17 

18 This component provides standardized input/output hooks. 

19 """ 

20 

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

22 

23 def __init__( 

24 self, 

25 name: str, 

26 config: Optional[Any] = None, 

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

28 ): 

29 """Initialize the unembedding bridge. 

30 

31 Args: 

32 name: The name of this component 

33 config: Optional configuration (unused for UnembeddingBridge) 

34 submodules: Dictionary of GeneralizedComponent submodules to register 

35 """ 

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

37 

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

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

40 

41 Args: 

42 original_component: The original transformer component to wrap 

43 """ 

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

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

46 # Get the output features (vocab size) 

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

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

49 

50 # .device is NOT safe the same way: under Accelerate offload this 

51 # weight is a meta placeholder outside of a materialized window, so 

52 # reading .device unguarded would build a meta-device (i.e. fake, 

53 # data-less) bias instead of real zeros. align_module_device gives 

54 # the real execution device for this one read. 

55 with align_module_device(original_component): 

56 device = original_component.weight.device 

57 

58 original_component.bias = torch.nn.Parameter( 

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

60 ) 

61 

62 super().set_original_component(original_component) 

63 

64 @property 

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

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

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

68 processed_W_U = self._parameters["_processed_W_U"] 

69 if processed_W_U is not None: 

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

71 # Transpose to TL format [d_model, d_vocab] 

72 return processed_W_U.T 

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

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

75 assert hasattr( 

76 self.original_component, "weight" 

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

78 weight = self.original_component.weight 

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

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

81 return weight.T 

82 

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

84 """Forward pass through the unembedding bridge. 

85 

86 Args: 

87 hidden_states: Input hidden states 

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

89 

90 Returns: 

91 Unembedded output (logits) 

92 """ 

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

94 raise RuntimeError( 

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

96 ) 

97 hidden_states = self.hook_in(hidden_states) 

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

99 

100 output = self.hook_out(output) 

101 return output 

102 

103 @property 

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

105 """Access the unembedding bias vector.""" 

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

107 param = self._parameters["_b_U"] 

108 if param is not None: 

109 return param 

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

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

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

113 bias = self.original_component.bias 

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

115 return bias 

116 else: 

117 assert hasattr( 

118 self.original_component, "weight" 

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

120 weight = self.original_component.weight 

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

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

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

124 # .device is not safe the same way under offload - see 

125 # set_original_component's identical guard above. 

126 with align_module_device(self.original_component): 

127 device = weight.device 

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