| """Senryu-2: a Gated DeltaNet-2 + Parallax hybrid on a 4-branch gated residual stream. |
| |
| Every sublayer reads its input from the residual branches through Qwen's |
| Gated Residual (per-branch RMSNorm, elementwise sigmoid gate from a low-rank |
| bottleneck over all branches) and writes its output back to every branch with |
| one data-dependent scalar per branch. A hashed n-gram memory, whose tables live |
| in host memory during training, is injected before block ``ngram_layer_index`` |
| (docs/architecture.md, "Senryu-2"). |
| |
| Training and scoring run packed: the batch is flattened into one row and split |
| into documents at every BOS token, so GDN-2, Parallax, the n-gram hashes and |
| the n-gram convolution never look across a document boundary. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| import torch.nn.functional as F |
| from fla.layers.gdn2 import GatedDeltaNet2 |
| from fla.models.utils import Cache as FlaCache |
| from fla.modules.rotary import rotary_embedding |
| from safetensors.torch import save_model |
| from torch import nn |
| from torch.utils.checkpoint import checkpoint |
|
|
| from tiny_gdn.config import TinyGDNConfig |
| from tiny_gdn.gated_residual_ops import fused_gated_residual_read, fused_gated_residual_write |
| from tiny_gdn.haiku_layers import SiTUGLU |
| from tiny_gdn.model import TinyGDNOutput, linear_cross_entropy |
| from tiny_gdn.ngram import ( |
| NGRAM_TABLES_EMA_FILE, |
| NGRAM_TABLES_FILE, |
| NGramLayout, |
| document_distance, |
| load_table_file, |
| ngram_row_indices, |
| ) |
| from tiny_gdn.nn_common import RMSNorm |
| from tiny_gdn.parallax_ops import parallax_attention, parallax_decode_cached |
|
|
| GATED_RESIDUAL_CHUNK_ROWS = 4096 |
| KV_CACHE_GRANULE = 256 |
|
|
|
|
| def mark_muon_row_blocks(parameter: nn.Parameter, blocks: int) -> None: |
| """Orthogonalize ``parameter`` as ``blocks`` equal row blocks (heads, fused halves).""" |
|
|
| if parameter.shape[0] % blocks: |
| raise ValueError(f"{blocks} row blocks do not divide {parameter.shape[0]} rows") |
| parameter._muon_row_blocks = blocks |
|
|
|
|
| def mark_adamw_only(*parameters: nn.Parameter) -> None: |
| """Keep 2D weights that are gates or elongated low-rank maps on AdamW.""" |
|
|
| for parameter in parameters: |
| parameter._adamw_only = True |
|
|
|
|
| def mark_gain(*parameters: nn.Parameter) -> None: |
| """2D per-branch norm gains and filters: AdamW without weight decay, like 1D gains.""" |
|
|
| mark_adamw_only(*parameters) |
| for parameter in parameters: |
| parameter._no_weight_decay = True |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Senryu2Context: |
| """Per-forward state shared by every block (passes through activation checkpointing).""" |
|
|
| batch: int |
| length: int |
| cu_seqlens: torch.Tensor | None |
| distance: torch.Tensor | None |
| rotary: tuple[torch.Tensor, torch.Tensor] |
| position_offset: int = 0 |
| cache: Senryu2Cache | None = None |
|
|
|
|
| class RotaryTables: |
| """Partial-RoPE cos/sin rows ``[length, rotary_dim / 2]``, built in FP32 and cached.""" |
|
|
| def __init__(self, rotary_dim: int, theta: float) -> None: |
| self.rotary_dim = rotary_dim |
| self.theta = theta |
| self._cached: tuple[torch.Tensor, torch.Tensor] | None = None |
| self._key: tuple[torch.device, torch.dtype, bool] | None = None |
|
|
| def get(self, length: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: |
| key = (device, dtype, torch.is_inference_mode_enabled()) |
| cached = self._cached |
| if cached is None or self._key != key or cached[0].shape[0] < length: |
| rows = max(length, 4096 if cached is None else 2 * cached[0].shape[0]) |
| inverse_frequency = 1.0 / ( |
| self.theta |
| ** (torch.arange(0, self.rotary_dim, 2, device=device, dtype=torch.float32) / self.rotary_dim) |
| ) |
| angles = torch.outer(torch.arange(rows, device=device, dtype=torch.float32), inverse_frequency) |
| cached = (angles.cos().to(dtype), angles.sin().to(dtype)) |
| self._cached = cached |
| self._key = key |
| return cached |
|
|
|
|
| def parallax_reference( |
| q: torch.Tensor, |
| r: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| *, |
| scale: float, |
| cu_seqlens: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """FP32 Parallax for queries at the last ``Tq`` of ``Tk`` key positions. |
| |
| ``out = O1 / d1 * (1 + d2 / d1) - O2 / d1`` with ``p1 = exp(scale q·k)``, |
| ``p2 = p1 * (r·k)``; ``cu_seqlens`` (batch 1, ``Tq == Tk``) masks across |
| packed documents. |
| """ |
|
|
| batch, queries, query_heads, head_dim = q.shape |
| keys, kv_heads = k.shape[1], k.shape[2] |
| groups = query_heads // kv_heads |
| qf = q.float().reshape(batch, queries, kv_heads, groups, head_dim) |
| rf = r.float().reshape(batch, queries, kv_heads, groups, head_dim) |
| kf, vf = k.float(), v.float() |
| s1 = torch.einsum("bqhgd,bkhd->bhgqk", qf, kf) * scale |
| s2 = torch.einsum("bqhgd,bkhd->bhgqk", rf, kf) |
| query_positions = torch.arange(keys - queries, keys, device=q.device)[:, None] |
| key_positions = torch.arange(keys, device=q.device)[None, :] |
| masked = key_positions > query_positions |
| if cu_seqlens is not None: |
| segment = torch.zeros(keys, dtype=torch.long, device=q.device) |
| segment[cu_seqlens[1:-1].long()] = 1 |
| segment = segment.cumsum(0) |
| masked = masked | (segment[keys - queries :, None] != segment[None, :]) |
| s1 = s1.masked_fill(masked, float("-inf")) |
| p1 = (s1 - s1.amax(dim=-1, keepdim=True)).exp() |
| d1 = p1.sum(dim=-1) |
| p2 = p1 * s2 |
| d2 = p2.sum(dim=-1) |
| o1 = torch.einsum("bhgqk,bkhd->bqhgd", p1, vf) |
| o2 = torch.einsum("bhgqk,bkhd->bqhgd", p2, vf) |
| inverse_d1 = (1.0 / d1).permute(0, 3, 1, 2)[..., None] |
| correction = (d2 / d1).permute(0, 3, 1, 2)[..., None] |
| out = o1 * inverse_d1 * (1.0 + correction) - o2 * inverse_d1 |
| return out.reshape(batch, queries, query_heads, head_dim).to(q.dtype) |
|
|
|
|
| |
|
|
|
|
| class _GatedResidualRead(torch.autograd.Function): |
| """``x = (1/n) Σ σ(W_u SiLU(W_d vec(R̂) / n))_i ⊙ R̂_i`` and ``s = 2σ(W_w vec(R̂) / n)``. |
| |
| Saves only the branches, the per-branch inverse RMS and the bottleneck |
| pre-activation; the normalized branches and gates are recomputed in |
| chunks during backward, so a 32k-token read keeps ~0.4 GB instead of ~4 GB. |
| CUDA runs the fused kernels in ``gated_residual_ops`` instead. |
| """ |
|
|
| @staticmethod |
| def forward(ctx, residual, gamma, down, up, write, eps): |
| rows, branches, hidden_size = residual.shape |
| dtype = down.dtype |
| gain = 1.0 + gamma.float() |
| out = residual.new_empty(rows, hidden_size) |
| rstd = residual.new_empty(rows, branches, dtype=torch.float32) |
| hidden = residual.new_empty(rows, down.shape[0]) |
| scales = residual.new_empty(rows, 0 if write is None else branches, dtype=torch.float32) |
| for start in range(0, rows, GATED_RESIDUAL_CHUNK_ROWS): |
| stop = min(start + GATED_RESIDUAL_CHUNK_ROWS, rows) |
| values = residual[start:stop].float() |
| chunk_rstd = torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps) |
| hat = values * chunk_rstd * gain |
| flat = hat.to(dtype).flatten(1) |
| chunk_hidden = (flat @ down.t()) / branches |
| activated = F.silu(chunk_hidden.float()).to(dtype) |
| gates = torch.sigmoid((activated @ up.t()).float()).view(stop - start, branches, hidden_size) |
| out[start:stop] = ((gates * hat).sum(dim=1) / branches).to(out.dtype) |
| rstd[start:stop] = chunk_rstd.squeeze(-1) |
| hidden[start:stop] = chunk_hidden |
| if write is not None: |
| scales[start:stop] = 2.0 * torch.sigmoid((flat @ write.t()).float() / branches) |
| ctx.save_for_backward(residual, gamma, down, up, write, rstd, hidden, scales) |
| scales_out = scales.to(residual.dtype, copy=True) |
| if write is None: |
| ctx.mark_non_differentiable(scales_out) |
| return out, scales_out |
|
|
| @staticmethod |
| def backward(ctx, grad_out, grad_scales): |
| residual, gamma, down, up, write, rstd, hidden, scales = ctx.saved_tensors |
| rows, branches, hidden_size = residual.shape |
| dtype = down.dtype |
| gain = 1.0 + gamma.float() |
| grad_residual = torch.empty(residual.shape, dtype=residual.dtype, device=residual.device) |
| grad_gamma = torch.zeros(branches, hidden_size, dtype=torch.float32, device=residual.device) |
| grad_down = torch.zeros(down.shape, dtype=torch.float32, device=down.device) |
| grad_up = torch.zeros(up.shape, dtype=torch.float32, device=up.device) |
| grad_write = None if write is None else torch.zeros(write.shape, dtype=torch.float32, device=write.device) |
| use_scales = write is not None and grad_scales is not None |
| for start in range(0, rows, GATED_RESIDUAL_CHUNK_ROWS): |
| stop = min(start + GATED_RESIDUAL_CHUNK_ROWS, rows) |
| count = stop - start |
| normalized = residual[start:stop].float() * rstd[start:stop, :, None] |
| hat = normalized * gain |
| flat = hat.to(dtype).flatten(1) |
| pre = hidden[start:stop].float() |
| sigmoid_pre = torch.sigmoid(pre) |
| activated = (pre * sigmoid_pre).to(dtype) |
| gates = torch.sigmoid((activated @ up.t()).float()).view(count, branches, hidden_size) |
| grad_x = grad_out[start:stop].float().unsqueeze(1) / branches |
| grad_hat = gates * grad_x |
| grad_gate_logits = (grad_x * hat * gates * (1.0 - gates)).flatten(1).to(dtype) |
| grad_up += (grad_gate_logits.t() @ activated).float() |
| grad_activated = (grad_gate_logits @ up).float() |
| grad_pre = grad_activated * sigmoid_pre * (1.0 + pre * (1.0 - sigmoid_pre)) |
| grad_pre = (grad_pre / branches).to(dtype) |
| grad_down += (grad_pre.t() @ flat).float() |
| grad_flat = (grad_pre @ down).float() |
| if use_scales: |
| scale_values = scales[start:stop] |
| grad_scale_logits = grad_scales[start:stop].float() * scale_values * (1.0 - scale_values / 2.0) |
| grad_scale_logits = (grad_scale_logits / branches).to(dtype) |
| grad_write += (grad_scale_logits.t() @ flat).float() |
| grad_flat += (grad_scale_logits @ write).float() |
| grad_hat = grad_hat + grad_flat.view(count, branches, hidden_size) |
| grad_gamma += (grad_hat * normalized).sum(dim=0) |
| grad_normalized = grad_hat * gain |
| grad_residual[start:stop] = ( |
| rstd[start:stop, :, None] |
| * (grad_normalized - normalized * (grad_normalized * normalized).mean(dim=-1, keepdim=True)) |
| ).to(grad_residual.dtype) |
| return ( |
| grad_residual, |
| grad_gamma.to(gamma.dtype), |
| grad_down.to(down.dtype), |
| grad_up.to(up.dtype), |
| None if grad_write is None else grad_write.to(write.dtype), |
| None, |
| ) |
|
|
|
|
| def gated_residual_read_reference( |
| residual: torch.Tensor, |
| gamma: torch.Tensor, |
| down: torch.Tensor, |
| up: torch.Tensor, |
| write: torch.Tensor | None, |
| eps: float, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Eager autograd version of :class:`_GatedResidualRead`, for tests.""" |
|
|
| rows, branches, hidden_size = residual.shape |
| gain = 1.0 + gamma.float() |
| values = residual.float() |
| hat = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + eps) * gain |
| flat = hat.to(down.dtype).flatten(1) |
| hidden = (flat @ down.t()) / branches |
| gates = torch.sigmoid((F.silu(hidden.float()).to(up.dtype) @ up.t()).float()).view(rows, branches, hidden_size) |
| out = ((gates * hat).sum(dim=1) / branches).to(residual.dtype) |
| if write is None: |
| return out, residual.new_empty(rows, 0) |
| scales = 2.0 * torch.sigmoid((flat @ write.t()).float() / branches) |
| return out, scales.to(residual.dtype) |
|
|
|
|
| class GatedResidual(nn.Module): |
| """Qwen's Gated Residual read (and write scales) over ``branches`` residual branches.""" |
|
|
| def __init__(self, branches: int, hidden_size: int, rank: int, eps: float, *, write: bool) -> None: |
| super().__init__() |
| self.branches = branches |
| self.eps = eps |
| self.weight = nn.Parameter(torch.zeros(branches, hidden_size)) |
| self.down_proj = nn.Linear(branches * hidden_size, rank, bias=False) |
| self.up_proj = nn.Linear(rank, branches * hidden_size, bias=False) |
| self.write_proj = nn.Linear(branches * hidden_size, branches, bias=False) if write else None |
| mark_gain(self.weight) |
| mark_adamw_only(self.down_proj.weight, self.up_proj.weight) |
| if self.write_proj is not None: |
| mark_adamw_only(self.write_proj.weight) |
|
|
| def read(self, residual: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """``(x [N, d], scales [N, n])`` from branches ``[N, n, d]``.""" |
|
|
| write = None if self.write_proj is None else self.write_proj.weight |
| read = fused_gated_residual_read if residual.is_cuda else _GatedResidualRead.apply |
| return read(residual, self.weight, self.down_proj.weight, self.up_proj.weight, write, self.eps) |
|
|
| @staticmethod |
| def write(residual: torch.Tensor, scales: torch.Tensor, update: torch.Tensor) -> torch.Tensor: |
| if residual.is_cuda: |
| return fused_gated_residual_write(residual, scales, update) |
| return residual + scales.unsqueeze(-1) * update.unsqueeze(-2) |
|
|
|
|
| |
|
|
|
|
| class ParallaxAttention(nn.Module): |
| """Parallax attention (softmax plus a learned first-order correction from ``r``). |
| |
| Grouped query/key/value heads with per-head zero-centered RMSNorm on q, r |
| and k, partial RoPE on q, r and k, and a full-rank sigmoid output gate. |
| """ |
|
|
| def __init__(self, config: TinyGDNConfig, layer_index: int) -> None: |
| super().__init__() |
| self.layer_index = layer_index |
| self.num_heads = config.num_attention_heads |
| self.num_kv_heads = config.num_key_value_heads |
| self.head_dim = config.attention_head_dim |
| self.scale = self.head_dim ** -0.5 |
| query_width = self.num_heads * self.head_dim |
| kv_width = self.num_kv_heads * self.head_dim |
| hidden = config.hidden_size |
| self.q_proj = nn.Linear(hidden, query_width, bias=False) |
| self.r_proj = nn.Linear(hidden, query_width, bias=False) |
| self.k_proj = nn.Linear(hidden, kv_width, bias=False) |
| self.v_proj = nn.Linear(hidden, kv_width, bias=False) |
| self.g_proj = nn.Linear(hidden, query_width, bias=True) |
| self.o_proj = nn.Linear(query_width, hidden, bias=False) |
| self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) |
| self.r_norm = RMSNorm(self.head_dim, config.rms_norm_eps) |
| self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) |
| for projection in (self.q_proj, self.r_proj): |
| mark_muon_row_blocks(projection.weight, self.num_heads) |
| for projection in (self.k_proj, self.v_proj): |
| mark_muon_row_blocks(projection.weight, self.num_kv_heads) |
| mark_adamw_only(self.g_proj.weight) |
|
|
| def forward(self, hidden_states: torch.Tensor, context: Senryu2Context) -> torch.Tensor: |
| batch, length, _ = hidden_states.shape |
| q = self.q_norm(self.q_proj(hidden_states).view(batch, length, self.num_heads, self.head_dim)) |
| r = self.r_norm(self.r_proj(hidden_states).view(batch, length, self.num_heads, self.head_dim)) |
| k = self.k_norm(self.k_proj(hidden_states).view(batch, length, self.num_kv_heads, self.head_dim)) |
| v = self.v_proj(hidden_states).view(batch, length, self.num_kv_heads, self.head_dim) |
| cosine, sine = context.rotary |
| q, r, k = ( |
| rotary_embedding( |
| tensor, |
| cosine, |
| sine, |
| seqlen_offsets=context.position_offset, |
| cu_seqlens=context.cu_seqlens, |
| ) |
| for tensor in (q, r, k) |
| ) |
| use_kernel = q.is_cuda and q.dtype in (torch.bfloat16, torch.float16) |
| cache = context.cache |
| if cache is not None: |
| keys, values = cache.append_attention(self.layer_index, k, v, context.position_offset) |
| if not use_kernel: |
| out = parallax_reference(q, r, keys, values, scale=self.scale) |
| elif keys.shape[1] == length: |
| out = parallax_attention(q, r, keys, values, scale=self.scale) |
| else: |
| out = parallax_decode_cached(q, r, keys, values, scale=self.scale) |
| elif use_kernel: |
| out = parallax_attention(q, r, k, v, scale=self.scale, cu_seqlens=context.cu_seqlens) |
| else: |
| out = parallax_reference(q, r, k, v, scale=self.scale, cu_seqlens=context.cu_seqlens) |
| gate = torch.sigmoid(self.g_proj(hidden_states)).view(batch, length, self.num_heads, self.head_dim) |
| return self.o_proj((out * gate).reshape(batch, length, self.num_heads * self.head_dim)) |
|
|
|
|
| def build_gated_deltanet2(config: TinyGDNConfig, layer_index: int) -> GatedDeltaNet2: |
| mixer = GatedDeltaNet2( |
| hidden_size=config.hidden_size, |
| expand_v=config.linear_expand_v, |
| head_dim=config.linear_head_dim, |
| num_heads=config.linear_num_heads, |
| num_v_heads=config.linear_num_value_heads, |
| mode="chunk", |
| use_short_conv=True, |
| allow_neg_eigval=config.allow_negative_eigenvalues, |
| conv_size=config.linear_conv_kernel_dim, |
| conv_bias=False, |
| layer_idx=layer_index, |
| norm_eps=config.rms_norm_eps, |
| ) |
| if config.linear_full_rank_gate: |
| |
| |
| mixer.g_proj = nn.Linear(config.hidden_size, mixer.value_dim, bias=True) |
| mark_adamw_only(mixer.g_proj.weight) |
| else: |
| mark_adamw_only(*(layer.weight for layer in mixer.g_proj)) |
| mark_adamw_only(*(layer.weight for layer in mixer.f_proj)) |
| for projection in (mixer.q_proj, mixer.k_proj, mixer.b_proj): |
| mark_muon_row_blocks(projection.weight, mixer.num_heads) |
| for projection in (mixer.v_proj, mixer.w_proj): |
| mark_muon_row_blocks(projection.weight, mixer.num_v_heads) |
| return mixer |
|
|
|
|
| |
|
|
|
|
| class BranchRMSNorm(nn.Module): |
| """Zero-centered RMSNorm over the last dim with one gain vector per branch.""" |
|
|
| def __init__(self, branches: int, hidden_size: int, eps: float) -> None: |
| super().__init__() |
| self.weight = nn.Parameter(torch.zeros(branches, hidden_size)) |
| self.eps = eps |
| mark_gain(self.weight) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| values = hidden_states.float() |
| values = values * torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + self.eps) |
| return values * (1.0 + self.weight.float()) |
|
|
|
|
| class NGramMemory(nn.Module): |
| """Engram-style memory: hashed n-gram rows, gated per residual branch, then a short conv. |
| |
| ``e`` concatenates one row per table. Each branch gates the shared value |
| ``W_V e`` with ``σ(RMSNorm(R_m) · RMSNorm(W_K,m e) / √d)``; the gated values |
| pass a depthwise causal convolution (dilation = largest order, zero init) |
| and are added to their branches. |
| """ |
|
|
| def __init__(self, config: TinyGDNConfig) -> None: |
| super().__init__() |
| self.layout = NGramLayout.from_config(config) |
| self.branches = config.residual_branches |
| self.hidden_size = config.hidden_size |
| self.kernel = config.ngram_conv_kernel |
| self.dilation = self.layout.max_order |
| memory_dim = self.layout.memory_dim |
| self.value_proj = nn.Linear(memory_dim, config.hidden_size, bias=False) |
| self.key_proj = nn.Linear(memory_dim, self.branches * config.hidden_size, bias=False) |
| self.hidden_norm = BranchRMSNorm(self.branches, config.hidden_size, config.rms_norm_eps) |
| self.key_norm = BranchRMSNorm(self.branches, config.hidden_size, config.rms_norm_eps) |
| self.conv_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.conv_weight = nn.Parameter(torch.zeros(config.hidden_size, self.kernel)) |
| mark_muon_row_blocks(self.key_proj.weight, self.branches) |
| mark_gain(self.conv_weight) |
| self.register_buffer("canonical_ids", torch.arange(config.vocab_size, dtype=torch.int64), persistent=False) |
| self.register_buffer("tables", None, persistent=False) |
|
|
| @property |
| def conv_history(self) -> int: |
| return (self.kernel - 1) * self.dilation |
|
|
| def set_canonical_ids(self, canonical: torch.Tensor) -> None: |
| if canonical.shape != self.canonical_ids.shape: |
| raise ValueError("canonical map does not match the vocabulary") |
| self.canonical_ids = canonical.to(device=self.canonical_ids.device, dtype=torch.int64) |
|
|
| def attach_tables(self, tables: torch.Tensor, canonical: torch.Tensor) -> None: |
| """Look rows up in ``tables [total_rows, embedding_dim]``. |
| |
| The tables may stay in host memory while the model runs on the GPU: |
| decoding reads only ``table_count`` rows per token. |
| """ |
|
|
| if tables.shape != (self.layout.total_rows, self.layout.embedding_dim): |
| raise ValueError(f"tables have shape {tuple(tables.shape)}") |
| self.set_canonical_ids(canonical) |
| self.tables = tables |
|
|
| def load_tables(self, path: Path, *, device: torch.device | str, dtype: torch.dtype) -> None: |
| tables, canonical = load_table_file(path, self.layout, device=device, dtype=dtype) |
| self.attach_tables(tables, canonical) |
|
|
| def lookup( |
| self, |
| input_ids: torch.Tensor, |
| *, |
| distance: torch.Tensor | None = None, |
| history: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Concatenated rows ``[..., memory_dim]`` for every position of ``input_ids``.""" |
|
|
| if self.tables is None: |
| raise RuntimeError( |
| "Senryu-2 needs n-gram embeddings: pass ngram_embeddings or attach tables " |
| "with model.ngram.load_tables(...)" |
| ) |
| canonical = self.canonical_ids[input_ids] |
| indices = ngram_row_indices(canonical, self.layout, distance=distance, history=history) |
| tables = self.tables |
| if tables.device == indices.device: |
| return F.embedding(indices, tables).flatten(-2) |
| rows = F.embedding(indices.to(tables.device), tables).flatten(-2) |
| return rows.to(indices.device, non_blocking=True) |
|
|
| def _convolve( |
| self, |
| values: torch.Tensor, |
| distance: torch.Tensor | None, |
| history: torch.Tensor | None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Causal dilated depthwise conv over ``[B, T, n, d]``; returns ``(out, new_history)``.""" |
|
|
| batch, length = values.shape[:2] |
| span = self.conv_history |
| if history is None: |
| history = values.new_zeros(batch, span, self.branches, self.hidden_size) |
| extended = torch.cat([history, values], dim=1) |
| weight = self.conv_weight.to(values.dtype) |
| out = values * weight[:, -1] |
| for tap in range(1, self.kernel): |
| shift = tap * self.dilation |
| shifted = extended[:, span - shift : span - shift + length] |
| if distance is not None: |
| shifted = shifted * (distance >= shift).to(values.dtype)[..., None, None] |
| out = out + shifted * weight[:, -1 - tap] |
| return out, extended[:, -span:] |
|
|
| def forward( |
| self, |
| residual: torch.Tensor, |
| embeddings: torch.Tensor, |
| *, |
| batch: int, |
| distance: torch.Tensor | None, |
| conv_state: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Branch updates ``[N, n, d]`` for branches ``[N, n, d]`` and rows ``[N, memory_dim]``.""" |
|
|
| rows = residual.shape[0] |
| value = self.value_proj(embeddings) |
| keys = self.key_proj(embeddings).view(rows, self.branches, self.hidden_size) |
| similarity = (self.hidden_norm(residual) * self.key_norm(keys)).sum(dim=-1) |
| alpha = torch.sigmoid(similarity / math.sqrt(self.hidden_size)).to(value.dtype) |
| gated = alpha.unsqueeze(-1) * value.unsqueeze(1) |
| normalized = self.conv_norm(gated).view(batch, rows // batch, self.branches, self.hidden_size) |
| convolved, new_state = self._convolve( |
| normalized, |
| None if distance is None else distance.view(batch, rows // batch), |
| conv_state, |
| ) |
| return F.silu(convolved).view(rows, self.branches, self.hidden_size) + gated, new_state |
|
|
|
|
| |
|
|
|
|
| class Senryu2Block(nn.Module): |
| def __init__(self, config: TinyGDNConfig, layer_index: int) -> None: |
| super().__init__() |
| self.layer_index = layer_index |
| self.layer_type = config.layer_types[layer_index] |
| residual = (config.residual_branches, config.hidden_size, config.residual_gate_rank, config.rms_norm_eps) |
| self.mixer_residual = GatedResidual(*residual, write=True) |
| self.mlp_residual = GatedResidual(*residual, write=True) |
| if self.layer_type == "gdn2": |
| self.token_mixer: nn.Module = build_gated_deltanet2(config, layer_index) |
| elif self.layer_type == "parallax": |
| self.token_mixer = ParallaxAttention(config, layer_index) |
| else: |
| raise ValueError(f"Senryu-2 has no {self.layer_type} layers") |
| self.mlp = SiTUGLU(config) |
| mark_muon_row_blocks(self.mlp.gate_up_proj.weight, 2) |
|
|
| def _mix(self, hidden_states: torch.Tensor, context: Senryu2Context) -> torch.Tensor: |
| batched = hidden_states.view(context.batch, context.length, -1) |
| if self.layer_type == "parallax": |
| mixed = self.token_mixer(batched, context) |
| else: |
| cache = context.cache |
| mixed, _, _ = self.token_mixer( |
| batched, |
| past_key_values=None if cache is None else cache.linear, |
| use_cache=cache is not None, |
| cu_seqlens=context.cu_seqlens, |
| ) |
| return mixed.reshape(hidden_states.shape) |
|
|
| def forward(self, residual: torch.Tensor, context: Senryu2Context) -> torch.Tensor: |
| hidden_states, scales = self.mixer_residual.read(residual) |
| residual = GatedResidual.write(residual, scales, self._mix(hidden_states, context)) |
| hidden_states, scales = self.mlp_residual.read(residual) |
| return GatedResidual.write(residual, scales, self.mlp(hidden_states)) |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Senryu2Cache: |
| """Decode state: GDN-2 states, Parallax keys/values, and the n-gram tails.""" |
|
|
| linear: Any = field(default_factory=FlaCache) |
| keys: dict[int, torch.Tensor] = field(default_factory=dict) |
| values: dict[int, torch.Tensor] = field(default_factory=dict) |
| length: int = 0 |
| ngram_tokens: torch.Tensor | None = None |
| ngram_conv: torch.Tensor | None = None |
|
|
| def append_attention( |
| self, |
| layer_index: int, |
| keys: torch.Tensor, |
| values: torch.Tensor, |
| start: int, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Write ``[B, T, H, D]`` keys/values at ``start``; return views of the whole prefix.""" |
|
|
| stop = start + keys.shape[1] |
| stored_keys = self.keys.get(layer_index) |
| if stored_keys is None or stored_keys.shape[1] < stop: |
| capacity = max(stop, 0 if stored_keys is None else 2 * stored_keys.shape[1]) |
| capacity = -(-capacity // KV_CACHE_GRANULE) * KV_CACHE_GRANULE |
| grown_keys = keys.new_empty(keys.shape[0], capacity, *keys.shape[2:]) |
| grown_values = values.new_empty(values.shape[0], capacity, *values.shape[2:]) |
| if stored_keys is not None: |
| grown_keys[:, :start] = stored_keys[:, :start] |
| grown_values[:, :start] = self.values[layer_index][:, :start] |
| self.keys[layer_index] = stored_keys = grown_keys |
| self.values[layer_index] = grown_values |
| stored_values = self.values[layer_index] |
| stored_keys[:, start:stop] = keys |
| stored_values[:, start:stop] = values |
| return stored_keys[:, :stop], stored_values[:, :stop] |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class _Packing: |
| """How a ``[B, T]`` batch maps onto one packed row of ``valid`` tokens.""" |
|
|
| batch: int |
| length: int |
| index: torch.Tensor | None |
| row_starts: torch.Tensor |
|
|
| def pack(self, tensor: torch.Tensor) -> torch.Tensor: |
| flat = tensor.reshape(self.batch * self.length, *tensor.shape[2:]) |
| return flat if self.index is None else flat.index_select(0, self.index) |
|
|
| def unpack(self, tensor: torch.Tensor) -> torch.Tensor: |
| if self.index is None: |
| return tensor.view(self.batch, self.length, *tensor.shape[1:]) |
| out = tensor.new_zeros(self.batch * self.length, *tensor.shape[1:]) |
| out.index_copy_(0, self.index, tensor) |
| return out.view(self.batch, self.length, *tensor.shape[1:]) |
|
|
|
|
| class Senryu2ForCausalLM(nn.Module): |
| def __init__(self, config: TinyGDNConfig) -> None: |
| super().__init__() |
| if config.residual_mode != "gated_residual": |
| raise ValueError("Senryu2ForCausalLM needs residual_mode 'gated_residual'") |
| self.config = config |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList(Senryu2Block(config, index) for index in range(config.num_hidden_layers)) |
| self.final_residual = GatedResidual( |
| config.residual_branches, |
| config.hidden_size, |
| config.residual_gate_rank, |
| config.rms_norm_eps, |
| write=False, |
| ) |
| self.ngram = NGramMemory(config) if config.ngram_orders else None |
| self.rotary = RotaryTables(config.rotary_dim, config.rope_theta) |
| self.checkpointed_layers: frozenset[int] = frozenset() |
| self.apply(self._initialize_module) |
| self._initialize_residual_projections() |
|
|
| |
| def _initialize_module(self, module: nn.Module) -> None: |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) |
| 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=self.config.initializer_range) |
|
|
| def _initialize_residual_projections(self) -> None: |
| residual_std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers) |
| for layer in self.layers: |
| nn.init.normal_(layer.token_mixer.o_proj.weight, mean=0.0, std=residual_std) |
| nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=residual_std) |
|
|
| def enable_gradient_checkpointing(self, enabled: bool = True, *, layers: int | None = None) -> None: |
| """Recompute block activations in backward; ``layers`` picks GDN-2 blocks first.""" |
|
|
| if not enabled: |
| if layers is not None: |
| raise ValueError("layers needs gradient checkpointing enabled") |
| self.checkpointed_layers = frozenset() |
| return |
| count = len(self.layers) if layers is None else layers |
| if not 1 <= count <= len(self.layers): |
| raise ValueError(f"layers must be between 1 and {len(self.layers)}") |
| cheapest_replay_first = sorted( |
| range(len(self.layers)), |
| key=lambda index: self.layers[index].layer_type != "gdn2", |
| ) |
| self.checkpointed_layers = frozenset(cheapest_replay_first[:count]) |
|
|
| def describe_gradient_checkpointing(self) -> str: |
| if not self.checkpointed_layers: |
| return "off" |
| layer_types = sorted({self.layers[index].layer_type for index in self.checkpointed_layers}) |
| return f"{len(self.checkpointed_layers)}/{len(self.layers)} blocks ({', '.join(layer_types)})" |
|
|
| def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| return F.linear(hidden_states, self.embed_tokens.weight) |
|
|
| |
| def _packing(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None) -> _Packing: |
| batch, length = input_ids.shape |
| device = input_ids.device |
| if attention_mask is None or bool(attention_mask.all()): |
| row_starts = torch.arange(0, batch * length, length, device=device) |
| return _Packing(batch, length, None, row_starts) |
| if attention_mask.shape != input_ids.shape: |
| raise ValueError("attention_mask must match input_ids") |
| valid = attention_mask.reshape(-1).bool() |
| index = valid.nonzero().flatten() |
| lengths = attention_mask.bool().sum(dim=1) |
| if bool((lengths == 0).any()): |
| raise ValueError("every row needs at least one unmasked token") |
| row_starts = torch.cumsum(lengths, dim=0) - lengths |
| return _Packing(batch, length, index, row_starts) |
|
|
| def packed_layout( |
| self, |
| tokens: torch.Tensor, |
| row_starts: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """``(cu_seqlens int32, distance int64)`` for a flat token stream split at BOS.""" |
|
|
| distance = document_distance(tokens, self.config.bos_token_id, row_starts) |
| starts = (distance == 0).nonzero().flatten() |
| cu_seqlens = torch.cat([starts, starts.new_tensor([tokens.numel()])]).to(torch.int32) |
| return cu_seqlens, distance |
|
|
| |
| def _run_block(self, index: int, residual: torch.Tensor, context: Senryu2Context) -> torch.Tensor: |
| layer = self.layers[index] |
| if self.training and index in self.checkpointed_layers: |
| return checkpoint(layer, residual, context, use_reentrant=False) |
| return layer(residual, context) |
|
|
| def _inject_ngrams( |
| self, |
| residual: torch.Tensor, |
| embeddings: torch.Tensor, |
| context: Senryu2Context, |
| ) -> torch.Tensor: |
| assert self.ngram is not None |
| cache = context.cache |
| conv_state = None if cache is None else cache.ngram_conv |
| distance = context.distance if cache is None else None |
| if self.training and self.checkpointed_layers: |
| update, new_state = checkpoint( |
| self.ngram, |
| residual, |
| embeddings, |
| batch=context.batch, |
| distance=distance, |
| conv_state=conv_state, |
| use_reentrant=False, |
| ) |
| else: |
| update, new_state = self.ngram( |
| residual, |
| embeddings, |
| batch=context.batch, |
| distance=distance, |
| conv_state=conv_state, |
| ) |
| if cache is not None: |
| cache.ngram_conv = new_state |
| return residual + update |
|
|
| def _ngram_rows( |
| self, |
| input_ids: torch.Tensor, |
| packing: _Packing | None, |
| distance: torch.Tensor | None, |
| ngram_embeddings: torch.Tensor | None, |
| cache: Senryu2Cache | None, |
| ) -> torch.Tensor: |
| assert self.ngram is not None |
| layout = self.ngram.layout |
| if cache is not None: |
| history = cache.ngram_tokens |
| if history is None: |
| history = input_ids.new_full((input_ids.shape[0], layout.max_order - 1), layout.pad_token_id) |
| canonical = self.ngram.canonical_ids[input_ids] |
| cache.ngram_tokens = torch.cat([history, canonical], dim=1)[:, -(layout.max_order - 1) :] |
| if ngram_embeddings is not None: |
| return ngram_embeddings.reshape(-1, layout.memory_dim) |
| return self.ngram.lookup(input_ids, history=history).reshape(-1, layout.memory_dim) |
| assert packing is not None |
| if ngram_embeddings is not None: |
| if ngram_embeddings.shape != (*input_ids.shape, layout.memory_dim): |
| raise ValueError( |
| f"ngram_embeddings must have shape {(*input_ids.shape, layout.memory_dim)}" |
| ) |
| return packing.pack(ngram_embeddings) |
| return self.ngram.lookup(packing.pack(input_ids), distance=distance) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| labels: torch.Tensor | None = None, |
| attention_mask: torch.Tensor | None = None, |
| past_key_values: Senryu2Cache | None = None, |
| *, |
| use_cache: bool = False, |
| return_logits: bool = True, |
| return_hidden_states: bool = False, |
| labels_are_shifted: bool = False, |
| include_mtp_loss: bool = True, |
| mtp_loss_weight: float | None = None, |
| z_loss_coefficient: float = 0.0, |
| logits_to_keep: int | None = None, |
| ngram_embeddings: torch.Tensor | None = None, |
| ) -> TinyGDNOutput: |
| del include_mtp_loss, mtp_loss_weight |
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape [batch, sequence]") |
| if labels is not None and labels.shape != input_ids.shape: |
| raise ValueError("labels must have the same shape as input_ids") |
| if z_loss_coefficient < 0.0: |
| raise ValueError("z_loss_coefficient cannot be negative") |
| if logits_to_keep is not None and logits_to_keep <= 0: |
| raise ValueError("logits_to_keep must be positive") |
| if use_cache and labels is not None: |
| raise ValueError("use_cache is not supported with labels") |
| if past_key_values is not None and not isinstance(past_key_values, Senryu2Cache): |
| raise TypeError("past_key_values must be a Senryu2Cache or None") |
| cached = use_cache or past_key_values is not None |
| batch, length = input_ids.shape |
| offset = 0 if past_key_values is None else past_key_values.length |
| if offset + length > self.config.max_position_embeddings: |
| raise ValueError("Input exceeds max_position_embeddings") |
|
|
| embeddings = self.embed_tokens(input_ids) |
| device, dtype = embeddings.device, embeddings.dtype |
| packing: _Packing | None = None |
| distance: torch.Tensor | None = None |
| if cached: |
| if attention_mask is not None and not bool(attention_mask.all()): |
| raise ValueError("Senryu-2 decodes unpadded rows; drop the padding instead of masking it") |
| cache = past_key_values if past_key_values is not None else Senryu2Cache() |
| context = Senryu2Context( |
| batch=batch, |
| length=length, |
| cu_seqlens=None, |
| distance=None, |
| rotary=self.rotary.get(offset + length, device, dtype), |
| position_offset=offset, |
| cache=cache, |
| ) |
| hidden = embeddings.reshape(batch * length, -1) |
| else: |
| cache = None |
| packing = self._packing(input_ids, attention_mask) |
| tokens = packing.pack(input_ids) |
| cu_seqlens, distance = self.packed_layout(tokens, packing.row_starts) |
| context = Senryu2Context( |
| batch=1, |
| length=tokens.numel(), |
| cu_seqlens=cu_seqlens, |
| distance=distance, |
| rotary=self.rotary.get(tokens.numel(), device, dtype), |
| ) |
| hidden = packing.pack(embeddings) |
| residual = hidden.unsqueeze(1).expand(-1, self.config.residual_branches, -1) |
|
|
| for index in range(len(self.layers)): |
| if self.ngram is not None and index == self.config.ngram_layer_index: |
| rows = self._ngram_rows(input_ids, packing, distance, ngram_embeddings, cache) |
| residual = self._inject_ngrams(residual, rows.to(dtype), context) |
| residual = self._run_block(index, residual, context) |
| hidden_states, _ = self.final_residual.read(residual) |
| if cache is not None: |
| cache.length = offset + length |
|
|
| main_loss = None |
| z_loss = None |
| total_loss = None |
| if labels is not None: |
| assert packing is not None |
| if labels_are_shifted: |
| targets = labels |
| else: |
| targets = torch.cat([labels[:, 1:], labels.new_full((batch, 1), -100)], dim=1) |
| total_loss, main_loss, z_loss = linear_cross_entropy( |
| hidden_states, |
| self.embed_tokens.weight, |
| packing.pack(targets).reshape(-1), |
| z_loss_coefficient=z_loss_coefficient, |
| ) |
| if z_loss_coefficient == 0.0: |
| z_loss = None |
|
|
| batched_states = ( |
| hidden_states.view(batch, length, -1) if packing is None else packing.unpack(hidden_states) |
| ) |
| output_states = batched_states if logits_to_keep is None else batched_states[:, -logits_to_keep:] |
| logits = self.project_to_vocabulary(output_states) if return_logits else None |
| return TinyGDNOutput( |
| loss=total_loss, |
| logits=logits, |
| main_loss=main_loss, |
| mtp_loss=None, |
| z_loss=z_loss, |
| hidden_states=batched_states if return_hidden_states else None, |
| past_key_values=cache if use_cache else None, |
| ) |
|
|
| |
| def parameter_report(self) -> dict[str, int]: |
| total = sum(parameter.numel() for parameter in self.parameters()) |
| embeddings = self.embed_tokens.weight.numel() |
| tables = 0 if self.ngram is None else self.ngram.layout.total_rows * self.ngram.layout.embedding_dim |
| return { |
| "deployable_core": total, |
| "training_total": total, |
| "embedding": embeddings, |
| "mtp_auxiliary": 0, |
| "non_embedding_core": total - embeddings, |
| "ngram_table": tables, |
| } |
|
|
| def save_checkpoint(self, output_dir: Path) -> None: |
| output_dir.mkdir(parents=True, exist_ok=True) |
| self.config.save_json(output_dir / "config.json") |
| save_model(self, output_dir / "model.safetensors") |
|
|
| def load_ngram_tables( |
| self, |
| checkpoint_dir: Path, |
| *, |
| weights: str, |
| device: str | torch.device, |
| dtype: torch.dtype, |
| ) -> None: |
| """Attach the checkpoint's tables that match ``weights`` ("normal" or "ema").""" |
|
|
| if self.ngram is None: |
| return |
| table_file = checkpoint_dir / (NGRAM_TABLES_EMA_FILE if weights == "ema" else NGRAM_TABLES_FILE) |
| if not table_file.is_file(): |
| raise FileNotFoundError(f"Senryu-2 checkpoint has no n-gram tables: {table_file}") |
| self.ngram.load_tables(table_file, device=device, dtype=dtype) |
|
|
| def extra_repr(self) -> str: |
| report = self.parameter_report() |
| return f"core_parameters={report['deployable_core']:,}, ngram_table_parameters={report['ngram_table']:,}" |
|
|
| def get_architecture_metadata(self) -> dict[str, Any]: |
| return { |
| "architecture": self.config.architecture, |
| "layer_types": list(self.config.layer_types), |
| "effective_num_layers": self.config.effective_num_layers, |
| "shared_layer_indices": [], |
| "parameter_report": self.parameter_report(), |
| "features": [ |
| "3:1 Gated DeltaNet-2 + Parallax attention hybrid", |
| "4-branch Gated Residual stream (gated read, per-branch write scales)", |
| "hashed n-gram memory (orders 2-3) with per-branch context gating", |
| "full-rank sigmoid output gates", |
| "partial RoPE, per-head QK/R normalization", |
| "SiTU-GLU", |
| "document-masked packed training", |
| "tied input-output embeddings", |
| ], |
| } |
|
|
