""" model_v2.py -- SpikeWhaleLM v2: optimized base architecture. Changes vs model.py (v1): PERFORMANCE - SparseMoEFFN: sort-based expert dispatch (one contiguous slice per expert, index_add_ scatter-back) replaces per-expert boolean masking. Far fewer kernel launches, torch.compile-friendly (no data-dependent boolean indexing in the hot path). - Shared experts fused into ONE ExpertFFN with n_shared * intermediate width (mathematically equivalent to the averaged sum, 1 matmul set instead of N). QUALITY / STABILITY - QK-Norm: per-head RMSNorm on Q and K before RoPE (Gemma2/OLMo2-style). Stabilizes attention logits, tolerates higher LR. (cfg.use_qk_norm, default ON) - z-loss on lm_head logits: zloss_coef * mean(log^2 Z). Prevents logit drift. (cfg.zloss_coef, default 1e-4; set 0 to disable) - MTP heads REDESIGNED: instead of K independent full H x V matrices (which at 50M params dwarfed the model), each MTP head is now a small zero-init H x H projection feeding the SHARED lm_head. Param cost per head: H^2 instead of H*V. MTP loss is down-weighted by cfg.mtp_loss_weight (default 0.3). - HC output: learned softmax mix over streams (HCOutputMix) instead of mean(). - Value-embedding residual (nanoGPT-speedrun style): per-layer learned gate (zero-init => exact no-op at init) adds a projection of the token embedding into each block's input. (cfg.use_value_embed, default OFF = opt-in) All new config keys are read with getattr(cfg, key, default) so your existing config.py works unmodified. NOTE: QK-Norm and HCOutputMix add parameters, so v1 checkpoints need load_state_dict(strict=False) (new params keep init; QK-Norm at init is NOT identity -- prefer training v2 from scratch, or set use_qk_norm=False to stay v1-loadable). XSA is kept byte-identical to v1 but read the note in MLADerfXSAAttention: with num_kv_heads == 1 it removes the SAME rank-1 value subspace from every head. A/B it at 50M before keeping it in the final base. """ import math import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional, Tuple, List from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast from torch.utils.checkpoint import checkpoint as gradient_checkpoint from config import SpikeWhaleConfig from fractal import fractal_rope_inv_freq # --------------------------------------------------------------------------- # Primitives # --------------------------------------------------------------------------- class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight class RotaryEmbedding(nn.Module): """RoPE for the rope partition of Q and K (qk_rope_head_dim dims only).""" def __init__(self, dim: int, max_positions: int = 4096, theta: float = 10000.0, inv_freq: Optional[torch.Tensor] = None): super().__init__() if inv_freq is None: inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) t = torch.arange(max_positions).float() freqs = torch.outer(t, inv_freq) self.register_buffer("cos_cache", freqs.cos()) self.register_buffer("sin_cache", freqs.sin()) def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: cos = self.cos_cache[position_ids].unsqueeze(1) # [B, 1, S, rope_dim//2] sin = self.sin_cache[position_ids].unsqueeze(1) d = cos.shape[-1] x1, x2 = x[..., :d], x[..., d:] return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) # --------------------------------------------------------------------------- # Engram: N-gram hash lookup + DERF gate (unchanged from v1) # --------------------------------------------------------------------------- class TokenCompressor(nn.Module): """Frozen random projection into the Engram hash space. The output is L2-NORMALIZED (v3). The downstream hash is a sign-bit LSH, which is scale-invariant on its own, but normalizing here means the hash input distribution no longer depends on `hidden_size`, `initializer_range`, or how far training has drifted the embedding norms -- so bucket occupancy stays uniform for the whole run instead of only at init. """ def __init__(self, embed_dim: int, compress_dim: int): super().__init__() self.proj = nn.Linear(embed_dim, compress_dim, bias=False) # 1/sqrt(fan_in) so the projection preserves scale rather than shrinking # it by ~0.02*sqrt(embed_dim); the normalize below makes this cosmetic, # but it keeps the pre-norm magnitude O(1) and readable in diagnostics. nn.init.normal_(self.proj.weight, std=embed_dim ** -0.5) # Frozen LSH-style projection: gradient never reaches it through the # integer hash cast, so a fixed random projection is correct (see v1). self.proj.weight.requires_grad_(False) def forward(self, x: torch.Tensor) -> torch.Tensor: z = self.proj(x) return z / (z.norm(dim=-1, keepdim=True) + 1e-6) class MultiHeadHashLookup(nn.Module): """Multi-head n-gram hash memory. HASH (v3, fixed). The v1/v2 hash was `h.abs().long() % table_size` where `h` was a dot product of magnitude ~0.1. `.long()` truncates toward zero, so EVERY token in every n-gram order hashed to bucket 0: the table collapsed to a single row, only that row ever received gradient, and Engram degenerated into a constant additive bias. Rescaling `h` would paper over it -- the truncation threshold is an absolute constant while `h`'s scale tracks hidden_size, initializer_range, and training drift, so the same collapse returns whenever those change. The fix is a real sign-bit LSH: hash by the SIGN of `hash_bits` random hyperplane projections, packed into an integer. Signs are invariant to any positive rescaling of the input, so the hash cannot silently collapse again no matter how the upstream scale moves. Occupancy is verifiable at any time via `bucket_occupancy()`. Set `legacy_trunc=True` ONLY to load a pre-v3 checkpoint bit-identically. """ def __init__(self, num_heads: int, table_size: int, compress_dim: int, out_dim: int, max_ngram: int = 3, legacy_trunc: bool = False): super().__init__() self.num_heads = num_heads self.table_size = table_size self.max_ngram = max_ngram self.out_dim = out_dim self.legacy_trunc = legacy_trunc # Hash width. Addressing the table needs ceil(log2(table_size)) bits, # but using EXACTLY that many makes `% table_size` non-uniform whenever # table_size is not a power of two: with table_size=300 the 9-bit hash # spans 512 values, so buckets 0-211 are reachable twice and 212-299 # once -- a measured 1.85x occupancy bias. 8 spare bits cut the worst-case # bias to 1 + table_size/2^bits (<0.4%), and for a power-of-two table # 2^bits divides table_size exactly, so it stays perfectly uniform. # Capped at 24: the bits are packed via a float32 accumulation, and # integers are only exact up to 2^24 there. self.hash_bits = min(24, max(1, math.ceil(math.log2(max(table_size, 2))) + 8)) self.tables = nn.ModuleList([ nn.Embedding(table_size, out_dim) for _ in range(num_heads) ]) for t in self.tables: nn.init.normal_(t.weight, std=0.01) # Bit weights 2^b for packing sign bits into a bucket index. self.register_buffer( "bit_weights", (2.0 ** torch.arange(self.hash_bits, dtype=torch.float32)), persistent=False, ) for n in range(1, max_ngram + 1): for k in range(n): if legacy_trunc: proj = torch.randn(num_heads, compress_dim) proj = proj / (proj.norm(dim=1, keepdim=True) + 1e-8) else: # [num_heads, hash_bits, compress_dim] -- one hyperplane per # (head, bit). Distinct per position k so the hash is # ORDER-SENSITIVE: "a b" and "b a" land in different buckets. proj = torch.randn(num_heads, self.hash_bits, compress_dim) proj = proj / (proj.norm(dim=-1, keepdim=True) + 1e-8) self.register_buffer(f"hash_proj_n{n}_p{k}", proj) def _bucket_ids(self, compressed: torch.Tensor, n: int, valid_len: int) -> torch.Tensor: """-> [B, valid_len, num_heads] int64 bucket indices for order-n grams.""" if self.legacy_trunc: h = None for k in range(n): proj = getattr(self, f"hash_proj_n{n}_p{k}") term = torch.matmul(compressed[:, k:k + valid_len, :].float(), proj.t()) h = term if h is None else h + term return h.abs().long() % self.table_size # Accumulate the hyperplane projections across the n gram positions # BEFORE taking the sign, so the bit pattern depends on the whole gram. h = None for k in range(n): proj = getattr(self, f"hash_proj_n{n}_p{k}") # [H, bits, D] # [B, valid_len, D] x [D, H*bits] -> [B, valid_len, H, bits] term = torch.matmul( compressed[:, k:k + valid_len, :].float(), proj.reshape(self.num_heads * self.hash_bits, -1).t(), ).view(-1, valid_len, self.num_heads, self.hash_bits) h = term if h is None else h + term bits = (h > 0).float() # scale-invariant idx = (bits * self.bit_weights).sum(dim=-1).long() return idx % self.table_size @torch.no_grad() def bucket_occupancy(self, compressed: torch.Tensor) -> dict: """Diagnostic: fraction of the table actually addressed, per n-gram order. A healthy hash gives `distinct / min(tokens, table_size)` near 1.0. The v1/v2 collapse showed up here as distinct == 1. Used by the regression test in tests/test_engram_hash.py and cheap enough to call during training warmup. """ S = compressed.shape[1] stats = {} for n in range(1, self.max_ngram + 1): if S < n: continue idx = self._bucket_ids(compressed, n, S - n + 1) distinct = int(idx.unique().numel()) stats[f"n{n}"] = { "distinct_buckets": distinct, "table_size": self.table_size, "occupancy": distinct / min(idx.numel(), self.table_size), } return stats def forward(self, compressed: torch.Tensor, prefix: Optional[torch.Tensor] = None) -> torch.Tensor: """`prefix` is the compressed tail of the preceding tokens (see EngramModule.retrieve). Output is only for the `compressed` positions. Without it, incremental decoding passes S=1 and every n-gram order with n > 1 is skipped by the `S < n` guard -- so a generated token sees only its own unigram while the same token in a full forward sees its 2- and 3-grams. That is a silent train/decode mismatch; it was invisible only because the Engram gate used to be a hard zero. """ if prefix is not None and prefix.shape[1] > 0: P = prefix.shape[1] compressed = torch.cat([prefix.to(compressed.dtype), compressed], dim=1) else: P = 0 B, S, _ = compressed.shape device = compressed.device out = torch.zeros(B, S, self.out_dim, device=device, dtype=compressed.dtype) norm = torch.zeros(S, device=device) for n in range(1, self.max_ngram + 1): if S < n: continue valid_len = S - n + 1 start = n - 1 idx = self._bucket_ids(compressed, n, valid_len) for head_idx, table in enumerate(self.tables): out[:, start:, :] = out[:, start:, :] + table(idx[:, :, head_idx]).to(out.dtype) norm[start:] += self.num_heads out = out / norm.view(1, -1, 1).clamp(min=1) return out[:, P:, :].to(compressed.dtype) class DERFContextGate(nn.Module): def __init__(self, obs_size: int, init_bias: float = -4.0): super().__init__() self.proj = nn.Linear(obs_size * 2, obs_size) self.alpha = nn.Parameter(torch.ones(obs_size)) self.bias = nn.Parameter(torch.full((obs_size,), init_bias)) self.gamma = nn.Parameter(torch.ones(obs_size)) def forward(self, retrieved: torch.Tensor, x: torch.Tensor) -> torch.Tensor: logits = self.proj(torch.cat([retrieved, x], dim=-1)) gate = self.gamma * ((torch.erf(self.alpha * logits + self.bias) + 1.0) / 2.0) return retrieved * gate class EngramModule(nn.Module): def __init__(self, cfg: SpikeWhaleConfig): super().__init__() self.compressor = TokenCompressor(cfg.hidden_size, cfg.engram_compress_dim) self.lookup = MultiHeadHashLookup( cfg.engram_num_heads, cfg.engram_table_size, cfg.engram_compress_dim, cfg.hidden_size, cfg.engram_max_ngram, legacy_trunc=getattr(cfg, "engram_legacy_hash", False), ) self.gate = DERFContextGate(cfg.hidden_size, cfg.engram_gate_init_bias) def compress(self, x: torch.Tensor) -> torch.Tensor: """Compressed hash-space vectors, one per position. These depend ONLY on the token embedding (x is detached and the projection is frozen), never on context -- which is what makes it valid to cache the tail of them and reuse it as the n-gram prefix on the next decode step. """ return self.compressor(x.detach()) def retrieve(self, x: torch.Tensor, prefix: Optional[torch.Tensor] = None) -> torch.Tensor: """Raw (ungated) n-gram memory read. Exposed so mid-network fusion layers can reuse the SAME retrieval instead of hashing again.""" return self.lookup(self.compress(x), prefix=prefix) def forward(self, x: torch.Tensor, retrieved: Optional[torch.Tensor] = None) -> torch.Tensor: if retrieved is None: retrieved = self.retrieve(x) return self.gate(retrieved, x) class EngramLayerFusion(nn.Module): """Mid-stack re-injection of the Engram n-gram read (Nanbeige NgramLayerFusion). Engram injects n-gram evidence once, at the embedding, where it competes with everything else the first layers need to do. Nanbeige re-injects the same n-gram features again partway up the stack, gated by the AGREEMENT between an n-gram-derived key and the current hidden state -- so a layer only pulls surface-form evidence back in when its own representation is already pointing that way. Bottlenecked to `fusion_dim` so the cost is 3*H*fusion_dim per fusion layer rather than 3*H^2. `value_proj` is zero-init => exact no-op at init. """ def __init__(self, hidden_size: int, fusion_dim: int, eps: float = 1e-6): super().__init__() self.fusion_dim = fusion_dim self.hidden_down = nn.Linear(hidden_size, fusion_dim, bias=False) self.key_proj = nn.Linear(hidden_size, fusion_dim, bias=False) self.value_proj = nn.Linear(hidden_size, fusion_dim, bias=False) self.out_proj = nn.Linear(fusion_dim, hidden_size, bias=False) self.hidden_norm = RMSNorm(fusion_dim, eps) self.ngram_norm = RMSNorm(fusion_dim, eps) for m in (self.hidden_down, self.key_proj, self.value_proj): nn.init.normal_(m.weight, std=0.02) nn.init.zeros_(self.out_proj.weight) def forward(self, hidden: torch.Tensor, ngram: torch.Tensor) -> torch.Tensor: key = self.ngram_norm(self.key_proj(ngram)) h = self.hidden_norm(self.hidden_down(hidden)) gate = (h * key).sum(-1, keepdim=True) / math.sqrt(self.fusion_dim) # Signed sqrt compresses the dynamic range of the agreement score before # the sigmoid, so the gate does not saturate on a few high-norm tokens. gate = torch.sigmoid(gate.abs().clamp_min(1e-6).sqrt() * torch.sign(gate)) return hidden + self.out_proj(gate * self.value_proj(ngram)) # --------------------------------------------------------------------------- # Hyper-Connections # --------------------------------------------------------------------------- class HyperConnectionLayer(nn.Module): """Simplified HC: softmax pre-mix / post-distribute over hc_mult streams. Asymmetric init (v1 bugfix) so streams diverge and gradients flow.""" def __init__(self, hidden_size: int, hc_mult: int, sinkhorn_iters: int = 20, eps: float = 1e-6): super().__init__() self.hc_mult = hc_mult self.pre_weight = nn.Parameter( torch.linspace(0.5, -0.5, hc_mult) / max(hc_mult, 1) ) self.post_weight = nn.Parameter( torch.linspace(-0.5, 0.5, hc_mult) / max(hc_mult, 1) ) def pre_op(self, copies: torch.Tensor) -> torch.Tensor: w = F.softmax(self.pre_weight, dim=0) return (copies * w.view(1, -1, 1, 1)).sum(dim=1) def post_op(self, copies: torch.Tensor, delta: torch.Tensor) -> torch.Tensor: w = F.softmax(self.post_weight, dim=0) return copies + delta.unsqueeze(1) * w.view(1, -1, 1, 1) class HCOutputMix(nn.Module): """ NEW (v2): learned combination of the hc_mult streams at the model output, replacing the v1 mean(dim=1). Mean forces the streams toward redundancy at exactly the point where you want them specialized. Initialized uniform so it starts identical to mean() -- a strict generalization, zero risk. """ def __init__(self, hc_mult: int): super().__init__() self.weight = nn.Parameter(torch.zeros(hc_mult)) # softmax(0)=uniform=mean def forward(self, copies: torch.Tensor) -> torch.Tensor: w = F.softmax(self.weight, dim=0) return (copies * w.view(1, -1, 1, 1)).sum(dim=1) # --------------------------------------------------------------------------- # MLA + (DERF) + XSA Attention, now with QK-Norm # --------------------------------------------------------------------------- class MLADerfXSAAttention(nn.Module): """ v2 additions: - QK-Norm (cfg.use_qk_norm, default True): per-head RMSNorm applied to Q and K BEFORE the rope/nope split. Bounds attention logits, the standard modern stability fix; composes cleanly with SDPA and partial RoPE. XSA NOTE (unchanged mechanics, important caveat): with num_kv_heads == 1 (MQA) every query head shares the same value vector, so the self-projection subtraction removes the SAME rank-1 value subspace from all heads -- much more aggressive than per-head XSA. Ablate use_xsa on/off at 50M before locking the base config. """ def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int = 0): super().__init__() self.layer_idx = layer_idx self.num_heads = cfg.num_attention_heads self.num_kv_heads = cfg.num_key_value_heads self.head_dim = cfg.head_dim self.qk_rope_head_dim = cfg.qk_rope_head_dim self.nope_head_dim = cfg.nope_head_dim self.hidden_size = cfg.hidden_size self.use_derf = cfg.use_derf self.use_xsa = cfg.use_xsa self.use_elo = getattr(cfg, "use_elo", False) self.dropout_p = cfg.attention_dropout self.kv_groups = self.num_heads // self.num_kv_heads self.use_qk_norm = getattr(cfg, "use_qk_norm", True) self.q_a_proj = nn.Linear(cfg.hidden_size, cfg.q_lora_rank, bias=False) self.q_a_norm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps) self.q_b_proj = nn.Linear(cfg.q_lora_rank, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.o_a_proj = nn.Linear(self.num_heads * self.head_dim, cfg.o_lora_rank, bias=False) self.o_b_proj = nn.Linear(cfg.o_lora_rank, cfg.hidden_size, bias=False) # QK-Norm: one RMSNorm over head_dim, shared across heads (Gemma-2 style). if self.use_qk_norm: self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) # Fractal RoPE: same endpoints/band count as standard RoPE, but the # exponents sit on a Cantor spectrum (fractal.py). gamma=0 == standard. rope_inv_freq = None if getattr(cfg, "use_fractal_rope", False): rope_inv_freq = fractal_rope_inv_freq( self.qk_rope_head_dim, cfg.rope_theta, cfg.fractal_rope_gamma) self.rope = RotaryEmbedding( self.qk_rope_head_dim, max_positions=cfg.max_position_embeddings, theta=cfg.rope_theta, inv_freq=rope_inv_freq, ) if self.use_derf: self.derf_alpha = nn.Parameter(torch.ones(self.num_heads)) self.derf_bias = nn.Parameter(torch.zeros(self.num_heads)) self.derf_gamma = nn.Parameter(torch.ones(self.num_heads)) # --- Depth Attention (Nanbeige-4.2 port) --- # Mixes this layer's V with the V of earlier checkpointed layers, using # softmax over the DEPTH axis of q.k_layer. Parameter-free except the # zero-init gate, which makes it an exact no-op at initialization. self.use_depth_attention = getattr(cfg, "use_depth_attention", False) self.depth_attention_stride = max(1, int(getattr(cfg, "depth_attention_stride", 4))) # Layer 0 has no earlier layer to mix from, so its gate could never # receive gradient. Don't allocate one rather than ship a parameter that # is dead by construction. self.depth_gate = ( nn.Parameter(torch.zeros(1)) if (self.use_depth_attention and layer_idx > 0) else None ) if self.use_elo: # per-head Elo K-factor (softplus>0) and rating->logit gain (tanh, # zero-init => no-op at start, so the tournament reduces to plain # softmax attention). self.elo_log_k = nn.Parameter(torch.full( (self.num_heads,), math.log(math.expm1(cfg.elo_k_init)))) self.elo_gain = nn.Parameter(torch.zeros(self.num_heads)) for m in (self.q_a_proj, self.q_b_proj, self.k_proj, self.v_proj, self.o_a_proj, self.o_b_proj): nn.init.normal_(m.weight, std=cfg.initializer_range) def _depth_mix_value(self, q, k, v, depth_kv): """Cross-layer value mixing (Nanbeige `_apply_depth_attention`). q: [B, num_heads, S, hd] (this layer's queries, pre-cache) k, v: [B, num_kv_heads, S, hd] (this layer's fresh K/V, post-RoPE) depth_kv: tuple of (k_l, v_l) from earlier checkpointed layers, same shape. For each (batch, kv-head, position) we softmax q.k over the LAYER axis -- current layer plus each source layer -- and take the corresponding convex combination of values. So a token can pull forward the value it had built at a shallower depth when that depth's key is a better match: a content-addressed skip along depth, complementing the sequence-axis attention. Returns the mixed V, shape unchanged. Queries are mean-pooled down to the KV-head grouping first, because K/V live in num_kv_heads (=2 under GQA) while Q lives in num_heads (=16). """ if not depth_kv: return v # [B, num_heads, S, hd] -> [B, num_kv_heads, S, hd] if self.num_heads != self.num_kv_heads: q_g = q.view(q.shape[0], self.num_kv_heads, self.kv_groups, q.shape[2], q.shape[3]).mean(dim=2) else: q_g = q keys = torch.stack([kk for kk, _ in depth_kv] + [k], dim=0) # [L+1,B,kvH,S,hd] vals = torch.stack([vv for _, vv in depth_kv] + [v], dim=0) logits = (q_g.unsqueeze(0).float() * keys.float()).sum(-1) # [L+1,B,kvH,S] probs = torch.softmax(logits * (self.head_dim ** -0.5), dim=0).to(vals.dtype) mixed = (probs.unsqueeze(-1) * vals).sum(dim=0) # Zero-init gate => `mixed` contributes nothing at step 0, so turning # depth attention on never perturbs an existing run's starting point. return v + torch.tanh(self.depth_gate) * (mixed - v) def _elo_attention(self, q, k, v, is_masked, past_rating): """Elo/Bradley-Terry tournament attention (ported from wheelerv2). Softmax attention already gives each key its win-probability against the field; Elo adds a PERSISTENT rating r_j that accumulates across queries and biases future logits (winners get boosted), like an AlphaFold-style branch-until-winner reinforcement. e_ij = q_i.k_j/sqrt(d) + gain * rating_j (rating = prior + intra-call) match : S_ij = softmax weight key j won at query i ; E_i = 1/n_i baseline update : dR_ij = K * (S_ij - E_i) rating_j(query i) = sum over earlier queries t pass B == pass A == plain softmax, so enabling Elo is an exact no-op at init. Returns (y, rating_out) with rating_out: [B, num_heads, N].""" B, H, S, _ = q.shape N = k.shape[2] device = q.device scale = 1.0 / math.sqrt(self.head_dim) scores = torch.matmul(q, k.transpose(-2, -1)) * scale # [B,H,S,N] # Take the working dtype from `scores`, NOT from q: under autocast the # matmul downcasts to bf16/fp16 while q stays fp32, and fp32's finfo.min # (-3.40282e38) overflows bf16's max finite value -- masked_fill then # raises "value cannot be converted to BFloat16 without overflow". dtype = scores.dtype keep = ~is_masked # broadcastable to [.,.,S,N] keepf = keep.to(dtype) neg = torch.finfo(dtype).min K = F.softplus(self.elo_log_k).view(1, H, 1, 1) # per-head K-factor >0 gain = torch.tanh(self.elo_gain).view(1, H, 1, 1) # per-head, 0 at init # prior rating of each key = what it accumulated BEFORE this forward # (from the cache) + 0 for the keys introduced this call. if past_rating is not None: prior = torch.cat( [past_rating, past_rating.new_zeros(B, H, S)], dim=-1) # [B,H,N] else: prior = scores.new_zeros(B, H, N) prior_bias = gain * prior.unsqueeze(2) # [B,H,1,N] masked = scores.masked_fill(is_masked, neg) # Match outcome = innate q.k compatibility (the "game result"), NOT the # accumulated rating. Keeping matches rating-INDEPENDENT is what makes # the parallel prefix-sum (tril @ dR, training) bit-identical to the # incremental running sum (decoding); rating is the REPUTATION that # biases attention, not what defines the match. a_base = torch.softmax(masked, dim=-1) # [B,H,S,N] n_i = keepf.sum(-1, keepdim=True).clamp(min=1.0) # valid keys per query dR = (K * (a_base - 1.0 / n_i)) * keepf # rating seen by query i = prior + causal sum of dR from earlier queries t 0 and self.training: a = F.dropout(a, p=self.dropout_p) y = torch.matmul(a, v) # [B,H,S,hd] # rating carried forward = prior + everything this call's queries contributed rating_out = prior + dR.sum(2) # [B,H,N] return y, rating_out def forward( self, x: torch.Tensor, position_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, depth_kv: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple]]: B, S, _ = x.shape q = self.q_a_norm(self.q_a_proj(x)) q = self.q_b_proj(q).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2) # QK-Norm before RoPE (v2). Cache stores the NORMALIZED k so prefill and # incremental decode agree. if self.use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) q_nope = q[..., :self.nope_head_dim] q_rope = q[..., self.nope_head_dim:] k_nope = k[..., :self.nope_head_dim] k_rope = k[..., self.nope_head_dim:] q_rope = self.rope(q_rope, position_ids) k_rope = self.rope(k_rope, position_ids) q = torch.cat([q_nope, q_rope], dim=-1) k = torch.cat([k_nope, k_rope], dim=-1) # Depth attention operates on the FRESH K/V for this position block, # before the KV cache is concatenated -- every layer's contribution is # then the same [B, kvH, S, hd] shape, and the mixed V is what gets # cached, so incremental decode inherits the mix (Nanbeige does the same # ordering: mix, then `past_key_value.update`). depth_entry = None if self.use_depth_attention: # Layer 0 CONSUMES nothing (no earlier layer, hence no gate) but # still PRODUCES a source for the layers above it. if self.depth_gate is not None: v = self._depth_mix_value(q, k, v, depth_kv or ()) if self.layer_idx % self.depth_attention_stride == 0: depth_entry = (k, v) past_rating = None if past_key_value is not None: if len(past_key_value) == 3: # (k, v, elo_rating) cache past_rating = past_key_value[2] k = torch.cat([past_key_value[0], k], dim=2) v = torch.cat([past_key_value[1], v], dim=2) present = (k, v) if use_cache else None N = k.shape[2] if self.kv_groups > 1: k = k.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape( B, self.num_heads, N, self.head_dim) v = v.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape( B, self.num_heads, N, self.head_dim) if self.use_elo: # Elo tournament over keys. Needs explicit scores, so it takes the # materialized-mask path rather than SDPA; it precedes DERF. if attention_mask is None and past_key_value is None: is_masked = torch.triu( torch.ones(S, N, dtype=torch.bool, device=q.device), diagonal=N - S + 1, ).unsqueeze(0).unsqueeze(0) else: is_masked = (attention_mask < -1.0) if attention_mask is not None \ else torch.triu( torch.ones(S, N, dtype=torch.bool, device=q.device), diagonal=N - S + 1, ).unsqueeze(0).unsqueeze(0) y, rating_out = self._elo_attention(q, k, v, is_masked, past_rating) if use_cache: present = (present[0], present[1], rating_out) elif self.use_derf: scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) if attention_mask is None and past_key_value is None: is_masked = torch.triu( torch.ones(S, N, dtype=torch.bool, device=scores.device), diagonal=N - S + 1, ).unsqueeze(0).unsqueeze(0) else: is_masked = (attention_mask < -1.0) if attention_mask is not None \ else torch.zeros_like(scores, dtype=torch.bool) safe_scores = scores.masked_fill(is_masked, -10000.0) a = self.derf_alpha.view(1, -1, 1, 1) b = self.derf_bias.view(1, -1, 1, 1) g = self.derf_gamma.view(1, -1, 1, 1) attn_weights = g * torch.erf(a * safe_scores + b) attn_weights = (attn_weights + g) / 2.0 attn_weights = attn_weights.masked_fill(is_masked, 0.0) attn_weights = attn_weights / (attn_weights.sum(dim=-1, keepdim=True) + 1e-8) if self.dropout_p > 0 and self.training: attn_weights = F.dropout(attn_weights, p=self.dropout_p) y = torch.matmul(attn_weights, v) else: q = q.contiguous() k = k.contiguous() v = v.contiguous() drop = self.dropout_p if self.training else 0.0 if past_key_value is None and attention_mask is None: y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop) else: if attention_mask is not None: is_masked = (attention_mask < -1.0) else: is_masked = torch.triu( torch.ones(S, N, dtype=torch.bool, device=q.device), diagonal=N - S + 1, ).unsqueeze(0).unsqueeze(0) y = F.scaled_dot_product_attention( q, k, v, attn_mask=~is_masked, dropout_p=drop) if self.use_xsa: past_len = N - S v_self = v[:, :, past_len:past_len + S, :] vn = v_self / (v_self.norm(dim=-1, keepdim=True) + 1e-8) projection = (y * vn).sum(dim=-1, keepdim=True) * vn y = y - projection y = y.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim) y = self.o_b_proj(self.o_a_proj(y)) return y, present, depth_entry # --------------------------------------------------------------------------- # MoE FFN -- v2: sort-based dispatch + fused shared expert # --------------------------------------------------------------------------- class ExpertFFN(nn.Module): """Single SwiGLU expert.""" def __init__(self, hidden_size: int, intermediate_size: int): super().__init__() self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) def sqrtsoftplus(x: torch.Tensor) -> torch.Tensor: return torch.sqrt(F.softplus(x) + 1e-8) class SparseMoEFFN(nn.Module): """ v2 changes: - FUSED shared expert: one ExpertFFN with width n_shared * intermediate, scaled by 1/n_shared on output -- equivalent to v1's averaged Python loop, one fused matmul set. (state-dict key changes: shared_expert.*) - SORT-BASED dispatch for routed experts: flatten (token, slot) pairs, argsort by expert id, run each expert on ONE contiguous slice, weighted index_add_ back. No boolean masks, no nonzero(), no per-expert scatter. Routing logic (hash routing, sqrtsoftplus, aux loss) is unchanged. """ def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int = 0): super().__init__() self.n_routed_experts = cfg.n_routed_experts self.n_shared_experts = cfg.n_shared_experts self.num_experts_per_tok = cfg.num_experts_per_tok self.norm_topk_prob = cfg.norm_topk_prob self.scoring_func = cfg.scoring_func self.routed_scaling_factor = cfg.routed_scaling_factor self.use_hash_routing = layer_idx < cfg.num_hash_layers self.aux_loss_coef = cfg.moe_aux_loss_coef self.router = nn.Linear(cfg.hidden_size, cfg.n_routed_experts, bias=False) if self.use_hash_routing: # The first `num_hash_layers` layers route by POSITION, so this # Linear is never called and would sit in the optimizer forever with # a zero gradient. Kept allocated (state-dict keys stay compatible # with existing checkpoints) but frozen, so it is honestly excluded # from training rather than silently dead. self.router.weight.requires_grad_(False) self.experts = nn.ModuleList([ ExpertFFN(cfg.hidden_size, cfg.moe_intermediate_size) for _ in range(cfg.n_routed_experts) ]) # Fused shared expert (v2) self.shared_expert = ( ExpertFFN(cfg.hidden_size, cfg.moe_intermediate_size * cfg.n_shared_experts) if cfg.n_shared_experts > 0 else None ) self._last_aux_loss: Optional[torch.Tensor] = None def forward(self, x: torch.Tensor, position_ids: Optional[torch.Tensor] = None) -> torch.Tensor: B, S, H = x.shape x_flat = x.view(B * S, H) T = B * S K = self.num_experts_per_tok # Shared expert: always active, single fused pass. if self.shared_expert is not None: shared_out = self.shared_expert(x_flat) if self.n_shared_experts > 1: shared_out = shared_out / self.n_shared_experts else: shared_out = None # ---- Routing (unchanged logic) ---- if self.use_hash_routing: if position_ids is not None: base = (position_ids.reshape(T, 1) % self.n_routed_experts).long() else: base = (torch.arange(T, device=x.device) % self.n_routed_experts).unsqueeze(1) offsets = torch.arange(K, device=x.device) top_k_indices = (base + offsets.unsqueeze(0)) % self.n_routed_experts # [T, K] top_k_weights = torch.full((T, K), 1.0 / K, device=x.device, dtype=x_flat.dtype) self._last_aux_loss = None else: router_logits = self.router(x_flat) if self.scoring_func == "sqrtsoftplus": routing_scores = sqrtsoftplus(router_logits) else: routing_scores = F.softmax(router_logits, dim=-1) top_k_scores, top_k_indices = torch.topk(routing_scores, K, dim=-1) if self.norm_topk_prob: top_k_weights = top_k_scores / (top_k_scores.sum(dim=-1, keepdim=True) + 1e-8) else: top_k_weights = top_k_scores top_k_weights = top_k_weights * self.routed_scaling_factor softmax_probs = F.softmax(router_logits, dim=-1) expert_mask = torch.zeros_like(softmax_probs) expert_mask.scatter_(1, top_k_indices, 1.0) f_e = expert_mask.mean(0) p_e = softmax_probs.mean(0) self._last_aux_loss = self.n_routed_experts * (f_e * p_e).sum() * self.aux_loss_coef # ---- Sort-based dispatch (v2) ---- # Flatten the (token, slot) assignment: T*K rows total. flat_expert = top_k_indices.reshape(-1) # [T*K] flat_weight = top_k_weights.reshape(-1, 1) # [T*K, 1] flat_token = torch.arange(T, device=x.device).repeat_interleave(K) # [T*K] order = torch.argsort(flat_expert, stable=True) # group by expert sorted_expert = flat_expert[order] sorted_token = flat_token[order] sorted_weight = flat_weight[order] counts = torch.bincount(sorted_expert, minlength=self.n_routed_experts) # boundaries per expert in the sorted order (CPU sync once per forward; # unavoidable without grouped-GEMM, still vastly cheaper than v1's # per-expert nonzero/masking) counts_list = counts.tolist() gathered = x_flat[sorted_token] # [T*K, H] out_flat = torch.zeros_like(x_flat) start = 0 for expert_idx, cnt in enumerate(counts_list): if cnt == 0: continue end = start + cnt seg = gathered[start:end] seg_out = self.experts[expert_idx](seg) * sorted_weight[start:end] out_flat.index_add_(0, sorted_token[start:end], seg_out.to(out_flat.dtype)) start = end if shared_out is not None: out_flat = out_flat + shared_out return out_flat.view(B, S, H) def get_aux_loss(self) -> Optional[torch.Tensor]: return self._last_aux_loss class DenseFFN(nn.Module): def __init__(self, cfg: SpikeWhaleConfig): super().__init__() self.gate_proj = nn.Linear(cfg.hidden_size, cfg.moe_intermediate_size, bias=False) self.up_proj = nn.Linear(cfg.hidden_size, cfg.moe_intermediate_size, bias=False) self.down_proj = nn.Linear(cfg.moe_intermediate_size, cfg.hidden_size, bias=False) def forward(self, x: torch.Tensor, position_ids: Optional[torch.Tensor] = None) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) def get_aux_loss(self) -> Optional[torch.Tensor]: return None # --------------------------------------------------------------------------- # Transformer block # --------------------------------------------------------------------------- class MemoryCacheBranch(nn.Module): """Memory Caching (arXiv 2602.24281): a linear-attention memory branch with growing cross-segment memory, adapted to sit PARALLEL to softmax MLA. The sequence is split into N segments of length L. For a token t in segment s the Gated Residual Memory readout is y_t = gamma_{t,s} * M_t^(s)(q_t) # own segment, CAUSAL + sum_{i ), i <= s. phi = elu(.)+1 (Katharopoulos linear-transformer feature map). Heads are small (mc_num_heads x mc_head_dim) so the branch is a cheap add-on. Cross-segment readout is O(N^2 L) -- fine at the sequence lengths used here; it is not a long-context speed win, it is the faithful gated-residual formulation. """ def __init__(self, cfg: SpikeWhaleConfig): super().__init__() self.h = cfg.mc_num_heads self.d = cfg.mc_head_dim self.L = cfg.mc_segment_len self.dg = cfg.mc_gate_dim self.eps = 1e-6 H, inner = cfg.hidden_size, self.h * self.d self.q_proj = nn.Linear(H, inner, bias=False) self.k_proj = nn.Linear(H, inner, bias=False) self.v_proj = nn.Linear(H, inner, bias=False) self.o_proj = nn.Linear(inner, H, bias=False) self.u_proj = nn.Linear(H, self.dg, bias=False) # per-token gate query self.p_proj = nn.Linear(H, self.dg, bias=False) # per-segment gate key self._st = {} # cache_idx -> incremental segment state @staticmethod def _phi(x): return F.elu(x) + 1.0 def forward(self, x, position_ids=None, attention_mask=None, state_key=0): # The linear-attention denominator underflows in bf16 as the branch's # projections grow, so the branch runs in fp32 with autocast disabled. # # state_key is the plan slot (cache_idx): unique per (layer, loop pass). # A looped stack calls this module once per pass on the same positions with # different activations, so each pass needs its own segment state. orig = x.dtype pos0 = None if position_ids is not None and position_ids.numel(): pos0 = int(position_ids.reshape(position_ids.shape[0], -1)[0, 0]) if not hasattr(self, "_st") or self._st is None: self._st = {} st = self._st.get(state_key) with torch.autocast(device_type=x.device.type, enabled=False): xf = x.float() if (x.shape[1] == 1 and pos0 is not None and st is not None and st["npos"] == pos0 and st["B"] == x.shape[0]): y = self._step(xf, state_key) else: y = self._mc_forward(xf, position_ids, attention_mask) if pos0 == 0 or st is None or st.get("B") != x.shape[0]: self._capture(xf, state_key) else: self._st.pop(state_key, None) # non-contiguous: state is stale return y.to(orig) def _mc_forward(self, x, position_ids=None, attention_mask=None): B, S, H = x.shape h, d, L = self.h, self.d, self.L pad = (L - S % L) % L if pad: x = F.pad(x, (0, 0, 0, pad)) Sp = S + pad N = Sp // L def heads(proj): return proj(x).view(B, N, L, h, d).permute(0, 3, 1, 2, 4) # [B,h,N,L,d] q = self._phi(heads(self.q_proj)) k = self._phi(heads(self.k_proj)) v = heads(self.v_proj) # [B,h,N,L,d] # --- own segment, causal within-segment linear attention (UNnormalized) --- A = torch.einsum('bhnld,bhnmd->bhnlm', q, k) # [B,h,N,L,L] causal = torch.tril(torch.ones(L, L, device=x.device, dtype=A.dtype)) A = A * causal num_intra = torch.einsum('bhnlm,bhnmd->bhnld', A, v) # [B,h,N,L,d] (no /Z) den_intra = A.sum(-1) # [B,h,N,L] # --- cached full-segment memories, cross readout for every (query,key) seg --- M = torch.einsum('bhnmd,bhnme->bhnde', k, v) # [B,h,N,d,d] z = k.sum(3) # [B,h,N,d] num = torch.einsum('bhqld,bhide->bhqile', q, M) # [B,h,Nq,Ni,L,d] UNnormalized den = torch.einsum('bhqld,bhid->bhqil', q, z) # [B,h,Nq,Ni,L] # diagonal (query seg == key seg) uses the CAUSAL running memory, not the # full-segment memory. Keep numerators/denominators UNnormalized here. idx = torch.arange(N, device=x.device) num = num.clone(); den = den.clone() num[:, :, idx, idx] = num_intra.to(num.dtype) # [B,h,N,L,d] den[:, :, idx, idx] = den_intra.to(den.dtype) # [B,h,N,L] # --- gates: gamma_{t,i} = softmax_i , masked to i<=s --- u = self.u_proj(x).view(B, N, L, self.dg) # [B,Nq,L,dg] pp = self.p_proj(x).view(B, N, L, self.dg) # [B,Ni,L,dg] pool_full = pp.mean(2) # [B,Ni,dg] full-seg mean score = torch.einsum('bqlg,big->bqil', u, pool_full) # [B,Nq,Ni,L] # CAUSALITY: the current-segment (diagonal) gate must pool ONLY tokens # 0..l of the segment -- a full-segment mean leaks that segment's future # tokens into token l's gate. Past segments (ibql', u, cpool) # [B,Nq,L] score = score.clone() score[:, idx, idx] = score_diag.to(score.dtype) # diagonal i==q -> causal seg_ok = (idx[None, :] <= idx[:, None]) # [Nq,Ni] i<=q score = score.masked_fill(~seg_ok[None, :, :, None], float('-inf')) gamma = torch.softmax(score, dim=2) # over key seg i [B,Nq,Ni,L] # Gated aggregation with a SINGLE GLOBAL normalizer. The paper's Gated # Residual weights the UNnormalized per-segment memories by gamma and does # NOT normalize each segment -- normalizing per-segment (the earlier bug) # erased the query-key match magnitude, so only the coarse mean-pool gate # could pick a segment. Here the numerators keep that magnitude; the gate # weights segments; one global Z normalizes the aggregate. g = gamma[:, None] # [B,1,Nq,Ni,L] y_num = (num * g.unsqueeze(-1)).sum(3) # [B,h,Nq,L,d] y_den = (den * g).sum(3) # [B,h,Nq,L] y = y_num / (y_den.unsqueeze(-1) + self.eps) y = y.permute(0, 2, 3, 1, 4).reshape(B, Sp, h * d) y = self.o_proj(y) return y[:, :S] # ---- incremental (KV-cache) support, loop-aware ------------------------ # State is keyed by the plan slot (cache_idx), because a looped stack calls # this same module once per pass with different activations. Sharing one slot # across passes folds each token in loop_count times and is worse than no fix. # # Exact, not approximate: M^(i) = sum_m phi(k_m) v_m^T is a plain sum and the # own-segment term is tril-masked, so a running sum over m<=l reproduces the # full path's einsum for the current token. def reset_state(self): self._st = {} def _tok(self, proj, x_t): B = x_t.shape[0] return proj(x_t).view(B, self.h, self.d) @torch.no_grad() def _capture(self, x, key): B, S, H = x.shape h, d, L = self.h, self.d, self.L k = self._phi(self.k_proj(x).view(B, S, h, d)).permute(0, 2, 1, 3) v = self.v_proj(x).view(B, S, h, d).permute(0, 2, 1, 3) pp = self.p_proj(x) nc = S // L dev, dt = x.device, x.dtype if nc: ks = k[:, :, :nc * L].reshape(B, h, nc, L, d) vs = v[:, :, :nc * L].reshape(B, h, nc, L, d) Ms = torch.einsum('bhnmd,bhnme->bhnde', ks, vs) zs = ks.sum(3) pools = pp[:, :nc * L].reshape(B, nc, L, self.dg).mean(2) else: Ms = torch.zeros(B, h, 0, d, d, device=dev, dtype=dt) zs = torch.zeros(B, h, 0, d, device=dev, dtype=dt) pools = torch.zeros(B, 0, self.dg, device=dev, dtype=dt) r = S - nc * L if r: kr, vr = k[:, :, nc * L:], v[:, :, nc * L:] M_run = torch.einsum('bhmd,bhme->bhde', kr, vr) z_run = kr.sum(2) pp_sum = pp[:, nc * L:].sum(1) else: M_run = torch.zeros(B, h, d, d, device=dev, dtype=dt) z_run = torch.zeros(B, h, d, device=dev, dtype=dt) pp_sum = torch.zeros(B, self.dg, device=dev, dtype=dt) self._st[key] = dict(Ms=Ms, zs=zs, pools=pools, M_run=M_run, z_run=z_run, pp_sum=pp_sum, n=r, npos=S, B=B) def _step(self, x_t, key): st = self._st[key] B = x_t.shape[0] h, d, L = self.h, self.d, self.L q = self._phi(self._tok(self.q_proj, x_t)) k = self._phi(self._tok(self.k_proj, x_t)) v = self._tok(self.v_proj, x_t) u = self.u_proj(x_t).view(B, self.dg) pp = self.p_proj(x_t).view(B, self.dg) st["M_run"] = st["M_run"] + torch.einsum('bhd,bhe->bhde', k, v) st["z_run"] = st["z_run"] + k st["pp_sum"] = st["pp_sum"] + pp st["n"] += 1 num_d = torch.einsum('bhd,bhde->bhe', q, st["M_run"]) den_d = torch.einsum('bhd,bhd->bh', q, st["z_run"]) nc = st["Ms"].shape[2] if nc: num_c = torch.einsum('bhd,bhide->bhie', q, st["Ms"]) den_c = torch.einsum('bhd,bhid->bhi', q, st["zs"]) score_c = torch.einsum('bg,big->bi', u, st["pools"]) else: num_c = num_d.new_zeros(B, h, 0, d) den_c = den_d.new_zeros(B, h, 0) score_c = u.new_zeros(B, 0) score_d = (u * (st["pp_sum"] / st["n"])).sum(-1, keepdim=True) gamma = torch.softmax(torch.cat([score_c, score_d], -1), dim=-1) num_all = torch.cat([num_c, num_d.unsqueeze(2)], dim=2) den_all = torch.cat([den_c, den_d.unsqueeze(2)], dim=2) g = gamma[:, None] y_num = (num_all * g.unsqueeze(-1)).sum(2) y_den = (den_all * g).sum(2) y = y_num / (y_den.unsqueeze(-1) + self.eps) y = self.o_proj(y.reshape(B, 1, h * d)) st["npos"] += 1 if st["n"] == L: st["Ms"] = torch.cat([st["Ms"], st["M_run"].unsqueeze(2)], dim=2) st["zs"] = torch.cat([st["zs"], st["z_run"].unsqueeze(2)], dim=2) st["pools"] = torch.cat([st["pools"], (st["pp_sum"] / L).unsqueeze(1)], dim=1) st["M_run"] = torch.zeros_like(st["M_run"]) st["z_run"] = torch.zeros_like(st["z_run"]) st["pp_sum"] = torch.zeros_like(st["pp_sum"]) st["n"] = 0 return y class TransformerBlock(nn.Module): def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int): super().__init__() self.use_hc = cfg.use_hyper_connections self.hidden_dropout = cfg.hidden_dropout self.use_value_embed = getattr(cfg, "use_value_embed", False) self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.attn = MLADerfXSAAttention(cfg, layer_idx=layer_idx) self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) # Memory Caching branch (arXiv 2602.24281), parallel to attention. Its own # norm + zero-init gate => exact no-op at init. self.use_mc = getattr(cfg, "use_memory_cache", False) if self.use_mc: self.mc_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.mc = MemoryCacheBranch(cfg) self.mc_gate = nn.Parameter(torch.zeros(1)) # Mid-stack Engram re-injection, only on the configured layers. fusion_layers = getattr(cfg, "engram_fusion_layers", None) or [] self.engram_fusion = ( EngramLayerFusion(cfg.hidden_size, getattr(cfg, "engram_fusion_dim", 256), cfg.rms_norm_eps) if (cfg.use_engram and layer_idx in fusion_layers) else None ) if cfg.use_moe and layer_idx in cfg.moe_layers: self.ffn = SparseMoEFFN(cfg, layer_idx) self.is_moe = True else: self.ffn = DenseFFN(cfg) self.is_moe = False if self.use_hc: self.hc_attn = HyperConnectionLayer(cfg.hidden_size, cfg.hc_mult, cfg.hc_sinkhorn_iters, cfg.hc_eps) self.hc_ffn = HyperConnectionLayer(cfg.hidden_size, cfg.hc_mult, cfg.hc_sinkhorn_iters, cfg.hc_eps) # NEW (v2, opt-in): value-embedding residual. Zero-init gate -> exact # no-op at init; learns to mix raw token-embedding signal into each # block's input (nanoGPT-speedrun "value embedding"/U-net skip family; # consistent wins at the 50-500M scale). if self.use_value_embed: self.ve_gate = nn.Parameter(torch.zeros(1)) def forward( self, x: torch.Tensor, # [B, hc_mult, S, H] if HC else [B, S, H] position_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple] = None, use_cache: bool = False, token_embed: Optional[torch.Tensor] = None, # [B, S, H] (value-embed) ngram_embed: Optional[torch.Tensor] = None, # [B, S, H] (Engram fusion) depth_kv: Optional[Tuple] = None, # ((k,v), ...) earlier layers mc_key: int = 0, # plan slot: per (layer, pass) ) -> Tuple[torch.Tensor, Optional[Tuple], Optional[torch.Tensor], Optional[Tuple]]: # --- Attention sub-layer --- if self.use_hc: h = self.hc_attn.pre_op(x) else: h = x if self.use_value_embed and token_embed is not None: h = h + torch.tanh(self.ve_gate) * token_embed attn_out, present, depth_entry = self.attn( self.attn_norm(h), position_ids, attention_mask, past_key_value, use_cache, depth_kv, ) attn_out = F.dropout(attn_out, p=self.hidden_dropout, training=self.training) # Memory Caching: fold the gated linear-memory readout into the attention # sub-layer contribution (zero-init gate => no-op at start). h is [B,S,H]. if self.use_mc: mc_out = self.mc(self.mc_norm(h), position_ids, attention_mask, mc_key) mc_out = F.dropout(mc_out, p=self.hidden_dropout, training=self.training) attn_out = attn_out + torch.tanh(self.mc_gate) * mc_out if self.use_hc: x = self.hc_attn.post_op(x, attn_out) h = self.hc_ffn.pre_op(x) else: h = h + attn_out # --- FFN sub-layer --- ffn_out = self.ffn(self.ffn_norm(h), position_ids) ffn_out = F.dropout(ffn_out, p=self.hidden_dropout, training=self.training) if self.use_hc: x = self.hc_ffn.post_op(x, ffn_out) else: x = h + ffn_out # Engram re-injection sits AFTER the block, on the residual stream, so # it feeds the next layer rather than competing inside this one. Under # hyper-connections x is [B, hc_mult, S, H]; broadcast the fusion across # streams by folding the stream axis into the batch. if self.engram_fusion is not None and ngram_embed is not None: if x.dim() == 4: B, M, S, H = x.shape fused = self.engram_fusion( x.reshape(B * M, S, H), ngram_embed.unsqueeze(1).expand(B, M, S, H).reshape(B * M, S, H), ) x = fused.view(B, M, S, H) else: x = self.engram_fusion(x, ngram_embed) return x, present, self.ffn.get_aux_loss(), depth_entry # --------------------------------------------------------------------------- # HRM refinement (unchanged) # --------------------------------------------------------------------------- class HRMRefinementBlock(nn.Module): """Iterative refinement of the final hidden state. DEAD-MODULE FIX. `gate` and `up.weight` were BOTH zero-initialized, and they sit on the same multiplicative path: h = h + tanh(gate[t]) * up(silu(down(...))) so dL/d(up) is proportional to tanh(gate) = 0, dL/d(gate) is proportional to up(...) = 0, and dL/d(down) and dL/d(norm) are proportional to up.weight = 0. Every tensor in the block received exactly zero gradient on every step: HRM refinement was frozen at the identity for the entire run, and deep supervision was supervising `hrm_refine_steps` identical copies of x. No error, no NaN, just a headline feature that silently did nothing -- the same failure mode as the Engram hash collapse. Only ONE endpoint of a multiplicative path may be zero-initialized. `up` keeps its zero (that is what makes the block an exact no-op at step 0) and `gate` now starts at `gate_init`, so gradient reaches `up` immediately and the rest of the block unfreezes as soon as `up` leaves zero. """ def __init__(self, hidden_size: int, refine_dim: int, steps: int, eps: float = 1e-6, gate_init: float = 1.0): super().__init__() self.steps = steps self.norm = RMSNorm(hidden_size, eps) self.down = nn.Linear(hidden_size * 2, refine_dim, bias=False) self.up = nn.Linear(refine_dim, hidden_size, bias=False) self.gate = nn.Parameter(torch.full((steps,), float(gate_init))) nn.init.normal_(self.down.weight, std=0.02) nn.init.zeros_(self.up.weight) def forward(self, x: torch.Tensor, return_all: bool = False): """return_all=True additionally returns the per-step states [h_1..h_T], so every refinement iteration can be supervised (LDT eq. 1). Default False keeps the original single-tensor return for existing callers.""" anchor = x h = x states = [] if return_all else None for t in range(self.steps): inp = torch.cat([self.norm(h), anchor], dim=-1) update = self.up(F.silu(self.down(inp))) h = h + torch.tanh(self.gate[t]) * update if return_all: states.append(h) return (h, states) if return_all else h # --------------------------------------------------------------------------- # Full model # --------------------------------------------------------------------------- def loop_layer_plan(num_layers: int, loop_count: int, loop_mode: str = "full", loop_middle_layers: Optional[int] = None): """Execution order for the looped stack. Returns (plan, num_pass_slots) where plan is a list of (layer_idx, pass_slot) and pass_slot indexes the per-pass embedding (None = not in a looped region). The list position IS the KV cache slot, which keeps cache indexing correct for both modes without either one having to know the other's layout. "full" -- the original Byrne-700M-Looped behaviour: the whole stack, `loop_count` times. Effective depth L*loop_count. "middle_split" -- Nanbeige-4.2 `enable_double_loop_split`: run a fixed prefix unlooped, repeat the middle block `loop_count` times, then a fixed suffix unlooped. Effective depth (L - M) + M*loop_count for a middle of width M, at correspondingly lower compute than "full". """ if loop_count <= 1 or loop_mode == "full": return ([(l, p if loop_count > 1 else None) for p in range(max(1, loop_count)) for l in range(num_layers)], max(1, loop_count)) if loop_mode != "middle_split": raise ValueError(f"unknown loop_mode {loop_mode!r} (expected 'full' or 'middle_split')") M = loop_middle_layers if loop_middle_layers is not None else num_layers // 2 if M <= 0 or num_layers % M != 0: raise ValueError( f"loop_middle_layers={M} must be positive and divide " f"num_hidden_layers={num_layers}" ) start = (num_layers - M) // 2 end = start + M plan = [(l, None) for l in range(start)] plan += [(l, p) for p in range(loop_count) for l in range(start, end)] plan += [(l, None) for l in range(end, num_layers)] return plan, loop_count class SpikeWhaleModel(nn.Module): """Decoder stack without LM head.""" def __init__(self, cfg: SpikeWhaleConfig): super().__init__() self.cfg = cfg self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) nn.init.normal_(self.embed_tokens.weight, std=cfg.initializer_range) self.engram = EngramModule(cfg) if cfg.use_engram else None self.layers = nn.ModuleList([ TransformerBlock(cfg, layer_idx=i) for i in range(cfg.num_hidden_layers) ]) self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) self.hc_out_mix = ( HCOutputMix(cfg.hc_mult) if cfg.use_hyper_connections else None ) self.hrm_refine = ( HRMRefinementBlock(cfg.hidden_size, cfg.hrm_refine_dim, cfg.hrm_refine_steps, cfg.rms_norm_eps, gate_init=getattr(cfg, "hrm_gate_init", 1.0)) if getattr(cfg, "use_hrm_refine", False) else None ) self.use_value_embed = getattr(cfg, "use_value_embed", False) self.deep_supervision = ( getattr(cfg, "hrm_deep_supervision", False) and self.hrm_refine is not None ) # --- looped transformer --- self.loop_count = max(1, int(getattr(cfg, "loop_count", 1))) # Per-pass embedding so the shared weights can condition on which # iteration they're in. Zero-init => loop_count=1 is bit-identical to # the dense baseline, and a fresh looped run starts from the same point. self.loop_mode = getattr(cfg, "loop_mode", "full") self.loop_plan, n_pass_slots = loop_layer_plan( cfg.num_hidden_layers, self.loop_count, self.loop_mode, getattr(cfg, "loop_middle_layers", None), ) # How many router invocations the plan makes, relative to one pass over # the stack -- used to keep the MoE aux-loss coefficient comparable # across loop modes (see the normalization in forward()). self.loop_router_scale = len(self.loop_plan) / cfg.num_hidden_layers # Cache layout: slots [0, _n_cache_slots) are the per-(pass, layer) K/V # entries; if Engram is on, one extra trailing slot carries its n-gram # context (see forward()). self._n_cache_slots = len(self.loop_plan) self.loop_pass_embed = None if self.loop_count > 1 and getattr(cfg, "loop_pass_embed", True): self.loop_pass_embed = nn.Parameter( torch.zeros(n_pass_slots, cfg.hidden_size)) self.use_depth_attention = getattr(cfg, "use_depth_attention", False) # The hc_mult residual streams start as IDENTICAL copies of x, and # HyperConnectionLayer.pre_op is a softmax-weighted sum over them -- on # identical streams that sum is invariant to the weights, so the very # first pre-mix in the network has an exactly-zero gradient. Under # loop_mode="full" layer 0 runs again on later passes (streams have # diverged by then) and it recovers; under "middle_split" layer 0 runs # exactly ONCE, so its pre_weight would be dead for the whole run. # Freeze it rather than ship a parameter that cannot train. if cfg.use_hyper_connections and self.loop_plan[0][0] == 0: runs_layer0_once = sum(1 for l, _ in self.loop_plan if l == 0) == 1 if runs_layer0_once: self.layers[0].hc_attn.pre_weight.requires_grad_(False) self.gradient_checkpointing = False def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_values: Optional[List[Tuple]] = None, use_cache: bool = False, ) -> Tuple[torch.Tensor, Optional[List[Tuple]], torch.Tensor]: B, S = input_ids.shape device = input_ids.device if position_ids is None: past_len = past_key_values[0][0].shape[2] if past_key_values else 0 position_ids = torch.arange( past_len, past_len + S, device=device ).unsqueeze(0).expand(B, -1) x = self.embed_tokens(input_ids) token_embed = x if self.use_value_embed else None # Read the n-gram memory ONCE and reuse it for both the embedding-level # injection and every mid-stack fusion layer (hashing is not free, and # re-hashing a drifted hidden state would address different buckets). ngram_embed = None engram_tail = None if self.engram is not None: # Incremental decode passes S=1, which would drop every n-gram order # above 1. The cache carries the last (max_ngram - 1) compressed # vectors of the prefix so the n-grams straddling the boundary are # still computed. See MultiHeadHashLookup.forward. prefix = None if past_key_values is not None and len(past_key_values) > self._n_cache_slots: prefix = past_key_values[self._n_cache_slots] compressed = self.engram.compress(x) ngram_embed = self.engram.lookup(compressed, prefix=prefix) x = x + self.engram(x, retrieved=ngram_embed) if use_cache: keep = max(0, self.cfg.engram_max_ngram - 1) full = compressed if prefix is None else \ torch.cat([prefix.to(compressed.dtype), compressed], dim=1) # max(0, ...) rather than relying on Python's negative-index # clamping: a prompt shorter than the n-gram window would # otherwise index with a negative start and only work by accident. start = max(0, full.shape[1] - keep) engram_tail = full[:, start:, :].detach() if keep else None if self.cfg.use_hyper_connections: x = x.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1, -1).clone() present_key_values = [] if use_cache else None total_aux_loss = torch.tensor(0.0, device=device) # Gradient checkpointing is incompatible with use_cache (the cache from # the discarded forward would be silently wrong on recompute). assert not (self.gradient_checkpointing and self.training and use_cache), \ "use_cache=True is not supported with gradient checkpointing" # LOOPED TRANSFORMER: run the whole stack `loop_count` times through the # SAME weights. Effective depth = loop_count * num_hidden_layers at the # parameter cost of num_hidden_layers. loop_count=1 reduces exactly to # the dense baseline (the pass embedding is skipped entirely). # # KV cache: every (pass, layer) pair produces its own K/V, because the # same weights see different activations on each pass. The cache is laid # out flat as [pass0_layer0 .. pass0_layerN, pass1_layer0 ..], indexed by # `pass_idx * num_layers + layer_idx`. Getting this wrong would silently # feed pass 1's keys to pass 0 and corrupt every cached decode. # # The plan (see loop_layer_plan) also covers loop_mode="middle_split", # where only the middle block repeats. The plan's list position is the # cache slot, so both modes get correct, non-overlapping cache indexing. # # Depth attention accumulates (k, v) from checkpointed layers as the # plan advances. It is RESET at the start of every looped pass: within a # pass, depth means "earlier layer"; carrying the list across passes # would grow it by loop_count and make pass 3 attend to a stale pass-1 # value of the same layer, which is a different (and unvalidated) # mechanism. The list is threaded through as an explicit argument rather # than mutated inside the layer, so gradient checkpointing's recompute # cannot double-append. prev_slot = None depth_kv: List[Tuple[torch.Tensor, torch.Tensor]] = [] depth_prefix: List[Tuple[torch.Tensor, torch.Tensor]] = [] for cache_idx, (layer_idx, pass_slot) in enumerate(self.loop_plan): layer = self.layers[layer_idx] if pass_slot is not None and pass_slot != prev_slot: if self.loop_pass_embed is not None: # [H] broadcasts over both [B,S,H] and the hyper-connection # [B,hc_mult,S,H] layout. x = x + self.loop_pass_embed[pass_slot] # Entering a new pass: keep whatever the UNLOOPED prefix # contributed (it ran once and is still "earlier depth"), drop # the previous pass's own entries. if prev_slot is None: depth_prefix = list(depth_kv) depth_kv = list(depth_prefix) prev_slot = pass_slot pkv = past_key_values[cache_idx] if past_key_values else None dkv = tuple(depth_kv) if self.use_depth_attention else None if self.gradient_checkpointing and self.training: x, present, aux_loss, depth_entry = gradient_checkpoint( layer, x, position_ids, attention_mask, None, False, token_embed, ngram_embed, dkv, cache_idx, use_reentrant=False, ) else: x, present, aux_loss, depth_entry = layer( x, position_ids, attention_mask, pkv, use_cache, token_embed, ngram_embed, dkv, cache_idx) if depth_entry is not None: depth_kv.append(depth_entry) if use_cache: present_key_values.append(present) if aux_loss is not None: total_aux_loss = total_aux_loss + aux_loss # Normalize the MoE aux loss by how many times the plan invoked a # router relative to a single pass over the stack: looping produces # proportionally more router invocations, so without this the aux # coefficient would effectively be scaled up in the looped variant and # the two runs would not be comparing like with like. (For loop_mode # "middle_split" the factor is not loop_count -- hence the plan-derived # ratio rather than a hardcoded divide.) if self.loop_router_scale > 1.0: total_aux_loss = total_aux_loss / self.loop_router_scale # Trailing cache slot for the Engram n-gram context (see above). if use_cache and engram_tail is not None: present_key_values.append(engram_tail) if self.cfg.use_hyper_connections: x = self.hc_out_mix(x) # v2: learned mix (init == mean) # Deep supervision (LDT): keep each refinement step's state so the LM # head can be applied to all of them. Only collected while training -- # at inference we read the final iteration only, as the paper does. refine_states = None if self.hrm_refine is not None: if self.deep_supervision and self.training: x, states = self.hrm_refine(x, return_all=True) refine_states = [self.norm(s) for s in states] else: x = self.hrm_refine(x) x = self.norm(x) return x, present_key_values, total_aux_loss, refine_states class MTPHead(nn.Module): """ v2 MTP head: small zero-init H x H projection feeding the SHARED lm_head. Cost per head: H^2 params (e.g. 1M at H=1024) instead of H*V (e.g. 50M+). Zero-init means at step 0 the head predicts exactly what lm_head predicts for the residual path = 0, i.e. uniform-ish gradient pressure; the residual form (x + proj(x)) keeps it anchored to the trunk representation. """ def __init__(self, hidden_size: int): super().__init__() self.proj = nn.Linear(hidden_size, hidden_size, bias=False) nn.init.zeros_(self.proj.weight) def forward(self, hidden: torch.Tensor) -> torch.Tensor: return hidden + self.proj(hidden) class SpikeWhaleLM(PreTrainedModel): """ v2 loss = CE + zloss_coef * z-loss + mtp_loss_weight * mean(MTP CE) + MoE aux loss """ config_class = SpikeWhaleConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["TransformerBlock"] def __init__(self, cfg: SpikeWhaleConfig): super().__init__(cfg) self.model = SpikeWhaleModel(cfg) self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) nn.init.normal_(self.lm_head.weight, std=cfg.initializer_range) self.zloss_coef = getattr(cfg, "zloss_coef", 1e-4) self.mtp_loss_weight = getattr(cfg, "mtp_loss_weight", 0.3) # v2 MTP: H x H residual projections sharing lm_head (see MTPHead). self.mtp_heads = nn.ModuleList([ MTPHead(cfg.hidden_size) for _ in range(cfg.num_nextn_predict_layers) ]) if cfg.num_nextn_predict_layers > 0 else None # LDT-adapted abstention head (see config.py / docs/LDT-notes.md). # Zero-init => logit 0 => p(abstain) = 0.5 everywhere at step 0, and it # contributes no gradient to the trunk until it learns something. self.use_abstain_head = getattr(cfg, "use_abstain_head", False) self.abstain_loss_weight = getattr(cfg, "abstain_loss_weight", 0.1) self.abstain_pos_weight = getattr(cfg, "abstain_pos_weight", 8.0) self.abstain_threshold = getattr(cfg, "abstain_threshold", 0.6) self.abstain_head = ( nn.Linear(cfg.hidden_size, 1, bias=True) if self.use_abstain_head else None ) self.post_init() # Zero-init AFTER post_init: HF's _init_weights runs inside post_init and # would otherwise overwrite these, leaving the head with a random opinion # about its own correctness before it has learned anything. if self.abstain_head is not None: nn.init.zeros_(self.abstain_head.weight) nn.init.zeros_(self.abstain_head.bias) # Same reason: EngramLayerFusion.out_proj, MTPHead.proj and # HRMRefinementBlock.up are all zero-init BY DESIGN (each makes its # module an exact no-op at step 0), and post_init() overwrites all three # with normal(0, initializer_range). Restore them here. for layer in self.model.layers: if getattr(layer, "engram_fusion", None) is not None: nn.init.zeros_(layer.engram_fusion.out_proj.weight) if self.model.hrm_refine is not None: nn.init.zeros_(self.model.hrm_refine.up.weight) for head in (self.mtp_heads or ()): # None when MTP is disabled nn.init.zeros_(head.proj.weight) def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def tie_weights(self, **kwargs): if self.config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight def save_pretrained(self, *args, **kwargs): tied = ( self.config.tie_word_embeddings and self.lm_head.weight.data_ptr() == self.model.embed_tokens.weight.data_ptr() ) if tied: self.lm_head.weight = nn.Parameter(self.model.embed_tokens.weight.detach().clone()) try: super().save_pretrained(*args, **kwargs) finally: if tied: self.lm_head.weight = self.model.embed_tokens.weight def _set_gradient_checkpointing(self, module, value=False): if isinstance(module, SpikeWhaleModel): module.gradient_checkpointing = value def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_values: Optional[List[Tuple]] = None, labels: Optional[torch.Tensor] = None, use_cache: bool = False, **kwargs, ) -> CausalLMOutputWithPast: hidden, present_kvs, aux_loss, refine_states = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, ) logits = self.lm_head(hidden) loss = None # LDT abstention signal: p(this prediction is wrong). Exposed on the # output so decoding can gate on it; `abstain_threshold` is deliberately # above 0.5 (paper raises theta_CLS at inference vs training). abstain_logit = None if self.abstain_head is not None: abstain_logit = self.abstain_head(hidden).squeeze(-1) # [B, S] if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() flat_logits = shift_logits.view(-1, shift_logits.size(-1)) flat_labels = shift_labels.view(-1) loss = F.cross_entropy(flat_logits, flat_labels, ignore_index=-100) # Deep supervision (LDT eq. 1): every refinement iteration gets its # own CE, all weighted equally, averaged with the final-state CE. # refine_states is populated only while training with the flag on. if refine_states: step_losses = [loss] for s in refine_states[:-1]: # last state == `hidden` s_logits = self.lm_head(s)[..., :-1, :].contiguous() step_losses.append(F.cross_entropy( s_logits.view(-1, s_logits.size(-1)), flat_labels, ignore_index=-100)) loss = torch.stack(step_losses).mean() # Abstention head: asymmetric BCE against "top-1 was wrong here". # pos_weight>1 makes missing a genuine error costlier than a false # alarm -- the paper's w+/w- = 8 asymmetry, pointed at the error # direction that actually matters for abstention. if abstain_logit is not None and self.abstain_loss_weight > 0: valid = flat_labels != -100 if valid.any(): with torch.no_grad(): wrong = (flat_logits.argmax(-1) != flat_labels).float() a_flat = abstain_logit[..., :-1].contiguous().view(-1) abstain_loss = F.binary_cross_entropy_with_logits( a_flat[valid].float(), wrong[valid], pos_weight=torch.tensor( self.abstain_pos_weight, device=loss.device)) loss = loss + self.abstain_loss_weight * abstain_loss # z-loss (v2): penalize log^2 of the partition function on valid # positions. Keeps logits from drifting; pairs well with Muon. if self.zloss_coef > 0: valid = flat_labels != -100 if valid.any(): log_z = torch.logsumexp(flat_logits[valid].float(), dim=-1) loss = loss + self.zloss_coef * (log_z ** 2).mean() # MTP (v2): residual H x H head -> shared lm_head, down-weighted. if self.mtp_heads is not None and self.mtp_loss_weight > 0: mtp_total = torch.tensor(0.0, device=loss.device) n_active = 0 for k, head in enumerate(self.mtp_heads, start=1): offset = k + 1 if hidden.size(1) > offset: mtp_hidden = head(hidden[..., :-offset, :]) mtp_logits = self.lm_head(mtp_hidden) mtp_labels = labels[..., offset:].contiguous() mtp_total = mtp_total + F.cross_entropy( mtp_logits.reshape(-1, mtp_logits.size(-1)), mtp_labels.reshape(-1), ignore_index=-100, ) n_active += 1 if n_active > 0: loss = loss + self.mtp_loss_weight * mtp_total / n_active loss = loss + aux_loss out = CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=present_kvs, ) # Attach the abstention signal without breaking the HF output contract. # p_abstain[b, t] = P(the token predicted at position t is wrong). if abstain_logit is not None: out.abstain_logit = abstain_logit out.p_abstain = torch.sigmoid(abstain_logit) return out def count_parameters(self) -> int: return sum(p.numel() for p in self.parameters()) def reset_memory_cache(model): """Clear incremental MemoryCacheBranch state (all loop passes) on every layer. Call between independent generations that reuse one model object.""" n = 0 for m in model.modules(): if isinstance(m, MemoryCacheBranch): m.reset_state(); n += 1 return n