0%
Applied AI Systems
Retrieval and Knowledge
Agents and Reasoning
Evaluation, Interpretability, and Safety
Production Operations and Safety
Multi-Agent Systems and Case Studies
Distributed Training
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 just for the weights, already exceeding the 80 GB capacity of a single H100 GPU. During training, the rough memory budget is . 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.
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 GPUs, you can train on times the data per step (or use the same effective batch size with times the throughput). The throughput gain is close to linear if the gradient communication is fast enough.
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:
- 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.
- 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 . 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.
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.