model/model.py

"""
Ivme-Coder-v1 (codename Otter 1) modeling file.
RoPE + SwiGLU + RMSNorm (pre-norm) decoder-only transformer, tied embeddings.
Load with: AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
from transformers.generation import GenerationMixin


class IvmeCoderConfig(PretrainedConfig):
    model_type = "ivme_coder"

    def __init__(
        self,
        vocab_size=16000,
        d_model=512,
        n_layer=12,
        n_head=8,
        context_len=1024,
        ffn_mult=8/3,
        rope_theta=10000.0,
        tie_word_embeddings=True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layer = n_layer
        self.n_head = n_head
        self.context_len = context_len
        self.ffn_mult = ffn_mult
        self.rope_theta = rope_theta
        # Standard HF attribute names, aliased to our own names. Recent `transformers`
        # generation/caching code (DynamicCache etc.) reads these directly off the
        # config regardless of architecture, even for trust_remote_code models.
        self.num_hidden_layers = n_layer
        self.num_attention_heads = n_head
        self.hidden_size = d_model
        self.use_cache = False  # this architecture does not implement KV caching -
                                 # forward() always recomputes the full sequence, see note below
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight


def precompute_rope(dim, max_len, theta=10000.0, device="cpu"):
    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(max_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    return torch.cos(freqs), torch.sin(freqs)


def apply_rope(x, cos, sin):
    T = x.size(2)
    # Cast cos/sin to match x's dtype. precompute_rope always builds these tables in
    # fp32 for precision, but if q/k arrive in bf16 (as they do when the model is
    # loaded with dtype=torch.bfloat16), multiplying against fp32 cos/sin silently
    # promotes the result back to fp32 - which then doesn't match `v` (which never
    # passes through this function and stays in bf16), and scaled_dot_product_attention
    # requires q, k, v to share one dtype.
    cos = cos[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
    sin = sin[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    rot1 = x1 * cos - x2 * sin
    rot2 = x1 * sin + x2 * cos
    return torch.stack([rot1, rot2], dim=-1).flatten(-2)


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.n_head = n_head
        self.d_head = d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, cos, sin):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)


class SwiGLU(nn.Module):
    def __init__(self, d_model, mult=8/3):
        super().__init__()
        hidden = int(d_model * mult)
        hidden = (hidden + 7) // 8 * 8
        self.w1 = nn.Linear(d_model, hidden, bias=False)
        self.w2 = nn.Linear(d_model, hidden, bias=False)
        self.w3 = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))


class Block(nn.Module):
    def __init__(self, d_model, n_head, ffn_mult):
        super().__init__()
        self.ln1 = RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head)
        self.ln2 = RMSNorm(d_model)
        self.mlp = SwiGLU(d_model, ffn_mult)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.ln1(x), cos, sin)
        x = x + self.mlp(self.ln2(x))
        return x


class IvmeCoderV1ForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = IvmeCoderConfig
    # Required so from_pretrained's tie-recovery logic knows how to reconnect
    # head.weight to tok_emb.weight when head.weight is absent from the checkpoint
    # (it's deliberately excluded from the safetensors file, since it's tied storage,
    # not distinct data). Without this, HF's loader treats the missing key as needing
    # fresh random initialization instead of re-tying it - which silently produces a
    # working-looking model with a completely untrained output head.
    _tied_weights_keys = {"head.weight": "tok_emb.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.context_len = config.context_len
        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList(
            [Block(config.d_model, config.n_head, config.ffn_mult) for _ in range(config.n_layer)]
        )
        self.ln_f = RMSNorm(config.d_model)
        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight  # tied embeddings

        # NOTE: RoPE cos/sin tables are intentionally NOT stored as a persistent=False
        # buffer here. transformers v5's from_pretrained() has a known bug where
        # persistent=False buffers get overwritten with uninitialized/garbage memory
        # after loading (https://github.com/huggingface/transformers/issues/44534),
        # even though they're computed correctly at __init__ time. That garbage then
        # silently produces huge (but finite) values through the RoPE rotation, which
        # overflow to NaN inside scaled_dot_product_attention. Recomputing the tables
        # fresh on every forward() call sidesteps this entirely - it's cheap (a cos/sin
        # over d_head/2 * context_len elements) relative to the rest of the forward pass.
        self.d_head = config.d_model // config.n_head
        self.rope_theta = config.rope_theta

        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def get_input_embeddings(self):
        return self.tok_emb

    def set_input_embeddings(self, value):
        self.tok_emb = value

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, value):
        self.head = value

    def forward(self, input_ids, labels=None, use_cache=False, past_key_values=None,
                attention_mask=None, **kwargs):
        # attention_mask is accepted for API compatibility with tokenizer output /
        # generate()'s kwarg validation, but intentionally unused: this model only
        # supports left-padding for batched generation (see model card), and single-
        # sequence causal generation (the common case) needs no mask at all.
        # Recompute RoPE tables fresh each call - see note in __init__ for why this
        # isn't cached in a buffer.
        rope_cos, rope_sin = precompute_rope(
            self.d_head, self.context_len, theta=self.rope_theta, device=input_ids.device
        )
        x = self.tok_emb(input_ids)
        for block in self.blocks:
            x = block(x, rope_cos, rope_sin)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        # This architecture has no KV cache implementation - every generation step
        # recomputes attention over the full (truncated) sequence from scratch.
        # This is correct but slower than a cached model; fine at 50M params / 1024 ctx.
        if input_ids.size(1) > self.context_len:
            input_ids = input_ids[:, -self.context_len:]
        return {"input_ids": input_ids, "use_cache": False}