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

diff --git a/src/zbp_scaling/zbp/wp.py b/src/zbp_scaling/zbp/wp.py
new file mode 100644
index 0000000..ff2511d
--- /dev/null
+++ b/src/zbp_scaling/zbp/wp.py
@@ -0,0 +1,28 @@
+"""Weight perturbation (parameter-space zeroth order) baseline: MeZO / CD-RGE / weight-space forward
+gradient.  All parameters are perturbed by eps*u (Rademacher), the loss is measured at +/- and the
+gradient estimate is (1/n) sum_i u_i (L(theta+eps u_i) - L(theta-eps u_i)) / (2 eps).
+2n forward passes of the whole network per step, no backward pass, variance ~ P/n with P = #params."""
+import torch
+
+
+def wp_step(model, loss_fn, x, y, n, eps, gen):
+    params = [p for p in model.parameters() if p.requires_grad]
+    grads = [torch.zeros_like(p) for p in params]
+    with torch.no_grad():
+        for i in range(n):
+            us = [(torch.randint(0, 2, p.shape, generator=gen, device=p.device).to(p.dtype) * 2 - 1) for p in params]
+            for p, u in zip(params, us):
+                p.add_(eps * u)
+            lp = loss_fn(model(x), y).item()
+            for p, u in zip(params, us):
+                p.sub_(2 * eps * u)
+            lm = loss_fn(model(x), y).item()
+            for p, u in zip(params, us):
+                p.add_(eps * u)
+            D = (lp - lm) / (2 * eps)
+            for g, u in zip(grads, us):
+                g.add_(u, alpha=D / n)
+        for p, g in zip(params, grads):
+            p.grad = g
+        loss = loss_fn(model(x), y)
+    return loss