Coverage for transformer_lens/components/t5_attention.py: 94%

1import math 

2from typing import Dict, Optional, Union 

3 

4import torch 

5import torch.nn as nn 

6from jaxtyping import Float, Int 

7 

8from transformer_lens.components.abstract_attention import AbstractAttention 

9from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

10from transformer_lens.hook_points import HookPoint 

11 

12 

13class T5Attention(AbstractAttention): 

14 r""" 

15 T5 attention - with relative attention bias and cross-attention support 

16 This realisation expects you to precompute relative positional bias, and then feed it to forward 

17 like 

18 ```python 

19 attn = T5Attention(cfg, has_relative_attention_bias=True) 

20 positional_bias = attn.compute_relative_attention_bias(query_len, key_len, device=device) 

21 result = attn(query, key, value, position_bias=positional_bias) 

22 ``` 

23 """ 

24 

25 def __init__( 

26 self, 

27 cfg: Union[Dict, HookedTransformerConfig], 

28 has_relative_attention_bias: bool = False, 

29 attn_type: str = "global", 

30 layer_id: Optional[int] = None, 

31 is_decoder: bool = False, 

32 ): 

33 super().__init__(cfg, attn_type, layer_id) 

34 if isinstance(cfg, Dict): 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true

35 cfg = HookedTransformerConfig.from_dict(cfg) 

36 self.cfg = cfg 

37 self.has_relative_attention_bias: bool = has_relative_attention_bias 

38 # Decoder buckets are unidirectional: every key sits at a non-positive 

39 # offset, so bidirectional bucketing halves the usable resolution. 

40 self.is_decoder: bool = is_decoder 

41 if is_decoder: 

42 # T5's cfg carries attention_dir="bidirectional" for the encoder; 

43 # decoder self-attention must still mask the future. 

44 self._attention_dir_override = "causal" 

45 

46 if self.has_relative_attention_bias: 

47 if ( 47 ↛ 51line 47 didn't jump to line 51 because the condition on line 47 was never true

48 cfg.relative_attention_num_buckets is None 

49 or cfg.relative_attention_max_distance is None 

50 ): 

51 raise ValueError( 

52 "You need to specify relative_attention_num_buckets and relative_attention_max_distance in config to use relative attention bias" 

53 ) 

54 

55 self.relative_attention_num_buckets = cfg.relative_attention_num_buckets 

56 self.relative_attention_max_distance = cfg.relative_attention_max_distance 

57 self.rel_pos_bias = nn.Embedding(self.relative_attention_num_buckets, self.cfg.n_heads) 

58 self.rel_pos_hook = HookPoint() 

59 

60 self.W_K = nn.Parameter( 

61 torch.empty(self.cfg.n_heads, self.cfg.d_model, self.cfg.d_head, dtype=cfg.dtype) 

62 ) 

63 self.W_V = nn.Parameter( 

64 torch.empty(self.cfg.n_heads, self.cfg.d_model, self.cfg.d_head, dtype=cfg.dtype) 

65 ) 

66 self.b_K = nn.Parameter(torch.zeros(self.cfg.n_heads, self.cfg.d_head, dtype=cfg.dtype)) 

67 self.b_V = nn.Parameter(torch.zeros(self.cfg.n_heads, self.cfg.d_head, dtype=cfg.dtype)) 

68 

69 @staticmethod 

70 def _relative_position_bucket( 

71 relative_position: Int[torch.Tensor, "query_pos kv_pos"], 

72 bidirectional=True, 

73 num_buckets=32, 

74 max_distance=128, 

75 ) -> Int[torch.Tensor, "query_pos kv_pos"]: 

76 """ 

77 added from 

78 https://github.com/huggingface/transformers/blob/e0c3cee17085914bbe505c159beeb8ae39bc37dd/src/transformers/models/t5/modeling_t5.py#L382 

79 which is adapted from 

80 https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 

81 

82 

83 Translate relative position to a bucket number for relative attention. The relative position is defined as 

84 memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to 

85 position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for 

86 small absolute relative_position and larger buckets for larger absolute relative_positions. All relative 

87 positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. 

88 This should allow for more graceful generalization to longer sequences than the model has been trained on 

89 

90 Args: 

91 relative_position: an int32 Tensor 

92 bidirectional: a boolean - whether the attention is bidirectional 

93 num_buckets: an integer 

94 max_distance: an integer 

95 

96 Returns: 

97 a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) 

98 """ 

99 relative_buckets = torch.zeros_like(relative_position) 

100 

101 if bidirectional: 

102 num_buckets //= 2 

103 relative_buckets += (relative_position > 0).to(torch.long) * num_buckets 

104 relative_position = torch.abs(relative_position) 

105 else: 

106 relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) 

107 # now relative_position is in the range [0, inf) 

108 

109 # half of the buckets are for exact increments in positions 

110 max_exact = num_buckets // 2 

111 is_small = relative_position < max_exact 

112 

113 # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance 

114 relative_position_if_large = max_exact + ( 

115 torch.log(relative_position.float() / max_exact) 

116 / math.log(max_distance / max_exact) 

117 * (num_buckets - max_exact) 

118 ).to(torch.long) 

119 relative_position_if_large = torch.min( 

120 relative_position_if_large, 

121 torch.full_like(relative_position_if_large, num_buckets - 1), 

122 ) 

123 

124 relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) 

125 return relative_buckets 

126 

127 def compute_relative_attention_bias( 

128 self, query_length: int, key_length: int, device=None 

129 ) -> Float[torch.Tensor, "1 head_index pos kv_pos"]: 

130 """Compute binned relative position bias""" 

131 if device is None: 

132 device = self.rel_pos_bias.weight.device 

133 context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] 

134 memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] 

135 relative_position = memory_position - context_position # shape (query_length, key_length) 

136 relative_position_bucket = self._relative_position_bucket( 

137 relative_position, # shape (query_length, key_length) 

138 bidirectional=not self.is_decoder, 

139 num_buckets=self.relative_attention_num_buckets, 

140 max_distance=self.relative_attention_max_distance, 

141 ) 

142 values = self.rel_pos_bias( 

143 relative_position_bucket 

144 ) # shape (query_length, key_length, num_heads) 

145 values = values.permute([2, 0, 1]).unsqueeze( 

146 0 

147 ) # shape (1, num_heads, query_length, key_length) 

148 return values