Changes in modelling_RW.py to be able to handle past_key_values for faster model generations

modelling_RW.py

@@ -87,10 +87,18 @@ class RotaryEmbedding(torch.nn.Module):
 
         return self.cos_cached, self.sin_cached
 
-    def forward(self, q, k):
-        batch, seq_len, head_dim = q.shape
+    def forward(self, q, k, past_seq_length=None):
+        if past_seq_length == None : 
+            batch, seq_len, head_dim = q.shape
+        else : 
+            # print("past_seq_length", past_seq_length)
+            batch, input_seq_len, head_dim = q.shape 
+            seq_len = past_seq_length + input_seq_len
         cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
-        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+        if past_seq_length != None : 
+            return (q * cos[:, past_seq_length:, :]) + (rotate_half(q) * sin[:, past_seq_length:, :]), (k * cos[:, past_seq_length:, :]) + (rotate_half(k) * sin[:, past_seq_length:, :])
+        else : 
+            return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
 
 
 def _make_causal_mask(
@@ -100,10 +108,10 @@ def _make_causal_mask(
     mask = torch.empty((target_length, target_length + past_key_values_length), dtype=torch.bool, device=device)
     # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround
     seq_ids = torch.arange(target_length, device=device)
-    mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]
+    mask[:, past_key_values_length:] = seq_ids[:, None] >= seq_ids[None, :]
 
     if past_key_values_length > 0:
-        mask[:, :past_key_values_length] = False
+        mask[:, :past_key_values_length] = True
 
     expanded_mask = mask[None, None, :, :].expand(batch_size, 1, target_length, target_length + past_key_values_length)
     return expanded_mask
@@ -248,6 +256,7 @@ class Attention(nn.Module):
         head_mask: Optional[torch.Tensor] = None,
         use_cache: bool = False,
         output_attentions: bool = False,
+        layer_number = None
     ):
         fused_qkv = self.query_key_value(hidden_states)  # [batch_size, seq_length, 3 x hidden_size]
 
@@ -264,31 +273,43 @@ class Attention(nn.Module):
         )
         value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)
 
-        query_layer, key_layer = self.maybe_rotary(query_layer, key_layer)
+        if layer_past is not None : 
+            past_key, past_value = layer_past
+            past_kv_length = past_key.shape[2]
+            query_layer, key_layer = self.maybe_rotary(query_layer, key_layer, past_kv_length)
+        else : 
+            query_layer, key_layer = self.maybe_rotary(query_layer, key_layer)
 
         if layer_past is not None:
             past_key, past_value = layer_past
             # concatenate along seq_length dimension:
             #  - key: [batch_size * self.num_heads, head_dim, kv_length]
             #  - value: [batch_size * self.num_heads, kv_length, head_dim]
+            past_key = past_key.permute(0, 2, 1)
             key_layer = torch.cat((past_key, key_layer), dim=1)
             value_layer = torch.cat((past_value, value_layer), dim=1)
 
         _, kv_length, _ = key_layer.shape
 
         if use_cache is True:
-            present = (key_layer, value_layer)
+            key_layer_permute = key_layer.permute(0, 2, 1)
+            present = (key_layer_permute, value_layer)
         else:
             present = None
-
+        
         if alibi is None:
             query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
             key_layer_ = key_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
             value_layer_ = value_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)
-
-            attn_output = F.scaled_dot_product_attention(
-                query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
-            )
+            
+            if attention_mask is not None : 
+                attn_output = F.scaled_dot_product_attention(
+                    query_layer_, key_layer_, value_layer_, attention_mask, 0.0, is_causal=False
+                )
+            else :
+                attn_output = F.scaled_dot_product_attention(
+                    query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
+                ) 
 
             x = attn_output.view(batch_size, self.num_heads, q_length, self.head_dim)
             x = x.permute(0, 2, 1, 3)
@@ -385,6 +406,7 @@ class DecoderLayer(nn.Module):
         head_mask: Optional[torch.Tensor] = None,
         use_cache: bool = False,
         output_attentions: bool = False,
+        layer_number = None
     ):
 
         ln_attn = self.ln_attn(hidden_states)
@@ -401,6 +423,7 @@ class DecoderLayer(nn.Module):
             head_mask=head_mask,
             use_cache=use_cache,
             output_attentions=output_attentions,
+            layer_number=layer_number
         )
 
         attention_output = attn_outputs[0]
@@ -528,10 +551,10 @@ class RWModel(RWPreTrainedModel):
         device = attention_mask.device
         _, src_length = input_shape
 
-        if src_length > 1:
-            combined_attention_mask = _make_causal_mask(
-                input_shape, device=device, past_key_values_length=past_key_values_length
-            )
+        # if src_length > 1:
+        combined_attention_mask = _make_causal_mask(
+            input_shape, device=device, past_key_values_length=past_key_values_length
+        )
 
         # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
         expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
@@ -651,15 +674,28 @@ class RWModel(RWPreTrainedModel):
                     head_mask[i],
                 )
             else:
-                outputs = block(
-                    hidden_states,
-                    layer_past=layer_past,
-                    attention_mask=causal_mask,
-                    head_mask=head_mask[i],
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    alibi=alibi,
-                )
+                if  i==0 : 
+                    outputs = block(
+                        hidden_states,
+                        layer_past=layer_past,
+                        attention_mask=causal_mask,
+                        head_mask=head_mask[i],
+                        use_cache=use_cache,
+                        output_attentions=output_attentions,
+                        alibi=alibi,
+                        layer_number=0
+                    )
+                else : 
+                    outputs = block(
+                        hidden_states,
+                        layer_past=layer_past,
+                        attention_mask=causal_mask,
+                        head_mask=head_mask[i],
+                        use_cache=use_cache,
+                        output_attentions=output_attentions,
+                        alibi=alibi,
+                    )
+
 
             hidden_states = outputs[0]
             if use_cache is True:
@@ -710,16 +746,19 @@ class RWForCausalLM(RWPreTrainedModel):
         **kwargs,
     ) -> dict:
         # only last token for input_ids if past is not None
-        if past:
+        if kwargs.get("past_key_values", None) : 
             input_ids = input_ids[:, -1].unsqueeze(-1)
-
+            past_key_values = kwargs["past_key_values"]
             # the cache may be in the stardard format (e.g. in contrastive search), convert to our's format if needed
-            if past[0][0].shape[0] == input_ids.shape[0]:
-                past = self._convert_to_rw_cache(past)
+            # if kwargs["past_key_values"][0][0].shape[0] == input_ids.shape[0]:
+            #     past_key_values = self._convert_to_rw_cache(kwargs["past_key_values"])
+                # past_key_values = kwargs["past_key_values"]
+        else : 
+            past_key_values = None
 
         return {
             "input_ids": input_ids,
-            "past_key_values": past,
+            "past_key_values": past_key_values,
             "use_cache": kwargs.get("use_cache"),
             "attention_mask": attention_mask,
         }