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

1 files changed, 35 insertions, 0 deletions

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