#!/bin/bash
#SBATCH --job-name=ar-ckpt-test
#SBATCH --partition=dev
#SBATCH --nodes=1
#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=00:20:00
#SBATCH --output=logs/ckpt_test_%j.log
#SBATCH --error=logs/ckpt_test_%j.err

set -euo 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 MASTER_PORT=$((29500 + SLURM_JOB_ID % 1000))
export FI_PROVIDER=efa
export NCCL_SOCKET_IFNAME=enp73s0
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

: "${MAX_STEPS:=60}"

echo "MASTER_ADDR=$MASTER_ADDR MAX_STEPS=$MAX_STEPS port=$MASTER_PORT"

srun bash -c '
torchrun \
  --nnodes=1 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  -m ar.train.train_fsdp \
    --bin /fsx/ajay-stu/ar-data/bin-val/fineweb.bin \
    --model_size small \
    --batch_size 2 \
    --seq_len 1024 \
    --max_steps '"$MAX_STEPS"' \
    --warmup_steps 5 \
    --ckpt_every 20 \
    --ckpt_dir /fsx/ajay-stu/ar-data/checkpoints-test
'