Skip to main content

Chapter 2: Distributed Training

Parallelism Strategies​

PyTorch DDP Example​

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train(rank, world_size):
    setup(rank, world_size)
    
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(model.parameters())
    
    for epoch in range(num_epochs):
        for batch in dataloader:
            outputs = model(batch)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    dist.destroy_process_group()

# Launch with torchrun
# torchrun --nproc_per_node=4 train.py

Key Takeaways​

Summary
  1. Data parallelism works for most use cases
  2. Model parallelism needed for very large models
  3. DeepSpeed/FSDP enable training models larger than GPU memory

Next: Chapter 3: Cloud Training