From 6a544fabfc2af22e4d5823410dd2387b5af89ea9 Mon Sep 17 00:00:00 2001 From: yurenh Date: Mon, 31 Aug 2026 18:14:09 -0500 Subject: scaffold: model (OLMo2-ish + ZBP partition), trainer (DDP/config), data shards, bench Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe --- src/zbp_scaling/zbp/wp.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 src/zbp_scaling/zbp/wp.py (limited to 'src/zbp_scaling/zbp/wp.py') 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 -- cgit v1.2.3