Update meshconfig.py

meshconfig.py

@@ -1,3 +1,64 @@
-# Source code for MeshConfig from cell VExhmWA0lXA_
-# Please replace this with the actual code from the notebook cell.
+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