Update modeling_auristream.py

modeling_auristream.py

@@ -424,7 +424,7 @@ class CausalSelfAttention(nn.Module):
         if hasattr(config, 'use_rope') and not config.use_rope:
             self.rotary = None
 
-    def forward(self, x, return_kv=False, ):
+    def forward(self, x, return_kv=False, return_attn_maps=False):
 
         B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
         # calculate query, key, values for all heads in batch and move head forward to be the batch dim
@@ -439,13 +439,35 @@ class CausalSelfAttention(nn.Module):
             q = apply_rotary_emb(q, cos, sin)
             k = apply_rotary_emb(k, cos, sin)
 
-        y = F.scaled_dot_product_attention(
-            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), 
-            is_causal=True)
+        if not return_kv and not return_attn_maps:
+            y = F.scaled_dot_product_attention(
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), 
+                is_causal=True)
+        else:
+            # manual implementation of attention
+            att = torch.einsum('bnsh,bnkh->bnsk', q, k) * (1.0 / math.sqrt(k.size(-1)))
+            mask = torch.triu(torch.ones(T, T), diagonal=1).to(dtype=torch.bool).to(x.device)
+            mask = mask.view(1, 1, T, T)
+            masked_att = att.masked_fill(mask, float('-inf'))
+            # upcast to float32 for numerical stability, as per llama implementation
+            masked_att = F.softmax(masked_att, dim=-1, dtype=torch.float32).to(q.dtype)
+            masked_att = self.attn_dropout(masked_att)
+            # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+            y = torch.einsum('bnsk,bnkh->bnsh', masked_att, v)
 
         y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
         # output projection
-        y = self.c_proj(y)
+        y = self.resid_dropout(self.c_proj(y))
+    
+        # return attention maps if requested
+        if return_attn_maps:
+            return y, F.softmax(att, dim=-1)
+
+        # return key and value caches if requested
+        if return_kv:
+            return y, k, v
+
         return y