multinode_trainer.slurm

#!/bin/bash
#SBATCH -D .
#SBATCH --ntasks=4
#SBATCH --nodes=4
#SBATCH --cpus-per-task=20
#SBATCH --time=00:30:00
#SBATCH --gres=gpu:4

#SBATCH --job-name=torchtitan_multi_node
#SBATCH --partition=acc
#SBATCH --mail-type=all
#SBATCH --mail-user=Tomas.Ockier@autonoma.cat

#SBATCH --output=/gpfs/scratch/ehpc391/titan_runs/outputs_4/run_5.out
#SBATCH --error=/gpfs/scratch/ehpc391/titan_runs/outputs_4/run_5.err

nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

echo Node IP: $head_node_ip
export LOGLEVEL=INFO
# Enable for A100
export FI_PROVIDER="efa"
# Ensure that P2P is available
# export NCCL_P2P_DISABLE=0
# export NCCL_IB_DISABLE=0

# debugging flags (optional)
export NCCL_DEBUG=WARN
export PYTHONFAULTHANDLER=1
# optional debug settings
# export NCCL_DEBUG=INFO
# NCCL_DEBUG_SUBSYS=INIT,GRAPH,ENV
#export LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH
#export LD_LIBRARY_PATH=/usr/local/lib/:$LD_LIBRARY_PATH
#export CUDA_LAUNCH_BLOCKING=0

# on your cluster you might need these:
# set the network interface
export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
export NCCL_BUFFSIZE=2097152
export TORCH_DIST_INIT_BARRIER=1
export FI_EFA_SET_CUDA_SYNC_MEMOPS=0

CONFIG_FILE=${CONFIG_FILE:-"./torchtitan/models/llama3/train_configs/llama3_1b_multinode.toml"}

# load env
source /gpfs/projects/ehpc391/envs/torch2_10/bin/activate

srun torchrun --nnodes 4 --nproc_per_node 4 --rdzv_id 101 --local-ranks-filter 0 --rdzv_backend c10d --rdzv_endpoint "$head_node_ip:29500" -m torchtitan.train --job.config_file ${CONFIG_FILE} "$@" --job.dump_folder /gpfs/scratch/ehpc391/outputs_4