Skip to content

Run multi-node training

View as Markdown

Run one training job across several nodes of your cluster. A single srun starts your container on every node of the allocation, and the nodes talk to each other over the cluster’s RDMA fabric. Getting the ranks to communicate efficiently is the part that needs care: shared storage, PMIx, and NCCL.

  • Shared storage mounted on the cluster. The image, the code, the dataset, and the checkpoints must be identical and reachable from every rank. A Parallel File System or Shared File System appears at the same path on the login node and every worker. These commands use /shared as the mount path; use the one you set.
  • Your container image imported once as a squash file on shared storage, so the nodes read one copy instead of each pulling it and hitting registry rate limits mid-run. See Cache container images.
  • --mpi=pmix on your srun command, which wires up the ranks across nodes inside containers.
  • Healthy nodes in one partition. A drained or unregistered node shrinks your allocation silently. Check sinfo and the Nodes tab. See Check cluster and node health.

Multi-node always means shared storage. Anything under /tmp, /dev/shm, or another node-local path belongs to one node. A container name exists only on the node that created it, and a checkpoint written to /tmp on rank 0 stays on rank 0’s node.

For a container job, srun starts one container per task, and Slurm’s environment tells each rank who it is. How you launch depends on the tool that starts your processes.

For torchrun, DeepSpeed, or any launcher that starts its own per-GPU processes, use one container per node and let the launcher fan out inside it:

#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1 # one container per node
#SBATCH --gpus-per-node=8 # which owns all 8 GPUs on it

For MPI ranks, use one task per GPU and let PMIx wire them up:

#SBATCH --nodes=4
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-task=1
srun --mpi=pmix --container-image=/shared/images/hpc.sqsh ./my_mpi_app

Every rank gets these variables from Slurm. Use them to build output paths and to work out each rank’s place in the job:

SLURM_JOB_ID The job's ID. Use it in output paths.
SLURM_JOB_NUM_NODES The number of nodes in the allocation.
SLURM_NODEID This node's index, starting at 0.
SLURM_PROCID This task's global rank.
SLURM_LOCALID This task's rank within its node.
SLURM_NTASKS The total number of tasks, which is the world size with one task per rank.
SLURM_JOB_NODELIST The compact node list, such as slinky-[0-3].

Derive the rendezvous host from the job’s node list instead of hard-coding an address, because Slurm picks the allocation’s nodes when the job starts:

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

The platform leaves NCCL and UCX tuning to you, and the defaults come from your image, so set the variables yourself on a cluster with an RDMA fabric. Mount /dev/infiniband into the container with --container-mounts too, because a container gets device nodes only when you mount them.

The InfiniBand device names differ by hardware, and they must be identical on every participating node. Wrong names are the most common cause of a hang. List them on one node:

srun --nodes=1 --container-image=/shared/images/hpc.sqsh \
--container-mounts=/dev/infiniband:/dev/infiniband \
bash -lc 'ls /sys/class/infiniband; ibdev2netdev 2>/dev/null'

Then confirm that more than one node reports the same names:

srun --nodes=4 --ntasks-per-node=1 bash -lc 'echo "$(hostname): $(ls /sys/class/infiniband | tr "\n" " ")"'

Start from this block and replace the device lists with what the commands above reported. Each line says what its variable controls:

export NCCL_DEBUG=WARN # how much NCCL prints: INFO shows the transport it chose, WARN suits production
export NCCL_IB_DISABLE=0 # 0 allows InfiniBand, 1 forces TCP: use 1 only to prove a fabric problem, never for real training
export NCCL_NET=IB
export NCCL_IB_HCA=<device-list> # which InfiniBand devices and ports NCCL may use
export UCX_NET_DEVICES=<device-list> # the same selection for UCX-based transports such as OpenMPI
export UCX_TLS=rc,cuda
export NCCL_SOCKET_IFNAME=eth0 # the interface NCCL uses to start up; it must exist on every node

Replace <device-list> with the device names and ports from the commands above, separated by commas.

To confirm that InfiniBand is in use, run once with the debug line set to INFO and look for a line naming [send] via NET/IB/.... A line naming NET/Socket means NCCL fell back to TCP, and your throughput will be far lower. Check the device names, and check that /dev/infiniband is mounted into the container. Every rank writes its topology decisions at INFO, so on a 32-rank job it adds megabytes of log per run. Use it to diagnose, then return to WARN.

