AR / train_real.sbatch

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

2ae6691 verified 18 days ago

3.8 kB

#!/bin/bash
#SBATCH --job-name=ar-real
#SBATCH --partition=dev
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=96
#SBATCH --exclusive
#SBATCH --exclude=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
#SBATCH --time=30-00:00:00
#SBATCH --signal=B:USR1@4200
#SBATCH --output=logs/real_%j.log
#SBATCH --error=logs/real_%j.err

# Training runs in a loop inside this one job, so the nodes are kept:
#   logs/RELOAD -> stop after the next checkpoint, re-read run_args.sh, continue
#   logs/STOP   -> stop after the next checkpoint and end the job cleanly
#   ~70 min before a time limit Slurm sends USR1 -> we request STOP (nothing lost)
#   crash -> retry up to 3 times here before giving the nodes back
set -uo pipefail
cd /fsx/ajay-stu/AR
source .venv/bin/activate
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export FI_PROVIDER=efa
export NCCL_PROTO=simple
export NCCL_SOCKET_IFNAME=enp73s0
export NCCL_IB_DISABLE=0
export FI_EFA_USE_DEVICE_RDMA=1
export LD_LIBRARY_PATH=/fsx/ajay-stu/AR/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH:-}
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TOKENIZERS_PARALLELISM=false
# keep ~524K tokens per step for any node count
case $SLURM_NNODES in
  1) BS=4; GA=4 ;;
  2) BS=4; GA=2 ;;
  3|4) BS=4; GA=1 ;;
  5|6|7|8) BS=2; GA=1 ;;     # 8 nodes: 64 GPUs x 2 seqs x 4096 = 524K tokens
  *) BS=1; GA=1 ;;
esac
#       # 4 nodes -> 1 micro-batch, 2 nodes -> 2 (same step size)

# handover from a smaller job: wait (max 30 min) until our other training job has ended
for i in $(seq 90); do
  squeue -u $USER -h -n ar-real -t RUNNING -o %i | grep -vxq "$SLURM_JOB_ID" || break
  sleep 20
done
if squeue -u $USER -h -n ar-real -t RUNNING -o %i | grep -vxq "$SLURM_JOB_ID"; then
  echo "=== other training job still running, giving nodes back ==="; exit 4
fi

rm -f logs/STOP logs/RELOAD    # clear leftovers from an earlier job
trap 'echo "=== time limit near: stopping after the next checkpoint ==="; touch logs/STOP logs/RELOAD' USR1

round=0; fails=0
while true; do
  round=$((round + 1))
  source run_args.sh                                   # latest settings every round
  ./preflight.sh || { echo "PREFLIGHT FAILED: busy GPUs"; exit 3; }
  export MASTER_PORT=$((29500 + (SLURM_JOB_ID + round) % 1000))
  echo "=== round $round: nodes=$SLURM_NNODES stop_at=$STOP_AT lr_horizon=$LR_HORIZON ==="
  echo "=== free space: $(df -h /fsx | tail -1) ==="

  # run in the background so the time-limit warning can be handled mid-training
  srun bash -c "torchrun --nnodes=\$SLURM_NNODES --nproc_per_node=8 --node_rank=\$SLURM_NODEID \
    --rdzv_id=${SLURM_JOB_ID}_$round --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    -m ar.train.train_fsdp --bin $BIN --weights $WEIGHTS --model_size full \
    --batch_size $BS --seq_len 4096 --max_steps $STOP_AT --lr_horizon $LR_HORIZON \
    --warmup_steps 2000 --peak_lr $PEAK_LR --grad_accum $GA \
    --ckpt_every 2000 --ckpt_dir /fsx/ajay-stu/ar-data/checkpoints" &
  pid=$!
  while true; do
    wait $pid; rc=$?
    kill -0 $pid 2>/dev/null || break                  # training ended; rc is its exit code
  done

  if [ -f logs/STOP ]; then
    rm -f logs/STOP logs/RELOAD
    echo "=== stopped cleanly after a checkpoint ==="; exit 0
  fi
  if [ -f logs/RELOAD ]; then
    rm -f logs/RELOAD; fails=0
    echo "=== reload requested: restarting with new settings, same nodes ==="
    continue
  fi
  [ $rc -eq 0 ] && { echo "=== training finished ==="; break; }
  fails=$((fails + 1))
  echo "=== training crashed (rc=$rc), retry $fails/3 on the same nodes ==="
  [ $fails -ge 3 ] && exit $rc
  sleep 60
