1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
|
# zbp-scaling
Scaling study for **ZBP (zeroth-order backpropagation)** vs BP on decoder-only language models
(part 2 of the ZBP paper). Every nonlinear block is treated as a physical black box trained from
forward queries only (score-space attention core + per-token SwiGLU FFN); linear maps and the score
product are digital. See the main zobp repo for the method, theory (NSR ≈ c·d/n) and part-1/3 results.
## Install
```
pip install -e . # or: pip install -e .[dev] && pytest tests/
```
## Layout
- `src/zbp_scaling/zbp/` — vendored ZBP package (probes, estimators, ZBPBlock autograd)
- `src/zbp_scaling/model.py` — OLMo2-style transformer (RMSNorm, SwiGLU, untied embeddings; learned
positions in v1) with the ZBP physical/digital partition
- `src/zbp_scaling/data.py` — FineWeb-Edu → GPT-2-BPE uint16 shards (WikiText-103 fallback for smoke)
- `src/zbp_scaling/trainer.py` — DDP (torchrun, 4 or 8 GPUs), bf16 autocast with fp32 measurement
accumulation, gradient accumulation to a fixed global batch, cosine + warmup, resume, JSONL logs
- `src/zbp_scaling/diagnostics.py` — per-block branch-gain rho_k profile and the NSR constant c(scale)
- `configs/` — model sizes (60m/124m/350m/1b) x training arms (bp / zbp_n16 / zbp_n4)
## Run (one click, on the H200 node)
```
./scripts/run_ladder.sh # env check+install -> data prep -> all sizes x arms (resume-safe)
VENV=1 ./scripts/run_ladder.sh # fresh node: bootstrap ./.venv first
SIZES="s60" ARMS="bp zbp_n16" TOKENS=3e8 ./scripts/run_ladder.sh # any subset / budget
DRY=1 ./scripts/run_ladder.sh # print the plan
```
Data prep runs **on the training node** (auto-invoked if `data/` is empty); `data/` and `runs/` are gitignored.
```
python scripts/prepare_data.py --tokens 3e9 --out data/fineweb # manual form
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.
## Citation
Paper in preparation ("Backpropagation Without Jacobians"); see the main zobp research repo for the
method, theory and part-1/part-3 results. License: Apache-2.0.
## After the runs (results flow)
Checkpoints stay on the node (`runs/` is gitignored). Collect the small JSONL logs and push them back:
```
python scripts/collect.py --runs runs --out results/h200node1
git add results && git commit -m "ladder results" && git push
```
Analysis (anywhere): `python scripts/plot_ladder.py --results results/h200node1` -> per-run table, loss
curves, and the gap-vs-scale figure (the paper's part-2 headline). If a specific checkpoint is needed for
the estimator audits, scp just that `runs/<name>/ckpt.pt`.
## Collaborator quickstart (zero tokens on the node)
You receive ONE file: the deploy key `zbp_scaling_deploy` (scoped to this repo only, revocable). Then:
```
install -m 600 zbp_scaling_deploy ~/.ssh/zbp_scaling_deploy
git clone -c core.sshCommand="ssh -i ~/.ssh/zbp_scaling_deploy -o IdentitiesOnly=yes" \
git@github.com:YurenHao0426/zbp-scaling.git
cd zbp-scaling && ./scripts/run_ladder.sh # env check -> data prep -> ladder -> results auto-pushed back
```
The `-c` persists `core.sshCommand` inside the clone, so the auto-push at the end works with no env setup
(nohup-safe; set `PUSH_RESULTS=0` to disable). No GitHub account, no HF token on the node: results JSONL
flow back through the deploy key; checkpoints stay on the node (scp on request) and HF uploads happen on
the maintainer's machine.
## HF upload & security (shared nodes)
Results (and optionally checkpoints) can go to a **private** HF repo: `HF_UPLOAD=1 [HF_CKPT=1] ./scripts/run_ladder.sh`
or manually `python scripts/upload_hf.py --results results/<tag> [--with-ckpt runs]` (default repo
`<whoami>/zbp-scaling-runs`, created private if missing).
Uploads authenticate ONLY via the `HF_TOKEN` environment variable or a standard `hf auth login`; tokens are
never CLI arguments (argv is world-readable via /proc on shared machines), never written by our scripts, and
`.gitignore` excludes token-like files. On a shared node, mint a **fine-grained HF token scoped to the single
private repo** (write permission only), `export HF_TOKEN=...` per session, and revoke it after the campaign.
Zero-token alternative: push only the small JSONL results to GitHub (a repo-scoped deploy key suffices) and
upload checkpoints from a trusted machine.
|
