| """
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 == " ":
|
| 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)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
| self.pad_cid = self.compressed_vocab
|
|
|
| if lookup is None:
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
| seen: set = set()
|
| self.head_sizes: List[List[int]] = []
|
| 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()]
|
|
|
| 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 = 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)
|
| 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]
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
| 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)
|
| 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)
|
|
|
|
|
|
|
| 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)
|
| key = self.key_proj(emb)
|
| nk = self.norm1(key)
|
| nq = self.norm2(hidden_states)
|
|
|
|
|
| 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)
|
| value = gate * self.value_proj(emb)
|
| return value + self.short_conv(value, doc_ids, cache=cache) |
