File size: 1,912 Bytes
2a2540a | 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 | #!/bin/bash
# Idempotent: safe to re-run after a stop/wake. Resumes training from ~/checkpoints/stage1/latest.pt.
set -uo pipefail
cd "$HOME/atlas"
export PYTHONUNBUFFERED=1
export HF_HOME="$HOME/.cache/huggingface"
export HF_HUB_ENABLE_HF_TRANSFER=1
echo "=== [1/4] python packages ==="
PIPFLAGS="--user --break-system-packages"
python3 -c "import sys; sys.exit(0 if sys.prefix != sys.base_prefix else 1)" && PIPFLAGS=""
python3 -m pip install -q $PIPFLAGS -U "transformers>=4.56" accelerate "huggingface_hub>=0.34" hf_transfer pillow sentencepiece protobuf || exit 10
export PATH="$HOME/.local/bin:$PATH"
python3 -c "import torch, transformers; print('torch', torch.__version__, '| transformers', transformers.__version__, '| cuda', torch.cuda.is_available())"
nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv
echo "=== [2/4] downloads ==="
python3 prepare.py || exit 11
df -h "$HOME" | tail -1
export HF_HUB_OFFLINE=1 # everything is cached now; training never needs the token
GC="${GRAD_CKPT:-0}"
if [ "${SKIP_SMOKE:-0}" != "1" ]; then
echo "=== [3/4] smoke test (20 steps) ==="
rm -rf "$HOME/checkpoints/smoke"
MAX_STEPS=20 LOG_EVERY=5 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/smoke" GRAD_CKPT=$GC python3 train.py
rc=$?
if [ $rc -eq 3 ] && [ "$GC" = "0" ]; then
echo "OOM without gradient checkpointing -> retrying with GRAD_CKPT=1"
GC=1
rm -rf "$HOME/checkpoints/smoke"
MAX_STEPS=20 LOG_EVERY=5 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/smoke" GRAD_CKPT=1 python3 train.py
rc=$?
fi
if [ $rc -ne 0 ]; then echo "SMOKE TEST FAILED (exit $rc)"; exit 12; fi
fi
echo "=== [4/4] full stage-1 run (GRAD_CKPT=$GC) ==="
GRAD_CKPT=$GC python3 train.py
rc=$?
if [ $rc -eq 3 ] && [ "$GC" = "0" ]; then
echo "OOM in full run -> resuming with GRAD_CKPT=1"
GRAD_CKPT=1 python3 train.py
rc=$?
fi
echo "=== run.sh finished (exit $rc) ==="
exit $rc
|