Update modeling_smartcoder_moe.py

modeling_smartcoder_moe.py

@@ -1,4 +1,4 @@
-"""
+""
 modeling_smartcoder_moe.py
 Custom model class for SmartCoderMoE.
 
@@ -13,18 +13,7 @@ Architecture (from tensor inspection):
     router:       [32, 2048]      router logits
 - LayerNorm: weight+bias (input_layernorm, post_attention_layernorm)
 - Final norm: model.norm.weight/bias
-
-NOTE ON THE EXPERT KEYS:
-The checkpoint stores the batched expert tensors as `experts_fc.weight` /
-`experts_proj.weight` (i.e. they were saved as a module holding a `.weight`
-Parameter). The previous version of this file declared them as bare
-`nn.Parameter` (`experts_fc`, no `.weight`) and tried to paper over the mismatch
-with a `_load_from_state_dict` remap hook — which does NOT fire correctly under
-`from_pretrained`, producing UNEXPECTED `experts_fc.weight` / MISSING
-`experts_fc`. The fix here is structural: wrap each batched expert tensor in a
-tiny `ExpertWeight` module so its natural state_dict key is `experts_fc.weight`,
-matching the checkpoint exactly. No load hook required.
-"""
+""
 
 import math
 import torch
@@ -82,23 +71,18 @@ class SmartCoderMoEConfig(PretrainedConfig):
 
 # ── RoPE ──────────────────────────────────────────────────────────────────────
 def rotate_half(x):
-    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
     return torch.cat([-x2, x1], dim=-1)
 
-
 def apply_rotary_emb(q, k, cos, sin):
     return (q * cos) + (rotate_half(q) * sin), \
            (k * cos) + (rotate_half(k) * sin)
 
-
 class RotaryEmbedding(nn.Module):
     def __init__(self, dim, max_pos=16384, base=10000.0):
         super().__init__()
         inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        # persistent=False: inv_freq is a computed constant, NOT a trained weight.
-        # The checkpoint correctly omits it; marking it non-persistent means it is
-        # not expected in the state_dict, killing the spurious MISSING warning.
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("inv_freq", inv_freq)
         self._cached_len = 0
 
     def _build_cache(self, seq_len, device):
@@ -121,32 +105,20 @@ class LayerNormWithBias(nn.Module):
     def __init__(self, hidden_size, eps=1e-5):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.bias   = nn.Parameter(torch.zeros(hidden_size))
         self.eps = eps
 
     def forward(self, x):
         return F.layer_norm(x, x.shape[-1:], self.weight, self.bias, self.eps)
 
 
-# ── Expert weight holder ──────────────────────────────────────────────────────
-class ExpertWeight(nn.Module):
-    """
-    Thin wrapper so a batched expert tensor registers under the key
-    `<name>.weight`, matching how the checkpoint saved it. The forward indexes
-    `.weight` directly; this module exists purely to produce the right key.
-    """
-    def __init__(self, *shape):
-        super().__init__()
-        self.weight = nn.Parameter(torch.empty(*shape))
-
-
 # ── Attention ─────────────────────────────────────────────────────────────────
 class SmartCoderAttention(nn.Module):
     def __init__(self, config: SmartCoderMoEConfig):
         super().__init__()
-        self.num_heads = config.num_attention_heads
+        self.num_heads    = config.num_attention_heads
         self.num_kv_heads = config.num_key_value_heads
-        self.head_dim = config.head_dim
+        self.head_dim     = config.head_dim
         self.num_kv_groups = self.num_heads // self.num_kv_heads
 
         self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * config.head_dim, bias=True)
@@ -163,17 +135,15 @@ class SmartCoderAttention(nn.Module):
         v = self.v_proj(hidden_states).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
 
         cos, sin = self.rotary_emb(T, hidden_states.device)