done

	#!/bin/bash
	#SBATCH --job-name=ar-real
	#SBATCH --partition=dev
	#SBATCH --nodes=4
	#SBATCH --ntasks-per-node=1
	#SBATCH --gres=gpu:8
	#SBATCH --cpus-per-task=96
	#SBATCH --exclusive
	#SBATCH --exclude=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
	#SBATCH --time=30-00:00:00
	#SBATCH --signal=B:USR1@4200
	#SBATCH --output=logs/real_%j.log
	#SBATCH --error=logs/real_%j.err

	# Training runs in a loop inside this one job, so the nodes are kept:
	# logs/RELOAD -> stop after the next checkpoint, re-read run_args.sh, continue
	# logs/STOP -> stop after the next checkpoint and end the job cleanly
	# ~70 min before a time limit Slurm sends USR1 -> we request STOP (nothing lost)
	# crash -> retry up to 3 times here before giving the nodes back
	set -uo pipefail
	cd /fsx/ajay-stu/AR
	source .venv/bin/activate
	export PYTHONPATH=$PWD:${PYTHONPATH:-}
	export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" \| head -n1)
	export FI_PROVIDER=efa
	export NCCL_PROTO=simple
	export NCCL_SOCKET_IFNAME=enp73s0
	export NCCL_IB_DISABLE=0
	export FI_EFA_USE_DEVICE_RDMA=1
	export LD_LIBRARY_PATH=/fsx/ajay-stu/AR/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH:-}
	export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
	export TOKENIZERS_PARALLELISM=false
	# keep ~524K tokens per step for any node count
	case $SLURM_NNODES in
	1) BS=4; GA=4 ;;
	2) BS=4; GA=2 ;;
	3\|4) BS=4; GA=1 ;;
	5\|6\|7\|8) BS=2; GA=1 ;; # 8 nodes: 64 GPUs x 2 seqs x 4096 = 524K tokens
	*) BS=1; GA=1 ;;
	esac
	# # 4 nodes -> 1 micro-batch, 2 nodes -> 2 (same step size)

	# handover from a smaller job: wait (max 30 min) until our other training job has ended
	for i in $(seq 90); do
	squeue -u $USER -h -n ar-real -t RUNNING -o %i \| grep -vxq "$SLURM_JOB_ID" \|\| break
	sleep 20
	done
	if squeue -u $USER -h -n ar-real -t RUNNING -o %i \| grep -vxq "$SLURM_JOB_ID"; then
	echo "=== other training job still running, giving nodes back ==="; exit 4
	fi

	rm -f logs/STOP logs/RELOAD # clear leftovers from an earlier job
	trap 'echo "=== time limit near: stopping after the next checkpoint ==="; touch logs/STOP logs/RELOAD' USR1

	round=0; fails=0
	while true; do
	round=$((round + 1))
	source run_args.sh # latest settings every round
	./preflight.sh \|\| { echo "PREFLIGHT FAILED: busy GPUs"; exit 3; }
	export MASTER_PORT=$((29500 + (SLURM_JOB_ID + round) % 1000))
	echo "=== round $round: nodes=$SLURM_NNODES stop_at=$STOP_AT lr_horizon=$LR_HORIZON ==="
	echo "=== free space: $(df -h /fsx \| tail -1) ==="

	# run in the background so the time-limit warning can be handled mid-training
	srun bash -c "torchrun --nnodes=\$SLURM_NNODES --nproc_per_node=8 --node_rank=\$SLURM_NODEID \
	--rdzv_id=${SLURM_JOB_ID}_$round --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
	-m ar.train.train_fsdp --bin $BIN --weights $WEIGHTS --model_size full \
	--batch_size $BS --seq_len 4096 --max_steps $STOP_AT --lr_horizon $LR_HORIZON \
	--warmup_steps 2000 --peak_lr $PEAK_LR --grad_accum $GA \
	--ckpt_every 2000 --ckpt_dir /fsx/ajay-stu/ar-data/checkpoints" &
	pid=$!
	while true; do
	wait $pid; rc=$?
	kill -0 $pid 2>/dev/null \|\| break # training ended; rc is its exit code
	done

	if [ -f logs/STOP ]; then
	rm -f logs/STOP logs/RELOAD
	echo "=== stopped cleanly after a checkpoint ==="; exit 0
	fi
	if [ -f logs/RELOAD ]; then
	rm -f logs/RELOAD; fails=0
	echo "=== reload requested: restarting with new settings, same nodes ==="
	continue
	fi
	[ $rc -eq 0 ] && { echo "=== training finished ==="; break; }
	fails=$((fails + 1))
	echo "=== training crashed (rc=$rc), retry $fails/3 on the same nodes ==="
	[ $fails -ge 3 ] && exit $rc
	sleep 60
	done