File size: 3,688 Bytes
b9301f1 | 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 | #!/usr/bin/env bash
# Resumable batched ATOM training driver (docs/BATCHED_TRAINING.md §Runbook). Re-running resumes from the
# newest checkpoint in $OUT (exact resume: weights, optimizer, RNG, B stream offsets, per-stream field state).
# Launch (detached):
# nohup env DATA_GLOB='data/fr/train_*.txt' VAL_GLOB='data/fr/val.txt' B=256 DEVICE=cuda \
# bash scripts/train_batched.sh >> logs/batched/master.log 2>&1 < /dev/null &
# Stop: kill $(cat $OUT/driver.pid) Progress: tail -f $OUT/progress.tsv (one row per checkpoint:
# train loss, byte CE, val CE normal / field-alone, learned leak, bytes/s, 3 chat samples)
# Env (defaults = scripts/train_bptt.sh recipe):
# DATA_GLOB (required) VAL_GLOB="" B=256 DEVICE=cuda THREADS=4 OUT=checkpoints/batched_b$B SUF=batched_b$B
# INIT="" (ckpt to start from, e.g. a single-stream checkpoint; empty = fresh init) TOTAL=4000000 (ticks)
# CKPT_EVERY=8000 LOG_EVERY=500 K=128 LEAK=0.05 LEARN=1 LR=2e-3 SLR=$LR LADROP=0.5 EPISODE=8000
# FLUSH=64 FIELD_LOSS=0.05 D=32 SEED=20260913 OFFSET_SEED=-1 SAMPLES=1 RESTARTS=3 EXTRA_ARGS=""
# Bytes consumed = ticks x B (B=256, 4M ticks = 1.02 GB). Checkpoints are never overwritten.
set -u
cd "$(dirname "$0")/.."
DATA_GLOB=${DATA_GLOB:?DATA_GLOB corpus glob}; VAL_GLOB=${VAL_GLOB:-}; B=${B:-256}; DEVICE=${DEVICE:-cuda}
THREADS=${THREADS:-4}; OUT=${OUT:-checkpoints/batched_b$B}; SUF=${SUF:-batched_b$B}; INIT=${INIT:-}
TOTAL=${TOTAL:-4000000}; CKPT_EVERY=${CKPT_EVERY:-8000}; LOG_EVERY=${LOG_EVERY:-500}; K=${K:-128}
LEAK=${LEAK:-0.05}; LEARN=${LEARN:-1}; LR=${LR:-2e-3}; SLR=${SLR:-$LR}; LADROP=${LADROP:-0.5}
EPISODE=${EPISODE:-8000}; FLUSH=${FLUSH:-64}; FIELD_LOSS=${FIELD_LOSS:-0.05}; D=${D:-32}; SEED=${SEED:-20260913}
OFFSET_SEED=${OFFSET_SEED:--1}; SAMPLES=${SAMPLES:-1}; RESTARTS=${RESTARTS:-3}; EXTRA_ARGS=${EXTRA_ARGS:-}
PY=${PY:-.venv/bin/python}
[ -x "$PY" ] || { echo "$(date -Is) $PY missing: run scripts/setup_gpu_box.sh"; exit 1; }
export PYTHONPATH=. OMP_NUM_THREADS=$THREADS MKL_NUM_THREADS=$THREADS
mkdir -p "$OUT"
exec 9>"$OUT/driver.lock"; flock -n 9 || { echo "$(date -Is) another driver holds $OUT/driver.lock"; exit 1; }
echo $$ > "$OUT/driver.pid"
LEARN_FLAG=--field-leak-learnable; [ "$LEARN" = "1" ] || LEARN_FLAG=--no-field-leak-learnable
SAMPLE_FLAG=--samples; [ "$SAMPLES" = "1" ] || SAMPLE_FLAG=--no-samples
VAL_ARGS=(); [ -n "$VAL_GLOB" ] && VAL_ARGS=(--val-glob "$VAL_GLOB")
INIT_ARGS=(); [ -n "$INIT" ] && INIT_ARGS=(--init "$INIT")
echo "=== $(date -Is) batched driver pid=$$ B=$B DEVICE=$DEVICE K=$K LR=$LR OUT=$OUT TOTAL=$TOTAL data=$DATA_GLOB ==="
[ "$DEVICE" = "cuda" ] && command -v nvidia-smi >/dev/null && nvidia-smi --query-gpu=name,memory.used,memory.total --format=csv,noheader
TRY=0
while :; do
"$PY" -u tools/train_batched.py --data-glob "$DATA_GLOB" "${VAL_ARGS[@]}" "${INIT_ARGS[@]}" --resume auto \
--batch-size "$B" --device "$DEVICE" --threads "$THREADS" --output-dir "$OUT" --suffix "$SUF" \
--total-ticks "$TOTAL" --ckpt-every "$CKPT_EVERY" --log-every "$LOG_EVERY" $SAMPLE_FLAG \
--d-model "$D" --n-modes "$D" --seed "$SEED" --offset-seed "$OFFSET_SEED" \
--learning-rate "$LR" --surface-learning-rate "$SLR" --field-bptt "$K" --field-leak "$LEAK" $LEARN_FLAG \
--last-atom-dropout "$LADROP" --episode-length "$EPISODE" --atom-flush-every "$FLUSH" \
--field-loss-weight "$FIELD_LOSS" $EXTRA_ARGS
RC=$?
if [ "$RC" -eq 0 ]; then echo "$(date -Is) DONE (rc=0)"; break; fi
TRY=$((TRY + 1)); echo "$(date -Is) trainer exited rc=$RC (restart $TRY/$RESTARTS, resumes from newest ckpt)"
[ "$TRY" -ge "$RESTARTS" ] && { echo "$(date -Is) giving up"; exit "$RC"; }
sleep 30
done
|