Distributed Training

Overview

Training large models requires distributing computation across GPUs/TPUs:

• DataParallel: Simple, single-machine multi-GPU
• DistributedDataParallel (DDP): Multi-machine, production-grade
• Gradient accumulation: Simulating large batches
• Synchronization: Ensuring all processes agree

Understanding these patterns is essential for scaling ML systems.

---

Single-GPU Training Baseline

Standard Training Loop

import torch
import torch.nn as nn
import torch.optim as optim

# Model and optimizer
model = nn.Linear(10, 5)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# Single GPU training
for epoch in range(10):
 for batch_x, batch_y in train_loader:
 # Forward
 output = model(batch_x)
 loss = criterion(output, batch_y)

 # Backward
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

# Characteristics:
# - Model on single GPU
# - Batch processed by single GPU
# - Simple, but limited by single GPU memory

-

DataParallel: Simple Multi-GPU

How It Works

import torch
import torch.nn as nn

# Single GPU model
model = nn.Linear(10, 5)

# Wrap with DataParallel
device_ids = [0, 1, 2, 3] # GPUs 0, 1, 2, 3
model = nn.DataParallel(model, device_ids=device_ids)

# Model automatically splits batch across GPUs
batch_size = 32 # Total batch
# GPU 0
# GPU 1
# GPU 2
# GPU 3

# Forward/backward works same as single GPU
for epoch in range(10):
 for batch_x, batch_y in train_loader:
 output = model(batch_x) # Auto-splits batch
 loss = criterion(output, batch_y)

 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

DataParallel Mechanics

# Simplified DataParallel forward pass:

def forward(self, input):
 # Scatter input across GPUs
 inputs = scatter(input, self.device_ids) # Split batch
 # inputs = [batch[0:8].to(gpu0), batch[8:16].to(gpu1),...]

 # Run replicas in parallel
 outputs = []
 for i, (input, device) in enumerate(zip(inputs, self.device_ids)):
 replica = self.module.to(device)
 output = replica(input)
 outputs.append(output)

 # Gather outputs back to device 0
 output = gather(outputs, self.device_ids[0])
 return output

Limitations of DataParallel

# DataParallel issues:

# 1. Bottleneck
# Other GPUs wait for master to finish

# 2. Single-process
# Can't utilize multi-core CPU for data loading

# 3. Slower than DistributedDataParallel for large models

# Example
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])

# GPU 0
# GPU 1, 2, 3
# CPU

# → Not recommended for production training

---

DistributedDataParallel: Production-Grade

Setup

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

def setup_ddp(rank, world_size):
 """Initialize distributed training."""
 os.environ['MASTER_ADDR'] = 'localhost'
 os.environ['MASTER_PORT'] = '12355'

 # Initialize process group
 dist.init_process_group(
 backend='nccl', # NVIDIA collective communication library
 rank=rank,
 world_size=world_size,
)

def cleanup_ddp():
 """Clean up distributed training."""
 dist.destroy_process_group()

def train_ddp(rank, world_size, model_class, train_loader):
 """Training function for single process."""
 # Setup
 setup_ddp(rank, world_size)

 # Create model and wrap with DDP
 model = model_class()
 model.to(rank) # Move to current GPU
 model = DDP(model, device_ids=[rank])

 optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

 # Training loop
 for epoch in range(10):
 # Important: shuffle differently on each rank
 train_loader.sampler.set_epoch(epoch)

 for batch_x, batch_y in train_loader:
 batch_x = batch_x.to(rank)
 batch_y = batch_y.to(rank)

 # Forward
 output = model(batch_x)
 loss = criterion(output, batch_y)

 # Backward
 optimizer.zero_grad()
 loss.backward()

 # DDP automatically synchronizes gradients
 optimizer.step()

 cleanup_ddp()

# Launch training on 4 GPUs
if __name__ == '__main__':
 world_size = 4
 mp.spawn(train_ddp, args=(world_size, MyModel, train_loader), nprocs=world_size)

DataLoader Setup for DDP

import torch
from torch.utils.data import DataLoader, DistributedSampler

