Distributed Training — Module 6.2

How This Lesson Fits the Module

Single-GPU training hits wall-clock limits on large datasets and models. Distributed training splits work across devices—most commonly data parallelism (DDP), where each GPU runs the same model on different batches and synchronizes gradients.

This is the capstone of Module 6.2: every prior lesson (loops, AMP, clipping) still applies inside each rank’s process.

Module 6.2 Recap Training loop → validation → test → checkpoints → early stopping → schedulers → warmup → clipping → AMP. Distributed wraps the same loop across processes.

Learning Objectives

By the end of this lesson, students should be able to:

Explain data parallelism vs model parallelism at a high level.
Initialize torch.distributed and wrap models with DDP.
Use DistributedSampler so each rank sees unique batches.
Restrict checkpointing and logging to rank 0.
Launch jobs with torchrun (multi-process per node).
Estimate near-linear speedup limits from communication overhead.

Data Parallel vs Model Parallel

Strategy	Idea	Typical use
Data Parallel (DDP)	Replica per GPU, split batches, sync gradients	Most CV/NLP training
Model Parallel	Split layers across devices	Models too large for one GPU
Pipeline / FSDP	Shard parameters across ranks	Very large LLMs (advanced)

DDP Setup Pattern

One process per GPU. Initialize process group, bind device, wrap model with DDP, use distributed sampler.

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

def setup():
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

local_rank = setup()
device = torch.device("cuda", local_rank)

model = MyModel().to(device)
model = DDP(model, device_ids=[local_rank])

train_sampler = DistributedSampler(train_dataset, shuffle=True)
train_loader = DataLoader(train_dataset, batch_size=64, sampler=train_sampler)

for epoch in range(num_epochs):
    train_sampler.set_epoch(epoch)  # new shuffle each epoch
    train_one_epoch(model, train_loader, ...)

    if local_rank == 0:
        val_loss, _ = validate(model.module, val_loader, ...)
        save_checkpoint(...)

dist.destroy_process_group()

Launch with torchrun

torchrun spawns one process per GPU and sets RANK, LOCAL_RANK, and WORLD_SIZE environment variables.

# 4 GPUs on one machine:
torchrun --standalone --nproc_per_node=4 train.py

# Effective batch size = batch_size * nproc_per_node * grad_accum_steps
# Often scale learning rate linearly with global batch size

Critical Mistake — Forgetting set_epoch on the Sampler

Without train_sampler.set_epoch(epoch), shuffling repeats identically every epoch across ranks—hurting convergence. Always call it at the start of each epoch.

DDP vs Deprecated DataParallel

nn.DataParallel (single process, multiple GPUs) is simpler but slower due to GIL and gradient gather on one GPU. Use DDP for multi-GPU training in production.

Accessing the Underlying Model

DDP wraps the module. For validation/saving on rank 0, use model.module.state_dict(), not model.state_dict() (keys include module. prefix).

Engineering Habit — Global Batch Size Math

Document: per-GPU batch × GPUs × accumulation = global batch. Learning rate schedules often depend on global batch, not per-device batch.

Knowledge Check

Short Answer: What does DDP synchronize? Answer: Gradients across processes after backward.
True/False: Each DDP rank should save checkpoints independently. Answer: False—usually rank 0 only.
Multiple Choice: DistributedSampler ensures: (a) same batches on all GPUs, (b) unique shards per rank, (c) no shuffling. Answer: (b).
Short Answer: Why call set_epoch on the sampler? Answer: Different shuffle seed each epoch per rank.
Short Answer: How to launch 8 processes on one node? Answer: torchrun --nproc_per_node=8 train.py.
True/False: DDP is slower than DataParallel for training. Answer: False—DDP is generally faster and preferred.
Multiple Choice: Save weights from DDP model via: (a) model.module.state_dict(), (b) model.cuda(), (c) rank 3 only. Answer: (a) on rank 0.
Short Answer: Effective batch with 4 GPUs, batch 32 each? Answer: 128 (without gradient accumulation).
Short Answer: When is model parallelism needed? Answer: Model does not fit on one GPU’s memory.
Multiple Choice: Backend for NVIDIA GPU training: (a) gloo, (b) nccl, (c) mpi only. Answer: (b).

Key Takeaways

DDP: one process per GPU, DistributedSampler, gradient sync, rank-0 logging.
Launch with torchrun; call set_epoch each epoch.
Scale global batch and LR together; save model.module weights.
Next module: 6.3 CPU vs GPU—hardware foundations.

Trainer’s Guide

Hands-on idea: Train MNIST 1 vs 2 GPU; measure speedup and discuss overhead below 2×.

Discussion prompt: What breaks first when scaling to 64 GPUs—compute, communication, or data loading?

Module Complete You can build production training loops end-to-end. Continue to Module 6.3 CPU vs GPU for hardware and precision deep dives.