Coverage for transformer_lens/model_bridge/generalized_components/rotary_embedding.py: 67%

1"""Rotary embedding bridge component. 

2 

3This module contains the bridge component for rotary position embedding layers. 

4""" 

5from typing import Any, Dict, Optional, Tuple, Union 

6 

7import torch 

8 

9from transformer_lens.hook_points import HookPoint 

10from transformer_lens.model_bridge.generalized_components.base import ( 

11 GeneralizedComponent, 

12) 

13from transformer_lens.utilities.heterogeneous_config import safe_config_get 

14 

15 

16class RotaryEmbeddingBridge(GeneralizedComponent): 

17 """Rotary embedding bridge that wraps rotary position embedding layers. 

18 

19 Unlike regular embeddings, rotary embeddings return a tuple of (cos, sin) tensors. 

20 This component properly handles the tuple return value without unwrapping it. 

21 """ 

22 

23 def __init__( 

24 self, 

25 name: str, 

26 config: Optional[Any] = None, 

27 submodules: Optional[Dict[str, GeneralizedComponent]] = None, 

28 ): 

29 """Initialize the rotary embedding bridge. 

30 

31 Args: 

32 name: The name of this component 

33 config: Optional configuration (unused for RotaryEmbeddingBridge) 

34 submodules: Dictionary of GeneralizedComponent submodules to register 

35 """ 

36 super().__init__(name, config, submodules=submodules or {}) 

37 self.hook_cos = HookPoint() 

38 self.hook_sin = HookPoint() 

39 

40 def get_random_inputs( 

41 self, 

42 batch_size: int = 2, 

43 seq_len: int = 8, 

44 device: Optional[torch.device] = None, 

45 dtype: Optional[torch.dtype] = None, 

46 ) -> Dict[str, Any]: 

47 """Generate random inputs for rotary embedding testing. 

48 

49 Rotary embeddings for Gemma-3 expect (x, position_ids) where: 

50 - x: tensor with shape [batch, seq, num_heads, head_dim] 

51 - position_ids: position indices with shape [batch, seq] 

52 

53 Args: 

54 batch_size: Batch size for generated inputs 

55 seq_len: Sequence length for generated inputs 

56 device: Device to place tensors on 

57 dtype: Dtype for generated tensors 

58 

59 Returns: 

60 Dictionary with positional args as tuple under 'args' key 

61 """ 

62 if device is None: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 device = torch.device("cpu") 

64 if dtype is None: 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true

65 dtype = torch.float32 

66 num_heads = safe_config_get(self.config, "num_attention_heads", 4) if self.config else 4 

67 head_dim = safe_config_get(self.config, "head_dim", 256) if self.config else 256 

68 x = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) 

69 position_ids = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1) 

70 args: tuple = (x, position_ids) 

71 # Gemma3's rotary embedding requires a layer_type argument (e.g., "sliding_attention") 

72 # to select the correct inv_freq buffer. Without it, forward() tries to access 

73 # "None_inv_freq" which doesn't exist. 

74 if self.original_component is not None and hasattr(self.original_component, "layer_types"): 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true

75 layer_type = self.original_component.layer_types[0] # type: ignore[index] 

76 args = (x, position_ids, layer_type) 

77 return {"args": args} 

78 

79 def forward( 

80 self, *args: Any, **kwargs: Any 

81 ) -> Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: 

82 """Forward pass through the rotary embedding bridge. 

83 

84 Rotary embeddings typically take seq_len or position_ids and return (cos, sin) tensors. 

85 This method ensures that cos and sin are passed through their respective hooks 

86 (hook_cos and hook_sin) to match HookedTransformer's behavior. 

87 

88 Args: 

89 *args: Positional arguments to pass to the original component 

90 **kwargs: Keyword arguments to pass to the original component 

91 

92 Returns: 

93 Tuple of (cos, sin) tensors for rotary position embeddings, after being 

94 passed through hook_cos and hook_sin respectively. For DeepSeek-V2-style 

95 embeddings that return a single complex ``freqs_cis`` tensor, that tensor is 

96 passed through unchanged for downstream complex multiplication. 

97 """ 

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

99 raise RuntimeError( 

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

101 ) 

102 

103 # Apply input hook if first arg is a tensor 

104 if args and isinstance(args[0], torch.Tensor): 104 ↛ 109line 104 didn't jump to line 109 because the condition on line 104 was always true

105 hooked_input = self.hook_in(args[0]) 

106 args = (hooked_input,) + args[1:] 

107 

108 # Call original component to get (cos, sin) tuple 

109 output = self.original_component(*args, **kwargs) 

110 

111 # Ensure output is a tuple — or a complex tensor (DeepSeek-V2 freqs_cis style) 

112 if not isinstance(output, tuple): 

113 if isinstance(output, torch.Tensor) and output.is_complex(): 113 ↛ 118line 113 didn't jump to line 118 because the condition on line 113 was always true

114 # V2-style: freqs_cis complex tensor — pass through without cos/sin split. 

115 # hook_cos/hook_sin do not apply here; the complex form is consumed by 

116 # MLAAttentionBridge which detects it and uses complex multiplication. 

117 return output 

118 if hasattr(output, "__iter__") and (not isinstance(output, torch.Tensor)): 

119 output = tuple(output) 

120 else: 

121 raise RuntimeError( 

122 f"Rotary embedding {self.name} returned {type(output)} instead of tuple. Expected (cos, sin) tuple." 

123 ) 

124 

125 # Extract cos and sin, apply their respective hooks, and return 

126 if len(output) == 2: 126 ↛ 136line 126 didn't jump to line 136 because the condition on line 126 was always true

127 cos, sin = output 

128 # Apply hooks to match HookedTransformer's rotary_cos/rotary_sin pattern 

129 cos = self.hook_cos(cos) 

130 sin = self.hook_sin(sin) 

131 # Return the hooked cos and sin as a tuple 

132 # Note: Don't pass tuple through hook_out as it expects a tensor 

133 return (cos, sin) 

134 else: 

135 # For unexpected tuple lengths, just pass through 

136 return output 

137 

138 def get_dummy_inputs( 

139 self, test_input: torch.Tensor, **kwargs: Any 

140 ) -> tuple[tuple[Any, ...], dict[str, Any]]: 

141 """Generate dummy inputs for rotary embedding forward method. 

142 

143 Rotary embeddings typically expect (x, position_ids) where: 

144 - x: input tensor [batch, seq, d_model] 

145 - position_ids: position indices [batch, seq] 

146 

147 Args: 

148 test_input: Base test input tensor [batch, seq, d_model] 

149 **kwargs: Additional context including position_ids 

150 

151 Returns: 

152 Tuple of (args, kwargs) for the rotary embedding forward method 

153 """ 

154 batch, seq_len, _ = test_input.shape 

155 

156 # Get position_ids from kwargs, or generate default 

157 position_ids = kwargs.get("position_ids") 

158 if position_ids is None: 

159 position_ids = ( 

160 torch.arange(seq_len, device=test_input.device).unsqueeze(0).expand(batch, -1) 

161 ) 

162 

163 # Rotary embeddings expect (x, position_ids) 

164 return (test_input, position_ids), {}