modeling_smartcoder_moe.py

#modeling_smartcoder_moe.py

#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
import torch.utils.checkpoint
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast


# ── 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._cached_len = 0

    def _build_cache(self, seq_len, device):
        t = torch.arange(seq_len, device=device).float()
        freqs = torch.outer(t, self.inv_freq.to(device))
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
        self._cached_len = seq_len

    def forward(self, seq_len, device):
        if seq_len > self._cached_len:
            self._build_cache(seq_len, device)
        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.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, **kwargs):
        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, hidden_states.device)
        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)
        attn = attn + causal.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).transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)


# ── MoE MLP ───────────────────────────────────────────────────────────────────
class SmartCoderMoEMLP(nn.Module):
    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

        self.num_experts = NE
        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)

    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
        # Checkpoint stores expert weights with a '.weight' suffix (as if
        # experts_fc/experts_proj were nn.Linear submodules), but they're
        # raw nn.Parameter tensors here (no '.weight' child -- needed for
        # batched bmm across all experts at once, see forward() below).
        # PyTorch's load_state_dict() recursion calls _load_from_state_dict
        # on EVERY submodule in the tree directly (using each module's own
        # class method, not a parent class's override) -- so the remap has
        # to live on THIS class, not on SmartCoderMoEForCausalLM. The
        # previous override sat on the top-level CausalLM class and only
        # ever fired for its own direct params/buffers (it has none), never
        # for this module's recursive call -- silently skipping every
        # expert tensor. That's the actual bug.
        remapped = {}
        for k, v in state_dict.items():
            if k == prefix + "experts_fc.weight":
                remapped[prefix + "experts_fc"] = v
            elif k == prefix + "experts_proj.weight":
                remapped[prefix + "experts_proj"] = v
            else:
                remapped[k] = v
        super()._load_from_state_dict(remapped, prefix, *args, **kwargs)

    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_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)

        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)
            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)
            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, **kwargs):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask)
        hidden_states = residual + hidden_states

        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


# ── 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)
        # Required for transformers' _set_gradient_checkpointing() to have
        # something to toggle. Without this attribute + the checkpoint()
        # call in forward(), declaring supports_gradient_checkpointing=True
        # at the PreTrainedModel level is a lie transformers will catch and
        # raise on -- which is exactly the ValueError this fixes.
        self.gradient_checkpointing = True

    def forward(self, input_ids, attention_mask=None, **kwargs):
        hidden_states = self.embed_tokens(input_ids)
        for layer in self.layers:
            if self.gradient_checkpointing and self.training:
                hidden_states = torch.utils.checkpoint.checkpoint(
                    layer, hidden_states, attention_mask, use_reentrant=True
                )
            else:
                hidden_states = layer(hidden_states, attention_mask=attention_mask)
        return self.norm(hidden_states)


# ── CausalLM ──────────────────────────────────────────────────────────────────
class SmartCoderMoEForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = SmartCoderMoEConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True

    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

    # transformers' _set_gradient_checkpointing (called by Unsloth/Trainer)
    # looks for this attribute on the *PreTrainedModel* root, finds the
    # submodule that has it, and toggles it. Exposing it here as a property
    # delegating to self.model keeps both objects in sync.
    @property
    def gradient_checkpointing(self):
        return self.model.gradient_checkpointing

    @gradient_checkpointing.setter
    def gradient_checkpointing(self, value):
        self.model.gradient_checkpointing = value

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        past_key_values=None,
        inputs_embeds=None,
        labels=None,
        use_cache=None,
        **kwargs,
    ):
        hidden_states = self.model(input_ids, attention_mask=attention_mask)
        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=None)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}


# ── Loader ────────────────────────────────────────────────────────────────────
def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloat16):
    import os
    from huggingface_hub import snapshot_download
    from safetensors.torch import load_file
    from pathlib import Path

    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...")
    sf_files = sorted(Path(model_dir).glob("*.safetensors"))
    state_dict = {}
    for f in sf_files:
        state_dict.update(load_file(str(f))))

    model = model.to(dtype)
    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)