1 files changed, 57 insertions, 0 deletions
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"