-        cos = cos[:, :, :T, : self.head_dim]
-        sin = sin[:, :, :T, : self.head_dim]
+        cos = cos[:, :, :T, :self.head_dim]
+        sin = sin[:, :, :T, :self.head_dim]
         q, k = apply_rotary_emb(q, k, cos, sin)
 
         k = k.repeat_interleave(self.num_kv_groups, dim=1)
         v = v.repeat_interleave(self.num_kv_groups, dim=1)
 
         attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        causal = torch.triu(
-            torch.full((T, T), float("-inf"), device=q.device, dtype=q.dtype), diagonal=1
-        )
+        causal = torch.triu(torch.full((T, T), float("-inf"), device=q.device, dtype=q.dtype), diagonal=1)
         attn = attn + causal.unsqueeze(0).unsqueeze(0)
         if attention_mask is not None:
             attn = attn + attention_mask
@@ -186,27 +156,26 @@ class SmartCoderAttention(nn.Module):
 class SmartCoderMoEMLP(nn.Module):
     def __init__(self, config: SmartCoderMoEConfig):
         super().__init__()
-        H = config.hidden_size
+        H  = config.hidden_size
         DI = config.dense_intermediate_size
         NE = config.num_experts
         EI = config.expert_intermediate_size
 
         self.num_experts = NE
-        self.top_k = config.num_experts_per_tok
+        self.top_k       = config.num_experts_per_tok
 
-        self.dense_fc = nn.Linear(H, DI, bias=True)
-        self.dense_proj = nn.Linear(DI, H, bias=True)
-        # Wrapped so keys are experts_fc.weight / experts_proj.weight (matches ckpt)
-        self.experts_fc = ExpertWeight(NE, EI, H)
-        self.experts_proj = ExpertWeight(NE, H, EI)
-        self.router = nn.Linear(H, NE, bias=False)
+        self.dense_fc    = nn.Linear(H, DI, bias=True)
+        self.dense_proj  = nn.Linear(DI, H, bias=True)
+        self.experts_fc  = nn.Parameter(torch.empty(NE, EI, H))
+        self.experts_proj = nn.Parameter(torch.empty(NE, H, EI))
+        self.router      = nn.Linear(H, NE, bias=False)
 
     def forward(self, x):
         B, T, H = x.shape
 
         dense_out = self.dense_proj(F.gelu(self.dense_fc(x)))
 
-        router_logits = self.router(x)
+        router_logits  = self.router(x)
         router_weights = F.softmax(router_logits, dim=-1)
         top_weights, top_indices = router_weights.topk(self.top_k, dim=-1)
         top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)
@@ -214,16 +183,13 @@ class SmartCoderMoEMLP(nn.Module):
         expert_out = torch.zeros_like(x)
         x_flat = x.view(B * T, H)
 
-        fc_all = self.experts_fc.weight       # [NE, EI, H]
-        proj_all = self.experts_proj.weight   # [NE, H, EI]
-
         for k in range(self.top_k):
             expert_ids = top_indices[:, :, k].reshape(B * T)
-            weights = top_weights[:, :, k].reshape(B * T, 1)
-            fc_w = fc_all[expert_ids]                       # [BT, EI, H]
-            proj_w = proj_all[expert_ids]                   # [BT, H, EI]
+            weights    = top_weights[:, :, k].reshape(B * T, 1)
+            fc_w   = self.experts_fc[expert_ids]
+            proj_w = self.experts_proj[expert_ids]
             hidden = F.gelu(torch.bmm(fc_w, x_flat.unsqueeze(-1)).squeeze(-1))
-            out = torch.bmm(proj_w, hidden.unsqueeze(-1)).squeeze(-1)
+            out    = torch.bmm(proj_w, hidden.unsqueeze(-1)).squeeze(-1)
             expert_out = expert_out + (out * weights).view(B, T, H)
 
         return dense_out + expert_out
@@ -233,10 +199,10 @@ class SmartCoderMoEMLP(nn.Module):
 class SmartCoderDecoderLayer(nn.Module):
     def __init__(self, config: SmartCoderMoEConfig):
         super().__init__()
-        self.input_layernorm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
-        self.self_attn = SmartCoderAttention(config)
+        self.input_layernorm          = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
+        self.self_attn                = SmartCoderAttention(config)
         self.post_attention_layernorm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
