Learn AI Series (#126) - Distributed Training

avatar

Learn AI Series (#126) - Distributed Training

variant-a-07-purple.png

What will I learn

  • You will learn data parallelism with DistributedDataParallel -- the boring, brilliant workhorse that gets you from one GPU to hundreds by splitting the batch;
  • how GPUs actually talk to each other -- all-reduce, all-gather, ring topologies, and why NCCL matters more than you'd think;
  • model parallelism and tensor parallelism -- what you do when a single card physically cannot hold the model;
  • pipeline parallelism and the infamous "bubble" -- keeping every GPU busy instead of idle-waiting;
  • FSDP and DeepSpeed/ZeRO -- sharding parameters, gradients and optimizer states so a 7B model that "doesn't fit" suddenly does;
  • and the scaling laws (Kaplan, then Chinchilla) -- the empirical relationship between compute, data and model size that quietly rewrote how everyone trains.

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.10+) distribution with PyTorch (pip install torch) -- and again, you do NOT need a rack of GPUs to follow along: we're after the mental model here, and the mental model is the part that transfers whether you own two cards or two thousand;
  • You've been through episode #60 (training large language models, where we first hit the memory wall) and episode #125 (GPU programming basics, where we learned how a single card actually works). Today we take that single card and multiply it.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#126) - Distributed Training

I ended episode #125 with a deliberately annoying cliffhanger. We'd just spent a whole post learning how to squeeze every drop of performance out of a single GPU -- keep your data resident, respect the memory hierarchy, turn on mixed precision -- and then I pointed at the ceiling hiding in plain sight: one box. What happens when the model, or the dataset, is so enormous that even a machine crammed with eight GPUs isn't enough, and you have to make several separate computers cooperate on one training run? That's today. And I promise it's less scary than it sounds, because the whole field basically rests on four ideas, and once they click you'll see them everywhere.

Here's the framing I want you to keep. A single modern GPU can comfortably train a model with a few hundred million parameters. Current frontier models have hundreds of billions. That's a factor of a thousand, and you do not close a factor of a thousand by waiting patiently -- you close it by putting a lot of GPUs to work at the same time. But "a lot of GPUs" is not one technique, it's a small family of them, each solving a different problem. Data too big? Split the data. Model too big to fit? Split the model. Both? Combine the two. The beautiful part is that the exact same principles apply whether you're wiring up two RTX cards on your desk or two thousand H100s in a datacenter. Scale is just more of the same idea. Let's go.

Data parallelism: the one you almost always want

Start with the common case, because 90% of distributed training in the wild is this one thing: data parallelism. The model is small enough to fit on a single GPU, but training is slow because you have mountains of data to grind through. The fix is embarrassingly simple in concept -- put a full copy of the model on every GPU, hand each GPU a different slice of every batch, and after each one computes its gradients, average those gradients across all the cards so everybody steps in the same direction. In PyTorch the tool for this is DistributedDataParallel (DDP), and it's the good one -- not the older, slower DataParallel we warned you off in episode #125.

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

def setup_ddp(rank, world_size):
    """Initialize the distributed environment (one call per process)."""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # 'nccl' is the NVIDIA backend -- fast. Use 'gloo' on CPU-only.
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train_ddp(rank, world_size, dataset):
    setup_ddp(rank, world_size)

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

    # Wrap with DDP -- this is what handles gradient synchronization
    model = DDP(model, device_ids=[rank])

    # DistributedSampler makes sure each GPU gets DIFFERENT data
    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)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)   # reshuffle differently every 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()   # <-- DDP synchronizes gradients HERE, automatically
            optimizer.step()

        if rank == 0:          # only rank 0 prints, or you get N copies
            print(f"Epoch {epoch} complete")

    cleanup()

# You don't call train_ddp yourself -- you launch it, one process per GPU:
#   torchrun --nproc_per_node=4 train.py

There are three small details in there that trip up quit some beginners, so let me spell them out. First, one process per GPU -- you don't loop over devices inside a single Python program, you launch N independent processes (that's what torchrun does) and each one owns exactly one card. This sidesteps Python's GIL entirely, which is precisely why DDP beats the old DataParallel. Second, the DistributedSampler -- without it, every GPU would happily train on the same batch, which is pointless; the sampler partitions the dataset so each rank sees a disjoint slice. And third, that if rank == 0 guard around printing and (later) checkpointing -- every process runs the same code, so anything you want to happen once has to be gated to a single rank.

