Multi-GPU and Multi-Node Training¶

This recipe builds on the PyTorch recipe. If you haven't set up a GPU-enabled PyTorch project yet, start there.

Do You Actually Need This?¶

Multi-GPU training adds real complexity: harder debugging, inter-GPU communication overhead, trickier batch size math, and Slurm scripts that are easy to get wrong. Before scaling out, ask:

• Single GPU is enough? Many medium-sized models train perfectly well on one A100 (80 GB). If your model and a reasonable batch fit in VRAM, start there.
• Gradient accumulation first. If training is slow because you can't use a large batch, try gradient accumulation (loss / accumulation_steps per micro-step) before adding GPUs. It's free and keeps your code simple.
• Profile before scaling. A data pipeline bottleneck or inefficient attention implementation won't get faster with more GPUs.

Multi-GPU is worth it when:

Data Parallel vs. Model Parallel¶

Data Parallelism (DDP)¶

The full model is replicated on each GPU. Each GPU processes a different shard of the batch. After the backward pass, gradients are averaged across all GPUs (via NCCL all-reduce) before the optimizer step.

• Best for: models that fit on one GPU but training is slow.
• Scaling efficiency is high: communication is a single gradient sync per step.
• PyTorch's DistributedDataParallel is the standard implementation.

Model Parallelism (Tensor / Pipeline)¶

The model itself is split across GPUs. Each GPU holds only a subset of the parameters.

• Required when: model weights alone exceed one GPU's VRAM.
• More complex to implement; pipeline bubbles and load imbalance are real costs.
• FSDP and DeepSpeed ZeRO automate this; see FSDP and DeepSpeed below.

PyTorch DDP¶

DistributedDataParallel is the standard approach for data-parallel training. It works within a single node (multi-GPU) or across multiple nodes.

Key Code Pattern¶

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def main():
    # torchrun sets these environment variables automatically
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # Wrap model in DDP
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # DistributedSampler ensures each GPU sees a unique data shard
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # shuffle differently each epoch
        for batch in loader:
            ...

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

The key ingredients are:

1. dist.init_process_group(backend="nccl") — initializes the communication backend. NCCL is the right choice for GPU-to-GPU.
2. DDP(model, device_ids=[local_rank]) — wraps the model so gradients are automatically synchronized.
3. DistributedSampler — partitions the dataset across workers so no two GPUs see the same samples. Without this, all GPUs train on identical data.

Launching with torchrun¶

torchrun is the modern launcher (replaces the deprecated torch.distributed.launch). It sets LOCAL_RANK, RANK, WORLD_SIZE, and MASTER_ADDR automatically.

torchrun --nproc_per_node=$SLURM_GPUS_ON_NODE \
         --nnodes=$SLURM_NNODES \
         --node_rank=$SLURM_NODEID \
         --master_addr=$(scontrol show hostname $SLURM_NODELIST | head -n1) \
         --master_port=29500 \
         train.py

Single-Node Multi-GPU Slurm Script¶

#!/bin/bash
#SBATCH --job-name=ddp_train
#SBATCH --partition=gpu
#SBATCH --time=04:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
#SBATCH --gres=gpu:4

# Run your work
torchrun --nproc_per_node=$SLURM_GPUS_ON_NODE train.py

Multi-Node Slurm Script¶

For multi-node jobs, request one task per node (the task is torchrun, which then spawns one process per GPU on that node):

#!/bin/bash
#SBATCH --job-name=ddp_multinode
#SBATCH --partition=gpu
#SBATCH --time=08:00:00
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
#SBATCH --gres=gpu:4

# Run your work
torchrun \
  --nproc_per_node=$SLURM_GPUS_ON_NODE \
  --nnodes=$SLURM_NNODES \
  --node_rank=$SLURM_NODEID \
  --master_addr=$(scontrol show hostname $SLURM_NODELIST | head -n1) \
  --master_port=29500 \
  train.py

HuggingFace Accelerate¶

Accelerate is a thin abstraction layer that handles DDP, FSDP, and DeepSpeed configuration automatically. It requires only minimal changes to single-GPU training code and is the recommended approach for most users who don't need low-level DDP control.

Minimal Code Change¶

model = model.to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for batch in dataloader:
    inputs = {k: v.to("cuda") for k, v in batch.items()}
    outputs = model(**inputs)
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

from accelerate import Accelerator

accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(
    model, optimizer, dataloader
)

for batch in dataloader:
    outputs = model(**batch)  # batch already on the right device
    loss = outputs.loss
    accelerator.backward(loss)  # replaces loss.backward()
    optimizer.step()
    optimizer.zero_grad()

accelerator.prepare() handles device placement, DistributedSampler wrapping, and gradient synchronization. You don't call dist.init_process_group or DDP() manually.

Slurm Script with accelerate launch¶

#!/bin/bash
#SBATCH --job-name=accelerate_train
#SBATCH --partition=gpu
#SBATCH --time=04:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
#SBATCH --gres=gpu:4

# Run your work
accelerate launch --num_processes $SLURM_GPUS_ON_NODE train.py

For multi-node, add --num_machines $SLURM_NNODES --machine_rank $SLURM_NODEID --main_process_ip $(scontrol show hostname $SLURM_NODELIST | head -n1) to the accelerate launch command.

FSDP and DeepSpeed — When You Need Them¶

When your model is too large to fit in a single GPU's VRAM even during training (weights + gradients + optimizer states can be 4–8× the weights-only size), you need to shard across GPUs.

FSDP (Fully Sharded Data Parallel) — Built into PyTorch. Shards model parameters, gradients, and optimizer states across all GPUs. Each GPU holds only a fraction of the model at any time; parameters are gathered on-demand for forward/backward passes.

DeepSpeed ZeRO — Microsoft's library with similar sharding (ZeRO-1/2/3) plus additional optimizations like CPU offloading and mixed precision. ZeRO-3 approaches FSDP in capability; DeepSpeed also provides kernel fusions that can improve throughput.

Both are supported as Accelerate backends. You configure them via an accelerate config YAML rather than changing training code:

# accelerate_config.yaml (DeepSpeed ZeRO-3 example)
compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
deepspeed_config:
  zero_stage: 3
  offload_optimizer_device: none
  offload_param_device: none
num_processes: 4

accelerate launch --config_file accelerate_config.yaml train.py

For a full walkthrough, see the Accelerate FSDP docs and DeepSpeed integration docs.

Checking GPU Utilization Across Nodes¶

After submitting a job, verify that all GPUs are actually being used:

# From inside a running job (add to your script for logging)
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 nvidia-smi

# After the job completes
sacct -j <JOBID> --format=JobID,Elapsed,NCPUS,NNodes,MaxRSS

The most common silent failure in DDP is that only GPU 0 is doing any work: nvidia-smi shows GPU 0 at ~100% utilization and all others at ~0%. This means dist.init_process_group() was never called, or your model/data aren't being moved to the correct local_rank device.

Enable NCCL debug output to diagnose communication issues:

export NCCL_DEBUG=INFO
torchrun ... train.py

Common Pitfalls¶

See Also¶

• GPU Computing Fundamentals — GPU architecture, VRAM sizing, and choosing the right GPU partition
• PyTorch Recipe — single-GPU setup and CUDA index selection
• Transformers Recipe — device_map="auto" for multi-GPU inference with HuggingFace models
• Job Arrays — embarrassingly parallel hyperparameter sweeps without DDP complexity