Update modeling_smartcoder_moe.py

modeling_smartcoder_moe.py

@@ -13,6 +13,17 @@ 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
@@ -71,18 +82,23 @@ 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))
-        self.register_buffer("inv_freq", inv_freq)
+        # 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._cached_len = 0
 
     def _build_cache(self, seq_len, device):
@@ -105,20 +121,32 @@ 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)
@@ -135,15 +163,17 @@ 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
@@ -156,26 +186,27 @@ 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)
-        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)
+        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)
 
     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)
@@ -183,13 +214,16 @@ 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   = self.experts_fc[expert_ids]
-            proj_w = self.experts_proj[expert_ids]
+            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]
             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
@@ -199,10 +233,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
@@ -223,8 +257,10 @@ 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)
@@ -241,20 +277,25 @@ 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()
 
-    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)
+    # 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 get_input_embeddings(self): return self.model.embed_tokens
-    def get_output_embeddings(self): return self.lm_head
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
 
     def forward(
         self,
@@ -285,7 +326,7 @@ class SmartCoderMoEForCausalLM(PreTrainedModel, GenerationMixin):
         return {"input_ids": input_ids}
 
 
-# ── Loader ────────────────────────────────────────────────────────────────────
+# ── Loader (standalone helper) ────────────────────────────────────────────────
 def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloat16):
     import os
     from huggingface_hub import snapshot_download
@@ -307,27 +348,26 @@ def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloa
     for f in sf_files:
         state_dict.update(load_file(str(f)))
 
-    # 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
-
+    # Keys now match natively (experts_fc.weight / experts_proj.weight) — no remap.
     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: {missing[:3]}{'...' if len(missing)>3 else ''}")
+        print(f"Missing (unexpected!): {missing[:5]}{'...' if len(missing) > 5 else ''}")
+    else:
+        print("No unexpected missing keys.")
     if unexpected:
-        print(f"Unexpected: {unexpected[:3]}{'...' if len(unexpected)>3 else ''}")
+        print(f"Unexpected: {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}")
+    else:
+        print("No unexpected keys.")
 
     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)