ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | # 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 | |