This job uses 32 GPUs: one container per node, with torchrun fanning out to eight local ranks.

#!/bin/bash
#SBATCH --job-name=llm-pretrain
#SBATCH --partition=all
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=180
#SBATCH --exclusive
#SBATCH --time=24:00:00
#SBATCH --output=/shared/logs/%x-%j.out
#SBATCH --error=/shared/logs/%x-%j.err
set -euo pipefail
IMAGE=/shared/images/pytorch-25.09-py3.sqsh
WORKDIR=/shared/project
CKPT=/shared/checkpoints/${SLURM_JOB_NAME}-${SLURM_JOB_ID}
mkdir -p "$CKPT" /shared/logs
# Rendezvous
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=29500
# Fabric
export NCCL_DEBUG=WARN
export NCCL_IB_DISABLE=0
export NCCL_NET=IB
export NCCL_SOCKET_IFNAME=eth0
echo "Job $SLURM_JOB_ID on $SLURM_JOB_NUM_NODES nodes: $SLURM_JOB_NODELIST"
echo "Rendezvous at $MASTER_ADDR:$MASTER_PORT"
srun --mpi=pmix \
--container-image="$IMAGE" \
--container-name=pretrain \
--container-mounts=/shared:/shared,/dev/infiniband:/dev/infiniband \
--container-workdir="$WORKDIR" \
bash -lc '
torchrun \
--nnodes=$SLURM_JOB_NUM_NODES \
--node_rank=$SLURM_NODEID \
--nproc_per_node=8 \
--master_addr=$MASTER_ADDR \
--master_port=$MASTER_PORT \
train.py \
--data /shared/datasets/corpus \
--checkpoint-dir '"$CKPT"' \
--save-every 500
'

Submit the job and follow it:

sbatch llm-pretrain.sh
squeue -u $USER
tail -f /shared/logs/llm-pretrain-<jobid>.out

Replace <jobid> with the number sbatch printed.

--exclusive gives the job whole nodes, so it has the CPU, memory bandwidth, and network fabric to itself. For large distributed training this is usually worth it for consistent step times, because a neighbour saturating the network card shows up as unexplained variation in step time.

Checkpoint often on a long run. A node failure or a preemption ends the job, and your progress is only as good as your last checkpoint on shared storage. The cluster recovers by itself, but the job ends; see What high availability covers. --save-every 500 in the example is there for that reason.

Before you commit to a 24-hour run, check the interconnect at the scale you plan to use with an NCCL all-reduce benchmark:

#!/bin/bash
#SBATCH --job-name=nccl-check
#SBATCH --partition=all
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=20
#SBATCH --exclusive
#SBATCH --time=00:30:00
#SBATCH --output=/shared/logs/nccl-%j.out
IMAGE=nvcr.io/nvidia/hpc-benchmarks:25.09
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_NET=IB
export NCCL_SOCKET_IFNAME=eth0
srun --mpi=pmix \
--container-image="$IMAGE" \
--container-name=nccl-bench \
--container-mounts=/dev/infiniband:/dev/infiniband,/shared:/shared \
--container-workdir=/workspace \
bash -lc 'all_reduce_perf_mpi -b 8 -e 8G -f 2 -g 1'

Read the busbw column at the largest message size and compare it with the bandwidth you expect for your hardware. A result far below expectation, or a run that hangs at initialisation, points at the InfiniBand device names, a missing /dev/infiniband mount, or one unhealthy node. Check the Nodes tab for XID errors.

Also prove that storage is shared across every node, since a surprising number of multi-node failures are a mount that differs from what you assumed:

echo "written from $(hostname) at $(date)" > /shared/shared-check.txt
srun --nodes=4 --ntasks-per-node=1 bash -lc 'echo "$(hostname): $(cat /shared/shared-check.txt)"'

Every node must print the same line. If one node fails to read the file, the volume is missing on that node. Check the Volumes tab; see Mount and unmount volumes.

To run a different image on each job or cache images on shared storage, see Run containers in jobs. If a multi-node job hangs at NCCL initialisation, see Troubleshooting containers on TIR Slurm Cluster.