But the real magic is one line: loss.backward(). That's where DDP earns its keep. It hooks into the backward pass and overlaps gradient computation with gradient communication. While GPU 0 is still busy computing gradients for layer N, it's simultaneously shipping the already-finished gradients for layer N+1 off to the other GPUs. Communication hides behind computation. The old DataParallel waited for the entire backward pass to finish before it sent a single byte -- which is why it left the network idle half the time. Having said that, DDP isn't magic pixie dust: it works because the model is small enough to replicate. The moment that stops being true, we need a different trick. Hold that thought.

How GPUs actually talk: collective communication

That "average the gradients across all GPUs" step has a proper name: all-reduce. It belongs to a family of operations called collective communication, and they're the vocabulary of every distributed training system on the planet. Worth knowing the words, because they show up in error messages and profiler output constantly:

# The collective operations you'll actually meet:

# ALL-REDUCE : every GPU has a value; afterwards every GPU has the
#              SUM (or average) of all of them.  <-- this is DDP's core op
# ALL-GATHER : every GPU has a piece; afterwards every GPU holds ALL
#              the pieces, concatenated.
# REDUCE     : every GPU contributes; ONE GPU ends up with the result.
# BROADCAST  : one GPU has the data; afterwards everyone has a copy.
# SCATTER    : one GPU splits its data into pieces, one piece per GPU.

# In PyTorch these live in torch.distributed:
#   dist.all_reduce(tensor)      # in-place, every rank gets the sum
#   dist.all_gather(list, tensor)
#   dist.broadcast(tensor, src=0)

Now, how do you actually compute a sum across 8 (or 800) GPUs without one of them becoming a bottleneck that everybody waits on? The naive way -- send everything to GPU 0, let it add, send the result back -- means GPU 0's network link carries all the traffic and everyone else twiddles their thumbs. The clever answer, and the one that made large-scale training practical, is the ring all-reduce:

def all_reduce_ring(local_chunk, world_size):
    """Ring all-reduce -- the intuition, not a runnable kernel.

    Arrange the N GPUs in a logical ring: 0 -> 1 -> 2 -> ... -> 0.
    Split each GPU's gradient tensor into N chunks.
    """
    # PHASE 1 -- scatter-reduce (N-1 steps):
    #   On each step every GPU sends ONE chunk to its right neighbor
    #   and receives one from its left, then ADDS what it received.
    #   After N-1 steps: each GPU holds ONE fully-summed chunk.

    # PHASE 2 -- all-gather (N-1 steps):
    #   Now every GPU sends its completed chunk around the ring.
    #   After N-1 steps: every GPU holds ALL the summed chunks.

    # Total data moved per GPU is ~2x the gradient size -- and crucially
    # it's INDEPENDENT of N. Double the GPUs, same bandwidth per card.
    return "summed gradients, everywhere"

That bandwidth-independent-of-N property is the whole ballgame. It's why you can scale to thousands of GPUs without the network melting. You do NOT implement this yourself -- NCCL (NVIDIA Collective Communication Library, pronounced "nickel") does, and it does it optimally for the hardware you actually have. NCCL automatically picks the fastest interconnect available: NVLink between GPUs inside the same machine (600 GB/s and up on modern boxes), and InfiniBand across machines (200-400 Gb/s). This hardware topology is exactly why, as we'll see, you prefer to keep the chattiest kind of parallelism inside a single node and the more forgiving kind across nodes.

When the model itself won't fit: model and tensor parallelism

Data parallelism assumes every GPU can hold a full copy of the model. But a 175-billion-parameter model in 16-bit is roughly 350GB of weights alone -- no single card on Earth holds that. When the model is the thing that doesn't fit, you split the model instead of the data. The crude version is model parallelism: put layers 1-12 on GPU 0, layers 13-24 on GPU 1, and pass the activations across as they flow forward. Simple, but wasteful (only one GPU is busy at a time). The refined version is tensor parallelism, which splits individual operations -- like a single big matrix multiply -- across cards:

