File size: 1,625 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 | #!/bin/bash
# Stage 2: smoke test -> full LoRA instruction tuning -> evaluation. Safe to re-run (resumes).
set -uo pipefail
cd "$HOME/atlas"
export PYTHONUNBUFFERED=1 HF_HOME="$HOME/.cache/huggingface" HF_HUB_OFFLINE=1 PATH="$HOME/.local/bin:$PATH"
python3 -m pip install -q --user --break-system-packages peft pyarrow || exit 10
python3 -c "import peft, pyarrow; print('peft', peft.__version__, '| pyarrow', pyarrow.__version__)"
BS="${BATCH_SIZE:-8}"; GA="${GRAD_ACCUM:-4}"
if [ "${SKIP_SMOKE:-0}" != "1" ]; then
echo "=== [1/3] smoke test (10 steps) ==="
rm -rf "$HOME/checkpoints/stage2_smoke"
MAX_STEPS=10 LOG_EVERY=2 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/stage2_smoke" BATCH_SIZE=$BS GRAD_ACCUM=$GA python3 train_stage2.py
rc=$?
if [ $rc -eq 3 ]; then
echo "OOM -> retrying with BATCH_SIZE=4 GRAD_ACCUM=8"; BS=4; GA=8
rm -rf "$HOME/checkpoints/stage2_smoke"
MAX_STEPS=10 LOG_EVERY=2 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/stage2_smoke" BATCH_SIZE=$BS GRAD_ACCUM=$GA python3 train_stage2.py
rc=$?
fi
if [ $rc -ne 0 ]; then echo "SMOKE TEST FAILED (exit $rc)"; exit 12; fi
fi
echo "=== [2/3] full stage-2 run (BATCH_SIZE=$BS GRAD_ACCUM=$GA) ==="
BATCH_SIZE=$BS GRAD_ACCUM=$GA python3 train_stage2.py
rc=$?
if [ $rc -eq 3 ] && [ "$BS" -gt 4 ]; then
echo "OOM in full run -> continuing with BATCH_SIZE=4 GRAD_ACCUM=8"
BATCH_SIZE=4 GRAD_ACCUM=8 python3 train_stage2.py
rc=$?
fi
if [ $rc -ne 0 ]; then echo "TRAINING FAILED (exit $rc)"; exit $rc; fi
echo "=== [3/3] evaluation ==="
python3 eval_stage2.py
rc=$?
echo "=== run_stage2.sh finished (exit $rc) ==="
exit $rc
|