File size: 2,009 Bytes
1dbd336 711a5f0 1dbd336 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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 |