class TensorParallelLinear(nn.Module):
    """Split ONE linear layer's output dimension across GPUs."""
    def __init__(self, in_features, out_features, world_size, rank):
        super().__init__()
        # Each GPU owns a vertical SLICE of the output neurons
        self.slice_size = out_features // world_size
        self.linear = nn.Linear(in_features, self.slice_size).to(rank)
        self.rank = rank
        self.world_size = world_size

    def forward(self, x):
        # Each GPU computes only its slice of the output...
        local_output = self.linear(x.to(self.rank))

        # ...then we all-gather to reconstruct the full output tensor
        gathered = [torch.zeros_like(local_output) for _ in range(self.world_size)]
        dist.all_gather(gathered, local_output)
        return torch.cat(gathered, dim=-1)

Megatron-LM (from NVIDIA) pioneered this for transformers, cleverly splitting attention heads and the columns of the MLP layers across GPUs so that the communication lands in just a couple of well-chosen spots per layer. This is how the genuinely enormous models get trained -- no single GPU holds the full thing, but each holds a well-defined fraction. The catch is that tensor parallelism is chatty: it needs an all-gather (or all-reduce) inside essentially every layer, so it wants that fat NVLink bandwidth. That's the rule of thumb worth memorising: tensor parallelism stays inside a node, data parallelism spans across nodes. Match the communication pattern to the wire.

Pipeline parallelism and the dreaded bubble

There's a third axis. If you've split your model into sequential stages (GPU 0 does the first half, GPU 1 the second), the naive schedule wastes almost everyone: while GPU 0 works, GPU 1 sits idle waiting for GPU 0's output, and vice versa. Pipeline parallelism fixes this by chopping each batch into micro-batches and streaming them through the stages like an assembly line -- so GPU 0 can already start micro-batch #2 while GPU 1 is still chewing on micro-batch #1:

class PipelineStage(nn.Module):
    """One stage of a pipeline-parallel model, pinned to one device."""
    def __init__(self, layers, device):
        super().__init__()
        self.layers = nn.Sequential(*layers).to(device)
        self.device = device

    def forward(self, x):
        return self.layers(x.to(self.device))

class SimplePipeline:
    def __init__(self, stages):
        self.stages = stages   # a list of PipelineStage, one per GPU

    def forward_micro_batches(self, x, n_micro_batches=4):
        micro_batch_size = x.shape[0] // n_micro_batches
        micro_batches = x.split(micro_batch_size)

        # Naive (no pipelining): stage 0 finishes ALL data, then stage 1
        # starts. Every GPU idle most of the time -- terrible utilisation.
        #
        # Pipelined: as soon as stage 0 finishes micro-batch 1, it hands
        # it to stage 1 and immediately starts micro-batch 2. The stages
        # overlap, and utilisation climbs toward "everybody busy".
        outputs = []
        for mb in micro_batches:
            h = mb
            for stage in self.stages:
                h = stage(h)
            outputs.append(h)
        return torch.cat(outputs)

The wrinkle is the bubble. At the very start of a batch, the later stages have nothing to do yet (nothing has reached them), and at the very end the earlier stages are done while the last ones finish up. Those idle moments at the edges are wasted GPU time. More micro-batches shrink the bubble (the steady-state middle grows relative to the edges) but cost you more memory, because more activations are in flight at once. Smarter schedules like 1F1B ("one forward, one backward", interleaving the passes) squeeze the bubble down further. It's a genuine engineering trade-off, not a solved problem -- which is why the big training frameworks give you knobs for it.

FSDP: shard everything (this is the big one)

Now the technique that changed what's possible on modest hardware. Go back to plain DDP for a second: every GPU holds a complete copy of the model, its gradients, AND the optimizer states. For Adam, the optimizer states are twice the size of the model itself (it keeps two running averages per parameter). That's a colossal amount of redundant storage -- every card holding an identical copy of the same gigabytes. FSDP (Fully Sharded Data Parallel), PyTorch's implementation of the idea DeepSpeed calls ZeRO, asks the obvious question: why replicate when you can shard? Split the parameters, gradients and optimizer states across the GPUs, and gather each piece back only for the brief moment you actually need it:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

