From 98d05c2f52527d31cf1872f8eda19d1b6d91bc5b Mon Sep 17 00:00:00 2001
From: yurenh <blackhao0426@gmail.com>
Date: Mon, 31 Aug 2026 18:25:33 -0500
Subject: one-click resume-safe ladder runner (env-configurable)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe
---
 README.md             | 11 +++++++---
 scripts/run_ladder.sh | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 3 deletions(-)
 create mode 100755 scripts/run_ladder.sh

diff --git a/README.md b/README.md
index d42f713..e0a2446 100644
--- a/README.md
+++ b/README.md
@@ -20,10 +20,15 @@ pip install -e .            # or: pip install -e .[dev] && pytest tests/
 - `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
-Data prep runs **on the training node** (H200), not on a dev machine; `data/` and `runs/` are gitignored.
+## Run (one click, on the H200 node)
 ```
-python scripts/prepare_data.py --tokens 3e9 --out data/fineweb        # on the H200 node
+./scripts/run_ladder.sh                       # all sizes x arms, sequential, resume-safe
+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.
diff --git a/scripts/run_ladder.sh b/scripts/run_ladder.sh
new file mode 100755
index 0000000..4d253f1
--- /dev/null
+++ b/scripts/run_ladder.sh
@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# One-click ZBP scaling ladder. Clone -> ./scripts/run_ladder.sh   (run on the H200 node, nohup-able)
+#
+# Configurable via env vars (defaults in brackets):
+#   SIZES="s60 m124 l350"        model configs from configs/model/   [s60 m124 l350]
+#   ARMS="bp zbp_n16 zbp_n64"    train configs from configs/train/   [bp zbp_n16 zbp_n64]
+#   NPROC=8                      GPUs per run                        [auto: nvidia-smi -L]
+#   TOKENS=2e9                   override token budget per run       [per train config]
+#   DATA=data/fineweb            token shards (prepared if missing)  [data/fineweb]
+#   PREP_TOKENS=3e9              tokens to tokenize if DATA missing  [3e9]
+#   OUT=runs                     output root                        [runs]
+#   MICRO_BS=8                   per-GPU micro batch                [8]
+#   DRY=1                        print the plan and exit
+#
+# Runs are sequential (each takes the whole node), resume-safe: a finished run leaves OUT/<name>/DONE
+# and is skipped on re-invocation, so the script can be re-run after interruptions.
+set -euo pipefail
+cd "$(dirname "$0")/.."
+
+SIZES=${SIZES:-"s60 m124 l350"}
+ARMS=${ARMS:-"bp zbp_n16 zbp_n64"}
+NPROC=${NPROC:-$(nvidia-smi -L 2>/dev/null | wc -l)}
+DATA=${DATA:-data/fineweb}
+PREP_TOKENS=${PREP_TOKENS:-3e9}
+OUT=${OUT:-runs}
+MICRO_BS=${MICRO_BS:-8}
+EXTRA=${TOKENS:+--set tokens=$TOKENS}
+
+echo "== zbp-scaling ladder: sizes=[$SIZES] arms=[$ARMS] nproc=$NPROC data=$DATA out=$OUT ${TOKENS:+tokens=$TOKENS}"
+[ "$NPROC" -ge 1 ] || { echo "no GPUs detected; set NPROC"; exit 1; }
+python -c "import zbp_scaling" 2>/dev/null || pip install -e . -q
+
+if [ ! -f "$DATA/train.bin" ]; then
+  echo "== preparing data ($PREP_TOKENS tokens -> $DATA) — run this on the training node"
+  [ "${DRY:-0}" = 1 ] || python scripts/prepare_data.py --tokens "$PREP_TOKENS" --out "$DATA"
+fi
+
+for size in $SIZES; do
+  for arm in $ARMS; do
+    name="${size}_${arm}"
+    dir="$OUT/$name"
+    if [ -f "$dir/DONE" ]; then echo "== skip $name (done)"; continue; fi
+    echo "== run $name  (model=configs/model/$size.yaml train=configs/train/$arm.yaml)"
+    [ "${DRY:-0}" = 1 ] && continue
+    mkdir -p "$dir"
+    if [ "$NPROC" -gt 1 ]; then
+      torchrun --standalone --nproc_per_node="$NPROC" scripts/train.py \
+        --model "configs/model/$size.yaml" --train "configs/train/$arm.yaml" \
+        --data "$DATA" --out "$dir" --micro_bs "$MICRO_BS" $EXTRA 2>&1 | tee -a "$dir/stdout.log"
+    else
+      python scripts/train.py --model "configs/model/$size.yaml" --train "configs/train/$arm.yaml" \
+        --data "$DATA" --out "$dir" --micro_bs "$MICRO_BS" $EXTRA 2>&1 | tee -a "$dir/stdout.log"
+    fi
+    touch "$dir/DONE"
+  done
+done
+echo "== ladder complete: $OUT"
-- 
cgit v1.2.3