""" SabiYarn Model Implementation - Optimized Version Memory-efficient with performance optimizations for generation. Matches original implementation exactly but with memory optimizations. """ from transformers import PreTrainedModel, AutoConfig, AutoModel, AutoModelForCausalLM from transformers.modeling_outputs import CausalLMOutputWithPast # use package-relative import to avoid colliding with unrelated `model` packages from .configuration import GPTJXMoEConfig from typing import Optional, List, Tuple from torch import nn import torch import torch.nn.functional as F import math class LayerNorm(nn.Module): """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """ def __init__(self, ndim, bias): super().__init__() self.weight = nn.Parameter(torch.ones(ndim)) self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None def forward(self, input): return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5) class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_heads == 0 # key, query, value projections for all heads, but in a batch self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) # output projection self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) # regularization self.attn_dropout = nn.Dropout(config.dropout) self.resid_dropout = nn.Dropout(config.dropout) self.n_heads = config.n_heads self.n_embd = config.n_embd self.head_dim = config.n_embd // config.n_heads self.dropout = config.dropout # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0 self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') def forward(self, x, attn_mask=None, past_key_value=None, use_cache=False): """ Forward pass with optional KV cache support. Args: x: (B, T, C) input embeddings attn_mask: Optional attention mask past_key_value: Optional tuple of (past_k, past_v) each (B, nh, past_len, hs) use_cache: Whether to return cache for next step Returns: If use_cache: (output, (k, v)) where output is (B, T, C) and k, v are (B, nh, total_len, hs) Else: output (B, T, C) """ B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) # calculate query, key, values for all heads in batch and move head forward to be the batch dim q, k, v = self.c_attn(x).split(self.n_embd, dim=2) k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # (B, nh, T, hs) q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # (B, nh, T, hs) v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # (B, nh, T, hs) # Concatenate with past KV cache if provided if past_key_value is not None: past_k, past_v = past_key_value k = torch.cat([past_k, k], dim=2) # (B, nh, past_len + T, hs) v = torch.cat([past_v, v], dim=2) # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, total_len) -> (B, nh, T, total_len) total_len = k.size(2) if self.flash: if attn_mask is not None: # efficient attention using Flash Attention CUDA kernels attn_mask = attn_mask.to(torch.bool) # Handle different mask shapes and convert to (B, nh, T, total_len) if attn_mask.dim() == 2: # (B, S) - expand to cover full sequence if needed B_mask = attn_mask.size(0) S = attn_mask.size(1) if S == total_len: # Mask already covers full sequence pass elif S == T: # Mask only covers current tokens - expand with ones for past tokens if past_key_value is not None: past_len = total_len - T past_mask = torch.ones(B_mask, past_len, device=x.device, dtype=attn_mask.dtype) attn_mask = torch.cat([past_mask, attn_mask], dim=1) else: # No cache, mask is correct as-is pass else: raise ValueError(f"Unsupported attention_mask shape: {attn_mask.shape}, expected (B, {T}) or (B, {total_len})") # Reshape to (B, 1, T, total_len) for Flash Attention # Flash Attention expects mask shape (B, nh, T, S) where T is query length # First ensure we have the right length if attn_mask.size(1) != total_len: raise ValueError(f"Mask length mismatch: got {attn_mask.size(1)}, expected {total_len}") # Reshape: (B, total_len) -> (B, 1, 1, total_len) -> (B, 1, T, total_len) -> (B, nh, T, total_len) attn_mask = attn_mask.view(B_mask, 1, 1, total_len) # Expand to (B, 1, T, total_len) - repeat for each query position attn_mask = attn_mask.expand(B_mask, 1, T, total_len) # Expand to include head dimension: (B, nh, T, total_len) attn_mask = attn_mask.expand(-1, self.n_heads, -1, -1) # Verify final shape assert attn_mask.shape == (B_mask, self.n_heads, T, total_len), \ f"Mask shape mismatch: got {attn_mask.shape}, expected ({B_mask}, {self.n_heads}, {T}, {total_len})" elif attn_mask.dim() == 4: # Already 4D mask - ensure it's the right shape B_mask = attn_mask.size(0) if attn_mask.size(-2) != T: # Slice to match query length if needed attn_mask = attn_mask[..., -T:, :] # Ensure head dimension matches if attn_mask.size(1) == 1: attn_mask = attn_mask.expand(-1, self.n_heads, -1, -1) elif attn_mask.size(1) != self.n_heads: raise ValueError(f"Mask head dimension {attn_mask.size(1)} doesn't match n_heads {self.n_heads}") else: raise ValueError(f"Unsupported attention_mask dimension: {attn_mask.dim()}, expected 2 or 4") y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0, is_causal=False) else: # No explicit mask provided if past_key_value is None: # No cache: use is_causal for efficiency y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True) else: # With cache: create causal mask manually (can't use is_causal when q and k have different lengths) causal_mask = torch.tril(torch.ones(T, total_len, device=x.device, dtype=torch.bool)) y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=causal_mask.view(1, 1, T, total_len), dropout_p=self.dropout if self.training else 0, is_causal=False) else: # manual implementation of attention total_len = k.size(2) att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim)) if attn_mask is not None: attn_mask = attn_mask.to(torch.bool) # Handle different mask shapes and convert to (B, nh, T, total_len) if attn_mask.dim() == 2: # (B, S) - expand to cover full sequence if needed B_mask = attn_mask.size(0) S = attn_mask.size(1) if S == total_len: # Mask already covers full sequence pass elif S == T: # Mask only covers current tokens - expand with ones for past tokens if past_key_value is not None: past_len = total_len - T past_mask = torch.ones(B_mask, past_len, device=x.device, dtype=torch.bool) attn_mask = torch.cat([past_mask, attn_mask], dim=1) else: # No cache, mask is correct as-is pass else: raise ValueError(f"Unsupported attention_mask shape: {attn_mask.shape}, expected (B, {T}) or (B, {total_len})") # Reshape to (B, 1, T, total_len) then expand to (B, nh, T, total_len) attn_mask = attn_mask.view(B_mask, 1, 1, total_len) attn_mask = attn_mask.expand(B_mask, 1, T, total_len) attn_mask = attn_mask.expand(-1, self.n_heads, -1, -1) elif attn_mask.dim() == 4: # Already 4D mask - ensure it's the right shape B_mask = attn_mask.size(0) if attn_mask.size(-2) != T: # Slice to match query length if needed attn_mask = attn_mask[..., -T:, :] # Ensure head dimension matches if attn_mask.size(1) == 1: attn_mask = attn_mask.expand(-1, self.n_heads, -1, -1) elif attn_mask.size(1) != self.n_heads: raise ValueError(f"Mask head dimension {attn_mask.size(1)} doesn't match n_heads {self.n_heads}") else: raise ValueError(f"Unsupported attention_mask dimension: {attn_mask.dim()}, expected 2 or 4") att = att.masked_fill(~attn_mask, float('-inf')) else: # Apply causal mask - created on-the-fly (memory efficient, scales to any length) # torch.tril() is fast and doesn't require storing large buffers # This approach works for 32k, 1M, or any context length causal_mask = torch.tril(torch.ones(T, total_len, device=x.device, dtype=torch.bool)) att = att.masked_fill(~causal_mask.view(1, 1, T, total_len), float('-inf')) att = F.softmax(att, dim=-1) att = self.attn_dropout(att) y = att @ v # (B, nh, T, total_len) x (B, nh, total_len, hs) -> (B, nh, T, hs) y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side # output projection y = self.resid_dropout(self.c_proj(y)) # Return cache if requested if use_cache: return y, (k.detach(), v.detach()) return y class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) self.gelu = nn.GELU() self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) self.dropout = nn.Dropout(config.dropout) def forward(self, x): x = self.c_fc(x) x = self.gelu(x) x = self.c_proj(x) x = self.dropout(x) return x class BlockJ(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = LayerNorm(config.n_embd, bias=config.bias) self.j = LayerNorm(config.n_embd, config.n_embd) self.attn = CausalSelfAttention(config) self.ln_2 = LayerNorm(config.n_embd, bias=config.bias) # Use MoE if configured, otherwise use dense MLP if getattr(config, 'use_moe', False): self.mlp = MoE( num_experts_per_tok=config.num_experts_per_tok, num_experts=config.num_experts, emb_dim=config.n_embd, moe_dim=config.moe_dim, dropout=config.dropout ) self.use_moe = True else: self.mlp = MLP(config) self.use_moe = False def forward(self, x, attn_mask=None, past_key_value=None, use_cache=False): """ Forward pass with optional KV cache support. Args: x: (B, T, C) input embeddings attn_mask: Optional attention mask past_key_value: Optional tuple of (past_k, past_v) for attention layer use_cache: Whether to return cache for next step Returns: If use_cache: (output, (k, v)) where output is (B, T, C) Else: output (B, T, C) """ h = x x_ln = self.ln_1(x) # Attention with optional KV cache if use_cache: attn_out, new_past = self.attn(x_ln, attn_mask=attn_mask, past_key_value=past_key_value, use_cache=True) x = h + attn_out + self.j(x_ln) else: attn_out = self.attn(x_ln, attn_mask=attn_mask, past_key_value=past_key_value, use_cache=False) x = h + attn_out + self.j(x_ln) x = x + self.mlp(self.ln_2(x)) if use_cache: return x, new_past return x class MoE(nn.Module): """ An MoE layer with MLP block with swiglue activation function. Optimized for production workflows with proper initialization and dropout support. """ def __init__(self, num_experts_per_tok: int, num_experts: int, emb_dim: int, moe_dim: int, dropout: float = 0.0, dtype=torch.float32): super().__init__() self.k = int(num_experts_per_tok) self.E = int(num_experts) self.D = int(emb_dim) self.H = int(moe_dim) self.dropout = dropout self.gate = nn.Linear(self.D, self.E, bias=False, dtype=dtype) # use gate variable bcause couldnt load from checkpoint # Match MLP structure: c_fc -> GELU -> c_proj self.fc_bank = nn.Parameter(torch.empty(self.E, self.D, self.H, dtype=dtype)) # Equivalent to c_fc: (n_embd -> 4*n_embd) self.proj_bank = nn.Parameter(torch.empty(self.E, self.H, self.D, dtype=dtype)) # Equivalent to c_proj: (4*n_embd -> n_embd) self.gelu = nn.GELU() # Match MLP activation self.dropout_layer = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() # Initialize parameters self._init_parameters() def expert_utilization(self, logits): """ This function compute expert utilization per token and also compute load balancer loss. Details of this load balancer can be found in https://arxiv.org/abs/2101.03961 """ _, selected = logits.topk(self.k, dim=-1) selected = F.one_hot(selected, num_classes=self.E).sum(dim=2) # B, T, E load = torch.mean(selected.float(), dim=(0,1)) # average router probability per expert P = torch.softmax(logits, dim=-1).float().mean(dim=(0,1)) # [E] self._router_probs = P.detach() # per-expert avg prob self._aux_lb = self.E * torch.sum(load * P) self._expert_utilization = load def _init_parameters(self): """Initialize MoE parameters following standard practices.""" # Initialize gate with small values to start with uniform routing nn.init.normal_(self.gate.weight, mean=0.0, std=0.02) # Initialize expert banks to match MLP initialization # fc_bank: standard normal (like c_fc in MLP) nn.init.normal_(self.fc_bank, mean=0.0, std=0.02) # proj_bank: smaller initialization for stability (like c_proj in MLP) nn.init.normal_(self.proj_bank, mean=0.0, std=0.02 / math.sqrt(2)) def forward(self, x): B, T, D = x.shape assert D == self.D, f"Expected emb_dim={self.D}, got {D}" logits = self.gate(x) # B, T, E if self.training: logits = logits + torch.randn_like(logits) * 1e-1 topk_logits, selected = logits.topk(self.k, dim=-1) topk_probs = F.softmax(topk_logits, dim=-1) # Match MLP structure exactly: c_fc -> GELU -> c_proj # Step 1: c_fc equivalent: x @ fc_bank -> (B, T, E, H) h = torch.einsum("btd,edh->bteh", x, self.fc_bank) # B, T, E, H # Step 2: GELU activation (matching MLP) h = self.gelu(h) # B, T, E, H # Step 3: c_proj equivalent: h @ proj_bank -> (B, T, E, D) y = torch.einsum("bteh,ehd->bted", h, self.proj_bank) # B, T, E, D # Step 4: Select top-k experts and combine gather_idx = selected.view(B, T, -1, 1).expand(-1, -1, -1, self.D) # B, T, K, D y = torch.gather(y, dim=2, index=gather_idx) # B, T, K, D # Step 5: Weighted sum of selected experts y = (y * topk_probs.unsqueeze(-1)).sum(dim=2) # B, T, D # Step 6: Apply dropout like MLP y = self.dropout_layer(y) self.expert_utilization(logits) return y class GPTJXMoEForCausalLM(PreTrainedModel): config_class = GPTJXMoEConfig base_model_prefix = "transformer" is_parallelizable = True supports_gradient_checkpointing = True _no_split_modules = ["BlockJ"] # _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) assert config.vocab_size is not None assert config.block_size is not None self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), wpe = nn.Embedding(config.block_size, config.n_embd), drop = nn.Dropout(config.dropout), h = nn.ModuleList([BlockJ(config) for _ in range(config.n_layer)]), ln_f = LayerNorm(config.n_embd, bias=config.bias), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight # No need to store causal mask buffer - masks are created on-the-fly when needed # Flash Attention handles causality internally with is_causal=True # For manual attention, torch.tril() creates masks efficiently on-the-fly # This approach scales to any context length (1M+ tokens) without memory overhead self.apply(self._init_weights) for pn, p in self.named_parameters(): if pn.endswith('c_proj.weight'): torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer)) print("number of parameters: %.2fM" % (self.get_num_params()/1e6,)) def get_num_params(self, non_embedding=True): """ Return the number of parameters in the model. For non-embedding count (default), the position embeddings get subtracted. The token embeddings would too, except due to the parameter sharing these params are actually used as weights in the final layer, so we include them. """ n_params = sum(p.numel() for p in self.parameters()) if non_embedding: n_params -= self.transformer.wpe.weight.numel() return n_params def get_expert_utilization(self): """ Get expert utilization statistics for MoE layers. Returns expert utilization per layer and load balancing loss. Only works when use_moe=True in config. """ if not getattr(self.config, 'use_moe', False): return None, None lb_loss, expert_utilization_per_layer = 0, [] moe_layers = 0 for block in self.transformer.h: if hasattr(block, 'use_moe') and block.use_moe and hasattr(block.mlp, '_aux_lb'): lb_loss += block.mlp._aux_lb expert_utilization_per_layer.append(block.mlp._expert_utilization.detach().cpu()) moe_layers += 1 if moe_layers > 0: lb_loss = lb_loss / moe_layers return expert_utilization_per_layer, lb_loss def get_input_embeddings(self): return self.transformer.wte def set_input_embeddings(self, new_embeddings): self.transformer.wte = new_embeddings def forward( self, input_ids, targets=None, attn_mask=None, attention_mask=None, # HF standard name past_key_values=None, position_ids=None, use_cache=None, output_hidden_states: Optional[bool] = None, **kwargs ): """ Forward pass with KV cache support for efficient generation. Args: input_ids: (B, T) Token indices targets: Optional (B, T) target token indices for training attn_mask: Optional attention mask (legacy name) attention_mask: Optional attention mask (HF standard name, takes precedence) past_key_values: Optional list of (k, v) tuples from previous steps for KV cache position_ids: Optional (B, T) position indices (if None, computed from past_key_values) use_cache: Whether to return past_key_values for next step (defaults to config.use_kv_cache) output_hidden_states: Whether to return hidden states Returns: CausalLMOutputWithPast with logits and optionally past_key_values """ device = input_ids.device b, t = input_ids.size() # Use attention_mask if provided (HF standard), otherwise fall back to attn_mask if attention_mask is not None: attn_mask = attention_mask # Determine if we're using KV cache use_kv_cache = use_cache if use_cache is not None else getattr(self.config, 'use_kv_cache', False) # Compute past sequence length if using cache past_len = 0 if past_key_values is not None: past_len = past_key_values[0][0].size(2) if len(past_key_values) > 0 else 0 # Handle position_ids if position_ids is None: # Compute position IDs: from past_len to past_len + t pos = torch.arange(past_len, past_len + t, dtype=torch.long, device=device) else: pos = position_ids # Validate sequence length total_len = past_len + t assert total_len <= self.config.block_size, f"Cannot forward sequence of length {total_len}, block size is only {self.config.block_size}" # forward the GPT model itself tok_emb = self.transformer.wte(input_ids) # token embeddings of shape (b, t, n_embd) # Handle position embeddings: wpe expects 1D position indices if pos.dim() == 2: # If position_ids is 2D (B, T), extract first row (assuming all sequences have same positions) pos_1d = pos[0] if pos.size(0) > 0 else pos.squeeze(0) else: pos_1d = pos pos_emb = self.transformer.wpe(pos_1d) # position embeddings of shape (t, n_embd) if pos_emb.dim() == 2: pos_emb = pos_emb.unsqueeze(0).expand(b, -1, -1) # Expand to (b, t, n_embd) x = self.transformer.drop(tok_emb + pos_emb) # Expand attention_mask to cover full sequence (past + current) if needed # HF's generation API may provide mask only for current tokens if attn_mask is not None and past_key_values is not None and use_kv_cache: # Check if mask needs expansion if attn_mask.dim() == 2: mask_len = attn_mask.size(1) if mask_len == t and total_len > t: # Mask only covers current tokens, expand with ones for past tokens past_len = total_len - t past_mask = torch.ones(b, past_len, device=device, dtype=attn_mask.dtype) attn_mask = torch.cat([past_mask, attn_mask], dim=1) # Process through transformer layers with KV cache new_past_key_values = [] if use_kv_cache else None for i, block in enumerate(self.transformer.h): layer_past = past_key_values[i] if past_key_values is not None else None if use_kv_cache: x, new_past = block(x, attn_mask=attn_mask, past_key_value=layer_past, use_cache=True) new_past_key_values.append(new_past) else: x = block(x, attn_mask=attn_mask, past_key_value=layer_past, use_cache=False) x = self.transformer.ln_f(x) # Compute logits and loss if targets is not None: # Training: compute logits for all positions logits = self.lm_head(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100) else: # Inference: only compute logits for last position when using cache, all positions otherwise if use_kv_cache and past_key_values is not None: logits = self.lm_head(x[:, [-1], :]) # Only last token else: logits = self.lm_head(x) # All tokens loss = None return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=tuple(new_past_key_values) if use_kv_cache else None, hidden_states=x if output_hidden_states else None, attentions=None, ) def prepare_inputs_for_generation( self, input_ids, attention_mask=None, past_key_values=None, position_ids=None, use_cache=None, **kwargs ): """ Prepare inputs for generation with KV cache support. This method is called by HF's generation API. """ # Determine if we should use cache use_kv_cache = use_cache if use_cache is not None else getattr(self.config, 'use_kv_cache', False) # Base model inputs model_inputs = { "input_ids": input_ids, } # ---- 1. Handle KV cache (past_key_values) ---- if past_key_values is not None and use_kv_cache: # Only feed the last token when using cached keys/values model_inputs["input_ids"] = input_ids[:, -1:] model_inputs["past_key_values"] = past_key_values # ---- 2. Handle attention mask ---- if attention_mask is not None: # When using cache, attention_mask should cover the full sequence (past + current) if past_key_values is not None and use_kv_cache: # Extend attention mask to include past tokens # HF generation will handle this, but we ensure it's passed through pass model_inputs["attention_mask"] = attention_mask # ---- 3. Handle position_ids correctly ---- # HF relies on this for models like GPT-J, GPT-NeoX, Llama, etc. if position_ids is not None: if past_key_values is not None and use_kv_cache: # Only use the last position when using cache position_ids = position_ids[:, -1].unsqueeze(-1) model_inputs["position_ids"] = position_ids elif past_key_values is not None and use_kv_cache: # Compute position_ids from past_key_values length past_len = past_key_values[0][0].size(2) if len(past_key_values) > 0 else 0 model_inputs["position_ids"] = torch.tensor([[past_len]], device=input_ids.device, dtype=torch.long) # ---- 4. Forward arbitrary extra kwargs safely ---- # For example: use_cache, output_attentions, token_type_ids, etc. if use_cache is not None: model_inputs["use_cache"] = use_cache for k, v in kwargs.items(): if v is not None: model_inputs[k] = v return model_inputs def _reorder_cache( self, past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], beam_idx: torch.Tensor, ) -> List[Tuple[torch.Tensor, torch.Tensor]]: """ Reorder cache for beam search. Required by HF for beam search to work correctly. Selects which beam samples to keep based on beam_idx. Args: past_key_values: List of (k, v) tuples from previous steps beam_idx: (batch_size,) tensor indicating which beams to keep Returns: Reordered past_key_values """ reordered_past = [] for layer_past in past_key_values: k, v = layer_past device = k.device beam_idx_dev = beam_idx.to(device) reordered_past.append(( k.index_select(0, beam_idx_dev), v.index_select(0, beam_idx_dev) )) return reordered_past def crop_block_size(self, block_size): assert block_size <= self.config.block_size self.config.block_size = block_size self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size]) for block in self.transformer.h: if hasattr(block.attn, 'bias'): block.attn.bias = block.attn.bias[:,:,:block_size,:block_size] def load_dense_weights_into_moe(self, dense_state_dict, strict=False): """ Migrate Dense MLP weights to MoE experts. Ensures exact mathematical equivalence by cloning weights/biases to ALL experts. """ if not getattr(self.config, 'use_moe', False): return self.load_state_dict(dense_state_dict, strict=strict) print("Converting Dense Checkpoint -> MoE Checkpoint...") moe_state_dict = {} # Get config details num_experts = self.config.num_experts moe_dim = self.config.moe_dim for key, value in dense_state_dict.items(): # Identify MLP weights if 'mlp.c_fc' in key or 'mlp.c_proj' in key: # Extract layer index and type (weight/bias) # key format: transformer.h.{i}.mlp.c_fc.{weight/bias} parts = key.split('.') layer_idx = parts[2] layer_key_prefix = f"transformer.h.{layer_idx}.mlp" is_bias = 'bias' in key is_fc = 'c_fc' in key # --- Handle c_fc (Input -> Hidden) --- if is_fc: if not is_bias: # Weight: Dense is (H, D) -> MoE needs (E, D, H) # 1. Transpose to (D, H) w_T = value.t() # 2. Slice to moe_dim if necessary w_T = w_T[:, :moe_dim] # 3. Expand to (E, D, H) new_val = w_T.unsqueeze(0).expand(num_experts, -1, -1).clone() moe_state_dict[f"{layer_key_prefix}.fc_bank"] = new_val else: # Bias: Dense is (H) -> MoE needs (E, H) b = value[:moe_dim] new_val = b.unsqueeze(0).expand(num_experts, -1).clone() moe_state_dict[f"{layer_key_prefix}.fc_bias"] = new_val # --- Handle c_proj (Hidden -> Output) --- else: if not is_bias: # Weight: Dense is (D, H) -> MoE needs (E, H, D) # 1. Transpose to (H, D) w_T = value.t() # 2. Slice source dimension (H) if necessary w_T = w_T[:moe_dim, :] # 3. Expand to (E, H, D) new_val = w_T.unsqueeze(0).expand(num_experts, -1, -1).clone() moe_state_dict[f"{layer_key_prefix}.proj_bank"] = new_val else: # Bias: Dense is (D) -> MoE needs (E, D) # Bias is on the output, so dimension is D, usually doesn't need slicing new_val = value.unsqueeze(0).expand(num_experts, -1).clone() moe_state_dict[f"{layer_key_prefix}.proj_bias"] = new_val # --- Initialize Gate (if not yet initialized) --- # We initialize gate to zero to ensure uniform routing probability initially, # which guarantees average of identical experts == single expert. gate_key = f"{layer_key_prefix}.gate.weight" if gate_key not in moe_state_dict: # Zeros = equal probability for all experts moe_state_dict[gate_key] = torch.zeros(num_experts, self.config.n_embd) else: # Copy non-MLP keys directly (Attn, LayerNorm, Embeddings) moe_state_dict[key] = value print("Loading constructed state dict...") return self.load_state_dict(moe_state_dict, strict=strict) AutoConfig.register("sabiyarn", GPTJXMoEConfig) AutoModel.register(GPTJXMoEConfig,GPTJXMoEForCausalLM) AutoModelForCausalLM.register(GPTJXMoEConfig, GPTJXMoEForCausalLM)