Distributing Training

Multi-GPU Training with TRL

The trainers in TRL are launched with torchrun, PyTorch’s distributed launcher. It starts one process per GPU:

torchrun --nproc_per_node 8 train.py <SCRIPT_ARGS>

The trl CLI does the same under the hood, so trl sft ... runs on every GPU of the machine; pass --nproc_per_node to use fewer.

Under the hood, 🤗 Accelerate creates one model per GPU. Each process:

• Processes its own batch of data
• Computes the loss and gradients for that batch
• Shares gradient updates across all GPUs

The effective batch size is calculated as: $$\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$$

To maintain a consistent batch size when scaling to multiple GPUs, make sure to update per_device_train_batch_size and gradient_accumulation_steps accordingly.

Example, these configurations are equivalent, and should yield the same results:

Having one model per GPU can lead to high memory usage, which may not be feasible for large models or low-memory GPUs. In such cases, you can leverage DeepSpeed, which provides optimizations like model sharding, Zero Redundancy Optimizer, mixed precision training, and offloading to CPU or NVMe. Check out our DeepSpeed Integration guide for more details.

Training on very long sequences has its own guide: Training Beyond 1M Tokens.

Multi-Node Training

When a single machine doesn’t have enough GPUs, torchrun can scale training across multiple machines (nodes). Every node runs the same command with the total number of nodes, its own rank, and a rendezvous endpoint on the main node.

Option 1: Manual Launch (Non-HPC)

Run the following on each node manually:

# Node 0 (main node, IP 10.0.0.1)
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 0 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

# Node 1
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 1 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

Replace 10.0.0.1 with the actual IP address of the rank 0 (main) node.

Option 2: SLURM Launch (HPC Clusters)

For clusters using SLURM job scheduler, create a job script (e.g., slurm_job.sh):

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --job-name=trl_multi

MAIN_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
srun torchrun --nnodes 2 --nproc_per_node 8 --node_rank $SLURM_NODEID --rdzv_backend c10d --rdzv_endpoint $MAIN_NODE:29500 train.py

Then submit the job:

sbatch slurm_job.sh

srun starts one torchrun per node, and each of them starts one process per GPU.

Key SLURM directives:

• --nodes=2: Request 2 compute nodes
• --gpus-per-node=8: Allocate 8 GPUs per node (16 total)
• --job-name: Label for tracking in the job queue

You can combine multi-node with DeepSpeed by setting deepspeed in the training config. See the DeepSpeed integration guide.

Further Reading

• torchrun (Elastic Launch)
• SLURM Workload Manager Documentation - For cluster job scheduling