Practical Guide to PyTorch Distributed Training: DP, DDP, Groups, and IO Considerations

This article provides a comprehensive overview of PyTorch's distributed training capabilities, focusing on both DataParallel (DP) and DistributedDataParallel (DDP) and covering practical details that are often overlooked.

Key Concepts

DP (single‑node multi‑GPU) replicates the model on each GPU, splits the batch, and gathers outputs on a designated output_device (usually rank 0). This can cause memory imbalance because rank 0 must hold the full output tensor.

DDP (multi‑node or multi‑GPU) synchronizes parameters and buffers once during initialization and then averages gradients each iteration, eliminating the output‑gather bottleneck.

Important Parameters world_size: total number of parallel processes. rank: the index of the current process. group_size, group_ws, group_rank, local_group_rank, group_rank_base: definitions for fine‑grained group communication.

Typical DP Usage

When using 4 GPUs on a single node with a batch size of 128, DP splits the batch into 32 samples per GPU, copies the model to GPUs 1‑3, runs forward passes, and then gathers results on the output device. This extra copy step reduces efficiency.

DDP Advantages and Pitfalls

Do not modify model parameters or buffers after DDP initialization; otherwise new parameters will not be synchronized.

Set find_unused_parameters=True for models that contain unused parameters (e.g., NAS sub‑graphs) to avoid silent gradient errors.

Buffers are synchronized before each forward pass; if exact reproducibility is required, call DDP._sync_params_and_buffers() after the last iteration.

When loss is already reduced (e.g., reduce_mean), do not manually divide by world_size because DDP already averages gradients.

Code Example: Basic Distributed Training Loop

from torch import distributed as dist
from torch.utils.data.distributed import DistributedSampler
import torch.utils.data as Data

assert torch.cuda.is_available()
if not dist.is_initialized():
    dist.init_process_group(backend='nccl')

rank = dist.get_rank()
world_size = dist.get_world_size()

model = MyModel().cuda()
ddp_model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()]).cuda()

dataset = MyDataset()
sampler = DistributedSampler(dataset, rank, world_size, shuffle=True)

dataloader = Data.DataLoader(dataset, batch_size, drop_last=False, sampler=sampler, shuffle=False, num_workers=8, pin_memory=True)

# training loop ...

Group Communication

Groups allow you to partition the set of processes into sub‑sets for independent communication. Example:

ranks = [0, 1, 2, 3]
gp = dist.new_group(ranks, backend='nccl')

All processes must create every group they might belong to, even if they are not members of a particular group.

rank = dist.get_rank()
group_ranks = [[0,1,2,3], [4,5,6,7], [8,9,10,11]]
cur_gp = None
for g_ranks in group_ranks:
    gp = dist.new_groups(g_ranks)
    if rank in g_ranks:
        cur_gp = gp
# use cur_gp for subsequent ops

Data Loading with DistributedSampler

Each node runs its own DataLoader with a DistributedSampler. The sampler splits the dataset indices across num_replicas (usually world_size) and assigns each rank a distinct subset.

real_data_num = int(math.ceil(len(dataset) * 1.0 / world_size)) * world_size

When evaluating, be aware that padding may cause duplicate samples, which can bias accuracy on small test sets.

IO Bottlenecks in Multi‑Node Scenarios

For small datasets, consider converting data to LMDB or copying it into the container’s local filesystem. For large datasets, avoid all‑node contention by broadcasting data from a subset of nodes (e.g., rank 0 of each group) or by scattering from a single node and then broadcasting.

Broadcasting requires that all tensors have the same shape; non‑tensor metadata (e.g., strings) cannot be broadcast. Additionally, the broadcast groups may differ from the model‑parallel groups, so you may need to create separate groups such as [0,3], [1,4], [2,5] for a 6‑GPU setup with two groups of three.

Optimizer Gradient Filtering

for p in model.parameters():
    if p.grad is not None and (p.grad == 0).all():
        p.grad = None

This prevents optimizers from mistakenly updating parameters whose gradients are zero due to previous sub‑graph optimizations.

Summary

The article walks through the mechanics of DP and DDP, highlights subtle bugs (memory imbalance, unused parameters, buffer synchronization), explains how to configure fine‑grained groups, details data loading strategies with DistributedSampler, and offers practical solutions for IO bottlenecks in large‑scale training.