#!/bin/bash
#SBATCH --job-name=ar-sanity-4node
#SBATCH --partition=dev
#SBATCH --nodes=4
#SBATCH --nodelist=ip-10-0-154-148,ip-10-0-184-219,ip-10-0-233-76,ip-10-0-245-227
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=96
#SBATCH --exclusive
#SBATCH --output=sanity_4node_%j.log

set -euo pipefail

cd /fsx/ajay-stu/AR
source .venv/bin/activate

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=29517

export FI_PROVIDER=efa
export NCCL_PROTO=simple
export NCCL_DEBUG=INFO
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 LD_LIBRARY_PATH=/fsx/ajay-stu/AR/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH

echo "MASTER_ADDR=$MASTER_ADDR NODES=$SLURM_NNODES job=$SLURM_JOB_ID"

srun bash -c '
torchrun \
  --nnodes=4 \
  --nproc_per_node=8 \
  --node_rank=$SLURM_NODEID \
  --rdzv_id=$SLURM_JOB_ID \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  -m ar.train.train_fsdp \
    --bin /fsx/ajay-stu/ar-data/bin/fineweb.bin \
    --model_size full \
    --max_steps 20 \
    --batch_size 1 \
    --seq_len 2048 \
    --warmup_steps 5
'