""" Engram -- static n-gram memory (arXiv:2601.07372, DeepSeek engram_demo_v1.py), adapted for Veylon/Arya's standard transformer (hc_mult = 1, tiktoken). Fixes vs. the drafted port (all verified in test_engram.py): 1. Hashing ran on CPU/numpy INSIDE forward(): a GPU->CPU sync + torch.compile graph break on every Engram layer, every step (and fatal on XLA). Now the hash is pure int64 torch ops on the model device, compile-friendly. 2. Every Engram layer built its own NgramHashMapping (re-decoding the whole vocab) and each call hashed ALL layers, then kept one. Now ONE NgramHasher lives on GPT, hashes all layers once per forward. 3. pad_id=None crashed np.pad; pad_id=-1 silently indexed the LAST lookup entry. Left-context padding is now a dedicated reserved id (never collides with a real token class). 4. nn.RMSNorm (torch>=2.4 only, eps=finfo(fp16)=1e-3 default) replaced by a local _RMSNorm (eps=1e-5, weight cast to the input dtype) matching model.RMSNorm's eager path. It does NOT upcast to fp32. 5. Gate dot-product computed in fp32 (sum over D in fp16 can overflow/lose precision), cast back after the sigmoid. 6. The compressed-vocab lookup is now a persistent buffer, so a checkpoint carries it and inference needs no tokenizer object. 7. sympy dependency dropped (trial-division primality is plenty for table sizes). 8. Special-id discovery no longer assumes TokenizerWrapper._special_ids is a set of ints (dict / missing attr both handled). 9. embed_per_ngram % n_heads != 0 silently produced a wrong-width table -> hard assert. Same for len(engram_vocab_size) < max_ngram-1. """ from __future__ import annotations import math import re import unicodedata import warnings from typing import List, Sequence import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # ───────────────────────────────────────────────────────────────────────────── # Primes # ───────────────────────────────────────────────────────────────────────────── def _is_prime(n: int) -> bool: if n < 2: return False if n < 4: return True if n % 2 == 0 or n % 3 == 0: return False i = 5 while i * i <= n: if n % i == 0 or n % (i + 2) == 0: return False i += 6 return True def find_next_prime(start: int, seen_primes: set) -> int: candidate = start + 1 while True: if _is_prime(candidate) and candidate not in seen_primes: return candidate candidate += 1 # ───────────────────────────────────────────────────────────────────────────── # CompressedTokenizer -- build-time only (numpy). Output is a lookup table. # ───────────────────────────────────────────────────────────────────────────── class CompressedTokenizer: """Maps raw token IDs -> compressed IDs (NFKC / strip-accents / lowercase / whitespace-collapse of each token's surface string). Build-time only: the result (`lookup_table`) is handed to the model as a persistent buffer.""" _SENTINEL = "\uE000" _WS_RE = re.compile(r"[ \t\r\n]+") def __init__(self, tokenizer_wrapper): self.tokenizer = tokenizer_wrapper self._special_ids = self._collect_special_ids(tokenizer_wrapper) self.lookup_table, self.num_new_token = self._build_lookup_table() def __len__(self): return self.num_new_token @staticmethod def _collect_special_ids(tw) -> set: ids = set() raw = getattr(tw, "_special_ids", None) if raw is not None: if isinstance(raw, dict): for k, v in raw.items(): for cand in (k, v): if isinstance(cand, (int, np.integer)): ids.add(int(cand)) else: for v in raw: if isinstance(v, (int, np.integer)): ids.add(int(v)) vs = int(tw.vocab_size) for name in ("bos_id", "eos_id", "pad_id"): v = getattr(tw, name, None) if isinstance(v, (int, np.integer)) and 0 <= int(v) < vs: ids.add(int(v)) return ids @classmethod def _normalize(cls, text: str) -> str: text = unicodedata.normalize("NFKC", text) text = unicodedata.normalize("NFD", text) text = "".join(c for c in text if unicodedata.category(c) != "Mn") text = text.lower() text = cls._WS_RE.sub(" ", text) if text == " ": # lone space survives strip() text = cls._SENTINEL text = text.strip() return text.replace(cls._SENTINEL, " ") def _decode_id(self, tid: int) -> str: tw = self.tokenizer try: enc = getattr(tw, "enc", None) if enc is not None: return enc.decode([tid]) return tw.decode([tid]) except Exception: return "" def _build_lookup_table(self): key2new, new_tokens = {}, [] vocab_size = int(self.tokenizer.vocab_size) lookup = np.empty(vocab_size, dtype=np.int64) for tid in range(vocab_size): if tid in self._special_ids: key = f"__special_{tid}" else: text = self._decode_id(tid) if not text or "\ufffd" in text: key = f"__raw_{tid}" else: norm = self._normalize(text) key = norm if norm else f"__raw_{tid}" nid = key2new.get(key) if nid is None: nid = len(new_tokens) key2new[key] = nid new_tokens.append(key) lookup[tid] = nid return lookup, len(new_tokens) # ───────────────────────────────────────────────────────────────────────────── # NgramHasher -- static, deterministic, layer-specific, pure torch int64 # ───────────────────────────────────────────────────────────────────────────── class NgramHasher(nn.Module): """(B, T) raw ids -> per-layer embedding indices of shape (B, T, n_cols). No parameters, no gradients. `lookup` and the hash multipliers `mult` are persistent (state_dict); the prime moduli `mods` are re-derived from config. `mult` comes from NumPy's Generator, which is not guaranteed stable across NumPy releases, so it is saved with the checkpoint instead of re-derived. """ def __init__(self, layer_ids: Sequence[int], max_ngram: int, vocab_size_per_ngram: Sequence[int], n_heads: int, compressed_vocab: int, raw_vocab_size: int, seed: int, lookup: torch.Tensor | None = None): super().__init__() assert max_ngram >= 2, "engram_max_ngram must be >= 2" assert len(vocab_size_per_ngram) >= max_ngram - 1, ( f"engram_vocab_size needs {max_ngram - 1} entries (ngram 2..{max_ngram}), " f"got {len(vocab_size_per_ngram)}" ) assert compressed_vocab > 0, "engram_compressed_vocab must be set (>0)" self.layer_ids = tuple(int(l) for l in layer_ids) self.max_ngram = int(max_ngram) self.n_heads = int(n_heads) self.n_cols = (self.max_ngram - 1) * self.n_heads self.compressed_vocab = int(compressed_vocab) # Reserved id for left-context padding: distinct from every token class. self.pad_cid = self.compressed_vocab if lookup is None: # Every token would hash to class 0 and Engram would silently ignore its input. This is only # valid if a checkpoint (which carries `lookup`) is loaded right after construction. warnings.warn( "NgramHasher built without a lookup table: all tokens map to class 0 until a checkpoint " "containing `lookup` is loaded. Pass engram_lookup from CompressedTokenizer for training.", stacklevel=2) lookup = torch.zeros(int(raw_vocab_size), dtype=torch.int64) else: lookup = lookup.detach().to(torch.int64).clone() assert lookup.numel() == int(raw_vocab_size), ( f"lookup has {lookup.numel()} entries, vocab_size={raw_vocab_size}") assert int(lookup.max()) < self.compressed_vocab self.register_buffer("lookup", lookup, persistent=True) # Multipliers: r*2+1, bounded so tok*mult can never overflow int64. # (+1 vocab slot for the reserved pad id.) max_long = int(np.iinfo(np.int64).max) M_max = max_long // (self.compressed_vocab + 1) half_bound = max(1, M_max // 2) PRIME_1 = 10007 mult = np.empty((len(self.layer_ids), self.max_ngram), dtype=np.int64) for li, layer_id in enumerate(self.layer_ids): g = np.random.default_rng(int(seed + PRIME_1 * int(layer_id))) r = g.integers(low=0, high=half_bound, size=(self.max_ngram,), dtype=np.int64) mult[li] = r * 2 + 1 self.register_buffer("mult", torch.from_numpy(mult), persistent=True) # Distinct prime moduli per (layer, ngram order, head); unique globally. seen: set = set() self.head_sizes: List[List[int]] = [] # per layer: flat list, len n_cols for _ in self.layer_ids: flat = [] for ngram in range(2, self.max_ngram + 1): search_start = int(vocab_size_per_ngram[ngram - 2]) - 1 for _h in range(self.n_heads): found = find_next_prime(search_start, seen) seen.add(found) flat.append(found) search_start = found self.head_sizes.append(flat) self.register_buffer( "mods", torch.tensor(self.head_sizes, dtype=torch.int64), persistent=False) head_offsets = [] for sizes in self.head_sizes: offsets = [] offset = 0 for size in sizes: offsets.append(offset) offset += size head_offsets.append(offsets) self.register_buffer( "head_offsets", torch.tensor(head_offsets, dtype=torch.int64), persistent=False) def compress(self, idx: torch.Tensor) -> torch.Tensor: return self.lookup[idx.long()] # idx is never negative (the -1 ignore index lives in targets) def forward(self, idx: torch.Tensor, doc_ids: torch.Tensor | None = None, cache: dict | None = None): """Hash tokens, optionally using cached history for incremental decoding. doc_ids (B, T), when supplied, replaces cross-document history with the left-padding id so packed documents hash as if processed separately. """ ids = self.compress(idx) # (B, T_new) B, T_new = ids.shape K = self.max_ngram - 1 if cache is not None: hist_ids = cache.get("hash_hist_ids") hist_doc = cache.get("hash_hist_doc") if hist_ids is None: hist_ids = ids.new_full((B, K), self.pad_cid) hist_doc = ( torch.full((B, K), -1, dtype=doc_ids.dtype, device=doc_ids.device) if doc_ids is not None else None ) else: hist_ids = hist_ids.to(device=ids.device) if hist_doc is not None: hist_doc = hist_doc.to(device=ids.device) full_ids = torch.cat([hist_ids, ids], dim=1) # (B, K + T_new) if doc_ids is not None and hist_doc is not None: full_doc = torch.cat([hist_doc, doc_ids], dim=1) else: full_doc = doc_ids cache["hash_hist_ids"] = full_ids[:, -K:].detach() if full_doc is not None: cache["hash_hist_doc"] = full_doc[:, -K:].detach() else: full_ids, full_doc = ids, doc_ids T_full = full_ids.shape[1] # Keep only one running hash per layer and one shifted-id tensor per order; # broadcasting all layers and orders at once creates a large int64 product. mixes = [full_ids * self.mult[li, 0] for li in range(len(self.layer_ids))] out = { layer_id: torch.empty( (B, T_full, self.n_cols), dtype=torch.int64, device=full_ids.device) for layer_id in self.layer_ids } for n in range(2, self.max_ngram + 1): shifted = F.pad(full_ids, (n - 1, 0), value=self.pad_cid)[:, :T_full] if full_doc is not None: same = F.pad(full_doc, (n - 1, 0), value=-1)[:, :T_full] == full_doc shifted = torch.where(same, shifted, torch.full_like(shifted, self.pad_cid)) lo = (n - 2) * self.n_heads hi = lo + self.n_heads for li, layer_id in enumerate(self.layer_ids): mix = torch.bitwise_xor(mixes[li], shifted * self.mult[li, n - 1]) mixes[li] = mix col = mix.unsqueeze(-1) % self.mods[li, lo:hi] col += self.head_offsets[li, lo:hi] out[layer_id][..., lo:hi] = col if T_full > T_new: out = {layer_id: values[:, -T_new:].contiguous() for layer_id, values in out.items()} return out # ───────────────────────────────────────────────────────────────────────────── # Modules # ───────────────────────────────────────────────────────────────────────────── class _RMSNorm(nn.Module): """RMSNorm with fused accumulation for the input dtype.""" def __init__(self, dim: int, eps: float = 1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return F.rms_norm(x, (x.shape[-1],), self.weight.to(x.dtype), self.eps) class ShortConv(nn.Module): """(B, T, D) -> (B, T, D). Depthwise causal conv, dilation = max_ngram.""" def __init__(self, hidden_size: int, kernel_size: int = 4, dilation: int = 1, activation: bool = True): super().__init__() self.activation = activation self.kernel_size = kernel_size self.dilation = dilation self.conv = nn.Conv1d( hidden_size, hidden_size, kernel_size=kernel_size, groups=hidden_size, bias=False, padding=(kernel_size - 1) * dilation, dilation=dilation, ) self.norm = _RMSNorm(hidden_size) self.act_fn = nn.SiLU() if activation else None def forward(self, x: torch.Tensor, doc_ids: torch.Tensor | None = None, cache: dict | None = None) -> torch.Tensor: T = x.shape[1] K, d = self.kernel_size, self.dilation rf = (K - 1) * d B, _, D = x.shape xn = self.norm(x) # (B, T, D) if cache is not None: hist_xn = cache.get("sc_hist_x") hist_doc = cache.get("sc_hist_doc") if hist_xn is None: hist_xn = xn.new_zeros(B, rf, D) hist_doc = ( torch.full((B, rf), -1, dtype=doc_ids.dtype, device=doc_ids.device) if doc_ids is not None else None ) else: hist_xn = hist_xn.to(dtype=xn.dtype, device=xn.device) if hist_doc is not None: hist_doc = hist_doc.to(device=xn.device) full_xn = torch.cat([hist_xn, xn], dim=1) if rf > 0 else xn if doc_ids is not None and hist_doc is not None: full_doc = torch.cat([hist_doc, doc_ids], dim=1) if rf > 0 else doc_ids else: full_doc = doc_ids if rf > 0: cache["sc_hist_x"] = full_xn[:, -rf:].detach() if full_doc is not None: cache["sc_hist_doc"] = full_doc[:, -rf:].detach() else: full_xn, full_doc = xn, doc_ids T_full = full_xn.shape[1] w = self.conv.weight[:, 0, :].to(full_xn.dtype) # (D, K) xp = F.pad(full_xn, (0, 0, rf, 0)) tap = xp[:, :T_full] if full_doc is not None and K > 1: s = (K - 1) * d same = F.pad(full_doc, (s, 0), value=-1)[:, :T_full] == full_doc tap = tap * same.unsqueeze(-1) y = tap * w[:, 0] for j in range(1, K): tap = xp[:, j * d:j * d + T_full] if full_doc is not None: s = (K - 1 - j) * d same = F.pad(full_doc, (s, 0), value=-1)[:, :T_full] == full_doc tap = tap * same.unsqueeze(-1) y.addcmul_(tap, w[:, j]) if T_full > T: y = y[:, -T:].contiguous() if self.activation: y = self.act_fn(y) return y class MultiHeadEmbedding(nn.Module): def __init__(self, list_of_N: List[int], D: int): super().__init__() self.embedding = nn.Embedding(sum(list_of_N), D) self.reset_parameters() def reset_parameters(self): nn.init.normal_(self.embedding.weight, std=0.01) def forward(self, hash_ids: torch.Tensor) -> torch.Tensor: return self.embedding(hash_ids) class Engram(nn.Module): """Static n-gram memory with conditional gated injection (hc_mult = 1). forward() returns the DELTA to add to the residual stream: x = x + engram(x, hash_ids) """ def __init__(self, hidden_size: int, head_sizes: List[int], max_ngram: int, embed_per_ngram: int, n_heads: int, kernel_size: int): super().__init__() assert embed_per_ngram % n_heads == 0, ( f"engram_embed_per_ngram={embed_per_ngram} must be divisible by " f"engram_n_heads={n_heads}") self.hidden_size = hidden_size d_head = embed_per_ngram // n_heads engram_hidden = (max_ngram - 1) * embed_per_ngram self.multi_head_embedding = MultiHeadEmbedding(list(head_sizes), d_head) self.value_proj = nn.Linear(engram_hidden, hidden_size, bias=False) self.key_proj = nn.Linear(engram_hidden, hidden_size, bias=False) self.norm1 = _RMSNorm(hidden_size) self.norm2 = _RMSNorm(hidden_size) self.short_conv = ShortConv(hidden_size, kernel_size=kernel_size, dilation=max_ngram) self.reset_parameters() def reset_parameters(self): self.multi_head_embedding.reset_parameters() nn.init.normal_(self.value_proj.weight, std=0.02) nn.init.normal_(self.key_proj.weight, std=0.02) # Zero-init the conv: ShortConv RMS-normalises its input, so without this the gated value is # blown back up to ~1 rms and the initial delta is ~0.13 rms (larger than the residual stream at # D<=256). With zeros the branch starts at ~0 and the conv weight gets gradient on step 1. nn.init.zeros_(self.short_conv.conv.weight) def forward(self, hidden_states: torch.Tensor, hash_ids: torch.Tensor, doc_ids: torch.Tensor | None = None, cache: dict | None = None): emb = self.multi_head_embedding(hash_ids).flatten(start_dim=-2) # (B,T,engram_hidden) key = self.key_proj(emb) nk = self.norm1(key) nq = self.norm2(hidden_states) # Product AND sum in fp32: with bf16 inputs, rounding nk*nq to bf16 first leaves gate errors up to # ~0.02 (sqrt(|g|) amplifies small pre-activations); the casts fuse under torch.compile. gate = (nk.float() * nq.float()).sum(dim=-1) / math.sqrt(self.hidden_size) gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() gate = gate.sigmoid().unsqueeze(-1).to(hidden_states.dtype) # (B,T,1) value = gate * self.value_proj(emb) return value + self.short_conv(value, doc_ids, cache=cache)