""" modeling_ivme_conversate.py HuggingFace transformers-compatible model wrapper for Ivme-Conversate-v3. Wraps the exact model architecture used during training/inference (RoPE, QK-norm, XSA, SwiGLU, tied embeddings) so it can be loaded via AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True) and used by any tool expecting a standard transformers CausalLM interface - including external benchmark runners like ArithMark's official script, which load models this way rather than via custom code. The underlying architecture code (RMSNorm, Attention/XSA, RoPE, SwiGLU) is copied verbatim from the verified, tested implementation used for training and the standalone inference script - including its dtype-correctness fixes - so results are identical to what training and infer_ivme_v3.py already produce. This file only adds the transformers-facing wrapper (config plumbing, forward() output format, generate() compatibility) around that unchanged core. """ import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast try: from .configuration_ivme_conversate import IvmeConversateConfig except ImportError: # Fallback for direct/standalone imports (e.g. a verification script # importing this file directly rather than through transformers' # trust_remote_code package-loading machinery, which does support # the relative import above). from configuration_ivme_conversate import IvmeConversateConfig # ============================================================================ # Core architecture - copied verbatim from the verified implementation. # ============================================================================ class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): dtype = x.dtype x = x.float() norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return norm.to(dtype) * self.weight.to(dtype) def precompute_rope(head_dim, max_seq_len, theta=10000.0, device=None): inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) t = torch.arange(max_seq_len, device=device).float() freqs = torch.outer(t, inv_freq) cos = torch.cos(freqs) sin = torch.sin(freqs) return cos, sin def apply_rope(x, cos, sin): d = x.shape[-1] x1, x2 = x[..., : d // 2], x[..., d // 2:] cos = cos[None, None, :, :].to(x.dtype) sin = sin[None, None, :, :].to(x.dtype) rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) return rotated class Attention(nn.Module): """Multi-head self-attention with QK-Norm + RoPE + XSA.""" def __init__(self, cfg): super().__init__() self.n_heads = cfg.n_heads self.head_dim = cfg.d_model // cfg.n_heads assert self.head_dim * cfg.n_heads == cfg.d_model self.wq = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.wk = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.wv = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.wo = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim, eps=cfg.qknorm_eps) self.k_norm = RMSNorm(self.head_dim, eps=cfg.qknorm_eps) self.xsa_eps = 1e-6 def forward(self, x, cos, sin): B, T, D = x.shape q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) q = self.q_norm(q) k = self.k_norm(k) q = apply_rope(q, cos[:T], sin[:T]) k = apply_rope(k, cos[:T], sin[:T]) y = F.scaled_dot_product_attention(q, k, v, is_causal=True) # --- XSA: z_i = y_i - (y_i . v_i)/(||v_i||^2+eps) * v_i, per token/head --- dot = (y * v).sum(dim=-1, keepdim=True) vnorm_sq = (v * v).sum(dim=-1, keepdim=True) proj_coef = dot / (vnorm_sq + self.xsa_eps) z = y - proj_coef * v z = z.transpose(1, 2).contiguous().view(B, T, D) return self.wo(z) class SwiGLU(nn.Module): def __init__(self, d_model, d_ffn): super().__init__() self.w_gate = nn.Linear(d_model, d_ffn, bias=False) self.w_up = nn.Linear(d_model, d_ffn, bias=False) self.w_down = nn.Linear(d_ffn, 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, cfg): super().__init__() self.attn_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps) self.attn = Attention(cfg) self.ffn_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps) self.ffn = SwiGLU(cfg.d_model, cfg.d_ffn) def forward(self, x, cos, sin): x = x + self.attn(self.attn_norm(x), cos, sin) x = x + self.ffn(self.ffn_norm(x)) return x # ============================================================================ # transformers-facing wrapper # ============================================================================ class IvmeConversateModel(PreTrainedModel): """Base model (no LM head) - exposed mainly so AutoModel works too; most users will want IvmeConversateForCausalLM below.""" config_class = IvmeConversateConfig def __init__(self, config): super().__init__(config) self.tok_embed = nn.Embedding(config.vocab_size, config.d_model) self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)]) self.final_norm = RMSNorm(config.d_model, eps=config.rmsnorm_eps) # NOTE: RoPE cos/sin are deliberately NOT cached as a # register_buffer here. transformers v5.x has a known, currently # open bug (confirmed via reproduction and matching upstream # reports: huggingface/transformers#43950, #44534) where # non-persistent buffers (persistent=False) get corrupted with # uninitialized/meta-device garbage memory during # from_pretrained()'s fast weight-materialization path - this # silently produced NaN/garbage RoPE values and completely wrong # model output when loaded via AutoModelForCausalLM, even though # the actual model weights loaded correctly. Recomputing cos/sin # fresh in forward() is cheap (a handful of trig ops) and # sidesteps this bug class entirely rather than depending on a # library fix. self._head_dim = config.d_model // config.n_heads self._rope_theta = config.rope_theta self._context_len = config.context_len self.post_init() def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, input_ids, **kwargs): # Defensive truncation: keep only the most recent context_len # tokens. This is a second safety layer beyond exposing # max_position_embeddings in the config (see # IvmeConversateConfig) - if any caller still manages to pass a # longer sequence, truncate cleanly to the trailing window # rather than crash inside RoPE with a cryptic tensor-shape # mismatch (confirmed reproducible: a 2048-token input against # this model's 1024 context previously crashed deep inside # apply_rope with no indication of the real cause). if input_ids.shape[1] > self._context_len: input_ids = input_ids[:, -self._context_len:] x = self.tok_embed(input_ids) cos, sin = precompute_rope(self._head_dim, self._context_len, theta=self._rope_theta, device=x.device) for block in self.blocks: x = block(x, cos, sin) x = self.final_norm(x) return x class IvmeConversateForCausalLM(PreTrainedModel, GenerationMixin): """Standard transformers CausalLM interface: forward() returns CausalLMOutputWithPast(logits=...), supports labels for loss computation, and is compatible with model.generate() and any tool (like ArithMark's benchmark runner) expecting a normal AutoModelForCausalLM-loadable model.""" config_class = IvmeConversateConfig # Tells transformers' save_pretrained/tie_weights machinery about the # tied embedding explicitly (required in transformers v5+; without # this, save_pretrained raises "shared tensors... not properly # defined" even though the weights ARE correctly tied via the # assignment in __init__ below - confirmed by reproducing the error # and fixing it with this attribute). _tied_weights_keys = {"lm_head.weight": "model.tok_embed.weight"} def __init__(self, config): super().__init__(config) self.model = IvmeConversateModel(config) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) if config.tie_embeddings: self.lm_head.weight = self.model.tok_embed.weight self.post_init() def get_input_embeddings(self): return self.model.tok_embed def set_input_embeddings(self, value): self.model.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 forward(self, input_ids, attention_mask=None, labels=None, **kwargs): # attention_mask accepted for interface compatibility but not # used - this model is trained/used causal-only with no padding # within a sequence (see ShardedTokenDataset / infer scripts). # # If input_ids exceeds context_len, self.model() truncates it # internally (see IvmeConversateModel.forward). labels must be # truncated identically here, BEFORE the loss is computed - # otherwise a caller passing both an over-length input_ids and # matching labels (e.g. a perplexity/loglikelihood task) would # hit a shape mismatch between logits and labels, or worse, a # silent misalignment if the shapes happened to coincide. context_len = self.config.max_position_embeddings if input_ids.shape[1] > context_len and labels is not None: labels = labels[:, -context_len:] hidden = self.model(input_ids) logits = self.lm_head(hidden) loss = None if labels is not None: loss = F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100 ) return CausalLMOutputWithPast( loss=loss, logits=logits, ) def prepare_inputs_for_generation(self, input_ids, **kwargs): # No KV-cache implementation - recomputes the full sequence each # step. Correct but not optimized for long generations; adequate # for benchmark scoring (log-likelihood over short continuations) # and interactive testing. return {"input_ids": input_ids} def generate(self, *args, **kwargs): # Force use_cache=False unconditionally. This model has no cache # implementation, and transformers' generate() otherwise tries to # construct a DynamicCache using standard config field names # (num_hidden_layers, etc.) that this config doesn't define, # crashing regardless of the config-level use_cache default set # in IvmeConversateConfig (confirmed empirically: that default # does not reliably reach generate()'s cache-setup path in all # transformers versions). Overriding here guarantees correct # behavior for any caller - including external tools like # ArithMark's benchmark runner - with zero special knowledge # required on their end. kwargs["use_cache"] = False return super().generate(*args, **kwargs) def num_params(self, non_embedding=False): n = sum(p.numel() for p in self.parameters()) if non_embedding: n -= self.model.tok_embed.weight.numel() return n