lexframe 0.1.0__py3-none-any.whl

lexframe/__init__.py

@@ -0,0 +1,7 @@
+"""A toolkit for interpreting LM residual streams in their native semantics."""
+
+__version__ = "0.1.0"
+
+from ._frame import Delta, LexFrame, State, View
+
+__all__ = ["LexFrame", "View", "Delta", "State", "__version__"]

lexframe/_basis.py

@@ -0,0 +1,166 @@
+"""The lexical basis at a layer, s = J^+ y, solved by LSMR on Jacobian-vector products without forming J."""
+
+import math
+
+import torch
+from torch.autograd import forward_ad as fwAD
+from torch.utils.checkpoint import checkpoint
+
+from . import _transport as TR
+from ._model import block
+
+
+def references(m, cal, xs, l, order):
+    L = int(m.n_layers)
+    B, Tn, d = xs[0].shape
+    pos = torch.arange(Tn, device=xs[0].device)
+    fi = torch.arange(0, B * Tn, Tn, device=xs[0].device)
+    if order is None:
+        return [TR.ref_kv(m, k, TR.innorm(m, k)(xs[k])) for k in range(l, L)]
+    low = [xs[l]] * order
+    out = []
+    with torch.no_grad():
+        for k in range(l, L):
+            norm = TR.innorm(m, k)
+            kv = [TR.ref_kv(m, k, norm(h)) for h in low]
+            out.append(kv[order - 1])
+            att = [TR.attend(m, k, norm(low[j]), pos, kv[j - 1]) for j in range(1, order)]
+            new = [TR._step(m, cal, k, low[0].reshape(-1, d), fi).view(B, Tn, d)]
+            for j in range(1, order):
+                a = att[j - 1].reshape(-1, d)
+                new.append(
+                    block(m, k, low[j].reshape(-1, d), lambda u, s, a=a: a[s : s + u.shape[0]].to(u.dtype)).view(
+                        B, Tn, d
+                    )
+                )
+            low = new
+    return out
+
+
+def stream(m, l, kv, X, ckpt=False):
+    L = int(m.n_layers)
+    b, Tn, d = X.shape
+    pos = torch.arange(Tn, device=X.device)
+
+    def step(h, k):
+        K, V = kv[k - l]
+        K, V = (K.expand(b, *K.shape[1:]), V.expand(b, *V.shape[1:])) if K.shape[0] != b else (K, V)
+        a = TR.attend(m, k, TR.innorm(m, k)(h), pos, (K, V)).reshape(-1, d)
+        return block(m, k, h.reshape(-1, d), lambda u, s: a[s : s + u.shape[0]].to(u.dtype)).view(b, Tn, d)
+
+    h = X
+    for k in range(l, L):
+        h = checkpoint(step, h, k, use_reentrant=False) if ckpt else step(h, k)
+    return m.final_norm(h).float()
+
+
+def jvp(m, l, kv, X0, V):
+    dt = X0.dtype
+    with torch.no_grad(), fwAD.dual_level():
+        out = stream(m, l, kv, fwAD.make_dual(X0.expand(V.shape[0], -1, -1).contiguous(), V.to(dt)))
+        return fwAD.unpack_dual(out).tangent.float()
+
+
+def vjp(m, l, kv, X0, U):
+    X = X0.expand(U.shape[0], -1, -1).clone().requires_grad_(True)
+    with torch.enable_grad():
+        return torch.autograd.grad(stream(m, l, kv, X, ckpt=True), X, U)[0].float()
+
+
+def _ortho(a, b):
+    r = torch.sqrt(a * a + b * b)
+    z = r == 0
+    return (
+        torch.where(z, torch.ones_like(a), a / r.clamp_min(1e-30)),
+        torch.where(z, torch.zeros_like(a), b / r.clamp_min(1e-30)),
+        r,
+    )
+
+
+def lsmr(Av, Atu, b, x0, damp, steps):
+    nrm = lambda u: u.norm(dim=-1, keepdim=True)
+    A = lambda v: (Av(v), damp * v)
+    At = lambda u: Atu(u[0]) + damp * u[1]
+    u = b - Av(x0), -damp * x0
+    beta = torch.sqrt(nrm(u[0]) ** 2 + nrm(u[1]) ** 2)
+    u = tuple(p / beta.clamp_min(1e-30) for p in u)
+    v = At(u)
+    alpha = nrm(v)
+    v = v / alpha.clamp_min(1e-30)
+    zetabar, alphabar = alpha * beta, alpha
+    one, zero = torch.ones_like(alpha), torch.zeros_like(alpha)
+    rho, rhobar, cbar, sbar, zeta = one, one, one, zero, zero
+    h, hbar, dx = v.clone(), torch.zeros_like(v), torch.zeros_like(v)
+    solved = alpha * beta == 0
+    for _ in range(0 if bool(solved.all()) else steps):
+        Avv = A(v)
+        u = Avv[0] - alpha * u[0], Avv[1] - alpha * u[1]
+        beta = torch.sqrt(nrm(u[0]) ** 2 + nrm(u[1]) ** 2)
+        u = tuple(p / beta.clamp_min(1e-30) for p in u)
+        v = At(u) - beta * v
+        alpha = nrm(v)
+        v = v / alpha.clamp_min(1e-30)
+        rhoold, rhobarold = rho, rhobar
+        c, s, rho = _ortho(alphabar, beta)
+        thetanew, alphabar = s * alpha, c * alpha
+        thetabar = sbar * rho
+        cbar, sbar, rhobar = _ortho(cbar * rho, thetanew)
+        zeta, zetabar = cbar * zetabar, -sbar * zetabar
+        hbar = h - thetabar * rho / (rhoold * rhobarold).clamp_min(1e-30) * hbar
+        step = zeta / (rho * rhobar).clamp_min(1e-30) * hbar
+        h = v - thetanew / rho.clamp_min(1e-30) * h
+        dx = dx + torch.where(solved, torch.zeros_like(step), step)
+    return x0 + dx
+
+
+def auto_batch(m, probe, n, cap=64):
+    if m.device.type != "cuda":
+        return min(n, 8)
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    base = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+    probe(2)
+    per = max(1, (torch.cuda.max_memory_allocated() - base) // 2)
+    free = torch.cuda.get_device_properties(m.device).total_memory - base
+    return max(1, min(n, cap, int(0.8 * free / per)))
+
+
+def basis(m, cal, xs, LS, Y, order=None, ridge=0.001, steps=30, after=None, warm=False):
+    nr = Y.shape[0]
+    LS = sorted(LS, reverse=True)
+    X0s = {l: xs[l][:1] for l in LS}
+    K0 = references(m, cal, xs, min(LS), None) if order is None else None
+    kvof = lambda l: K0[l - min(LS) :] if order is None else references(m, cal, xs, l, order)
+
+    def run(rows, l, x0, kv):
+        X0 = X0s[l]
+        y = Y[rows].float()
+        Av = lambda v: jvp(m, l, kv, X0, v)
+        Atu = lambda u: vjp(m, l, kv, X0, u)
+        damp = math.sqrt(ridge) * Atu(y).norm(dim=-1, keepdim=True) / y.norm(dim=-1, keepdim=True).clamp_min(1e-30)
+        return lsmr(Av, Atu, y, x0, damp, steps)
+
+    def probe(n):
+        l = LS[-1]
+        v = Y[:1].float().expand(n, -1, -1)
+        jvp(m, l, kvof(l), X0s[l], v)
+        vjp(m, l, kvof(l), X0s[l], v)
+
+    b = auto_batch(m, probe, nr)
+    r0 = 0
+    while r0 < nr:
+        rows = slice(r0, min(nr, r0 + b))
+        try:
+            x = torch.zeros_like(Y[rows].float())
+            for l in LS:
+                kv = kvof(l)
+                x = run(rows, l, x if warm else torch.zeros_like(x), kv)
+                if after is not None:
+                    after(l, X0s[l], kv, x, rows)
+            r0 = rows.stop
+        except torch.cuda.OutOfMemoryError:
+            if b == 1:
+                raise
+            b = max(1, b // 2)
+            torch.cuda.empty_cache()

lexframe/_calibrate.py

@@ -0,0 +1,240 @@
+"""The typical attention output of every block: a ridge regression of the attention output on its input."""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+
+from ._model import Model, forward
+
+RHO = [10 ** (e / 8) for e in range(-56, 25)]
+
+
+class _Moments:
+
+    def __init__(self, m: Model, batches: Iterable[torch.Tensor], folds: int = 2, work=torch.float32):
+        L, d = m.n_layers, m.hidden_dim
+        self.L, self.d, self.K, self.dev = L, d, folds, m.device
+        self.n = 0
+        acc = torch.float64 if 2 * folds * L * d * d * 8 < 4000000000.0 else torch.float32
+        z = lambda *s: torch.zeros(*s, dtype=acc, device=m.device)
+        self.C = [[[z(d, d), z(d, d)] for _ in range(folds)] for _ in range(L)]
+        self.cnt = [0] * folds
+        at = [0]
+        self.mu = torch.zeros(L, folds, d, dtype=torch.float64, device=m.device)
+        self.ma = torch.zeros_like(self.mu)
+        self.tr = torch.zeros(L, folds, dtype=torch.float64)
+
+        def cap(k, u, a):
+            g = (at[0] + torch.arange(u.shape[0], device=u.device)) % folds
+            for f in range(folds):
+                if not bool((g == f).any()):
+                    continue
+                n2, mu2, ma2, C2, B2, t2 = _moments(
+                    u[g == f].reshape(-1, d).to(work), a[g == f].reshape(-1, d).to(work)
+                )
+                n1 = self.cnt[f]
+                w = n1 * n2 / (n1 + n2)
+                du, da = mu2 - self.mu[k, f], ma2 - self.ma[k, f]
+                Cf = self.C[k][f]
+                Cf[0] += (C2 + w * torch.outer(du, du)).to(acc)
+                Cf[1] += (B2 + w * torch.outer(du, da)).to(acc)
+                self.tr[k, f] += t2 + w * float(da @ da)
+                self.mu[k, f] += du * (n2 / (n1 + n2))
+                self.ma[k, f] += da * (n2 / (n1 + n2))
+                if k == L - 1:
+                    self.cnt[f] = n1 + n2
+
+        for ids in batches:
+            ids = m.ids(ids)
+            forward(m, ids, on_attention=cap)
+            at[0] += ids.shape[0]
+            self.n += ids.numel()
+        if at[0] < folds:
+            raise ValueError(f"calibration needs at least {folds} sequences")
+
+    def moments(self, k):
+        out = [
+            (
+                self.cnt[f],
+                self.mu[k, f],
+                self.ma[k, f],
+                self.C[k][f][0].double(),
+                self.C[k][f][1].double(),
+                float(self.tr[k, f]),
+            )
+            for f in range(self.K)
+        ]
+        self.C[k] = None
+        return out
+
+
+def _moments(U, A):
+    n = U.shape[0]
+    mu, ma = U.mean(0), A.mean(0)
+    Uc = U - mu
+    Ac = A - ma
+    return n, mu.double(), ma.double(), (Uc.T @ Uc).double(), (Uc.T @ Ac).double(), float((Ac.double() ** 2).sum())
+
+
+def _merge(ms):
+    n = sum((m[0] for m in ms))
+    mu = sum((m[0] * m[1] for m in ms)) / n
+    ma = sum((m[0] * m[2] for m in ms)) / n
+    C = sum((m[3] + m[0] * torch.outer(m[1] - mu, m[1] - mu) for m in ms))
+    B = sum((m[4] + m[0] * torch.outer(m[1] - mu, m[2] - ma) for m in ms))
+    return n, mu, ma, C, B, sum((m[5] + m[0] * float((m[2] - ma) @ (m[2] - ma)) for m in ms))
+
+
+class Calibration:
+
+    def __init__(self, ubar, abar, Q, lam, QA, rho, stats, model_id):
+        self.ubar, self.abar, self.Q, self.lam, self.QA, self.rho, self.stats, self.model_id = (
+            ubar,
+            abar,
+            Q,
+            lam,
+            QA,
+            rho,
+            stats,
+            model_id,
+        )
+
+    def __repr__(self) -> str:
+        return f"Calibration(model={self.model_id['architecture']}, layers={self.model_id['layers']})"
+
+    def typical(self, k: int):
+        ub, ab = self.ubar[k], self.abar[k]
+        dev = ub.device
+        Q, QA, lam, rid = self.Q[k].to(dev), self.QA[k].to(dev), self.lam[k], self.rho[k]
+
+        def fn(u, s):
+            uc = u.float() - ub
+            return (ab.expand_as(uc) + uc @ Q / (lam[None] + rid) @ QA).to(u.dtype)
+
+        return fn
+
+    def to(self, device):
+        mv = lambda t: t.to(device)
+        return Calibration(
+            mv(self.ubar),
+            mv(self.abar),
+            self.Q if self.Q.device.type == "cpu" else mv(self.Q),
+            mv(self.lam),
+            self.QA if self.QA.device.type == "cpu" else mv(self.QA),
+            mv(self.rho),
+            self.stats,
+            self.model_id,
+        )
+
+    def save(self, path: str) -> None:
+        torch.save(
+            {
+                "format": 1,
+                "model_id": self.model_id,
+                "stats": self.stats,
+                **{k: getattr(self, k).cpu() for k in ("ubar", "abar", "Q", "lam", "QA", "rho")},
+            },
+            path,
+        )
+
+    @staticmethod
+    def load(path: str, model: Model | None = None) -> Calibration:
+        z = torch.load(path, map_location="cpu", weights_only=True)
+        if z.get("format") != 1:
+            raise ValueError(f"{path}: not a lexframe calibration")
+        c = Calibration(z["ubar"], z["abar"], z["Q"], z["lam"], z["QA"], z["rho"], z["stats"], z["model_id"])
+        if model is not None:
+            if c.model_id != model_id(model):
+                raise ValueError(f"{path} was calibrated for {c.model_id}, not for {model_id(model)}")
+            c = c.to(model.device)
+        return c
+
+
+def model_id(m: Model) -> dict:
+    w = lambda p: float(sum((t.double().sum() for t in p.parameters())))
+    return {
+        "architecture": type(m.hf).__name__,
+        "layers": m.n_layers,
+        "hidden": m.hidden_dim,
+        "checksum": round(w(m.layers[0]) + w(m.layers[-1]), 3),
+    }
+
+
+def chunks(m: Model, texts: Iterable[str], length: int = 128, batch: int = 8, limit: int | None = None):
+    if m.tokenizer is None:
+        raise ValueError("this model has no tokenizer: pass token ids to calibrate")
+    eos = m.tokenizer.eos_token_id
+    buf, out, n = [], [], 0
+    for t in texts:
+        if not t or not t.strip():
+            continue
+        buf += m.tokenizer.encode(t, add_special_tokens=False) + ([eos] if eos is not None else [])
+        while len(buf) >= length:
+            out.append(buf[:length])
+            buf = buf[length:]
+            n += 1
+            if len(out) == batch:
+                yield torch.tensor(out)
+                out = []
+            if limit is not None and n >= limit:
+                if out:
+                    yield torch.tensor(out)
+                return
+    if out:
+        yield torch.tensor(out)
+
+
+@torch.no_grad()
+def calibrate(m: Model, batches: Iterable, folds: int = 2, work=torch.float32) -> Calibration:
+    bank = _Moments(m, batches, folds, work)
+    L, d, dev, K = bank.L, bank.d, bank.dev, bank.K
+    f64 = torch.float64
+    store = "cpu" if L * d * d * 8 > 1500000000.0 else dev
+    ubar = torch.empty(L, d, device=dev)
+    abar = torch.empty(L, d, device=dev)
+    stats = []
+    Q = torch.empty(L, d, d, device=store)
+    QA = torch.empty(L, d, d, device=store)
+    lam = torch.empty(L, d, device=dev)
+    rho = torch.empty(L, device=dev)
+    grid = torch.tensor(RHO, dtype=f64, device=dev)
+    for k in range(L):
+        mo = bank.moments(k)
+        n, mu, ma, Cuu, Cua, _ = _merge(mo)
+        ubar[k], abar[k] = mu.float(), ma.float()
+        rg = grid * (Cuu.diagonal().sum() / d)
+        err = torch.zeros_like(rg)
+        tts = 0.0
+        for g in range(K):
+            _, mug, mag, Cg, Bg, _ = _merge([mo[i] for i in range(K) if i != g])
+            nf, _, _, Cf, Bf, tf = mo[g]
+            eu, ea = mo[g][1] - mug, mo[g][2] - mag
+            Duu = Cf + nf * torch.outer(eu, eu)
+            Dua = Bf + nf * torch.outer(eu, ea)
+            tt = tf + nf * float(ea @ ea)
+            tts += tt
+            ev, Qg = torch.linalg.eigh(Cg.to(work))
+            ev = ev.double().clamp_min(0)
+            QAg = Qg.T @ Bg.to(work)
+            Dr = Qg.T @ Duu.to(work) @ Qg
+            p = (Qg.T @ Dua.to(work) * QAg).sum(1).double()
+            KS = (QAg @ QAg.T * Dr).double()
+            W = 1.0 / (ev[None] + rg[:, None])
+            err += tt - 2 * (W @ p) + (W @ KS * W).sum(1)
+        j = int(err.argmin())
+        ev, Qk = torch.linalg.eigh(Cuu.to(work))
+        ev = ev.double().clamp_min(0)
+        QAk = Qk.T @ Cua.to(work)
+        rk = float(rg[j])
+        stats.append(
+            {
+                "rows": n,
+                "ridge_relative": RHO[j],
+                "dof": float((ev / (ev + rk)).sum()),
+                "r2_heldout": 1 - float(err[j]) / tts,
+            }
+        )
+        Q[k], QA[k], lam[k], rho[k] = Qk.float().to(store), QAk.float().to(store), ev.float(), rk
+    return Calibration(ubar, abar, Q, lam, QA, rho, stats, model_id(m))