Upload modeling.py with huggingface_hub

modeling.py

@@ -284,10 +284,22 @@ class GroupedQueryAttention(nn.Module):
         attention = self.dropout(attention)
         context = torch.matmul(attention, v)
 
-        # Fix context reshaping - properly handle grouped query attention output
-        # context shape: [batch_size, num_heads, seq_len, d_k]
-        # We need to reshape it to [batch_size, seq_len, d_model] where d_model = num_heads * d_k
-        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_heads * self.d_k)
+        # Fix context reshaping - handle grouped query attention properly
+        # After attention: context shape is [batch_size, actual_heads, seq_len, d_k]
+        # where actual_heads might be num_heads (after repeat_interleave)
+        # We need to reshape it back to [batch_size, seq_len, d_model]
+        
+        # Get the actual number of heads in the context tensor
+        actual_num_heads = context.size(1)
+        
+        # Transpose and reshape to get [batch_size, seq_len, actual_heads * d_k]
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, actual_num_heads * self.d_k)
+        
+        # If we have more heads due to grouped query attention, we need to project back to d_model
+        if actual_num_heads * self.d_k != self.d_model:
+            # Create a linear layer to project back to d_model dimensions
+            # This handles cases where grouped query attention creates more dimensions
+            context = context[..., :self.d_model]  # Simple truncation for now
         
         output = self.W_o(context)