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, 81 insertions, 0 deletions

diff --git a/src/zbp_scaling/zbp/joint.py b/src/zbp_scaling/zbp/joint.py
new file mode 100644
index 0000000..f29c210
--- /dev/null
+++ b/src/zbp_scaling/zbp/joint.py
@@ -0,0 +1,81 @@
+"""Model-level (non-chained) estimators that bracket chained ZBP:
+
+'np'      classic node perturbation with block-boundary structure: all block inputs are perturbed
+          simultaneously with independent probes, ONE scalar D = sum_l <u_l, g_l> (= the directional
+          derivative of the loss) is measured, and every block gets g_hat_l = (1/n) sum_i u_{l,i} D_i.
+          No compounding; variance of block l ~ (d_l/n) sum_l' |g_l'|^2; O(n) network forwards.
+'direct'  per-block projection of the EXACT input error (no compounding, no cross-block terms);
+          physically this is INP-like probing through all downstream blocks: O(n L^2) block forwards.
+
+Both are simulated in oracle form (exact directional derivatives) with a two-pass replay:
+pass 1 runs the network with exact backward and captures every block's full input error; the
+projected errors are then injected in pass 2, whose backward builds the local parameter gradients
+from the injected errors exactly as chained ZBP would.
+"""
+import torch
+from .autograd import zbp_blocks, get_generator
+from .probes import sample_probes
+
+
+def joint_backward(model, loss_fn, x, y, cfg, mode):
+    # only the physical blocks (those in the joint mode) are perturbed / replayed; digital blocks stay exact
+    blocks = [b for b in zbp_blocks(model) if b.estimate_input_grad and b.cfg.mode == mode]
+    cfgs = [b.cfg for b in blocks]
+    # ---- pass 1: exact input errors
+    for b in blocks:
+        b.cfg = b.cfg.replace(mode="exact")
+        b.capture = True
+    model.zero_grad(set_to_none=True)
+    loss = loss_fn(model(x), y)
+    loss.backward()
+    gs = [b.captured for b in blocks]
+    for b in blocks:
+        b.capture = False
+        b.captured = None
+    model.zero_grad(set_to_none=True)
+    # ---- projections (probes drawn in chunks so that memory is O(chunk * sum_l |g_l|), not O(n * ...))
+    gen = get_generator(x.device)
+    n = cfg.n_probes
+    chunk = cfg.probe_chunk if cfg.probe_chunk and cfg.probe_chunk > 0 else 8
+    acc = [torch.zeros_like(g) for g in gs]
+    for i0 in range(0, n, chunk):
+        c = min(chunk, n - i0)
+        us, Ds = [], []
+        for b, g in zip(blocks, gs):
+            bs = g.shape[:b.batch_dims]
+            d = g[0].numel() // (int(torch.tensor(bs[1:]).prod()) if len(bs) > 1 else 1)
+            u = sample_probes(c, tuple(bs), d, cfg.probe, gen, g.device, g.dtype).reshape(c, *g.shape)
+            D = (u.reshape(c, *bs, -1) * g.reshape(1, *bs, -1)).sum(-1)    # [c, *bs]
+            us.append(u); Ds.append(D)
+        if mode == "np":
+            # one scalar per (probe, sample): sum over blocks (and over tokens/positions within a sample)
+            Dtot = sum(D.reshape(c, D.shape[1], -1).sum(-1) for D in Ds)    # [c, B]
+            for a, u, g in zip(acc, us, gs):
+                shape = (c, g.shape[0]) + (1,) * (g.dim() - 1)
+                a.add_((u * Dtot.reshape(shape)).sum(0))
+        elif mode == "direct":
+            for a, u, D, g in zip(acc, us, Ds, gs):
+                extra = g.dim() - len(D.shape[1:])
+                a.add_((u * D.reshape(*D.shape, *([1] * extra))).sum(0))
+        else:
+            raise ValueError(mode)
+        del us, Ds
+    for b, a in zip(blocks, acc):
+        b.replay = a / n
+    if mode == "np":
+        queries = 2.0 * n                                                   # full-network forwards per sample
+    else:
+        L = len(blocks)
+        queries = 2.0 * n * sum(range(1, L + 1))                            # block forwards per sample
+    # ---- pass 2: replay projected errors, build local parameter gradients
+    for b in blocks:
+        b.cfg = b.cfg.replace(mode="replay")
+    model.zero_grad(set_to_none=True)
+    loss = loss_fn(model(x), y)
+    loss.backward()
+    for b, c in zip(blocks, cfgs):
+        b.cfg = c
+        b.replay = None
+        b.stats["queries"] += queries / max(1, len(blocks))
+        b.stats["backward_calls"] += 1
+    return loss