GPU jobs and multi-node training

One compute node, one GPU

#!/bin/bash
#SBATCH --gres=gpu:1 -c 8 --mem=64G --time=30 -J train
source ~/venv/bin/activate
python train.py

Refer to Your first job for the steps.

What each process must know

Each process needs its rank, the number of processes and its place on its compute node. Slurm gives these as SLURM_PROCID, SLURM_NTASKS and SLURM_LOCALID. PyTorch reads RANK, WORLD_SIZE and LOCAL_RANK. The recipes below connect the two.

Start each job script with these lines:

source ~/venv/bin/activate
export MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -1) MASTER_PORT=29501

The recipes show 2 compute nodes with 2 GPUs each, so the world size is 4. The srun and torchrun forms also work with -N1. To use more GPUs, change --gres=gpu:2 and the same number in the recipe.

Training in a venv uses InfiniBand with no extra settings.

Plain PyTorch, HF Trainer or DeepSpeed

This form starts one process per GPU. Save this text as train.sbatch:

#!/bin/bash
#SBATCH -N2 --ntasks-per-node=2 --gres=gpu:2 -c8 --mem-per-gpu=48G --time=60
source ~/venv/bin/activate
export MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -1) MASTER_PORT=29501
srun --kill-on-bad-exit=1 bash -c 'RANK=$SLURM_PROCID WORLD_SIZE=$SLURM_NTASKS LOCAL_RANK=$SLURM_LOCALID exec python train.py'

Use this form for DeepSpeed too. DeepSpeed's own launchers need extra tools (pdsh, mpi4py) and fail to train across 2 compute nodes. DeepSpeed also needs the managed Python from Software environments.

torchrun

This form starts one launcher per compute node. The launcher starts one process per GPU.

#!/bin/bash
#SBATCH -N2 --ntasks-per-node=1 --gres=gpu:2 -c16 --mem-per-gpu=48G --time=60
source ~/venv/bin/activate
export MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -1) MASTER_PORT=29501
srun --kill-on-bad-exit=1 bash -c 'torchrun --nnodes=$SLURM_NNODES --nproc_per_node=2 --node_rank=$SLURM_NODEID --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT train.py'

Set --nproc_per_node to the same number as --gres=gpu:N. This recipe gives the job 8 CPUs per GPU.

accelerate

This form also starts one launcher per compute node.

#!/bin/bash
#SBATCH -N2 --ntasks-per-node=1 --gres=gpu:2 -c16 --mem-per-gpu=48G --time=60
source ~/venv/bin/activate
export MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -1) MASTER_PORT=29501
srun --kill-on-bad-exit=1 bash -c 'accelerate launch --multi_gpu --num_machines=$SLURM_NNODES --machine_rank=$SLURM_NODEID --num_processes=$((SLURM_NNODES*2)) --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT train.py'

In --num_processes, the 2 is the number of GPUs per compute node.

In a container

In a PyTorch image, the cluster's container setup sets RANK, WORLD_SIZE and LOCAL_RANK, so plain srun is enough:

srun --kill-on-bad-exit=1 --container-image=$HOME/ptib.sqsh --container-workdir=$PWD python train.py

You build ptib.sqsh in Containers. The stock PyTorch image does not use InfiniBand.

Check that it works

Do these checks on your first run:

1. Add this line to the job script, before srun.

   export NCCL_DEBUG=INFO

2. Submit the job.

3. Search the job's output for the network that NCCL uses.

   grep -E "NET/IB|Using network" slurm-<jobid>.out

4. Confirm that the output has Using network IB.

5. Confirm that your program reports the correct world size. The size is 4 in the recipes above, if your script prints it.

6. Confirm that both compute nodes show the job's GPUs.

   ssh node-1 nvidia-smi -L
   ssh node-2 nvidia-smi -L

The SSH session joins your job, so it shows your job's GPUs.

Log line	Meaning
Using network IB	NCCL uses InfiniBand.
Failed to open libibverbs.so, then Using network Socket	NCCL uses Ethernet. This happens in a container with no InfiniBand library.

Results that show a wrong form

These forms give no clear error, or an error that points to the wrong cause. They are normal Slurm and PyTorch behaviour.

Form	What happens
srun python train.py with a script that does not read the variables	Each process starts its own training, with no error. With HF Trainer on 2 compute nodes with 2 GPUs each, you get 4 separate trainings.
The same form with a plain PyTorch script, outside a container	The job stops with environment variable RANK expected, but not set.
torchrun without srun in a -N2 job	The training uses the first compute node only. The other compute node stays idle and you pay for it.
accelerate launch without --num_machines and --machine_rank	Each compute node trains alone.
--gpus-per-task or --ntasks-per-gpu with 2 ranks on a compute node	The job crashes with ncclUnhandledCudaError. Use --gres=gpu:N.
--signal without B:	The signal reaches every step, so a rank can die from it.

Things to do on each job

• Add --kill-on-bad-exit=1 to srun. Without it, a crashed process leaves the job RUNNING and the GPUs busy until the time limit. Refer to Job options.
• Keep srun as the last line of the script.
• Set --time.
• Save checkpoints in your home folder. Refer to Node failures and checkpoints.
• Add ulimit -n 65536 if you use many DataLoader workers.

First run with an empty cache

A first run of Hugging Face with an empty cache fails on 2 compute nodes at the same time, with FileNotFoundError or DatasetGenerationError. Fill the cache first with a 1-node job. After that, 2-node jobs work. Refer to Storage.

A public model and dataset from Hugging Face download with no token.

Speed that you can expect

Training scales well over InfiniBand. 4 GPUs across 2 compute nodes can run at about the same speed as 4 GPUs on one compute node. On Ethernet, the same job runs much slower. Your speed depends on how much your job sends per step.

Data that a compute node has read before reads faster. So the first epoch is slower than later epochs. Refer to Storage.