Upload 7 files

crossexpertattention.py

@@ -0,0 +1,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

meshconfig.py

@@ -0,0 +1,64 @@
+from transformers import PretrainedConfig, PreTrainedModel, AutoModelForCausalLM
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+class MeshConfig(PretrainedConfig):
+    model_type = "mesh"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=768,
+        intermediate_size=2048,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_key_value_heads=12,
+        max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        # Mesh specific configurations
+        mesh_grid_size=(2, 2), # 2x2 grid
+        expert_intermediate_size=256, # Example size for expert intermediate layer
+        routing_k=2, # Top-k routing
+        neighbor_exchange_enabled=True,
+        cross_expert_attention_enabled=True,
+        **kwargs
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            max_position_embeddings=max_position_embeddings,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            use_cache=use_cache,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.mesh_grid_size = mesh_grid_size
+        # Calculate expert_intermediate_size based on the shared and expert parameter split
+        # Total parameters = Shared (Embedding, Norm, LM Head) + Experts + Overhead
+        # This calculation is complex and depends on the specific layer mapping.
+        # For now, let's use a placeholder or calculate it based on the target parameter count.
+        # Target A242M (top-2): 100M shared + 135M (2 experts) + 7M overhead = 242M
+        # Let's assume the 135M for 2 experts is primarily in the intermediate size.
+        # We need to determine how Gemma's intermediate size maps to the expert intermediate size.
+        # For now, I will keep a placeholder or a simple ratio.
+        self.expert_intermediate_size = intermediate_size // (mesh_grid_size[0] * mesh_grid_size[1]) # Example: divide intermediate size by number of experts
+        self.routing_k = routing_k
+        self.neighbor_exchange_enabled = neighbor_exchange_enabled
+        self.cross_expert_attention_enabled = cross_expert_attention_enabled

meshexpert.py

@@ -0,0 +1,17 @@
+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 a single Expert within the Mesh
+class MeshExpert(nn.Module):
+    def __init__(self, config: MeshConfig):
+        super().__init__()
+        self.fc1 = nn.Linear(config.hidden_size, config.expert_intermediate_size)
+        self.gelu = nn.GELU() # Using GELU as an example activation
+        self.fc2 = nn.Linear(config.expert_intermediate_size, config.hidden_size)
+
+    def forward(self, x):
+        return self.fc2(self.gelu(self.fc1(x)))

meshlayer.py

@@ -0,0 +1,55 @@
+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 main Mesh Layer
+class MeshLayer(nn.Module):
+    def __init__(self, config: MeshConfig):
+        super().__init__()
+        self.config = config
+        self.router = MeshRouter(config)
+        self.experts = nn.ModuleList([MeshExpert(config) for _ in range(config.mesh_grid_size[0] * config.mesh_grid_size[1])])
+        self.neighbor_exchange = NeighborExchange(config)
+        self.cross_expert_attention = CrossExpertAttention(config)
+
+    def forward(self, hidden_states):
+        # hidden_states shape: (batch_size, sequence_length, hidden_size)
+
+        # 1. Routing
+        topk_weights, topk_indices = self.router(hidden_states)
+        # topk_weights shape: (batch_size, sequence_length, k)
+        # topk_indices shape: (batch_size, sequence_length, k)
+
+        # Prepare expert inputs: repeat hidden_states for each expert
+        # shape: (batch_size, sequence_length, num_experts, hidden_size)
+        expanded_hidden_states = hidden_states.unsqueeze(2).expand(-1, -1, self.config.mesh_grid_size[0] * self.config.mesh_grid_size[1], -1)
+
+        # 2. Expert Computation
+        # Compute output for all experts (can be optimized to only compute for selected experts)
+        expert_outputs = torch.stack([expert(expanded_hidden_states[:, :, i, :]) for i, expert in enumerate(self.experts)], dim=2)
+        # expert_outputs shape: (batch_size, sequence_length, num_experts, hidden_size)
+
+        # 3. Neighbor Exchange (conceptual implementation needed)
+        exchanged_expert_outputs = self.neighbor_exchange(expert_outputs, topk_indices)
+
+        # 4. Cross-Expert Attention (conceptual implementation needed)
+        cross_attned_expert_outputs = self.cross_expert_attention(exchanged_expert_outputs)
+
+        # 5. Combine expert outputs based on routing weights
+        # Create a tensor to gather the outputs of the selected experts
+        # shape: (batch_size, sequence_length, k, hidden_size)
+        gathered_outputs = torch.gather(
+            cross_attned_expert_outputs,
+            dim=2,
+            index=topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.config.hidden_size)
+        )
+
+        # Apply routing weights: (batch_size, sequence_length, k, 1) * (batch_size, sequence_length, k, hidden_size)
+        combined_output = (gathered_outputs * topk_weights.unsqueeze(-1)).sum(dim=2)
+        # combined_output shape: (batch_size, sequence_length, hidden_size)
+
+        # Return the combined output and the expert indices for potential visualization
+        return combined_output, topk_indices # Return combined output and expert indices

