Data Parallelism › DDP in Practice: PyTorch Implementation
Language AI HandbookAll chapters Data Parallelism › DDP in Practice: PyTorch ImplementationOne section · Sections in this chapter5 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us walk through a complete DDP training setup using PyTorch. Modern PyTorch DDP requires specific components to work correctly, and the setup is slightly more involved than single-GPU training.
Launching with torchrun
Before looking at the code, it is worth understanding how DDP processes are launched. DDP requires one process per GPU, and each process must know its rank (which GPU it is), the total number of processes (world size), and how to find the other processes (via a rendezvous point, typically a hostname and port).
PyTorch provides torchrun (replacing the older torch.distributed.launch) as the recommended launcher. You invoke it as:
# torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
# --master_addr=node0 --master_port=29500 train.pytorchrun starts the specified number of processes on each node, sets the environment variables RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, and MASTER_PORT, and waits for all processes to join the rendezvous before any of them proceed. Your training script reads these environment variables to configure its local rank and total world size.
For single-node training, the simplest invocation is:
# torchrun --nproc_per_node=8 train.pyThis starts 8 processes, one per GPU, with ranks 0 through 7.
Setting Up the Process Group
import os
import torch.distributed as dist
def setup_distributed():
"""Initialize the distributed process group."""
# In real training, these are set by the launcher (torchrun or SLURM)
# We simulate a single-process environment here
if not dist.is_initialized():
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "12355")
dist.init_process_group(
backend="gloo", # Use "nccl" for real GPU training
rank=0,
world_size=1,
)
def cleanup_distributed():
"""Clean up the distributed process group."""
if dist.is_initialized():
dist.destroy_process_group()Process group setup functions defined. In real multi-GPU training, use 'nccl' backend and launch with torchrun.
Setting up the distributed process group is the first step in any DDP training script. The process group coordinates all communication between ranks. The gloo backend used here works on CPU and is suitable for demonstration; real GPU training always uses the nccl backend. The init_process_group call blocks until all processes in the group have called it, making sure that communication infrastructure is fully initialized before training begins.
Defining a Model and Wrapping with DDP
import torch.nn as nn
class SimpleTransformerBlock(nn.Module):
"""A minimal transformer block for illustration."""
def __init__(self, d_model=256, nhead=4, dim_feedforward=512):
super().__init__()
self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
self.ff = nn.Sequential(
nn.Linear(d_model, dim_feedforward),
nn.ReLU(),
nn.Linear(dim_feedforward, d_model),
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x):
attn_out, _ = self.attn(x, x, x)
x = self.norm1(x + attn_out)
x = self.norm2(x + self.ff(x))
return x
class SimpleModel(nn.Module):
def __init__(self, vocab_size=1000, d_model=256, num_layers=2):
super().__init__()
self.embed = nn.Embedding(vocab_size, d_model)
self.layers = nn.ModuleList(
[SimpleTransformerBlock(d_model) for _ in range(num_layers)]
)
self.head = nn.Linear(d_model, vocab_size)
def forward(self, x):
x = self.embed(x)
for layer in self.layers:
x = layer(x)
return self.head(x)
# Count parameters
model = SimpleModel()
num_params = sum(p.numel() for p in model.parameters())
param_size_mb = sum(
p.numel() * p.element_size() for p in model.parameters()
) / (1024**2)Model parameters: 1,567,208 Model size (fp32): 6.0 MB In DDP with N GPUs, each GPU holds 6.0 MB for model parameters. Gradients require the same memory: 6.0 MB per GPU.
from torch.nn.parallel import DistributedDataParallel as DDP
# Wrapping a model with DDP is a single line change
# device = torch.device("cuda:0") # In real training
# model = model.to(device)
# ddp_model = DDP(model, device_ids=[local_rank])
# For demonstration (CPU / single-process):
setup_distributed()
ddp_model = DDP(model)Model wrapped with DDP. DDP hooks are now registered on all parameters. Gradient all-reduce will be triggered automatically during backward().
Wrapping a model with DDP is the only model-related change needed in the training code. DDP registers backward hooks on all parameters automatically when the wrapper is created. When a gradient is computed during backpropagation, the hook fires and places the gradient into the appropriate bucket. When the bucket is full (or when the last parameter in the bucket gets its gradient), DDP launches the asynchronous NCCL all-reduce operation on that bucket.
One important detail: when you first construct the DDP wrapper, it broadcasts parameters from rank 0 to all other ranks. This ensures that all replicas start from identical initial parameters, even if rank 0 loaded a checkpoint and other ranks initialized randomly. After this initial broadcast, the invariant holds: all ranks have identical parameters at the start of every step.
The DistributedSampler
import torch
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler
# Create synthetic training data
num_samples = 1024
seq_len = 32
vocab_size = 1000
X = torch.randint(0, vocab_size, (num_samples, seq_len))
y = torch.randint(0, vocab_size, (num_samples, seq_len))
dataset = TensorDataset(X, y)
# DistributedSampler ensures each rank processes a non-overlapping shard
# world_size is the number of GPUs; rank is the current GPU's index
sampler = DistributedSampler(
dataset,
num_replicas=1, # world_size in real training
rank=0, # local_rank in real training
shuffle=True,
seed=42,
)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
# Count batches per rank
num_batches = len(dataloader)
samples_per_rank = len(sampler)Total dataset size: 1,024 samples Samples assigned to this rank: 1,024 Batches per rank per epoch: 32 With N GPUs, effective global batch size = 32 x N per step.
The DistributedSampler is critical for correctness. It partitions the dataset indices across ranks, making sure that each rank sees a different non-overlapping subset of samples each epoch. Without it, every rank would process the same samples, and you would be performing redundant computation instead of parallel computation. The effective batch size would remain rather than , and you would get no speedup from additional GPUs.
The sampler also handles an important edge case: when the dataset size is not evenly divisible by the world size, some ranks would get slightly more samples than others. By default, DistributedSampler pads the dataset with replicated samples to ensure equal-length epochs across all ranks. This prevents DDP from hanging when one rank finishes its epoch and calls all-reduce while another rank still has samples to process.
The shuffle=True parameter causes the sampler to reshuffle the assignment of samples to ranks at the start of each epoch, using a shared seed based on the epoch number. This is why you must call sampler.set_epoch(epoch) before iterating the DataLoader each epoch: it updates the epoch counter used to compute the shuffled index assignments.
Training Loop with DDP
import time
optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
# Simulate one epoch
ddp_model.train()
total_loss = 0.0
num_steps = 0
step_times = []
sampler.set_epoch(0) # Must be called each epoch for correct shuffling
for batch_X, batch_y in dataloader:
t0 = time.perf_counter()
optimizer.zero_grad()
# Forward pass: local computation only
logits = ddp_model(batch_X) # (B, T, vocab_size)
loss = criterion(logits.view(-1, vocab_size), batch_y.view(-1))
# Backward pass: DDP automatically all-reduces gradients
loss.backward()
# Gradient clipping (standard practice)
nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=1.0)
# Optimizer step: identical update on all ranks
optimizer.step()
step_time = time.perf_counter() - t0
step_times.append(step_time)
total_loss += loss.item()
num_steps += 1
avg_loss = total_loss / num_steps
avg_step_ms = np.mean(step_times) * 1000Training steps completed: 32 Average loss: 7.0233 Average step time: 26.6 ms (CPU, single-process simulation) In real DDP training: - backward() triggers async all-reduce on each gradient bucket - Communication overlaps with computation for earlier layers - All ranks apply identical optimizer updates
# Cleanup
cleanup_distributed()The training loop with DDP looks almost identical to single-GPU training. The only changes are wrapping the model with DDP, using DistributedSampler, and calling sampler.set_epoch(epoch) at the start of each epoch. The all-reduce happens automatically inside loss.backward(). This transparency is one of DDP's great strengths: existing training code requires minimal modification to become distributed.
One subtlety in the loop above is gradient clipping. The call to clip_grad_norm_ computes the global gradient norm across all parameters and clips it if it exceeds max_norm. Because all ranks have identical gradients after all-reduce, they compute identical norms and apply identical clipping. This is correct behavior. However, if you call clip_grad_norm_ before the all-reduce completes (which could happen in unusual training loop configurations), different ranks would clip based on their local gradient norms and diverge. DDP ensures all-reduce completes before the optimizer step, so standard gradient clipping is safe.
Key Parameters
The key parameters when setting up DDP are:
backend: Communication backend. Use"nccl"for GPU training (fastest, CUDA-aware),"gloo"for CPU training or debugging.bucket_cap_mb: Maximum size of gradient buckets in megabytes (default 25). Larger buckets reduce communication launch overhead but increase memory usage and per-bucket latency before the first all-reduce starts.find_unused_parameters: Set toTrueif some parameters may not receive gradients during every forward pass, for example in models with conditional routing or dynamic computation graphs. This adds overhead by scanning for unused parameters after each forward pass. Leave itFalsewhen all parameters are used in every forward pass, which is the common case for standard transformers.gradient_as_bucket_view: WhenTrue, gradients are stored directly in the communication buffer rather than being copied into it. This saves a memory copy per bucket and reduces peak memory usage. It is recommended for memory-efficient training, though it slightly constrains how you can modify gradients between the backward pass and the optimizer step.static_graph: WhenTrue, DDP assumes the computational graph is identical every step. This enables further optimizations including more aggressive bucket overlap and elimination of the find-unused-parameters scan. Set it toTruefor standard transformer models where the graph does not change.device_ids: List of CUDA device IDs to use. Typically[local_rank]in multi-GPU setups where each process owns exactly one GPU.