Systems Saturday 18: An Introduction to Distributed Training
How modern ML scales (from someone who learned about it this week)
Forewarning: I am still learning about distributed training myself and am therefore far from an expert. That being said, if I get something wrong — please let me know! I am more excited to be wrong and learn something new rather than be confidently incorrect.
In the era of massive AI models, training a neural network has become more than simply clever algorithms. Nowadays, it sometimes seems like a battle against hardware limits. Single GPUs, once the workhorses of deep learning, now buckle under the weight of billion- or trillion-parameter behemoths. Distributed training encompasses the sophisticated orchestration of computation across multiple devices that turns impossible tasks into reality.
Training currently faces three key constraints:
Model size exceeds GPU memory
Training time becomes impossibly long with infinitely increasing parameters
Data throughput demands parallelism
These three unforgiving constraints demand we scale out and also inspire a suite of rather ingenious techniques that make it all work.
Data Parallelism
At the heart of distributed training lies data parallelism. This technique is deceptively simple and is yet incredibly powerful for scaling up training throughput. Imagine you’re training a large language model on a massive dataset. Instead of squeezing everything onto one GPU, parallelism essentially replicates the entire model across multiple GPUs. Each GPU then handles a different subset of the data batch, computing forward passes, losses, and gradients independently from the other batches.
Here’s how it unfolds step by step: During the forward pass, every GPU processes its mini-batch slice using its identical model copy. Backpropagation follows suit, yielding local gradients. The magic happens next with an “all-reduce” operation, where gradients from all GPUs are averaged and synchronized back to each device. This ensures the model stays consistent across the cluster, ready for the next optimizer step.
The appeal here is clear. Data parallelism is straightforward to implement and shines when you can crank up the global batch size without memory issues. It scales almost linearly with GPU count for compute-bound workloads, making it the go-to for many practitioners. Frameworks like PyTorch’s DistributedDataParallel (DDP), Microsoft’s DeepSpeed, Horovod from Uber, and NVIDIA’s Megatron-Core have optimized DP to minimize overhead, often integrating with high-speed interconnects like NVLink for faster communication.
But it’s not without trade-offs. Since each GPU holds a full model replica, you’re basically capped by the memory of a single device. This is generally fine for smaller models, but it’s another wall for giants like those with hundreds of billions of parameters. Communication costs also creep up as your cluster grows; that all-reduce step can become a bottleneck, especially on slower networks, leading to sub-linear scaling efficiency. Still, for many scenarios, DP remains the foundational workhorse, often combined with other strategies for even greater scale.
Model Parallelism
When a model’s sheer size — hundreds of billions or even trillions of parameters — outstrips the memory of a single GPU, data parallelism alone won’t cut it. That’s where model parallelism steps in, dividing the model itself across multiple devices to make training feasible. Rather than replicating the entire model, MP shards it strategically, allowing massive architectures to run on hardware that would otherwise be insufficient.
There are two primary flavors of model parallelism that each tackle the problem from a slightly different angle:
1. Tensor (Sharded) Parallelism
This approach slices large tensors, like weight matrices in transformer layers, across GPUs. Each device holds only a portion of the parameters, drastically cutting per-GPU memory use. During forward and backward passes, GPUs exchange data via collective operations (e.g., all-gather for activations, reduce-scatter for gradients) to compute as if the model were whole. It’s communication-intensive but scales well for wide models. Pioneered in frameworks like NVIDIA’s Megatron-LM, it’s a staple in training behemoths such as GPT and other similar LLMs, where layers are too large for one GPU.
2. Pipeline Parallelism
With this method, the model is split by layers or stages, somewhat like an assembly line across GPUs. Data flows through in micro-batches: one GPU handles early layers, then passes activations to the next, and so on. This minimizes memory per device (because each holds only a subset of layers) and keeps hardware utilized if stages are balanced. The forward pass pipelines micro-batches sequentially, with backward passes reversing the flow. However, pipeline bubbles, or the idle GPU time between batches, can reduce efficiency, often mitigated by techniques like interleaved scheduling or gradient accumulation. Tools like PyTorch’s Pipeline Parallelism or DeepSpeed’s implementation make this practical for deep networks.
Model parallelism excels in memory-bound scenarios but introduces some significant complexity: more communication overhead, potential imbalances, and the need for careful partitioning. It’s rarely used in isolation and is often layered with data parallelism for hybrid scaling, but it’s indispensable for pushing the boundaries of model size in today’s AI landscape.
Memory Optimization
Even with data and model parallelism, memory constraints can still throttle training for ultra-large-scale models. That’s where advanced memory optimization techniques like ZeRO (Zero Redundancy Optimizer) and sharding come into play. These methods mainly focus on partitioning not just the model but also the ephemeral state (optimizers, gradients, and parameters) across GPUs, slashing redundancy and enabling models to grow larger on the same hardware.
ZeRO partitions training components into stages:
ZeRO-1: Shards only the optimizer states (e.g., Adam’s momentum and variance) across devices, reducing memory by offloading these bulky elements while keeping parameters and gradients replicated.
ZeRO-2: Extends sharding to gradients as well, further cutting memory use during backpropagation, with minimal added communication.
ZeRO-3: The full monty. This shards parameters too, allowing truly massive models by dynamically gathering shards only when needed for computations. This introduces more all-gather/reduce-scatter ops but unlocks unprecedented scale.
Sharding, in a broader sense, is implemented in tools like PyTorch’s Fully Sharded Data Parallel, which mirrors ZeRO-3 by sharding parameters and reassembling them on-the-fly, or XLA’s Single Program Multiple Data for compiler-driven optimizations. These approaches minimize replication, overlapping communication with computation to hide latency.
The payoff is massive — lower per-GPU memory footprint means fitting bigger models or using larger batches for better convergence. However, they demand fast interconnects to handle the increased comms volume. In practice, ZeRO and sharding are often the glue in hybrid parallelism, making trillion-parameter training viable without extensive hardware upgrades.
Communication at Massive Scale
While parallelism strategies like data, model, and memory optimization enable scaling, possibly the most make-or-break factor in distributed training is communication. As clusters grow to hundreds or thousands of GPUs, the overhead of exchanging data can effectively eclipse computation time, turning a high-performance system into a sluggish one. Collective operations form the core of this bottleneck:
All-Reduce: Aggregates and distributes gradients across all devices, essential for synchronization in data parallelism.
All-Gather: Collects data from every GPU to each, often used in tensor parallelism for reassembling sharded tensors.
Reduce-Scatter: Reduces data and scatters results, common in sharding to distribute partial computations.
Broadcast: Sends data from one device to all others, handy for initializing or updating shared states.
These ops rely on high-bandwidth interconnects like NVIDIA’s NVLink for intra-node speed or InfiniBand/RoCE for inter-node transfers. Without them, latency understandably skyrockets. Optimization tricks include smarter scheduling to overlap communication with computation (hiding transfer time behind ongoing calculations) and compression techniques to shrink data volume. Frameworks like NCCL (NVIDIA Collective Communications Library) are well tuned for this and provide some efficient primitives.
A timeless axiom in systems engineering holds: As compute power becomes abundant and affordable, communication emerges as the dominant limiter. Ignoring it leads to poor scaling efficiency. Your 100-GPU cluster may perform like just 50 if comms are bottlenecked. Monitoring tools, such as NCCL traces or network profilers, are crucial for diagnosing and mitigating these issues in real-world deployments.
3-Dimensional Parallelism
In the quest for training ever-larger language models, no single parallelism strategy suffices. 3D parallelism composes a hybrid approach that combines data, model (tensor and pipeline), and memory optimization techniques into a generally cohesive framework. This multi-dimensional strategy is the powerhouse behind scaling to trillion-parameter models on massive GPU clusters, maximizing efficiency by addressing throughput, model size, and memory constraints simultaneously.
At its core, 3D parallelism layers the techniques orthogonally:
Data Parallelism handles batch scaling, distributing data across replicas for higher throughput.
Tensor Parallelism shards wide layers within nodes, reducing per-device memory for intra-layer computations.
Pipeline Parallelism divides deep models across devices or nodes, pipelining stages to keep hardware saturated.
Sharding (e.g., ZeRO/FSDP) overlays memory efficiency, partitioning states globally to minimize redundancy.
For instance, in setups like NVIDIA’s Megatron-Turing NLG or Meta’s OPT, you might see data parallelism across nodes, tensor parallelism within nodes for fast NVLink comms, and pipeline parallelism for depth, all with ZeRO-3 sharding to fit enormous parameters. This orchestration requires careful configuration in order to balance dimensions to avoid bottlenecks, like ensuring communication doesn’t overwhelm compute.
The result is near to linear scaling on thousands of GPUs, enabling feats like training models that would otherwise be impossible. Frameworks such as DeepSpeed, PyTorch FSDP, and Hugging Face’s Accelerate make 3D parallelism accessible, often with auto-tuning (!) for optimal configs.
Closing Thoughts
As AI models continue to grow in complexity and scale, distributed training has evolved from a niche optimization into an absolutely essential pillar of machine learning engineering. Mastering the art of partitioning data, models, and memory across vast GPU clusters is more than a technical, rote skill. It’s the key to unlocking breakthroughs in everything from natural language processing to generative AI. In this rapidly advancing field, staying ahead means embracing these distributed paradigms — often times in seemingly repackaged ways.
As aforementioned, I am still in the process of learning this myself. If you spot something that seems amiss here, please let me know!
I will also be taking a break from writing next week, but will resume on December 26!