scaffold: model (OLMo2-ish + ZBP partition), trainer (DDP/config), data shards, bench

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe

1 files changed, 421 insertions, 0 deletions

diff --git a/src/zbp_scaling/zbp/estimators.py b/src/zbp_scaling/zbp/estimators.py
new file mode 100644
index 0000000..81f29a3
--- /dev/null
+++ b/src/zbp_scaling/zbp/estimators.py
@@ -0,0 +1,421 @@
+"""Zeroth-order estimators of g = J^T v for a black-box block.
+
+Conventions
+-----------
+A block maps x (shape [*batch_shape, *feat]) to y.  A *measurement* function
+`h(x_pert)` returns v^T f(x_pert) per measurement unit:
+  * global units:  h -> [*batch_shape]              (sum over all feature dims)
+  * local units:   h -> [*batch_shape, H', W']      (sum over channels only; conv blocks)
+`h` accepts any integer multiple of the batch stacked along dim 0, so several
+probes can be evaluated in one call.
+
+For local units the directional derivatives D_p (one per output position p) are
+box-summed over the receptive field before multiplying the probe, which keeps the
+estimator unbiased while removing the variance contributed by far-away positions.
+
+Every estimator returns (g_hat, queries) where queries is the number of block
+evaluations *per sample unit* that a physical system would have to perform.
+"""
+import math
+import torch
+import torch.nn.functional as F
+
+from .probes import sample_probes
+
+
+class Local:
+    """Receptive-field geometry of a conv-like block: output position p depends on input
+    positions q with |q - stride*p|_inf <= radius (input coordinates)."""
+
+    def __init__(self, radius, stride=1):
+        self.radius, self.stride = int(radius), int(stride)
+
+    def gather(self, D, in_hw):
+        """D: [n, *bs, H', W'] per-output-position derivatives -> S: [n, *bs, H, W] box sums."""
+        lead = D.shape[:-2]
+        Hp, Wp = D.shape[-2:]
+        H, W = in_hw
+        k = 2 * self.radius + 1
+        Dm = D.reshape(-1, 1, Hp, Wp)
+        ones = torch.ones(1, 1, k, k, dtype=D.dtype, device=D.device)
+        oph = H - ((Hp - 1) * self.stride + 1)
+        opw = W - ((Wp - 1) * self.stride + 1)
+        S = F.conv_transpose2d(Dm, ones, stride=self.stride, padding=self.radius, output_padding=(oph, opw))
+        return S.reshape(*lead, H, W)
+
+
+class Local1D:
+    """Causal token-window locality for [B, T, d] blocks: output token t depends on input tokens
+    s in [t-window+1, t].  Measurement units = output tokens; gather = causal window box-sum."""
+    kind = "1d"
+
+    def __init__(self, window):
+        self.window = int(window)
+
+    def gather(self, D, in_shape):
+        """D: [n, *bs, T] per-output-token derivatives -> S: [n, *bs, T], S[s] = sum_{t=s}^{s+W-1} D[t]."""
+        T = D.shape[-1]
+        c = torch.cumsum(D, dim=-1)
+        pad = torch.zeros_like(c[..., :1])
+        c0 = torch.cat([pad, c], dim=-1)                                # c0[t] = sum_{<t}
+        hi = torch.clamp(torch.arange(T, device=D.device) + self.window, max=T)
+        lo = torch.arange(T, device=D.device)
+        return c0[..., hi] - c0[..., lo]
+
+
+def _probes_for(x, n_probes, family, gen, n_batch_dims, cycle=None):
+    bs = x.shape[:n_batch_dims]
+    feat = x.shape[n_batch_dims:]
+    d = math.prod(feat)
+    u = sample_probes(n_probes, tuple(bs), d, family, gen, x.device, x.dtype, cycle=cycle)
+    return u.reshape(n_probes, *bs, *feat), d
+
+
+def zo_vjp_cd_parts(h, x, cfg, gen, n_batch_dims=1, cycle=None):
+    """Central-difference measurements together with their probes (global units only):
+    returns (u [n,*bs,*feat], D [n,*bs] = u_i^T g + O(eps^2), g_hat, queries)."""
+    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle)
+    eps = _eps_for(cfg, d)
+    D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+    return u, D, _combine(u, D, n_batch_dims, None, x.shape), float(2 * cfg.n_probes)
+
+
+def _combine(u, D, n_batch_dims, local, in_shape):
+    """g_hat = (1/n) sum_i u_i * S_i, where S_i are the (gathered) directional derivatives."""
+    n = u.shape[0]
+    if local is None:
+        S = D.reshape(*D.shape, *([1] * (u.dim() - 1 - n_batch_dims)))
+    elif getattr(local, "kind", "2d") == "1d":
+        S = local.gather(D, in_shape).unsqueeze(-1)           # [n, *bs, T, 1] broadcast over features
+    else:
+        S = local.gather(D, in_shape[-2:])                   # [n, *bs, H, W]
+        S = S.unsqueeze(n_batch_dims + 1)                     # broadcast over channels: [n, *bs, 1, H, W]
+    return (u * S).mean(0)
+
+
+def _eps_for(cfg, d):
+    return cfg.eps if cfg.eps_mode == "coord" else cfg.eps / math.sqrt(d)
+
+
+def _eval_pairs(h, x, u, eps_list, n_batch_dims, batch_probes=True, chunk=0):
+    """D(eps, u_i) = (h(x+eps u_i) - h(x-eps u_i)) / (2 eps) for all probes/eps -> [len(eps), n, *bs, *meas]."""
+    n = u.shape[0]
+    B0 = x.shape[0]
+    outs = []
+    if batch_probes:
+        step = n if chunk <= 0 else max(1, min(chunk, n))
+        for eps in eps_list:
+            parts = []
+            for i0 in range(0, n, step):
+                xs = []
+                for i in range(i0, min(n, i0 + step)):
+                    xs.append(x + eps * u[i])
+                    xs.append(x - eps * u[i])
+                hv = h(torch.cat(xs, 0))
+                hv = hv.reshape(len(xs) // 2, 2, B0, *hv.shape[1:])
+                parts.append((hv[:, 0] - hv[:, 1]) / (2 * eps))
+            outs.append(torch.cat(parts, 0))
+    else:
+        for eps in eps_list:
+            outs.append(torch.stack([(h(x + eps * u[i]) - h(x - eps * u[i])) / (2 * eps) for i in range(n)]))
+    return torch.stack(outs)
+
+
+def exact_directional(f, x, v, u, n_batch_dims=1, local=None, chunk=16):
+    """Oracle: exact directional derivatives D_i = <v, J u_i> per measurement unit via forward-mode AD
+    (vmapped over probes in chunks; falls back to a loop if the block is not vmap-compatible)."""
+    def one(t):
+        return torch.func.jvp(f, (x,), (t,))[1]
+    outs = []
+    n = u.shape[0]
+    try:
+        for i in range(0, n, chunk):
+            outs.append(torch.func.vmap(one)(u[i:i + chunk]))
+        Ju = torch.cat(outs, 0)
+    except Exception:
+        Ju = torch.stack([one(u[i]) for i in range(n)])
+    prod = Ju * v.unsqueeze(0)
+    if local is None:
+        return prod.reshape(n, *x.shape[:n_batch_dims], -1).sum(-1)
+    if getattr(local, "kind", "2d") == "1d":
+        return prod.sum(-1)
+    return prod.sum(n_batch_dims + 1)                         # sum over channel dim only
+
+
+def oracle_vjp(f, x, v, cfg, gen, n_batch_dims=1, local=None):
+    """Stage A: (1/n) sum_i u_i <v, J u_i>  (== (1/n) sum_i u_i (u_i^T g)).  Conditionally unbiased."""
+    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims)
+    with torch.no_grad():
+        D = exact_directional(f, x, v, u, n_batch_dims, local)
+    return _combine(u, D, n_batch_dims, local, x.shape), 0.0
+
+
+def oracle_projection(g, cfg, gen, n_batch_dims=1):
+    """Projection form of the oracle for a given exact VJP g (global units only)."""
+    u, d = _probes_for(g, cfg.n_probes, cfg.probe, gen, n_batch_dims)
+    n = u.shape[0]
+    bs = g.shape[:n_batch_dims]
+    D = (u.reshape(n, *bs, -1) * g.reshape(1, *bs, -1)).sum(-1)
+    return _combine(u, D, n_batch_dims, None, g.shape), 0.0
+
+
+def zo_vjp(h, x, cfg, gen, n_batch_dims=1, local=None, cycle=None):
+    """Function-query estimate of J^T v.  Returns (g_hat, queries_per_sample_unit)."""
+    if cfg.mode == "cd" and cfg.probe == "coordinate_all" and local is None:
+        # analytic combine: u_i = sqrt(d) e_i  ->  g_hat[..., i] = D_i / sqrt(d); the probe tensor is a broadcast
+        # view of eye(d), so memory stays O(d^2 + measurements) instead of O(n * batch * d)
+        bs = x.shape[:n_batch_dims]
+        d = math.prod(x.shape[n_batch_dims:])
+        if cfg.n_probes != d:
+            raise ValueError(f"coordinate_all needs n == d (n={cfg.n_probes}, d={d})")
+        eps = _eps_for(cfg, d)
+        ident = math.sqrt(d) * torch.eye(d, device=x.device, dtype=x.dtype)
+        uv = ident.reshape(d, *([1] * len(bs)), d).expand(d, *bs, d)
+        D = _eval_pairs(h, x, uv, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+        return (D.movedim(0, -1) / math.sqrt(d)).reshape(x.shape), float(2 * d)
+    if cfg.mode == "cd" and cfg.probe_chunk and cfg.n_probes > 4 * cfg.probe_chunk and cfg.probe != "hadamard_cycle":
+        # probe-chunked sample -> evaluate -> accumulate: memory O(chunk * batch * d) for any n
+        bs = x.shape[:n_batch_dims]
+        d = math.prod(x.shape[n_batch_dims:])
+        eps = _eps_for(cfg, d)
+        g_acc = torch.zeros_like(x)
+        done = 0
+        while done < cfg.n_probes:
+            c = min(cfg.probe_chunk, cfg.n_probes - done)
+            uc, _ = _probes_for(x, c, cfg.probe, gen, n_batch_dims)
+            D = _eval_pairs(h, x, uc, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+            g_acc += _combine(uc, D, n_batch_dims, local, x.shape) * c
+            done += c
+        return g_acc / cfg.n_probes, float(2 * cfg.n_probes)
+    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle)
+    n = cfg.n_probes
+    eps = _eps_for(cfg, d)
+    bs = x.shape[:n_batch_dims]
+    B0 = x.shape[0]
+
+    if cfg.mode == "forward":
+        if cfg.batch_probes:
+            hv = h(torch.cat([x] + [x + eps * u[i] for i in range(n)], 0))
+            hv = hv.reshape(n + 1, B0, *hv.shape[1:])
+            D = (hv[1:] - hv[:1]) / eps
+        else:
+            h0 = h(x)
+            D = torch.stack([(h(x + eps * u[i]) - h0) / eps for i in range(n)])
+        return _combine(u, D, n_batch_dims, local, x.shape), float(n + 1)
+
+    if cfg.mode == "cd":
+        D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+        return _combine(u, D, n_batch_dims, local, x.shape), float(2 * n)
+
+    if cfg.mode == "richardson":
+        D = _eval_pairs(h, x, u, [eps, eps / 2], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)
+        R = (4 * D[1] - D[0]) / 3
+        return _combine(u, R, n_batch_dims, local, x.shape), float(4 * n)
+
+    if cfg.mode in ("ml", "ml_richardson"):
+        rich = cfg.mode == "ml_richardson"
+        if cfg.level_sampling == "per_sample":
+            U = torch.rand((n, *bs), generator=gen, device=x.device, dtype=torch.float64)
+        else:
+            U = torch.rand((n,), generator=gen, device=x.device, dtype=torch.float64)
+            U = U.reshape(n, *([1] * len(bs))).expand(n, *bs)
+        N = torch.floor(-torch.log2(U) / cfg.alpha).clamp(max=cfg.max_level).to(torch.int64)
+        if cfg.eps_min > 0:   # truncate the ladder where readout noise would dominate (bias O(eps_min^2))
+            kmax = max(0, int(math.floor(math.log2(eps / cfg.eps_min)))) if eps > cfg.eps_min else 0
+            N = N.clamp(max=kmax)
+        Nmax = int(N.max().item())
+        K = Nmax + (2 if rich else 1)
+        eps_list = [eps * 2.0 ** (-k) for k in range(K)]
+        D = _eval_pairs(h, x, u, eps_list, n_batch_dims, cfg.batch_probes, cfg.probe_chunk)   # [K, n, *bs, *meas]
+        S = (4 * D[1:] - D[:-1]) / 3 if rich else D
+        est = S[0].clone()
+        extra = S.dim() - 1 - N.dim()
+        for k in range(1, Nmax + 1):
+            mask = (N >= k).to(x.dtype).reshape(*N.shape, *([1] * extra))
+            est = est + mask * (S[k] - S[k - 1]) / (2.0 ** (-cfg.alpha * k))
+        per = 2.0 * (N.to(torch.float64) + 1 + (1 if rich else 0))
+        queries = float(per.sum().item() / (math.prod(bs) if len(bs) else 1))
+        return _combine(u, est, n_batch_dims, local, x.shape), queries
+
+    raise ValueError(f"unknown ZBP mode {cfg.mode!r}")
+
+
+def causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None):
+    """Estimate attention's K-side VJP by probing causal score rows.
+
+    The black-box interface here is ``softmax(scores) @ value`` rather than K itself.
+    One paired query perturbs one valid score coordinate in every (sample, head,
+    output-token) row, and the output supplies a separate scalar measurement for
+    every such row.  For a row with ``m`` valid keys, the estimator visits all
+    coordinates when ``m <= n``; otherwise it samples ``n`` coordinates without
+    replacement and applies the Horvitz--Thompson weight ``m / n``.  Thus the
+    row-gradient relative MSE is exactly ``max(m / n - 1, 0)`` in the directional-
+    derivative limit, instead of ``(m - 1) / n`` for i.i.d. dense probes.
+
+    The estimated score VJP is mapped to K through the digitally known bilinear
+    score map QK^T/sqrt(dh).  Q and V are held fixed during the function queries.
+    Returns ``(g_k, queries_per_sequence)``; full-attention K is [B,T,d], while
+    multi-query K is [B,T,dh].
+    """
+    if cfg.mode not in ("cd", "oracle"):
+        raise ValueError(f"causal score-coordinate probing supports cd/oracle, got {cfg.mode!r}")
+    B, T, _ = qkv.shape
+    H = int(n_heads)
+    dh = int(d) // H
+    dk = int(d if dk is None else dk)
+    dv = int(d if dv is None else dv)
+    if int(d) % H:
+        raise ValueError(f"attention width {d} is not divisible by {H} heads")
+    if dk not in (int(d), dh) or dv not in (int(d), dh):
+        raise ValueError("score-coordinate probing supports full or multi-query K/V layouts")
+
+    q0, k0, value0 = qkv.split([int(d), dk, dv], dim=-1)
+    q = q0.view(B, T, H, dh).transpose(1, 2)
+    if dk == dh:
+        k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
+    else:
+        k = k0.view(B, T, H, dh).transpose(1, 2)
+    if dv == dh:
+        value = value0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
+    else:
+        value = value0.view(B, T, H, dh).transpose(1, 2)
+    delta = v_out.view(B, T, H, dh).transpose(1, 2)
+
+    idx = torch.arange(T, device=qkv.device)
+    valid = idx[None, :] <= idx[:, None]
+    if window is not None:
+        valid = valid & (idx[None, :] > idx[:, None] - int(window))
+    counts = valid.sum(-1)                                      # [T], valid coordinates per row
+    scores = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
+    scores = scores.masked_fill(~valid, float("-inf"))
+
+    # Analytic score VJP is useful for projection-only audits.  The CD path below
+    # obtains the same quantities solely from paired softmax/value evaluations.
+    if cfg.mode == "oracle":
+        att = scores.softmax(-1)
+        g_att = delta @ value.transpose(-2, -1)
+        g_score_exact = att * (g_att - (att * g_att).sum(-1, keepdim=True))
+        g_score_exact = g_score_exact.masked_fill(~valid, 0)
+
+    n = int(cfg.n_probes)
+    if n <= 0:
+        raise ValueError("score-coordinate probing needs at least one probe")
+    slots = min(n, T)
+    # Independent random subsets for every sample/head/row.  Invalid positions
+    # sort last, so the first min(n,m) entries are a uniform subset of the m-prefix.
+    priorities = torch.rand((B, H, T, T), generator=gen, device=qkv.device, dtype=torch.float32)
+    priorities.masked_fill_(~valid.view(1, 1, T, T), 2.0)
+    chosen = priorities.argsort(dim=-1)[..., :slots]             # [B,H,T,slots]
+    chosen = chosen.permute(3, 0, 1, 2).contiguous()             # [slots,B,H,T]
+    slot_id = torch.arange(slots, device=qkv.device).view(slots, 1, 1, 1)
+    active = slot_id < counts.view(1, 1, 1, T)
+    active = active.expand(slots, B, H, T)
+    weights = torch.where(counts <= n, torch.ones_like(counts, dtype=qkv.dtype),
+                          counts.to(qkv.dtype) / n)               # [T]
+    g_score = torch.zeros_like(scores)
+
+    if cfg.mode == "oracle":
+        picked = torch.gather(g_score_exact.unsqueeze(0).expand(slots, -1, -1, -1, -1),
+                              -1, chosen.unsqueeze(-1)).squeeze(-1)
+        contrib = picked * active.to(qkv.dtype) * weights.view(1, 1, 1, T)
+        for j in range(slots):
+            g_score.scatter_add_(-1, chosen[j].unsqueeze(-1), contrib[j].unsqueeze(-1))
+        queries = 0.0
+    else:
+        eps = float(cfg.eps)  # directions are unit coordinates, so coord/norm conventions coincide
+        if eps <= 0:
+            raise ValueError("score-coordinate probing needs eps > 0")
+        chunk = slots if cfg.probe_chunk <= 0 else max(1, min(int(cfg.probe_chunk), slots))
+        for i0 in range(0, slots, chunk):
+            ids = chosen[i0:i0 + chunk]
+            act = active[i0:i0 + chunk]
+            c = ids.shape[0]
+            plus = scores.unsqueeze(0).expand(c, -1, -1, -1, -1).clone()
+            minus = plus.clone()
+            shift = act.to(qkv.dtype).unsqueeze(-1) * eps
+            plus.scatter_add_(-1, ids.unsqueeze(-1), shift)
+            minus.scatter_add_(-1, ids.unsqueeze(-1), -shift)
+            y_plus = plus.softmax(-1) @ value.unsqueeze(0)
+            y_minus = minus.softmax(-1) @ value.unsqueeze(0)
+            if cfg.readout_noise > 0:
+                y_plus = y_plus + cfg.readout_noise * torch.randn(
+                    y_plus.shape, generator=gen, device=y_plus.device, dtype=y_plus.dtype)
+                y_minus = y_minus + cfg.readout_noise * torch.randn(
+                    y_minus.shape, generator=gen, device=y_minus.device, dtype=y_minus.dtype)
+            deriv = ((y_plus - y_minus) * delta.unsqueeze(0)).sum(-1) / (2 * eps)
+            contrib = deriv * act.to(qkv.dtype) * weights.view(1, 1, 1, T)
+            for j in range(c):
+                g_score.scatter_add_(-1, ids[j].unsqueeze(-1), contrib[j].unsqueeze(-1))
+        queries = float(2 * n)
+
+    _LAST_SCORE_VJP['g'] = g_score
+    gk_heads = g_score.transpose(-2, -1) @ q / math.sqrt(dh)
+    if dk == dh:
+        gk = gk_heads.sum(1).transpose(1, 2)                      # shared K accumulates all heads
+    else:
+        gk = gk_heads.transpose(1, 2).reshape(B, T, int(d))
+    return gk, queries
+
+
+def make_h(f, v, n_batch_dims=1, local=None, readout_noise=0.0, gen=None):
+    """Measurement function h(x_pert) = <v, f(x_pert)> per unit, tolerant to stacked batches.
+    readout_noise > 0 adds i.i.d. Gaussian noise to every measured output coordinate of every query
+    (a physical readout-noise model; v is applied digitally afterwards)."""
+    B0 = v.shape[0]
+
+    def h(xp):
+        with torch.no_grad():
+            y = f(xp)
+            if readout_noise > 0:
+                y = y + readout_noise * torch.randn(y.shape, generator=gen, device=y.device, dtype=y.dtype)
+            k = y.shape[0] // B0
+            vv = v.repeat(k, *([1] * (v.dim() - 1))) if k > 1 else v
+            prod = y * vv
+            if local is None:
+                return prod.reshape(*prod.shape[:n_batch_dims], -1).sum(-1)
+            if getattr(local, "kind", "2d") == "1d":
+                return prod.sum(-1)                           # per output token
+            return prod.sum(n_batch_dims)                     # keep spatial dims
+    return h
+
+
+def control_vjp(g, cfg, gen, n_batch_dims=1):
+    """Control estimators.  'zero': no credit is propagated (downstream layers frozen).
+    'noise': isotropic Gaussian noise with the per-sample norm that the oracle estimate with
+    the same n would have (sqrt(1 + (d-1)/n) ||g||): same magnitude, zero information."""
+    if cfg.mode == "zero":
+        return torch.zeros_like(g), 0.0
+    if cfg.mode == "noise":
+        bs = g.shape[:n_batch_dims]
+        gf = g.reshape(*bs, -1)
+        d = gf.shape[-1]
+        z = torch.randn(gf.shape, generator=gen, device=g.device, dtype=g.dtype)
+        z = z / z.norm(dim=-1, keepdim=True).clamp_min(1e-30)
+        target = gf.norm(dim=-1, keepdim=True) * math.sqrt(1 + (d - 1) / cfg.n_probes)
+        return (z * target).reshape(g.shape), 0.0
+    raise ValueError(cfg.mode)
+
+
+def causal_score_coordinate_qk_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None):
+    """Same score-space probing as causal_score_coordinate_k_vjp, but returns the Q-side VJP as well
+    (both follow from the score-VJP through the known bilinear map S = Q K^T / sqrt(d_h))."""
+    B, T, _ = qkv.shape
+    H = int(n_heads); dh = int(d) // H
+    dk = int(d if dk is None else dk)
+    gk, queries = causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk, dv, window)
+    # recover g_score from gk is not possible in general; recompute the score VJP estimate via the same
+    # probes by calling the K routine on a transposed view is not exact either, so instead re-derive:
+    # g_K[s] = sum_t g_S[t,s] q_t / sqrt(dh)  and  g_Q[t] = sum_s g_S[t,s] k_s / sqrt(dh).
+    # We obtain g_S by solving nothing: the K routine exposes it through `_last_score_vjp`.
+    g_score = _LAST_SCORE_VJP["g"]
+    k0 = qkv[..., int(d):int(d) + dk]
+    if dk == dh:
+        k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
+    else:
+        k = k0.view(B, T, H, dh).transpose(1, 2)
+    gq = (g_score @ k) / math.sqrt(dh)                              # [B,H,T,dh]
+    gq = gq.transpose(1, 2).reshape(B, T, int(d))
+    return gq, gk, queries
+
+
+_LAST_SCORE_VJP = {"g": None}