diff options
| -rw-r--r-- | README.md | 11 | ||||
| -rwxr-xr-x | scripts/run_ladder.sh | 57 |
2 files changed, 65 insertions, 3 deletions
@@ -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" |