def train_with_fsdp(rank, world_size):
    setup_ddp(rank, world_size)

    model = LargeTransformer(n_layers=24, d_model=2048).to(rank)

    # Wrap with FSDP -- parameters are now SHARDED across GPUs, not copied
    model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        # FULL_SHARD   : shard params + grads + optimizer  (ZeRO-3)
        # SHARD_GRAD_OP: shard grads + optimizer only      (ZeRO-2)
        # NO_SHARD     : behaves like plain DDP
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for batch in dataloader:
        output = model(batch['input'].to(rank))
        loss = criterion(output, batch['target'].to(rank))
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    cleanup()

The memory arithmetic is where it stops being abstract and starts being the difference between "impossible" and "runs on the hardware I have". Watch what sharding does to a 7-billion-parameter model:

memory_comparison = """
Model: 7B parameters  (7e9 x 4 bytes = 28 GB in FP32)

DDP (4 GPUs) -- everything replicated on every card:
  Per GPU:  28GB model + 28GB grads + 56GB optimizer (Adam) = 112 GB
  ==> does NOT fit on an 80GB A100. Full stop.

FSDP FULL_SHARD (4 GPUs) -- everything split across cards:
  Per GPU:  7GB model + 7GB grads + 14GB optimizer = 28 GB
  ==> fits on an 80GB A100, with room to spare for activations.

Same 112 GB total. The difference is you no longer store 4 identical
copies -- you store ONE copy, spread across the 4 GPUs.
"""
print(memory_comparison)

That's the whole trick, and it's genuinely close to free lunch -- almost. The cost FSDP pays is communication: because no GPU permanently holds a full layer, it has to all-gather the parameters just before each forward and backward pass, use them, then throw the gathered copy away (re-shard) to reclaim the memory. More network traffic than DDP, yes. But for large models the memory savings are so enormous that the extra communication is a bargain -- and much of it overlaps with computation anyway, the same trick DDP uses. If your model won't fit under plain DDP, FSDP is very often the first thing to reach for.

DeepSpeed and CPU offload

DeepSpeed (from Microsoft) is the other big name here, and it gave us the ZeRO vocabulary that FSDP later adopted. It does the same sharding across three "stages" (ZeRO-1 shards optimizer states, ZeRO-2 adds gradients, ZeRO-3 adds the parameters themselves), plus one more trick that borders on cheating -- offloading cold data all the way out to CPU RAM, or even NVMe disk:

# ds_config.json -- a DeepSpeed configuration
deepspeed_config = {
    "train_batch_size": 256,
    "gradient_accumulation_steps": 4,
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,                         # ZeRO-3: shard everything
        "offload_param":     {"device": "cpu"},   # params -> CPU RAM
        "offload_optimizer": {"device": "cpu"},   # optimizer -> CPU RAM
    },
}

# Usage sketch:
#   model, optimizer, _, _ = deepspeed.initialize(
#       model=model, config=deepspeed_config)
#
# Launch:  deepspeed --num_gpus=4 train.py --deepspeed_config ds_config.json

