blob: 4d253f1405d84f802abfb66ed5c65ad2d94d0737 (
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
|
#!/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"
|
