data prep (FineWeb-Edu->GPT2 BPE), rho diagnostics, tests, measured-cost notes

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

5 files changed, 102 insertions, 0 deletions

diff --git a/README.md b/README.md
index c9f5821..9c038d5 100644
--- a/README.md
+++ b/README.md
@@ -21,3 +21,9 @@ python scripts/prepare_data.py --dataset fineweb-edu --tokens 3e9 --out data/fin
 torchrun --nproc_per_node=8 scripts/train.py --model configs/model/m124.yaml --train configs/train/zbp_n16.yaml
 ```
 Global batch is fixed in the train config; per-rank micro-batch and accumulation adapt to world size.
+
+## Measured cost (48 GB Ampere-class, d=512 L=8 seq 1024, shared GPU)
+BP 312 ms/step; ZBP n=16 **8.7x**, n=64 **29x** (FLOPs-bound: probe batching is already saturated at
+probe_chunk 8; the score core uses its own chunk <= 4 since its memory goes as chunk*B*H*T^2).
+Ladder arms: **bp / zbp_n16 / zbp_n64** per size (n=4 optional).  Headroom if needed: torch.compile on the
+query path; forward differences (n+1 instead of 2n queries) as a cheaper biased arm.
diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py
new file mode 100644
index 0000000..626e7a8
--- /dev/null
+++ b/scripts/prepare_data.py
@@ -0,0 +1,35 @@
+"""FineWeb-Edu -> GPT-2 BPE uint16 shards (train.bin / val.bin).
+  python scripts/prepare_data.py --tokens 3e9 --out data/fineweb [--dataset HuggingFaceFW/fineweb-edu --name sample-10BT]"""
+import os, sys, argparse
+import numpy as np
+
+p = argparse.ArgumentParser()
+p.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
+p.add_argument("--name", default="sample-10BT")
+p.add_argument("--tokens", type=float, default=3e9)
+p.add_argument("--val_tokens", type=float, default=5e6)
+p.add_argument("--out", default="data/fineweb")
+a = p.parse_args()
+os.makedirs(a.out, exist_ok=True)
+import tiktoken
+from datasets import load_dataset
+enc = tiktoken.get_encoding("gpt2")
+ds = load_dataset(a.dataset, name=a.name, split="train", streaming=True)
+train_path, val_path = os.path.join(a.out, "train.bin"), os.path.join(a.out, "val.bin")
+ftr, fva = open(train_path, "wb"), open(val_path, "wb")
+n_tr = n_va = 0
+target_tr, target_va = int(a.tokens), int(a.val_tokens)
+buf = []
+for i, ex in enumerate(ds):
+    ids = enc.encode_ordinary(ex["text"]) + [enc.eot_token]
+    arr = np.array(ids, dtype=np.uint16)
+    if n_va < target_va and i % 100 == 0:          # every 100th doc to val until filled
+        fva.write(arr.tobytes()); n_va += len(arr)
+    else:
+        ftr.write(arr.tobytes()); n_tr += len(arr)
+    if n_tr % 50_000_000 < len(arr):
+        print(f"train {n_tr/1e6:.0f}M  val {n_va/1e6:.1f}M tokens", flush=True)
+    if n_tr >= target_tr and n_va >= target_va:
+        break
+ftr.close(); fva.close()
+print(f"DONE train {n_tr/1e6:.1f}M val {n_va/1e6:.1f}M -> {a.out}")
diff --git a/src/zbp_scaling/diagnostics.py b/src/zbp_scaling/diagnostics.py
new file mode 100644
index 0000000..4fa3747
--- /dev/null
+++ b/src/zbp_scaling/diagnostics.py
@@ -0,0 +1,35 @@
+"""Branch-gain profile rho_k and the NSR constant c(scale) = L * mean(rho) (THEORY.md T2/corollary).
+Validation-only: uses the simulation autograd for exact J_F^T v."""
+import torch
+import torch.nn as nn
+from .zbp import zbp_blocks, ZBPConfig
+
+
+@torch.enable_grad()
+def rho_profile(model, x, y):
+    blocks = zbp_blocks(model)
+    cfgs = [b.cfg for b in blocks]
+    for b in blocks:
+        b.cfg = b.cfg.replace(mode="exact")
+        b.capture = True
+    model.zero_grad(set_to_none=True)
+    nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).backward()
+    out = {}
+    for b, c in zip(blocks, cfgs):
+        b.cfg = c
+        b.capture = False
+        b.captured = None
+    # rho via a recorded pass: |J_F^T v| / |v| per block from the Recorder
+    from .zbp import Recorder
+    for b in blocks:
+        b.cfg = b.cfg.replace(mode="zero")
+    with Recorder() as rec:
+        model.zero_grad(set_to_none=True)
+        nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).backward()
+    for b, c in zip(blocks, cfgs):
+        b.cfg = c
+    model.zero_grad(set_to_none=True)
+    rows = rec.summary()
+    for name, r in rows.items():
+        out[name] = (r["gnorm"] / max(r["vnorm"], 1e-30)) ** 2 if "vnorm" in r else None
+    return out
diff --git a/tests/__pycache__/test_scaling.cpython-313-pytest-9.0.2.pyc b/tests/__pycache__/test_scaling.cpython-313-pytest-9.0.2.pyc
new file mode 100644
index 0000000..a1801fd
--- /dev/null
+++ b/tests/__pycache__/test_scaling.cpython-313-pytest-9.0.2.pyc
diff --git a/tests/test_scaling.py b/tests/test_scaling.py
new file mode 100644
index 0000000..b66a05c
--- /dev/null
+++ b/tests/test_scaling.py
@@ -0,0 +1,26 @@
+import os, sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+import torch
+import torch.nn as nn
+from zbp_scaling.model import ScalingLM
+from zbp_scaling.zbp import ZBPConfig, zbp_blocks
+
+
+def test_exact_mode_matches_bp():
+    torch.manual_seed(0)
+    m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="bp"))
+    x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
+    g_bp = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
+    for b in zbp_blocks(m):
+        b.cfg = ZBPConfig(mode="exact")
+    g = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
+    err = max((a - b).norm().item() / (b.norm().item() + 1e-12) for a, b in zip(g, g_bp))
+    assert err < 1e-5, err
+
+
+def test_zbp_cd_trains_shape():
+    torch.manual_seed(0)
+    m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="cd", n_probes=4, eps=0.1, probe_chunk=4))
+    x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
+    nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()).backward()
+    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in m.parameters())