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