def create_ddp_dataloaders(train_dataset, val_dataset, batch_size, num_workers):
 """Create dataloaders for DDP."""

 # Sampler handles splitting data across ranks
 train_sampler = DistributedSampler(
 train_dataset,
 num_replicas=torch.distributed.get_world_size(),
 rank=torch.distributed.get_rank(),
 shuffle=True,
 drop_last=True, # Important: same batch size on all ranks
)

 train_loader = DataLoader(
 train_dataset,
 batch_size=batch_size,
 sampler=train_sampler,
 num_workers=num_workers,
 pin_memory=True, # Faster GPU transfer
)

 # Validation: typically single-process or synchronized
 val_loader = DataLoader(
 val_dataset,
 batch_size=batch_size,
 shuffle=False,
 num_workers=num_workers,
)

 return train_loader, val_loader

-

Gradient Synchronization

How DDP Synchronizes Gradients

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

# DDP automatically:
# 1. After backward(), collects gradients from all processes
# 2. Performs all_reduce() to average gradients
# 3. Each process gets same averaged gradients
# 4. All processes update parameters identically

# Example with 2 processes:

# Process 0 has gradients
# Process 1 has gradients

# After DDP all_reduce (average):
# Process 0
# Process 1

# Both processes update parameters identically

Manual Synchronization

import torch.distributed as dist

# Get own gradient
local_grad = model[0].weight.grad

# Create tensor to hold reduced gradient
reduced_grad = local_grad.clone()

# All-reduce
dist.all_reduce(reduced_grad, op=dist.ReduceOp.SUM)

# Average
reduced_grad /= torch.distributed.get_world_size()

# Update model
model[0].weight.grad = reduced_grad

-

Communication Patterns

All-Reduce (Most Common)

# Each process contributes, result goes to all processes
# 4 processes, each has local value:

# Process 0
# Process 1
# Process 2
# Process 3

# After all_reduce (SUM):
# Process 0
# Process 1
# Process 2
# Process 3

dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

Reduce (One Process Receives)

# Only process 0 gets result
dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)

# Used for collecting validation metrics from all processes

Broadcast (One Process Sends)

# Process 0 sends to all
dist.broadcast(tensor, src=0)

# Used for distributing model updates or hyperparameters

Gather/Scatter

# Gather
if rank == 0:
 gathered_list = [torch.zeros_like(tensor) for _ in range(world_size)]
else:
 gathered_list = None

dist.gather(tensor, gather_list=gathered_list, dst=0)

# Scatter
dist.scatter(tensor, scatter_list=scatter_list, src=0)

-

Advanced: Gradient Accumulation with DDP

Effective Batch Size > Physical Batch Size

import torch
import torch.nn as nn

# Training with gradient accumulation
model = nn.Linear(10, 5)
model = DDP(model)

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# Simulate batch size 128 with gradient accumulation
effective_batch_size = 128
actual_batch_size = 32
accumulation_steps = effective_batch_size // actual_batch_size

for epoch in range(10):
 train_loader.sampler.set_epoch(epoch)

 for batch_idx, (batch_x, batch_y) in enumerate(train_loader):
 # Forward
 output = model(batch_x)

 # Loss scaled by accumulation steps
 loss = criterion(output, batch_y) / accumulation_steps

 # Backward: gradients accumulate
 loss.backward()

 # Update after K steps
 if (batch_idx + 1) % accumulation_steps == 0:
 optimizer.step()
 optimizer.zero_grad()

Mixed Precision with DDP

from torch.cuda.amp import autocast, GradScaler

model = nn.Linear(10, 5)
model = DDP(model)

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = GradScaler() # Scales loss to prevent underflow

for epoch in range(10):
 for batch_x, batch_y in train_loader:
 # Forward in float16
 with autocast(device_type='cuda'):
 output = model(batch_x)
 loss = criterion(output, batch_y)

 # Backward with loss scaling
 scaler.scale(loss).backward()

 # Unscale before update
 scaler.unscale_(optimizer)

 # Gradient clipping (optional)
 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

 # Update
 scaler.step(optimizer)
 scaler.update()
 optimizer.zero_grad()

---

Multi-Node Training

Launch Across Multiple Machines

# Machine 1 (Master)
# python train.py --rank 0 --world_size 8 --master_addr 192.168.1.100

# Machine 2
# python train.py --rank 4 --world_size 8 --master_addr 192.168.1.100

import argparse
import os

def setup_ddp_multinode(rank, world_size, master_addr, master_port):
 """Setup multi-node DDP."""
 os.environ['MASTER_ADDR'] = master_addr
 os.environ['MASTER_PORT'] = str(master_port)

 dist.init_process_group(
 backend='nccl',
 rank=rank,
 world_size=world_size,
 timeout=timedelta(minutes=30),
)

