""" Modeling file for Ivme-Conversate-S-v2-Instruct. Standard decoder-only Transformer architecture, deliberately matching Ivme-Conversate-v2-Base's proven recipe (pulled directly from its real config.json): tied embeddings, standard multi-head attention (no GQA, no DIFF), RoPE, SwiGLU, RMSNorm, pre-norm. No architectural novelty by design -- this model tests a DATA strategy (instruct-heavy, single-epoch pretraining) in isolation, on infrastructure already proven stable. Trained on ~900M tokens, single epoch, instruct-heavy mix (UltraChat-200k as the dominant 45% share, plus SODA, UltraInteract, orca-math, dolly-15k, sql-create-context) -- all permissively licensed (MIT/CC-BY/CC-BY-SA), no CC-BY-NC sources, matching v2-Base's Apache-2.0 license. Uses standard HF tied-embedding conventions (get_output_embeddings / set_output_embeddings + config.tie_word_embeddings), so PreTrainedModel's own tie_weights() machinery handles the tie correctly through from_pretrained -- more robust than manual weight assignment, since it's re-applied automatically by HF's own loading path rather than needing to survive it. """ import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutput try: from .configuration_ivme_s_v2_instruct import IvmeConversateSV2InstructConfig except ImportError: from configuration_ivme_s_v2_instruct import IvmeConversateSV2InstructConfig def build_rope_cache(dim, max_seq_len, base=10000.0): assert dim % 2 == 0 inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) t = torch.arange(max_seq_len).float() freqs = torch.outer(t, inv_freq) emb = torch.cat([freqs, freqs], dim=-1) return emb.cos(), emb.sin() def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat([-x2, x1], dim=-1) def apply_rope(x, cos, sin): T = x.shape[-2] cos = cos[:T].unsqueeze(0).unsqueeze(0).to(x.dtype) sin = sin[:T].unsqueeze(0).unsqueeze(0).to(x.dtype) return x * cos + rotate_half(x) * sin class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-5): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): norm = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() return x * norm * self.weight class StandardAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() assert d_model % n_heads == 0 self.n_heads = n_heads self.head_dim = d_model // n_heads self.wqkv = nn.Linear(d_model, 3 * d_model, bias=False) self.wo = nn.Linear(d_model, d_model, bias=False) def forward(self, x, rope_cos, rope_sin): B, T, D = x.shape qkv = self.wqkv(x) q, k, v = qkv.split(D, dim=-1) q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) q = apply_rope(q, rope_cos, rope_sin) k = apply_rope(k, rope_cos, rope_sin) out = F.scaled_dot_product_attention(q, k, v, is_causal=True) out = out.transpose(1, 2).contiguous().view(B, T, D) return self.wo(out) class SwiGLU(nn.Module): def __init__(self, d_model, d_ff): super().__init__() self.w_gate = nn.Linear(d_model, d_ff, bias=False) self.w_up = nn.Linear(d_model, d_ff, bias=False) self.w_down = nn.Linear(d_ff, d_model, bias=False) def forward(self, x): return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) class Block(nn.Module): def __init__(self, d_model, n_heads, d_ff, eps=1e-5): super().__init__() self.norm1 = RMSNorm(d_model, eps) self.attn = StandardAttention(d_model, n_heads) self.norm2 = RMSNorm(d_model, eps) self.ffn = SwiGLU(d_model, d_ff) def forward(self, x, rope_cos, rope_sin): x = x + self.attn(self.norm1(x), rope_cos, rope_sin) x = x + self.ffn(self.norm2(x)) return x class IvmeConversateSV2InstructModel(PreTrainedModel): """HF-compatible wrapper. Load with: AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True) """ config_class = IvmeConversateSV2InstructConfig # Explicit declarative tied-weights mapping -- confirmed via direct # inspection of transformers' PreTrainedModel.get_expanded_tied_weights_keys # that get_input_embeddings()/get_output_embeddings() ALONE do not trigger # automatic tying in this version; the class needs _tied_weights_keys set # explicitly (same convention used by e.g. GPT2LMHeadModel: # {'lm_head.weight': 'transformer.wte.weight'}). Verified this actually # ties the weights via post_init() -> init_weights() -> tie_weights(): # an earlier version of this file relied on get_output_embeddings() alone # and the weights were NOT tied (model.tok_embed.weight is model.lm_head. # weight was False) despite tie_word_embeddings=True in config. _tied_weights_keys = {"lm_head.weight": "tok_embed.weight"} def __init__(self, config: IvmeConversateSV2InstructConfig): super().__init__(config) self.tok_embed = nn.Embedding(config.vocab_size, config.d_model) nn.init.normal_(self.tok_embed.weight, mean=0.0, std=0.02) self.blocks = nn.ModuleList([ Block(config.d_model, config.n_heads, config.d_ff, config.norm_eps) for _ in range(config.n_layers) ]) self.norm_f = RMSNorm(config.d_model, config.norm_eps) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) head_dim = config.d_model // config.n_heads cos, sin = build_rope_cache(head_dim, config.max_seq_len, config.rope_theta) self.register_buffer("rope_cos", cos, persistent=True) self.register_buffer("rope_sin", sin, persistent=True) self.post_init() def get_input_embeddings(self): return self.tok_embed def set_input_embeddings(self, value): self.tok_embed = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def can_generate(self): return True def forward(self, input_ids, labels=None, **kwargs): x = self.tok_embed(input_ids) for block in self.blocks: x = block(x, self.rope_cos, self.rope_sin) x = self.norm_f(x) logits = self.lm_head(x) loss = None if labels is not None: loss = F.cross_entropy( logits[:, :-1, :].reshape(-1, self.config.vocab_size), labels[:, 1:].reshape(-1), ) return CausalLMOutput(loss=loss, logits=logits)