Distributed Training: DataParallel to DistributedDataParallel

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: 8 samples
# GPU 1: 8 samples
# GPU 2: 8 samples
# GPU 3: 8 samples

# 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: Master GPU (0) does gathering/scattering
#    Other GPUs wait for master to finish

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

# 3. Slower than DistributedDataParallel for large models

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

# GPU 0: Master (gathers gradients)
# GPU 1, 2, 3: Compute (wait for master)
# CPU: Single process (bottleneck for data loading)

# → 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: [2.0, 3.0, 1.0]
# Process 1 has gradients: [4.0, 1.0, 5.0]

# After DDP all_reduce (average):
# Process 0: [3.0, 2.0, 3.0]  (sum / 2)
# Process 1: [3.0, 2.0, 3.0]  (sum / 2)

# 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: sum across all processes
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: [1, 2, 3]
# Process 1: [4, 5, 6]
# Process 2: [7, 8, 9]
# Process 3: [10, 11, 12]

# After all_reduce (SUM):
# Process 0: [22, 26, 30]  (sum of all)
# Process 1: [22, 26, 30]
# Process 2: [22, 26, 30]
# Process 3: [22, 26, 30]

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: collect from all processes
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: distribute to all processes
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: Only rank 0 updates learning rate
if rank == 0:
    adjust_learning_rate(optimizer, epoch)

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

# CORRECT: All ranks adjust identically
# (or broadcast from rank 0)

adjust_learning_rate(optimizer, epoch)

Avoiding Deadlocks

# WRONG: Different code paths on different ranks
if rank == 0:
    # Long validation
    validate(model, val_loader)

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

# CORRECT: All ranks follow same code path
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 - Profiling DDP overhead