CPU offloading is what lets a hobbyist with a single decent GPU fine-tune a model that has no business fitting in their VRAM. Adam's optimizer states are twice the model size and they're only touched once per step -- so DeepSpeed parks them in your (much larger, much cheaper) system RAM and pages them in when the optimizer actually needs them. The trade is speed: shuffling gigabytes across the PCIe bus every step makes that bus the bottleneck (remember from episode #125 how brutally slow PCIe is compared to on-device memory). You can end up training 4-8x bigger models at maybe half or a third the speed. For a lot of people -- especially anyone fine-tuning rather than training from scratch -- that's a trade they'll happily take. Slower-but-possible beats fast-but-impossible every single time.

Scaling laws: how big, and on how much data?

There's a question lurking under all of this. If I've built the machinery to train a bigger model, should I? Or would I be better off training a smaller model on more data? For years this was folklore and vibes. Then Kaplan et al. (2020) measured it, and found clean power-law relationships between compute, dataset size, model size and the final loss. Hoffmann et al. (2022) -- the "Chinchilla" paper -- refined it into a rule you can actually plan around:

def compute_optimal_model(compute_budget_flops):
    """Chinchilla-style estimate: given a compute budget, roughly how
    big a model and how many tokens should you aim for?"""
    # Chinchilla's key finding: for a fixed compute budget, model size N
    # and training tokens D should scale in roughly EQUAL proportion.
    # (These constants are deliberately rough -- order of magnitude only.)
    N_optimal = (compute_budget_flops / 6) ** 0.5   # parameters
    D_optimal = compute_budget_flops / (6 * N_optimal)  # tokens
    return N_optimal, D_optimal

budgets = {
    'A100_1day':     1e20,    # one A100 for a day
    'A100_8_1week':  5.6e21,  # eight A100s for a week
    'GPT3_scale':    3.6e23,  # GPT-3 training compute
    'GPT4_scale':    2e25,    # rough estimate for GPT-4
}

for name, flops in budgets.items():
    n, d = compute_optimal_model(flops)
    print(f"{name:14s}: ~{n/1e9:.0f}B params, ~{d/1e9:.0f}B tokens")

The practical punchline reshaped the whole open-source world. Chinchilla showed that most big models built before 2022 were undertrained -- they had loads of parameters but hadn't seen nearly enough data to justify them. GPT-3, at 175B parameters trained on 300B tokens, would have been better as a smaller model fed far more data. Meta's LLaMA models took the hint: 7B to 65B parameters, trained on 1.4 trillion tokens, and they matched or beat models many times their size. This is a huge part of why open models caught up so fast -- you didn't need a 500-billion-parameter monster, you needed a right-sized model and a mountain of tokens. Distributed training is what makes chewing through a mountain of tokens possible in the first place. It all connects ;-)

TL;DR - what we split up today

  • Data parallelism (DDP) replicates the model on every GPU, splits the batch, and averages gradients with an all-reduce -- it overlaps communication with the backward pass, and it's the strategy you want 90% of the time;
  • collective communication (all-reduce, all-gather, broadcast) is the shared language of distributed training, and ring all-reduce via NCCL keeps per-GPU bandwidth constant no matter how many cards you add;
  • tensor/model parallelism splits the model itself across GPUs for models too big to fit -- it's chatty, so you keep it inside a node on fast NVLink;
  • pipeline parallelism streams micro-batches through sequential stages to keep every GPU busy, at the cost of a start/end "bubble" that smarter schedules (1F1B) shrink;
  • FSDP / ZeRO shard parameters, gradients and optimizer states instead of replicating them, cutting per-GPU memory from O(model) down to O(model / N) -- this is what makes "it doesn't fit" turn into "it fits";
  • DeepSpeed adds CPU/NVMe offloading on top, trading speed for the ability to train models 4-8x bigger than raw VRAM allows;
  • Chinchilla scaling laws say model size and training data should grow together -- many old models were undertrained, and right-sizing is why smaller open models caught up so fast.

Exercises

Three to chew on before next time. No rack of GPUs required for the reasoning -- and as always I'll walk through full solutions in the next episode.

  1. Reason through the memory wall. Take a model with 3B parameters. Using the same accounting as the memory_comparison above (FP32: 4 bytes/param for weights, 4 for gradients, 8 for Adam's two optimizer states), compute the total memory needed under plain DDP per GPU. Then compute per-GPU memory under FSDP FULL_SHARD on 8 GPUs. Write one sentence on why the total across all cards is the same in both cases, but only one of them actually fits on a 24GB consumer card.

  2. Trace a ring all-reduce by hand. Take 4 GPUs, each holding a gradient vector split into 4 chunks. On paper, walk through the scatter-reduce phase (3 steps) and note, after each step, which chunk on which GPU is now fully summed. Then do the all-gather phase. Confirm that every GPU ends with the complete sum, and count the total number of send operations. How does that count change if you go from 4 GPUs to 8?

  3. Pick the right parallelism. For each of these three scenarios, decide which strategy (or combination) you'd use and write one sentence justifying it: (a) a 200M-parameter model, but 5TB of training data and you want it done this week; (b) a 70B-parameter model that cannot fit on any single GPU you own, but you have a node with 8 NVLinked cards; (c) fine-tuning a 13B model on a single 24GB GPU where it almost fits but not quite. Hint: match the bottleneck (data? model size? a little of both?) to the technique built for it.

We now have the full picture of how a model gets trained at scale -- one GPU, many GPUs in a box, many boxes across a network. But a trained model isn't the end of the story; it's a valuable, attackable artifact sitting in the world, and there are people who'd love to poison it, steal it, or trick it into misbehaving. That whole adversarial side of the coin is where we head next time ;-)

Bedankt en tot de volgende keer!

@scipio

stem stemsocial steemstem python programming
0
0
0.000
0 comments