File size: 1,307 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | #!/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
'
|