meshconfig.py

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