diff options
| -rw-r--r-- | README.md | 6 | ||||
| -rw-r--r-- | scripts/prepare_data.py | 35 | ||||
| -rw-r--r-- | src/zbp_scaling/diagnostics.py | 35 | ||||
| -rw-r--r-- | tests/__pycache__/test_scaling.cpython-313-pytest-9.0.2.pyc | bin | 0 -> 5063 bytes | |||
| -rw-r--r-- | tests/test_scaling.py | 26 |
5 files changed, 102 insertions, 0 deletions
@@ -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 Binary files differnew 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()) |