blob: 7402111fa3ddaf8476bc90ec2cf2a1c6a7a31ea1 (
plain)
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
81
82
83
84
85
86
87
|
#!/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]
# SET="k=v k2=v2" extra --set overrides for every run (e.g. vocab=8192 seq_len=256)
# DRY=1 print the plan and exit
# PUSH_RESULTS=0 skip the auto collect + git-push of results JSONL [on]
# HF_UPLOAD=1 [HF_CKPT=1] also upload to a private HF repo (needs HF_TOKEN; see README)
#
# 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="--set ${TOKENS:+tokens=$TOKENS} ${SET:-}"
[ "$EXTRA" = "--set " ] && EXTRA=""
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; }
# --- environment bootstrap: current env first; VENV=1 creates ./.venv (fresh nodes) ---
PY=${PY:-python3}
if [ "${VENV:-0}" = 1 ]; then
[ -d .venv ] || "$PY" -m venv .venv
. .venv/bin/activate
fi
python -c "import zbp_scaling" 2>/dev/null || pip install -e . -q || {
echo "pip install failed; on managed nodes run with VENV=1 (creates ./.venv) or install into your conda env"; exit 1; }
python scripts/env_check.py || exit 1
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
[ -f "configs/model/$size.yaml" ] || { echo "missing configs/model/$size.yaml"; exit 1; }
[ -f "configs/train/$arm.yaml" ] || { echo "missing configs/train/$arm.yaml"; exit 1; }
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
TAG=${HF_TAG:-$(hostname)-$(date +%Y%m%d)}
if [ "${PUSH_RESULTS:-1}" = 1 ] && [ "${DRY:-0}" != 1 ]; then
# zero-token default: collect the small JSONL results and push them back over this clone's git auth
if python scripts/collect.py --runs "$OUT" --out "results/$TAG"; then
git add results
git -c user.email=ladder@zbp -c user.name=ladder commit -m "results: $TAG" || true # nothing new is fine
git push || echo "!! git push failed — push manually later or send results/$TAG"
else
echo "!! collect found no finished runs — skipping push"
fi
fi
if [ "${HF_UPLOAD:-0}" = 1 ]; then
# optional direct-to-HF (needs HF_TOKEN in the environment; see README Security)
python scripts/collect.py --runs "$OUT" --out "results/$TAG"
python scripts/upload_hf.py --results "results/$TAG" ${HF_REPO:+--repo "$HF_REPO"} ${HF_CKPT:+--with-ckpt "$OUT"}
fi
echo "== ladder complete: $OUT"
|