if __name__ == '__main__':
 parser = argparse.ArgumentParser()
 parser.add_argument('--rank', type=int)
 parser.add_argument('--world_size', type=int)
 parser.add_argument('--master_addr', type=str)
 parser.add_argument('--master_port', type=int, default=29500)
 args = parser.parse_args()

 setup_ddp_multinode(args.rank, args.world_size, 
 args.master_addr, args.master_port)

 # Training code
 model = nn.Linear(10, 5).cuda(args.rank)
 model = DDP(model, device_ids=[args.rank])

 #... training loop...

---

Practical Example: Complete DDP Training

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch.multiprocessing as mp

class SimpleModel(nn.Module):
 def __init__(self):
 super().__init__()
 self.fc1 = nn.Linear(10, 20)
 self.fc2 = nn.Linear(20, 5)

 def forward(self, x):
 x = torch.relu(self.fc1(x))
 return self.fc2(x)

def train(rank, world_size):
 # Setup
 os.environ['MASTER_ADDR'] = 'localhost'
 os.environ['MASTER_PORT'] = '12355'
 dist.init_process_group("nccl", rank=rank, world_size=world_size)

 # Create model
 model = SimpleModel()
 model.to(rank)
 model = DDP(model, device_ids=[rank])

 # Data
 dataset = TensorDataset(
 torch.randn(1000, 10),
 torch.randint(0, 5, (1000,))
)
 sampler = DistributedSampler(dataset, rank=rank, world_size=world_size)
 loader = DataLoader(dataset, batch_size=32, sampler=sampler)

 # Optimizer
 optimizer = optim.SGD(model.parameters(), lr=0.01)
 criterion = nn.CrossEntropyLoss()

 # Training
 for epoch in range(10):
 sampler.set_epoch(epoch)

 for batch_x, batch_y in loader:
 batch_x, batch_y = batch_x.to(rank), batch_y.to(rank)

 output = model(batch_x)
 loss = criterion(output, batch_y)

 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

 if rank == 0:
 print(f"Epoch {epoch}: loss={loss.item():.4f}")

 dist.destroy_process_group()

if __name__ == '__main__':
 world_size = 4
 mp.spawn(train, args=(world_size,), nprocs=world_size)

---

Synchronization Issues

Common Pitfall: Rank 0 Behaves Differently

# WRONG
if rank == 0:
 adjust_learning_rate(optimizer, epoch)

# At next step, rank 0 has different lr than others
# Divergence in training!

# CORRECT
# (or broadcast from rank 0)

adjust_learning_rate(optimizer, epoch)

Avoiding Deadlocks

# WRONG
if rank == 0:
 # Long validation
 validate(model, val_loader)

# Rank 1, 2, 3 reach barrier and wait forever!

# CORRECT
validate(model, val_loader)

# Or synchronize explicitly
dist.barrier()

---

Debugging DDP

Check Process Group

if dist.is_available() and dist.is_initialized():
 print(f"Rank: {dist.get_rank()}")
 print(f"World size: {dist.get_world_size()}")
 print(f"Backend: {dist.get_backend()}")

Hang Detection

# Set timeout to detect hangs
dist.init_process_group(
 backend='nccl',
 timeout=timedelta(minutes=30),
)

# Hangs will raise exception after 30 minutes

Gradient Verification

# Check gradients match across processes
if dist.is_initialized():
 all_grads = []
 for param in model.parameters():
 local_grad = param.grad.clone()
 gathered = [torch.zeros_like(local_grad) for _ in range(world_size)]
 dist.all_gather(gathered, local_grad)
 all_grads.append(gathered)

 # Verify all ranks have identical gradients
 for rank in range(1, world_size):
 for i, grads in enumerate(all_grads):
 assert torch.allclose(grads[0], grads[rank])

---

Summary: DDP vs DataParallel

Feature	DataParallel	DDP
GPUs	Single machine	Multiple machines
Processes	Single	Multiple
Scalability	Poor	Excellent
Communication	No explicit sync	All-reduce
Complexity	Simple	Moderate
Production	Not recommended	Standard
Speed	Slower	Fastest

-

Related Topics

• 02 Autograd Implementation - Gradient computation
• 00 Readme - Communication overhead
• [04 Profiling & Performance Analysis](/05-py3/09-bytecode-and-execution/(04-profiling-performance-analysis/) - Profiling DDP overhead