Upload modeling_smartcoder.py

modeling_smartcoder.py

@@ -0,0 +1,413 @@
+"""
+modeling_smartcoder_moe.py
+Custom model class for SmartCoderMoE.
+
+Architecture (from tensor inspection):
+- vocab_size: 65536, hidden: 2048, layers: 40
+- Attention: q[2048,2048], k/v[512,2048] — 16 heads, 4 KV heads, head_dim=128
+- MLP (hybrid dense + MoE):
+    dense_fc:     [8192, 2048]    up
+    dense_proj:   [2048, 8192]    down
+    experts_fc:   [32, 512, 2048] expert up (batched)
+    experts_proj: [32, 2048, 512] expert down (batched)
+    router:       [32, 2048]      router logits
+- LayerNorm: weight+bias (input_layernorm, post_attention_layernorm)
+- Final norm: model.norm.weight/bias
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from typing import Optional, Tuple, List
+
+
+# ── Config ────────────────────────────────────────────────────────────────────
+class SmartCoderMoEConfig(PretrainedConfig):
+    model_type = "smartcoder_moe"
+
+    def __init__(
+        self,
+        vocab_size=65536,
+        hidden_size=2048,
+        num_hidden_layers=40,
+        num_attention_heads=16,
+        num_key_value_heads=4,
+        dense_intermediate_size=8192,
+        num_experts=32,
+        expert_intermediate_size=512,
+        num_experts_per_tok=2,
+        max_position_embeddings=16384,
+        rope_theta=10000.0,
+        rms_norm_eps=1e-5,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=0,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = hidden_size // num_attention_heads
+        self.dense_intermediate_size = dense_intermediate_size
+        self.num_experts = num_experts
+        self.expert_intermediate_size = expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        super().__init__(
+            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,
+        )
+
+
+# ── RoPE ──────────────────────────────────────────────────────────────────────
+def rotate_half(x):
+    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)
+        self.max_pos = max_pos
+        self._build_cache(max_pos)
+
+    def _build_cache(self, seq_len):
+        t = torch.arange(seq_len, device=self.inv_freq.device).float()
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :])
+
+    def forward(self, seq_len):
+        if seq_len > self.max_pos:
+            self._build_cache(seq_len)
+        return self.cos_cached[:, :, :seq_len, :], \
+               self.sin_cached[:, :, :seq_len, :]
+
+
+# ── LayerNorm (with bias) ─────────────────────────────────────────────────────
+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.eps = eps
+
+    def forward(self, x):
+        return F.layer_norm(x, x.shape[-1:], self.weight, self.bias, self.eps)
+
+
+# ── Attention ─────────────────────────────────────────────────────────────────
+class SmartCoderAttention(nn.Module):
+    def __init__(self, config: SmartCoderMoEConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads   = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        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)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=True)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=True)
+        self.o_proj = nn.Linear(config.num_attention_heads * config.head_dim, config.hidden_size, bias=True)
+
+        self.rotary_emb = RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta)
+
+    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
+        B, T, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        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)
+        cos = cos[:, :, :T, :self.head_dim]
+        sin = sin[:, :, :T, :self.head_dim]
+        q, k = apply_rotary_emb(q, k, cos, sin)
+
+        if past_key_value is not None:
+            k = torch.cat([past_key_value[0], k], dim=2)
+            v = torch.cat([past_key_value[1], v], dim=2)
+        present = (k, v) if use_cache else None
+
+        # Expand KV heads to match Q heads (GQA)
+        k = k.repeat_interleave(self.num_kv_groups, dim=1)
+        v = v.repeat_interleave(self.num_kv_groups, dim=1)
+
+        scale = math.sqrt(self.head_dim)
+        attn = torch.matmul(q, k.transpose(-2, -1)) / scale
+
+        kv_len = k.shape[2]
+        causal_mask = torch.triu(
+            torch.full((T, kv_len), float("-inf"), device=q.device, dtype=q.dtype),
+            diagonal=1 + kv_len - T
+        )
+        attn = attn + causal_mask.unsqueeze(0).unsqueeze(0)
+
+        if attention_mask is not None:
+            attn = attn + attention_mask
+
+        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.o_proj(out), present
+
+
+# ── MoE MLP ───────────────────────────────────────────────────────────────────
+class SmartCoderMoEMLP(nn.Module):
+    """
+    Hybrid Dense + MoE MLP.
+    dense path:   hidden -> dense_fc (8192) -> gelu -> dense_proj (2048)
+    expert path:  router picks top-k experts from experts_fc/experts_proj
+    output = dense_out + expert_out
+    """
+    def __init__(self, config: SmartCoderMoEConfig):
+        super().__init__()
+        H  = config.hidden_size
+        DI = config.dense_intermediate_size
+        NE = config.num_experts
+        EI = config.expert_intermediate_size
+        K  = config.num_experts_per_tok
+
+        self.num_experts     = NE
+        self.top_k           = K
+
+        # Dense residual path
+        self.dense_fc   = nn.Linear(H, DI, bias=True)
+        self.dense_proj = nn.Linear(DI, H, bias=True)
+
+        # MoE path — stored as batched weight matrices matching safetensors layout
+        # experts_fc:   [NE, EI, H]
+        # experts_proj: [NE, H, EI]
+        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 path
+        dense_out = self.dense_proj(F.gelu(self.dense_fc(x)))
+
+        # Router
+        router_logits = self.router(x)                          # [B, T, NE]
+        router_weights = F.softmax(router_logits, dim=-1)
+        top_weights, top_indices = router_weights.topk(self.top_k, dim=-1)  # [B, T, K]
+        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)   # normalize
+
+        # Expert computation — iterate over top-k (K is small so this is fine)
+        expert_out = torch.zeros_like(x)
+        x_flat = x.view(B * T, H)
+
+        for k in range(self.top_k):
+            expert_ids = top_indices[:, :, k].reshape(B * T)   # [B*T]
+            weights    = top_weights[:, :, k].reshape(B * T, 1) # [B*T, 1]
+
+            # Batched expert forward using einsum
+            # For each token, pick its expert's weights
+            # experts_fc: [NE, EI, H] → gather → [B*T, EI, H]
+            fc_w   = self.experts_fc[expert_ids]    # [B*T, EI, H]
+            proj_w = self.experts_proj[expert_ids]  # [B*T, H, EI]
+
+            # up: [B*T, EI]
+            hidden = F.gelu(torch.bmm(fc_w, x_flat.unsqueeze(-1)).squeeze(-1))
+            # down: [B*T, H]
+            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
+
+
+# ── Decoder Layer ─────────────────────────────────────────────────────────────
+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.post_attention_layernorm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
+        self.mlp                     = SmartCoderMoEMLP(config)
+
+    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
+        # Attention
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, present = self.self_attn(
+            hidden_states, attention_mask=attention_mask,
+            past_key_value=past_key_value, use_cache=use_cache
+        )
+        hidden_states = residual + hidden_states
+
+        # MLP
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states, present
+
+
+# ── Full Model ────────────────────────────────────────────────────────────────
+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)
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False):
+        hidden_states = self.embed_tokens(input_ids)
+        presents = [] if use_cache else None
+
+        for i, layer in enumerate(self.layers):
+            pkv = past_key_values[i] if past_key_values else None
+            hidden_states, present = layer(
+                hidden_states, attention_mask=attention_mask,
+                past_key_value=pkv, use_cache=use_cache
+            )
+            if use_cache:
+                presents.append(present)
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states, presents
+
+
+# ── CausalLM wrapper ──────────────────────────────────────────────────────────
+class SmartCoderMoEForCausalLM(PreTrainedModel):
+    config_class = SmartCoderMoEConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: SmartCoderMoEConfig):
+        super().__init__(config)
+        self.model   = SmartCoderMoEModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        hidden_states, presents = self.model(
+            input_ids, attention_mask=attention_mask,
+            past_key_values=past_key_values, use_cache=use_cache
+        )
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=presents,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+        return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}
+
+
+# ── Loader ────────────────────────────────────────────────────────────────────
+def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloat16):
+    """Load SmartCoderMoE with correct custom architecture."""
+    import os
+    from huggingface_hub import snapshot_download
+    from safetensors.torch import load_file
+
+    os.environ["HF_HUB_DISABLE_XET"] = "1"
+
+    print(f"Downloading {model_id}...")
+    model_dir = snapshot_download(model_id)
+
+    config = SmartCoderMoEConfig()
+    print("Initializing model...")
+    model = SmartCoderMoEForCausalLM(config)
+
+    print("Loading weights...")
+    from pathlib import Path
+    sf_files = sorted(Path(model_dir).glob("*.safetensors"))
+    state_dict = {}
+    for f in sf_files:
+        state_dict.update(load_file(str(f)))
+
+    # experts_fc in safetensors: [32, 512, 2048] — matches our [NE, EI, H] ✓
+    # experts_proj in safetensors: [32, 2048, 512] — matches our [NE, H, EI] ✓
+    # router in safetensors: [32, 2048] — stored as Linear weight [out, in] ✓
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        print(f"Missing keys: {missing[:5]}{'...' if len(missing)>5 else ''}")
+    if unexpected:
+        print(f"Unexpected keys: {unexpected[:5]}{'...' if len(unexpected)>5 else ''}")
+
+    model = model.to(dtype)
+    print(f"Loaded! Params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
+    return model, config
+
+
+if __name__ == "__main__":
+    from transformers import AutoTokenizer
+    import torch
+
+    model, config = load_smartcoder_moe()
+    model.eval()
+    model = model.cuda()
+
+    tokenizer = AutoTokenizer.from_pretrained("Johnblick187/SmartCoderMoE", trust_remote_code=True)
+
+    prompt = "def fibonacci(n):"
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    input_len = inputs["input_ids"].shape[-1]
+
+    with torch.no_grad():
+        out = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+            repetition_penalty=1.3,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    print(tokenizer.decode(out[0][input_len:], skip_special_tokens=True))