Phase 6 • EduArtha

Systems, Infrastructure & Research

Building a real AI model at scale requires distributed systems, massive compute, and research skills. This is what separates an ML practitioner from an AI engineer.

⏱ Ongoing | 14 Chapters | 50+ Exercises | Industry Problems

Part I

Distributed Training

Training models across multiple GPUs and nodes

Chapter 1

Data Parallelism (DDP)

Learning Objectives

Understand DistributedDataParallel — the workhorse of multi-GPU training
Implement DDP training loops from scratch
Know how gradient synchronization works via AllReduce
Scale from 1 GPU to 8 GPUs with minimal code changes

How Data Parallelism Works

Each GPU gets a copy of the model + a different mini-batch → Forward → Backward → AllReduce gradients → Update

Python
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

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

def train_ddp(rank, world_size):
    setup(rank, world_size)

    # Each GPU gets the SAME model
    model = nn.Sequential(
        nn.Linear(784, 512), nn.ReLU(),
        nn.Linear(512, 10)
    ).to(rank)

    # Wrap with DDP — handles gradient sync automatically
    model = DDP(model, device_ids=[rank])

    # DistributedSampler ensures each GPU gets DIFFERENT data
    dataset = torchvision.datasets.MNIST('./data', train=True)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(10):
        sampler.set_epoch(epoch)  # Shuffle differently each epoch
        for X, y in loader:
            X, y = X.to(rank), y.to(rank)
            loss = nn.functional.cross_entropy(model(X.view(-1,784)), y)
            optimizer.zero_grad()
            loss.backward()       # DDP: AllReduce gradients automatically
            optimizer.step()

    dist.destroy_process_group()

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

Industry Problem: Linear Scaling Efficiency

Problem: Going from 1 GPU to 8 GPUs should give 8× speedup — but communication overhead reduces this. With 8× A100 on NVLink you get ~7.5× speedup. Across nodes (InfiniBand), you might only get 6× for 8 GPUs.

Solutions: (1) Gradient compression — reduce communication volume. (2) Overlap communication with computation — DDP does this by default, syncing gradients of earlier layers while computing later ones. (3) Large batch training — LARS/LAMB optimizers scale learning rate with batch size. (4) Gradient accumulation — simulate larger batches without more GPUs.

Exercises

Exercise 1.1: Why must sampler.set_epoch(epoch) be called?

Without set_epoch(), DistributedSampler generates the same shuffled order every epoch (deterministic seed). Each GPU would see the same data in the same order — effectively no shuffling between epochs. set_epoch() changes the random seed, ensuring different shuffles each epoch while keeping GPUs synchronized.

Exercise 1.2: What is AllReduce and why is it used for gradient sync?

AllReduce sums tensors across all GPUs and distributes the result back to every GPU. After backward pass, each GPU has different gradients (from different data). AllReduce averages them, so all GPUs have identical averaged gradients → identical weight updates → models stay synchronized. NCCL provides hardware-optimized AllReduce on NVIDIA GPUs.

Exercise 1.3: How does effective batch size change with DDP?

Effective batch = per-GPU batch × num_GPUs. With batch=64 on 8 GPUs: effective batch = 512. This changes training dynamics — you may need to adjust learning rate (linear scaling rule: LR × num_GPUs) or use warmup. Very large batches (>8K) may hurt generalization.

Chapter Summary

DDP replicates the model on every GPU and synchronizes gradients via AllReduce
DistributedSampler ensures each GPU processes different data
Near-linear scaling (90%+) with proper overlap of communication and computation
Effective batch size = per-GPU batch × num_GPUs — adjust LR accordingly