File size: 1,515 Bytes
275cf6f | 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 | #!/usr/bin/env bash
# Stop the AdamW control arm cleanly once step 400 is checkpointed.
set -uo pipefail; cd "$(dirname "$0")"
say(){ echo "[$(date -u +%H:%M:%S)] $*"; }
LOG=p1_const_adamw_logs/train.log
CK=checkpoints_const_adamw/rank_0/blumodelcp_run2_step_400.ckpt
SCRIPT_PID=30792
say "waiting for step 400 checkpoint to be written safely..."
for i in $(seq 1 360); do
grep -qa 'Saved: .*step_400\.ckpt (safely)' "$LOG" && { say "step 400 checkpoint confirmed safe"; break; }
kill -0 "$SCRIPT_PID" 2>/dev/null || { say "chain already exited"; break; }
sleep 10
done
grep -qa 'Saved: .*step_400\.ckpt (safely)' "$LOG" || { say "ABORT: no safe step-400 ckpt, leaving run alone"; exit 1; }
sleep 5 # let the .zero_g1 shard flush
say "sizes: $(stat -Lc%s "$CK" 2>/dev/null) + $(stat -Lc%s "${CK%.ckpt}.zero_g1" 2>/dev/null)"
# parent script first so it cannot advance to the next stage, then the trainer
kill "$SCRIPT_PID" 2>/dev/null && say "stopped abl_phase1.sh ($SCRIPT_PID)"
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do
kill "$p" 2>/dev/null && say "TERM trainer $p"
done
for i in $(seq 1 30); do
[ -z "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader)" ] && break; sleep 2; done
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do
kill -9 "$p" 2>/dev/null && say "KILL -9 $p"; done
sleep 3
say "GPU now: $(nvidia-smi --query-gpu=memory.used --format=csv,noheader)"
say "last step: $(grep '^step' "$LOG" | tail -1)"
say "CUT_DONE"
|