| #!/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 | |
| ' | |