""" 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}