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

diff --git a/src/zbp_scaling/zbp/config.py b/src/zbp_scaling/zbp/config.py
new file mode 100644
index 0000000..e6269e2
--- /dev/null
+++ b/src/zbp_scaling/zbp/config.py
@@ -0,0 +1,78 @@
+from dataclasses import dataclass, asdict, replace
+
+
+@dataclass
+class ZBPConfig:
+    """Configuration of the zeroth-order VJP estimator used inside every ZBPBlock.
+
+    mode:
+        'bp'           exact reverse-mode autograd through the block (baseline)
+        'zero'         control: no activation error is propagated (upstream layers get no credit)
+        'noise'        control: random error with the norm an n-probe oracle estimate would have
+        'fa'           feedback alignment: fixed random linear map of the incoming error (skip path exact)
+        'dfa'          direct feedback alignment: fixed random projection of the output error (no chain)
+        'exact'        same as 'bp' but routed through the ZBP backward (debug/consistency check)
+        'oracle'       Stage A: exact g = J^T v computed internally, only its random
+                       projection (1/n) sum_i u_i (u_i^T g) is returned  (projection noise only)
+        'forward'      one-sided finite difference  (h(x+eps u) - h(x)) / eps
+        'cd'           Stage B: central difference   (h(x+eps u) - h(x-eps u)) / (2 eps)
+        'richardson'   Richardson-corrected central difference  (4 D_{eps/2} - D_eps) / 3
+        'ml'           Stage C: randomized multilevel debiased central difference (unbiased)
+        'ml_richardson'  randomized multilevel built on the Richardson sequence (unbiased)
+    n_probes:  number of random directions per sample unit per backward call.
+    probe:     'rademacher' | 'gaussian' | 'orthogonal' | 'hadamard' | 'coordinate'.
+               All families are normalized so that E[u u^T] = I_d (entries are O(1)).
+    eps:       finite-difference step.  With eps_mode='coord' the perturbation is x + eps*u
+               (each coordinate moves by ~eps, Euclidean norm eps*sqrt(d)); with
+               eps_mode='norm' it is x + eps*u/sqrt(d) (fixed Euclidean norm eps).
+    alpha:     survival exponent of the multilevel truncation level, P(N >= k) = 2^{-alpha k}.
+    max_level: hard cap on the multilevel index (P(N > max_level) = 2^{-alpha (max_level+1)}).
+    level_sampling: 'per_sample' (independent N for every (probe, sample) pair — the physical
+               cost model) or 'per_probe' (one N per probe shared by the batch — cheaper to simulate).
+    batch_probes: evaluate all probes of a backward call in one batched forward (simulation speed).
+    """
+    mode: str = "oracle"
+    n_probes: int = 4
+    probe: str = "rademacher"
+    eps: float = 1e-2
+    eps_mode: str = "coord"
+    alpha: float = 2.0
+    max_level: int = 12
+    level_sampling: str = "per_sample"
+    batch_probes: bool = True
+    probe_chunk: int = 0        # >0: evaluate at most this many probes per batched forward (bounds memory)
+    param_mode: str = "local"   # 'local': J_theta^T v via in-block autograd from the incoming error
+    readout_noise: float = 0.0  # std of additive Gaussian noise on every measured output coordinate (per query)
+    eps_min: float = 0.0        # multilevel ladder is truncated at eps_k >= eps_min (residual bias O(eps_min^2))
+    surrogate: bool = False     # learned linear control variate J^T ~ A per block (fitted online from the ZO estimates)
+    surrogate_ema: float = 0.9  # EMA factor of the regression statistics E[g v^T], E[v v^T]
+    surrogate_ridge: float = 1e-3
+
+    def replace(self, **kw):
+        return replace(self, **kw)
+
+    def asdict(self):
+        return asdict(self)
+
+    @property
+    def is_zo(self):
+        return self.mode in ("forward", "cd", "richardson", "ml", "ml_richardson")
+
+    def expected_queries_per_probe(self):
+        """Expected number of block evaluations (per sample unit) per probe direction."""
+        if self.mode == "oracle" or self.mode in ("bp", "exact"):
+            return 0.0
+        if self.mode == "forward":
+            return 1.0  # plus one shared base evaluation per backward call
+        if self.mode == "cd":
+            return 2.0
+        if self.mode == "richardson":
+            return 4.0
+        q = 2.0 ** (-self.alpha)
+        # E[N+1] = sum_{k>=0} Q_k = 1/(1-q)  (ignoring the cap)
+        en1 = 1.0 / (1.0 - q)
+        if self.mode == "ml":
+            return 2.0 * en1
+        if self.mode == "ml_richardson":
+            return 2.0 * (en1 + 1.0)  # levels 0..N+1
+        raise ValueError(self.mode)