AR / save_final.sh

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 18 days ago

2.53 kB

#!/bin/bash
# Waits for a final checkpoint, then:
#  1) keeps it forever as a hard-link copy (no extra disk space; survives deletion of the original)
#  2) exports a weights-only bf16 file (~8 GB) for downloading to a PC, plus its md5 fingerprint
# Settings (environment): NAME = checkpoint folder, OUT = target folder, FILE = export file name,
#   MAXSTEP (optional) = warn and stop if training passes this step without the checkpoint.
# Example, final Stage 2 model (login node):
#   NAME=ckpt_619699 OUT=/fsx/ajay-stu/ar-data/stage2_final FILE=model_stage2_final.pt \
#     nohup ./save_final.sh > logs/save_stage2.log 2>&1 &
set -uo pipefail
AR=${AR:-/fsx/ajay-stu/AR}
CKDIR=${CKDIR:-/fsx/ajay-stu/ar-data/checkpoints}
: "${NAME:?set NAME, e.g. ckpt_619699}" "${OUT:?set OUT, the target folder}" "${FILE:?set FILE, e.g. model_stage2_final.pt}"
MAXSTEP=${MAXSTEP:-0}
EXC="ip-10-0-136-14,ip-10-0-136-94,ip-10-0-136-253,ip-10-0-138-178,ip-10-0-139-113,ip-10-0-140-247"
cd "$AR" || exit 1
[ -f .venv/bin/activate ] && source .venv/bin/activate
export PYTHONPATH=$PWD

# latest training step from the running job's log (empty if no job is running)
latest_step() {
  local j
  j=$(squeue -u "${USER:-$(id -un)}" -h -n ar-real -t RUNNING -o %i 2>/dev/null | head -1)
  [ -n "$j" ] && grep "^step" "logs/real_$j.log" 2>/dev/null | tail -1 | awk '{print $2}' || true
}

# 1) wait until the checkpoint is completely written (.metadata is written last)
echo "$(date -u) waiting for $CKDIR/$NAME"
until [ -f "$CKDIR/$NAME/.metadata" ]; do
  s=$(latest_step)
  if [ "$MAXSTEP" -gt 0 ] && [ "${s:-0}" -gt "$MAXSTEP" ]; then
    echo "$(date -u) WARNING: training is at step $s but $NAME was never saved (disk full?)"
    exit 1
  fi
  sleep "${POLL:-300}"
done
sleep "${SETTLE:-60}"                                     # let the last file writes settle
echo "$(date -u) found complete $NAME"

# 2) permanent hard-link copy
mkdir -p "$OUT"
[ -d "$OUT/$NAME" ] || cp -al "$CKDIR/$NAME" "$OUT/$NAME" || { echo "hard-link copy failed"; exit 1; }
echo "$(date -u) permanent copy: $OUT/$NAME ($(ls "$OUT/$NAME" | wc -l) files)"

# 3) weights-only bf16 export for the PC (CPU job, ~5 min) and its fingerprint
M=$OUT/$FILE
if [ ! -f "$M" ]; then
  srun -p dev -N1 -n1 --gpus=0 --cpus-per-task=16 --exclude=$EXC \
    python -m ar.eval.export_ckpt --ckpt "$OUT/$NAME" --out "$M" || { echo "export failed"; exit 1; }
fi
( cd "$OUT" && md5sum "$FILE" > "${FILE%.pt}.md5" )
ls -lh "$M"
cat "$OUT/${FILE%.pt}.md5"
echo "$(date -u) DONE: $FILE saved"

	#!/bin/bash
	# Waits for a final checkpoint, then:
	# 1) keeps it forever as a hard-link copy (no extra disk space; survives deletion of the original)
	# 2) exports a weights-only bf16 file (~8 GB) for downloading to a PC, plus its md5 fingerprint
	# Settings (environment): NAME = checkpoint folder, OUT = target folder, FILE = export file name,
	# MAXSTEP (optional) = warn and stop if training passes this step without the checkpoint.
	# Example, final Stage 2 model (login node):
	# NAME=ckpt_619699 OUT=/fsx/ajay-stu/ar-data/stage2_final FILE=model_stage2_final.pt \
	# nohup ./save_final.sh > logs/save_stage2.log 2>&1 &
	set -uo pipefail
	AR=${AR:-/fsx/ajay-stu/AR}
	CKDIR=${CKDIR:-/fsx/ajay-stu/ar-data/checkpoints}
	: "${NAME:?set NAME, e.g. ckpt_619699}" "${OUT:?set OUT, the target folder}" "${FILE:?set FILE, e.g. model_stage2_final.pt}"
	MAXSTEP=${MAXSTEP:-0}
	EXC="ip-10-0-136-14,ip-10-0-136-94,ip-10-0-136-253,ip-10-0-138-178,ip-10-0-139-113,ip-10-0-140-247"
	cd "$AR" \|\| exit 1
	[ -f .venv/bin/activate ] && source .venv/bin/activate
	export PYTHONPATH=$PWD

	# latest training step from the running job's log (empty if no job is running)
	latest_step() {
	local j
	j=$(squeue -u "${USER:-$(id -un)}" -h -n ar-real -t RUNNING -o %i 2>/dev/null \| head -1)
	[ -n "$j" ] && grep "^step" "logs/real_$j.log" 2>/dev/null \| tail -1 \| awk '{print $2}' \|\| true
	}

	# 1) wait until the checkpoint is completely written (.metadata is written last)
	echo "$(date -u) waiting for $CKDIR/$NAME"
	until [ -f "$CKDIR/$NAME/.metadata" ]; do
	s=$(latest_step)
	if [ "$MAXSTEP" -gt 0 ] && [ "${s:-0}" -gt "$MAXSTEP" ]; then
	echo "$(date -u) WARNING: training is at step $s but $NAME was never saved (disk full?)"
	exit 1
	fi
	sleep "${POLL:-300}"
	done
	sleep "${SETTLE:-60}" # let the last file writes settle
	echo "$(date -u) found complete $NAME"

	# 2) permanent hard-link copy
	mkdir -p "$OUT"
	[ -d "$OUT/$NAME" ] \|\| cp -al "$CKDIR/$NAME" "$OUT/$NAME" \|\| { echo "hard-link copy failed"; exit 1; }
	echo "$(date -u) permanent copy: $OUT/$NAME ($(ls "$OUT/$NAME" \| wc -l) files)"

	# 3) weights-only bf16 export for the PC (CPU job, ~5 min) and its fingerprint
	M=$OUT/$FILE
	if [ ! -f "$M" ]; then
	srun -p dev -N1 -n1 --gpus=0 --cpus-per-task=16 --exclude=$EXC \
	python -m ar.eval.export_ckpt --ckpt "$OUT/$NAME" --out "$M" \|\| { echo "export failed"; exit 1; }
	fi
	( cd "$OUT" && md5sum "$FILE" > "${FILE%.pt}.md5" )
	ls -lh "$M"
	cat "$OUT/${FILE%.pt}.md5"
	echo "$(date -u) DONE: $FILE saved"