Update crossexpertattention.py

crossexpertattention.py

@@ -1,3 +1,40 @@
-# Source code for CrossExpertAttention from cell 39f2782d
-# Please replace this with the actual code from the notebook cell.
+from transformers import PretrainedConfig, PreTrainedModel, AutoModelForCausalLM # Import AutoModelForCausalLM
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from transformers.modeling_outputs import CausalLMOutputWithPast # Import the necessary output class
 
+# Define the Cross-Expert Attention mechanism
+class CrossExpertAttention(nn.Module):
+    def __init__(self, config: MeshConfig):
+        super().__init__()
+        self.config = config
+        # Define multi-head attention layers or similar for cross-expert communication
+        # This is a placeholder and needs detailed implementation
+        self.cross_attention = nn.MultiheadAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads, # Using model's attention heads for now
+            batch_first=True
+        )
+
+    def forward(self, expert_outputs):
+        # expert_outputs shape: (batch_size, sequence_length, num_experts, hidden_size)
+
+        if not self.config.cross_expert_attention_enabled:
+            return expert_outputs
+
+        # Reshape for attention: (batch_size * sequence_length, num_experts, hidden_size)
+        batch_seq_size = expert_outputs.shape[0] * expert_outputs.shape[1]
+        reshaped_outputs = expert_outputs.view(batch_seq_size, self.config.mesh_grid_size[0] * self.config.mesh_grid_size[1], self.config.hidden_size)
+
+        # Apply cross-expert attention. Query, Key, Value are the same here (self-attention across experts)
+        # Attention mask could be used to restrict communication if needed
+        cross_attn_output, _ = self.cross_attention(reshaped_outputs, reshaped_outputs, reshaped_outputs)
+
+        # Reshape back: (batch_size, sequence_length, num_experts, hidden_size)
+        cross_attn_output = cross_attn_output.view(
+            expert_outputs.shape[0], expert_outputs.shape[1], self.config.mesh_grid_size[0] * self.config.mesh_grid_size[1], self.config.hidden_size
+        )
+
+        return cross_attn_output