fuse qkv

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ba58d7bbc20355cd3083e789a88fa6b9016ec36ffaf113e94df03d1449ecadf6
-size 4903283827
+oid sha256:ecb114682d0f35efe851ed162500e9d7babdf7c8c008fdcf76ec679e9a788533
+size 4903278480

modeling_gpt2_mq.py

@@ -13,7 +13,6 @@ from transformers.modeling_outputs import (
 )
 from transformers.models.gpt2.modeling_gpt2 import GPT2Model, GPT2Block, GPT2PreTrainedModel, GPT2LMHeadModel
 from transformers.utils import logging
-
 from .configuration_gpt2_mq import GPT2CustomConfig, MULTI_QUERY
 
 logger = logging.get_logger(__name__)
@@ -130,10 +129,7 @@ class GPT2MQAttention(nn.Module):
         if self.is_cross_attention:
             raise NotImplementedError("Cross-attention not implemented for MQA")
         else:
-            # self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
-            self.q_attn = nn.Linear(self.embed_dim, self.embed_dim)
-            # Keys and values are shared across heads
-            self.kv_attn = nn.Linear(self.embed_dim, 2 * self.head_dim)
+            self.attn = nn.Linear(self.embed_dim, self.embed_dim + 2 * self.head_dim)
         self.c_proj = nn.Linear(self.embed_dim, self.embed_dim)
 
         self.attn_dropout = nn.Dropout(config.attn_pdrop)
@@ -143,13 +139,13 @@ class GPT2MQAttention(nn.Module):
         self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
 
     def _attn(self, query, key, value, attention_mask=None, head_mask=None):
-        # query: (b, num_heads * sq, head_dim)
+        # query: (b, sq * num_heads, head_dim)
         # key: (b, head_dim, sk)
         # value: (b, sk, head_dim)
         batch_size = query.size(0)
         query_length = query.size(1) // self.num_heads
         key_length = key.size(2)
-        # (b, num_heads * sq, head_dim) x (b, head_dim, sk) -> (b, num_heads * sq, sk)
+        # (b, sq * num_heads, head_dim) x (b, head_dim, sk) -> (b, sq * num_heads, sk)
 
         if self.scale_attn_weights:
             query *= self.inv_norm_factor
@@ -157,7 +153,7 @@ class GPT2MQAttention(nn.Module):
         attn_weights = torch.bmm(query, key)
 
         # -> (b, num_heads, sq, sk)
-        attn_weights = attn_weights.view(batch_size, self.num_heads, query_length, key_length)
+        attn_weights = attn_weights.view(batch_size, query_length, self.num_heads, key_length)
 
         # Layer-wise attention scaling
         if self.scale_attn_by_inverse_layer_idx:
@@ -174,13 +170,13 @@ class GPT2MQAttention(nn.Module):
 
         # Mask heads if we want to
         if head_mask is not None:
-            attn_weights = attn_weights * head_mask
+            raise NotImplementedError
 
         # (b, num_heads, sq, sk) -> (b, num_heads * sq, sk)
-        _attn_weights = attn_weights.view(batch_size, self.num_heads * query_length, key_length)
+        _attn_weights = attn_weights.view(batch_size, query_length * self.num_heads, key_length)
         # (b, num_heads * sq, sk) x (b, sk, head_dim) -> (b, num_heads * sq, head_dim)
         attn_output = torch.bmm(_attn_weights, value)
-        attn_output = attn_output.view(batch_size, self.num_heads, query_length, self.head_dim)
+        attn_output = attn_output.view(batch_size, query_length, self.num_heads, self.head_dim)
 
         return attn_output, attn_weights
 
@@ -188,10 +184,8 @@ class GPT2MQAttention(nn.Module):
         """
         Merges attn_head_size dim and num_attn_heads dim into hidden_size
         """
-        batch_size, num_heads, seq_length, head_dim = tensor.shape
-
-        tensor = tensor.permute(0, 2, 1, 3)
-        return tensor.reshape(batch_size, seq_length, num_heads * head_dim)
+        batch_size, seq_length, num_heads, head_dim = tensor.shape
+        return tensor.view(batch_size, seq_length, num_heads * head_dim)
 
     def forward(
             self,
@@ -207,17 +201,14 @@ class GPT2MQAttention(nn.Module):
         if encoder_hidden_states is not None:
             raise NotImplementedError("Cross-attention not implemented for MQA")
         else:
-            query = self.q_attn(hidden_states)
-            key, value = self.kv_attn(hidden_states).split(self.head_dim, dim=2)
+            qkv = self.attn(hidden_states)
+            query, key, value = qkv.split([self.embed_dim, self.head_dim, self.head_dim], dim=2)
 
         batch_size, seq_length = query.shape[:2]
-        # (query_length, batch, num_heads, head_dim)
-        # (batch, num_heads * query_length, head_dim)\
 
-        # (batch, query_length, hidden_size) -> (batch, num_heads, query_length, head_dim)
-        query = query.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
-        # -> (batch, num_heads * query_length, head_dim)
-        query = query.reshape(batch_size, self.num_heads * seq_length, self.head_dim)
+        # (batch, query_length, hidden_size) -> (batch, query_length * num_heads, head_dim)
+        # forced to reshape here
+        query = query.reshape(batch_size, seq_length * self.num_heads, self.head_dim)
 
         key = key.transpose(1, 2)  # (batch_size, head_dim, seq_length)
 
@@ -360,8 +351,7 @@ class GPT2CustomModel(GPT2Model):
             past_key_values_length=past_key_values_length,
         )
 
-        attention_mask = attention_mask.unsqueeze(1).expand(batch_size, self.config.num_attention_heads,
-                                                            *attention_mask.shape[1:])
+        attention_mask = attention_mask.unsqueeze(2).expand(batch_size, attention_mask.shape[1], self.config.num_attention_heads, attention_mask.shape[2])
 
         # If a 2D or 3D attention mask is provided for the cross-attention
         # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]