-        self.mlp = SmartCoderMoEMLP(config)
+        self.mlp                      = SmartCoderMoEMLP(config)
 
     def forward(self, hidden_states, attention_mask=None, **kwargs):
         residual = hidden_states
@@ -257,10 +223,8 @@ class SmartCoderMoEModel(nn.Module):
     def __init__(self, config: SmartCoderMoEConfig):
         super().__init__()
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList(
-            [SmartCoderDecoderLayer(config) for _ in range(config.num_hidden_layers)]
-        )
-        self.norm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
+        self.layers = nn.ModuleList([SmartCoderDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm   = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
 
     def forward(self, input_ids, attention_mask=None, **kwargs):
         hidden_states = self.embed_tokens(input_ids)
@@ -277,25 +241,20 @@ class SmartCoderMoEForCausalLM(PreTrainedModel, GenerationMixin):
 
     def __init__(self, config: SmartCoderMoEConfig):
         super().__init__(config)
-        self.model = SmartCoderMoEModel(config)
+        self.model   = SmartCoderMoEModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.post_init()
 
-    # No _load_from_state_dict / load_state_dict override needed:
-    # the module structure now produces experts_fc.weight / experts_proj.weight
-    # natively, matching the checkpoint, and inv_freq is non-persistent.
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        remapped = {}
+        for k, v in state_dict.items():
+            k = k.replace('experts_fc.weight', 'experts_fc')
+            k = k.replace('experts_proj.weight', 'experts_proj')
+            remapped[k] = v
+        super()._load_from_state_dict(remapped, prefix, *args, **kwargs)
 
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
+    def get_input_embeddings(self): return self.model.embed_tokens
+    def get_output_embeddings(self): return self.lm_head
 
     def forward(
         self,
@@ -326,7 +285,7 @@ class SmartCoderMoEForCausalLM(PreTrainedModel, GenerationMixin):
         return {"input_ids": input_ids}
 
 
-# ── Loader (standalone helper) ────────────────────────────────────────────────
+# ── Loader ────────────────────────────────────────────────────────────────────
 def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloat16):
     import os
     from huggingface_hub import snapshot_download
@@ -348,26 +307,27 @@ def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloa
     for f in sf_files:
         state_dict.update(load_file(str(f)))
 
-    # Keys now match natively (experts_fc.weight / experts_proj.weight) — no remap.
+    # Remap expert keys — safetensors has .weight suffix, our params don't
+    remapped = {}
+    for k, v in state_dict.items():
+        if 'experts_fc.weight' in k:
+            remapped[k.replace('experts_fc.weight', 'experts_fc')] = v
+        elif 'experts_proj.weight' in k:
+            remapped[k.replace('experts_proj.weight', 'experts_proj')] = v
+        else:
+            remapped[k] = v
+    state_dict = remapped
+
     missing, unexpected = model.load_state_dict(state_dict, strict=False)
-    # inv_freq is non-persistent so it will show as "missing"; that is expected
-    # and harmless (recomputed at init). Filter it from the report.
-    missing = [m for m in missing if "inv_freq" not in m]
     if missing:
-        print(f"Missing (unexpected!): {missing[:5]}{'...' if len(missing) > 5 else ''}")
-    else:
-        print("No unexpected missing keys.")
+        print(f"Missing: {missing[:3]}{'...' if len(missing)>3 else ''}")
     if unexpected:
-        print(f"Unexpected: {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}")
-    else:
-        print("No unexpected keys.")
+        print(f"Unexpected: {unexpected[:3]}{'...' if len(unexpected)>3 else ''}")
 
     model = model.to(dtype)
-    print(f"Loaded! Params: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B")
+    print(f"Loaded! Params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
     return model, config
 
-
 from transformers import AutoConfig, AutoModelForCausalLM
-
 AutoConfig.register("smartcoder_moe", SmartCoderMoEConfig)
 AutoModelForCausalLM.register(SmartCoderMoEConfig, SmartCoderMoEForCausalLM)