Distributed Training

Topics Covered

Data Parallelism

Why Single-GPU Is Not Enough

What Data Parallelism Is

How the Communication Works

Synchronous vs. Asynchronous Data Parallelism

The Effective Batch Size

Limits of Data Parallelism

Model Parallelism

Tensor Parallelism

Megatron-LM and Tensor Parallelism for Transformers

Pipeline Parallelism

GPipe and PipeDream

When to Use Which

Communication Patterns

The Practical Choice

ZeRO Stages

The Memory Waste in Standard Data Parallelism

ZeRO Stage 1: Optimizer State Sharding

ZeRO Stage 2: Gradient Sharding

ZeRO Stage 3: Parameter Sharding

ZeRO-Offload and ZeRO-Infinity

When to Use Which Stage

FSDP: PyTorch's ZeRO Equivalent

Comparison with Tensor Parallelism

Why ZeRO Matters

Gradient Checkpointing

The Problem It Solves

The Recomputation Trick

How Checkpointing Is Implemented

Granularity of Checkpointing

Selective Checkpointing

When Checkpointing Is the Right Choice

Combining with Other Techniques

Training large models requires more compute than any single GPU can provide. Distributed training spreads the work across many GPUs, but doing this well is one of the harder engineering problems in modern ML. This lesson covers the four main approaches: data parallelism, model parallelism, ZeRO optimization, and gradient checkpointing. Each addresses a different bottleneck and most production training stacks combine several.

Why Single-GPU Is Not Enough

Modern LLMs have hundreds of billions of parameters. A 70B parameter model in FP16 needs 70×109⋅2 bytes=140 GB70 \times 10^9 \cdot 2\ \text{bytes} = 140\ \text{GB} just for the weights, already exceeding the 80 GB capacity of a single H100 GPU. During training, the rough memory budget is params⋅2 (FP16)+grads⋅2+optimizer⋅8=12⋅params bytes\mathrm{params} \cdot 2\ (\text{FP16}) + \mathrm{grads} \cdot 2 + \mathrm{optimizer} \cdot 8 = 12 \cdot \mathrm{params}\ \text{bytes}. For 70B params that is 840 GB before counting activations. The total memory requirement for training a 70B model can easily exceed 1 TB.

Even if you could fit it in memory, training time matters. Training a frontier model on a single GPU would take years. To finish in weeks, you need thousands of GPUs working in parallel. Distributed training is the framework that makes this work.

Key Insight

The total memory for training a model is roughly 12 times the parameter count in bytes (weights + gradients + optimizer state). A 70B model needs over 840 GB before activations. This is why distributed training is not optional for large models.

What Data Parallelism Is

Data parallelism is the simplest and most common form of distributed training. The idea: replicate the entire model on every GPU, but give each GPU a different slice of the training data. Every GPU computes gradients on its slice, then the gradients are averaged across all GPUs and used to update the weights. After the update, all GPUs have identical weights again, and the next iteration begins.

This works because gradient descent is averaging-friendly. If you compute gradients on the full batch on one GPU, or gradients on different parts of the batch on multiple GPUs and average them, you get the same result. The mathematics are equivalent, the difference is just where the computation happens.

Data parallelism scales well as long as the model fits on each GPU. With NN GPUs, you can train on NN times the data per step (or use the same effective batch size with NN times the throughput). The throughput gain is close to linear if the gradient communication is fast enough.

Shard gradients averaged against the full batch gradient, with the ring all reduce cost and the strong scaling efficiency that ends at 58.0 pct by 512 GPUs.

How the Communication Works

The key operation in data parallelism is gradient averaging. After each GPU computes its local gradients, the gradients must be averaged across all GPUs. This is implemented as an "all-reduce" operation: every GPU sends its gradients to every other GPU and computes the sum (or average).

All-reduce has been heavily optimized over the years. Modern implementations like NCCL use ring algorithms or tree algorithms to minimize the time and bandwidth needed. On a high-speed interconnect like NVLink or InfiniBand, all-reduce can be fast enough that it does not bottleneck training.

The cost of all-reduce scales with the model size: a 70B model has 70B parameters of gradients that need to be communicated. Even at 100 Gbps, this takes meaningful time. Faster interconnects and more efficient algorithms reduce but cannot eliminate this overhead.

Synchronous vs. Asynchronous Data Parallelism

Two flavors of data parallelism exist:

  1. Synchronous: All GPUs wait for the gradient averaging to complete before moving to the next iteration. The result is mathematically equivalent to single-GPU training with a larger batch size.
  2. Asynchronous: GPUs do not wait, they update their weights as soon as they have new gradients. The result is faster per iteration but mathematically different (the updates use stale weights).

Synchronous is the default for modern LLM training because it gives reproducible results and works well with high-speed interconnects. Asynchronous was more popular when interconnects were slower and the wait time was significant. With modern hardware, asynchronous offers little benefit and adds complexity.

The Effective Batch Size

Data parallelism has a subtle effect: the effective batch size scales with the number of GPUs. If each GPU processes 32 examples and there are 256 GPUs, the effective batch size is 32⋅256=819232 \cdot 256 = 8192. This is much larger than what would fit on a single GPU.

Larger batch sizes affect training dynamics. The learning rate often needs to be scaled (the linear scaling rule says learning rate scales with batch size). The number of epochs to convergence may change. The generalization properties of the trained model can differ. Practitioners need to tune for their effective batch size, not their per-GPU batch size.

For very large effective batch sizes (millions of tokens per batch), the training dynamics can become unfavorable. There are diminishing returns to larger batches, and beyond a certain point, more GPUs do not help training speed because the additional batch size does not improve gradient quality enough.

Limits of Data Parallelism

Data parallelism has a fundamental limit: it requires the model to fit on each GPU. For models larger than a single GPU can hold, data parallelism alone is not sufficient. You need to combine it with other parallelism strategies that split the model across GPUs.

Data parallelism also has a communication cost that grows with model size. For very large models, the time to communicate gradients can become a significant fraction of total training time. Other techniques (like ZeRO, discussed later) address this by avoiding gradient duplication.

Despite these limits, data parallelism is the foundation of distributed training. Almost all real distributed training setups use data parallelism as one component, combined with other techniques to handle larger models.

Common Pitfall

When scaling data parallelism, do not forget to tune the learning rate for the effective batch size. Using a single-GPU learning rate with a 256-GPU effective batch often causes divergence in the first few steps.