Distributed Training: DDP, Gradient Accumulation, and Mixed Precision

🎯 Learning Objectives

• Understand the need for distributed training
• Implement DistributedDataParallel (DDP)
• Use gradient accumulation for large batch training
• Apply Automatic Mixed Precision (AMP)

📋 Prerequisites

• Multi-GPU concepts
• PyTorch training loop

---

1. 📖 Core Content

1.1 Why Distributed Training?

• GPT-3 (175B): ~350GB in FP16 — needs 5+ A100s (80GB each)
• Training on 1 GPU would take years Parallelism strategies:

1.2 DataParallel vs DistributedDataParallel

1.3 DDP Implementation

python
Run Inline
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
def setup(rank, world_size):
    """Initialize the process group"""
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
def cleanup():
    dist.destroy_process_group()
def train_ddp(rank, world_size):
    setup(rank, world_size)
    # Create model and move to GPU
    model = MyModel().to(rank)
    # Wrap in DDP
    ddp_model = DDP(model, device_ids=[rank])
    # Data loader (split across GPUs)
    dataset = MyDataset()
    sampler = torch.utils.data.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank
    )
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    optimizer = torch.optim.Adam(ddp_model.parameters())
    for epoch in range(10):
        sampler.set_epoch(epoch)  # Shuffle differently each epoch
        for batch in dataloader:
            inputs, labels = batch
            inputs, labels = inputs.to(rank), labels.to(rank)
            outputs = ddp_model(inputs)
            loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    cleanup()
# Launch: mp.spawn(train_ddp, args=(n_gpus,), nprocs=n_gpus)

1.4 Gradient Accumulation

python
Run Inline
def train_with_gradient_accumulation(model, dataloader, accumulation_steps=4):
    optimizer.zero_grad()
    for i, batch in enumerate(dataloader):
        outputs = model(batch)
        loss = criterion(outputs, batch['labels'])
        # Scale loss to account for accumulation
        loss = loss / accumulation_steps
        loss.backward()
        if (i + 1) % accumulation_steps == 0:
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            optimizer.zero_grad()
    return model

1.5 Automatic Mixed Precision (AMP)

python
Run Inline
def train_amp(model, dataloader):
    scaler = torch.cuda.amp.GradScaler()
    for epoch in range(num_epochs):
        for batch in dataloader:
            optimizer.zero_grad()
            # Auto-cast to appropriate precision
            with torch.cuda.amp.autocast():
                outputs = model(batch['input'])
                loss = criterion(outputs, batch['label'])
            # Scale gradients to prevent underflow
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
    return model

---

📝 Practice Questions

Q1

<strong>Q1</strong>: With gradient accumulation of 4 steps, effective batch size 256, and memory for batch size 64, what's the relationship?

Per-GPU batch: 64 Gradient accumulation: 4 steps Gradient updates: (256 / 64) / 4 = 1 update per 4 steps...

Actually: effective batch = per-GPU batch × accumulation_steps × n_gpus

With 4 GPUs, per-GPU batch=64, accumulation=4: effective batch = 64 × 4 × 4 = 1024

The model sees 1024 samples per optimizer step, enabling stable training with large batch sizes while only needing memory for batch 64. Q2

<strong>Q2

<strong>Q2</strong>: Why does AMP use a gradient scaler?

FP16 gradients can underflow (become 0 for small values). The GradScaler multiplies the loss by a scale factor before backward, making gradients larger and preventing underflow.

After backward, gradients are unscaled before the optimizer step. The scaler adjusts its scale factor dynamically — increasing if no overflow, decreasing if overflow detected.

Without scaling: small gradients → underflow to 0 in FP16 → no learning With scaling: small gradients → scaled up → represented in FP16 → unscaled after backward Q3

<strong>Q3

<strong>Q3

<strong>Q3

<strong>Q3

<strong>Q3</strong>: A model uses 40GB memory with batch size 32 on one GPU. How much memory with batch size 128 on 4 GPUs using DDP?

With DDP (data parallelism), each GPU processes batch/N GPUs:

• Each GPU: batch 128/4 = 32
• Memory per GPU: 40GB (same as before — each GPU sees the same per-GPU batch)

Total memory: 4 × 40GB = 160GB

DDP doesn't reduce per-GPU memory for the same per-GPU batch size. It enables larger total batch sizes. To reduce per-GPU memory, use gradient accumulation with smaller per-GPU batches. Q4

<strong>Q4

<strong>Q4

<strong>Q4

<strong>Q4

<strong>Q4

<strong>Q4</strong>: In DDP, gradients are synchronized across GPUs after each backward pass. How does this affect training?

After loss.backward(), DDP performs an all-reduce of gradients across all GPUs — each GPU's gradients are averaged:

g_i_avg = (g_1 + g_2 + ... + g_N) / N

This ensures that after optimizer.step(), all GPUs have identical model parameters. Without this sync, different GPUs would diverge.

The all-reduce adds communication overhead proportional to model size. This is why DDP is most efficient when computation time >> communication time (large models, large batches).

Modern DDP implementations overlap communication with computation (gradient computation + all-reduce happen simultaneously) to minimize overhead.

</details> * * * ## 🔗 Cross-References - **Next**: [Mixed Precision](/notes/04-degree-electives-bsda5013-dl-practice-week06-06-mixed-precision) - **Previous**: [Training & Debugging](/notes/04-degree-electives-bsda5013-dl-practice-week04-04-training-debugging) - **Video**: BSDA5013 Week 5 transcripts [Join Discord](https://discord.gg/gE2m4Qrdqv) [Previous**Training & Debugging**](/notes/04-degree-electives-bsda5013-dl-practice-week04-04-training-debugging)[Next**Mixed Precision Training**](/notes/04-degree-electives-bsda5013-dl-practice-week06-06-mixed-precision)