Coverage for transformer_lens/components/t5_block.py: 88%

1from typing import Optional 

2 

3import torch 

4import torch.nn as nn 

5from jaxtyping import Float 

6 

7from transformer_lens.cache.key_value_cache_entry import ( 

8 TransformerLensKeyValueCacheEntry, 

9) 

10from transformer_lens.components import RMSNorm, T5Attention 

11from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

12from transformer_lens.factories.mlp_factory import MLPFactory 

13from transformer_lens.hook_points import HookPoint 

14from transformer_lens.utilities import repeat_along_head_dimension 

15 

16 

17class T5Block(nn.Module): 

18 """ 

19 T5 decoder Block. Uses T5Layernorm, and T5attention instead of usual ones. 

20 Also uses cross attention if is_decoder is True. 

21 """ 

22 

23 def __init__(self, cfg: HookedTransformerConfig, block_index: int, is_decoder: bool): 

24 super().__init__() 

25 self.cfg = cfg 

26 self.is_decoder = is_decoder 

27 

28 self.ln1 = RMSNorm(cfg) 

29 self.attn = T5Attention( 

30 cfg, has_relative_attention_bias=block_index == 0, is_decoder=is_decoder 

31 ) 

32 self.ln2 = RMSNorm(cfg) 

33 if self.is_decoder: 

34 self.cross_attn = T5Attention(cfg) 

35 self.ln3 = RMSNorm(cfg) 

36 self.mlp = MLPFactory.create_mlp(self.cfg) # [batch, pos, n_heads] 

37 

38 self.hook_q_input = HookPoint() # [batch, pos, n_heads, d_model] 

39 self.hook_k_input = HookPoint() # [batch, pos, n_heads, d_model] 

40 self.hook_v_input = HookPoint() # [batch, pos, n_heads, d_model] 

41 

42 self.hook_attn_in = HookPoint() # [batch, pos, d_model] 

43 self.hook_attn_out = HookPoint() # [batch, pos, d_model] 

44 if self.is_decoder: 

45 self.hook_cross_attn_in = HookPoint() # [batch, pos, d_model] 

46 self.hook_cross_attn_out = HookPoint() # [batch, pos, d_model] 

47 self.hook_resid_mid_cross = HookPoint() # [batch, pos, d_model] 

48 

49 self.hook_mlp_in = HookPoint() # [batch, pos, d_model] 

50 self.hook_mlp_out = HookPoint() # [batch, pos, d_model] 

51 self.hook_resid_pre = HookPoint() # [batch, pos, d_model] 

52 self.hook_resid_mid = HookPoint() # [batch, pos, d_model] 

53 self.hook_resid_post = HookPoint() # [batch, pos, d_model] 

54 

55 def forward( 

56 self, 

57 resid_pre: Float[torch.Tensor, "batch pos d_model"], 

58 additive_attention_mask: Optional[Float[torch.Tensor, "batch 1 1 pos"]] = None, 

59 encoder_additive_attention_mask: Optional[ 

60 Float[torch.Tensor, "batch 1 1 encoder_pos"] 

61 ] = None, 

62 position_bias: Optional[Float[torch.Tensor, "1 head_index pos kv_pos"]] = None, 

63 encoder_hidden_states: Optional[Float[torch.Tensor, "batch encoder_pos d_model"]] = None, 

64 past_kv_cache_entry: Optional[TransformerLensKeyValueCacheEntry] = None, 

65 ) -> Float[torch.Tensor, "batch pos d_model"]: 

66 """A single Transformer block. 

67 

68 Args: 

69 resid_pre (torch.Tensor): The residual stream - shape [batch, pos, d_model] 

70 encoder_hidden_states (torch.Tensor): The hidden states of the encoder for cross attention - shape [batch, encoder_pos, d_model] 

71 cache (TransformerLensKeyValueCache): A cache of previous keys and values, used only when generating text. Defaults to None. 

72 attention_mask (torch.Tensor, optional): The attention mask for padded tokens. Defaults to None. 

73 

74 Returns: 

75 Float[torch.Tensor, "batch pos d_model"]: The block output residual stream. 

76 """ 

77 resid_pre = self.hook_resid_pre(resid_pre) # [batch, pos, d_model] 

78 

79 attn_in = resid_pre 

80 

81 if self.cfg.use_attn_in: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 attn_in = self.hook_attn_in( 

83 repeat_along_head_dimension(resid_pre, n_heads=self.cfg.n_heads) 

84 ) 

85 

86 if self.cfg.use_split_qkv_input: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 n_kv_heads = ( 

88 self.cfg.n_key_value_heads 

89 if self.cfg.n_key_value_heads is not None 

90 else self.cfg.n_heads 

91 ) 

92 query_input = self.hook_q_input( 

93 repeat_along_head_dimension(resid_pre, n_heads=self.cfg.n_heads) 

94 ) 

95 key_input = self.hook_k_input( 

96 repeat_along_head_dimension(resid_pre, n_heads=n_kv_heads) 

97 ) 

98 value_input = self.hook_v_input( 

99 repeat_along_head_dimension(resid_pre, n_heads=n_kv_heads) 

100 ) 

101 else: 

102 query_input = attn_in 

103 key_input = attn_in 

104 value_input = attn_in 

105 

106 attn_out = self.hook_attn_out( 

107 # hook the residual stream states that are used to calculate the 

108 # queries, keys and values, independently. 

109 # Then take the layer norm of these inputs, and pass these to the attention module. 

110 self.attn( 

111 query_input=self.ln1(query_input), 

112 key_input=self.ln1(key_input), 

113 value_input=self.ln1(value_input), 

114 past_kv_cache_entry=past_kv_cache_entry, 

115 additive_attention_mask=additive_attention_mask, 

116 position_bias=position_bias, 

117 ) 

118 ) 

119 

120 # [batch, pos, d_model] 

121 

122 resid_mid = self.hook_resid_mid(resid_pre + attn_out) # [batch, pos, d_model] 

123 

124 if self.is_decoder: 

125 cross_attn_in = ( 

126 resid_mid 

127 if not self.cfg.use_attn_in 

128 else self.hook_cross_attn_in(resid_mid.clone()) 

129 ) 

130 

131 if encoder_hidden_states is None: 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true

132 raise ValueError("Encoder hidden states must be provided for cross attention!") 

133 

134 cross_attn_out = self.hook_cross_attn_out( 

135 self.cross_attn( 

136 query_input=self.ln2(cross_attn_in), 

137 key_input=encoder_hidden_states, 

138 value_input=encoder_hidden_states, 

139 additive_attention_mask=encoder_additive_attention_mask, 

140 ) 

141 ) 

142 resid_mid_cross = self.hook_resid_mid_cross(resid_mid + cross_attn_out) 

143 

144 mlp_in = ( 

145 resid_mid_cross 

146 if not self.cfg.use_hook_mlp_in 

147 else self.hook_mlp_in(resid_mid_cross.clone()) 

148 ) 

149 

150 normalized_resid_mid = self.ln3(mlp_in) 

151 else: 

152 mlp_in = ( 

153 resid_mid if not self.cfg.use_hook_mlp_in else self.hook_mlp_in(resid_mid.clone()) 

154 ) 

155 normalized_resid_mid = self.ln2(mlp_in) 

156 

157 mlp_out = self.hook_mlp_out(self.mlp(normalized_resid_mid)) # [batch, pos, d_model] 

158 resid_post = self.hook_resid_post(mlp_in + mlp_out) # [batch, pos, d_model] 

159 

160 return resid_post