slurm_torchrun

#!/usr/bin/env bash

set -euo pipefail

torchrun \
    --nproc-per-node=$SLURM_GPUS_ON_NODE\
    --nnodes=$SLURM_JOB_NUM_NODES \
    --node_rank=$SLURM_NODEID \
    --master_addr=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n1) \
    --master_port=$(expr 10000 + $(echo -n $SLURM_JOBID | tail -c 4)) \
    "$@"