meshmodel.py

@@ -0,0 +1,88 @@
+from transformers import PretrainedConfig, PreTrainedModel, AutoModelForCausalLM
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+class MeshModel(PreTrainedModel):
+    config_class = MeshConfig
+
+    def __init__(self, config: MeshConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([MeshLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        # Ensure return_dict is set to True by default if not specified
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            inputs_embeds = self.embedding(input_ids)
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        hidden_states = inputs_embeds
+        expert_indices_list = [] # To collect expert indices from each layer
+
+        for i, layer in enumerate(self.layers):
+            hidden_states, expert_indices = layer(hidden_states)
+            expert_indices_list.append(expert_indices) # Collect indices
+
+        hidden_states = self.norm(hidden_states)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Compute loss (e.g., CrossEntropyLoss)
+            loss_fct = nn.CrossEntropyLoss()
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Calculate scalar loss
+            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+
+        # Return a CausalLMOutputWithPast object or a tuple
+        if return_dict:
+             return CausalLMOutputWithPast(
+                 loss=loss,
+                 logits=logits,
+                 past_key_values=None, # Need to implement caching
+                 hidden_states=hidden_states,
+                 attentions=None, # Need to implement attention handling
+             )
+        else:
+             # Return a tuple including loss, logits, and collected expert indices
+             # Ensure the order and content match what the Trainer expects or can handle
+             # Trainer expects (loss, logits, hidden_states, attentions) or similar.
+             # We can return (loss, logits) as the primary outputs for the Trainer
+             # and potentially include expert_indices as an additional output if needed
+             # by a custom callback or logging, but the default Trainer expects loss as the first element for backward.
+             return (loss, logits, hidden_states, expert_indices_list) # Include expert_indices_list

meshrouter.py

@@ -0,0 +1,27 @@
+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 Router for dynamic routing
+class MeshRouter(nn.Module):
+    def __init__(self, config: MeshConfig):
+        super().__init__()
+        self.gate = nn.Linear(config.hidden_size, config.mesh_grid_size[0] * config.mesh_grid_size[1])
+        self.softmax = nn.Softmax(dim=-1)
+        self.routing_k = config.routing_k
+
+    def forward(self, x):
+        # x shape: (batch_size, sequence_length, hidden_size)
+        gate_scores = self.gate(x) # shape: (batch_size, sequence_length, num_experts)
+        gate_weights = self.softmax(gate_scores)
+
+        # Select top-k experts
+        topk_weights, topk_indices = torch.topk(gate_weights, self.routing_k, dim=-1)
+
+        # Normalize top-k weights
+        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-6)
+
+        return topk_weights, topk_indices # shapes: (batch_size, sequence_length, k), (batch_size, sequence_length, k)

neighborexchange.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+
+class NeighborExchange(nn.Module):
+    def __init__(self, config: MeshConfig):
+        super().__init__()
+        self.config = config
+        self.num_experts_x = config.mesh_grid_size[0]
+        self.num_experts_y = config.mesh_grid_size[1]
+        self.num_experts = self.num_experts_x * self.num_experts_y
+
+        # Define parameters for neighbor communication.
+        # A simple approach: a learned linear combination of neighbor features.
+        # We can define a weight for each potential neighbor direction (e.g., up, down, left, right).
+        # For a 2x2 grid, each expert has 2 or 3 neighbors.
+        # A more general approach is a linear layer that takes concatenated neighbor features.
+        # Let's use a linear layer to transform the aggregated neighbor information.
+        # The input size to this layer will be the sum of hidden sizes of all potential neighbors
+        # multiplied by the hidden size, but that's too complex.
+        # A simpler approach: a linear layer per direction, or a single layer after aggregating.
+
+        # Let's define a linear layer to process the information received from neighbors.
+        # The input size is the hidden size (from neighbors), output size is hidden size
+        # This layer will transform the aggregated neighbor features before adding to the expert's own output.
+        self.exchange_projection = nn.Linear(config.hidden_size, config.hidden_size) # Projects aggregated neighbor info
+
+        # Optional: Learned weights for different neighbor directions
+        # self.neighbor_weights = nn.Parameter(torch.ones(4)) # Example for 4 directions (N, S, E, W)
+
+    def forward(self, expert_outputs, expert_indices=None):
+        # expert_outputs shape: (batch_size, sequence_length, num_experts, hidden_size)
+        # expert_indices shape: (batch_size, sequence_length, k) - indices of selected experts (not directly used for neighbor exchange in this simple model)
+
+        if not self.config.neighbor_exchange_enabled:
+            return expert_outputs
+
+        batch_size, seq_length, num_experts, hidden_size = expert_outputs.shape
+
+        # Reshape expert_outputs to reflect the grid structure (batch_size, seq_length, grid_x, grid_y, hidden_size)
+        reshaped_outputs = expert_outputs.view(batch_size, seq_length, self.num_experts_x, self.num_experts_y, hidden_size)
+
+        # Create a tensor to store the aggregated neighbor information for each expert
+        aggregated_neighbor_info = torch.zeros_like(reshaped_outputs)
+
+        # Implement neighbor exchange logic
+        # Iterate through each expert in the grid
+        for i in range(self.num_experts_x):
+            for j in range(self.num_experts_y):
+                current_expert_output = reshaped_outputs[:, :, i, j, :]
+                neighbor_info = torch.zeros_like(current_expert_output) # Accumulate info from neighbors
+
+                # Define neighbor directions (example: up, down, left, right)
+                neighbors = []
+                if i > 0: # Up neighbor
+                    neighbors.append(reshaped_outputs[:, :, i-1, j, :])
+                if i < self.num_experts_x - 1: # Down neighbor
+                    neighbors.append(reshaped_outputs[:, :, i+1, j, :])
+                if j > 0: # Left neighbor
+                    neighbors.append(reshaped_outputs[:, :, i, j-1, :])
+                if j < self.num_experts_y - 1: # Right neighbor
+                    neighbors.append(reshaped_outputs[:, :, i, j+1, :])
+
+                # Aggregate information from neighbors (simple average as an example)
+                if neighbors:
+                    # Stack neighbors along a new dimension and take the mean
+                    neighbor_stack = torch.stack(neighbors, dim=-2) # shape (batch, seq, num_neighbors, hidden)
+                    aggregated_info = torch.mean(neighbor_stack, dim=-2) # shape (batch, seq, hidden)
+                    neighbor_info = aggregated_info # Use the aggregated info
+
+                # Apply the exchange projection to the aggregated neighbor information
+                transformed_neighbor_info = self.exchange_projection(neighbor_info)
+
+                # Store the transformed neighbor info for the current expert's position
+                aggregated_neighbor_info[:, :, i, j, :] = transformed_neighbor_info
+
+        # Reshape aggregated_neighbor_info back to (batch_size, sequence_length, num_experts, hidden_size)
+        aggregated_neighbor_info = aggregated_neighbor_info.view(batch_size, seq_length, num_experts, hidden_size)
+
+        # Combine expert outputs with aggregated neighbor information (additive combination)
+        exchanged_expert_outputs = expert_outputs + aggregated_neighbor_info
+
+        return exchanged_expert_outputs