ZeRO Optimization › Practical Considerations
Language AI HandbookAll chapters ZeRO Optimization › Practical ConsiderationsOne section · Sections in this chapter10 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Choosing the Right Stage
The three stages present increasing memory savings at increasing communication cost. The right choice depends on your hardware configuration, model size, and network topology:
-
Stage 1: The recommended default for most training runs. It provides near-zero communication overhead increase over baseline data parallelism while delivering substantial optimizer state savings. For typical model sizes and reasonable GPU counts, Stage 1 often cuts the dominant memory component (optimizer states) by a factor of 8-64x, which is frequently enough to make a training job feasible. The implementation complexity is minimal because the communication pattern is nearly identical to standard data parallelism.
-
Stage 2: Suitable when optimizer state savings alone do not free enough memory, especially at small GPU counts where the term is still large. Stage 2 adds gradient partitioning with no increase in communication volume over Stage 1, making it a free upgrade in terms of communication cost. On a small cluster of 4-8 GPUs, Stage 2 may be the difference between fitting a model and not fitting it. The gradient bucketing implementation adds some complexity but is handled transparently by libraries like DeepSpeed and FSDP.
-
Stage 3: Use when Stage 2 memory is still insufficient, when training models that simply cannot fit in bytes regardless of how the optimizer states are handled, or when scaling to hundreds or thousands of GPUs where the communication overhead becomes proportionally smaller relative to compute time. At very large scale, the 3x communication volume of Stage 3 is spread over many GPUs, and the per-device communication fraction of total time decreases as the cluster grows.
Communication Topology Awareness
ZeRO's communication patterns are sensitive to network topology. Within a single node, NVLink provides 600 GB/s of bidirectional bandwidth on A100 DGX systems. Across nodes, InfiniBand typically offers 200-400 Gb/s per port. Stage 3 generates 3x more communication than Stages 1 and 2, which can become a bottleneck on inter-node links when the model is large and the nodes are connected by slower fabrics.
A practical heuristic: if you can fit the entire training run within a single node (using Stage 1 or Stage 2 and gradient checkpointing), the all-NVLink communication makes even Stage 3 fast. Once you go multi-node, the inter-node bandwidth becomes the governing constraint, and you should prefer the lowest ZeRO stage that fits the memory budget. ZeRO++ hierarchical partitioning (hpZ) is specifically designed to address this by keeping as much communication as possible within the fast NVLink domain.
Interaction with Gradient Checkpointing
ZeRO can be combined with gradient checkpointing (activation recomputation) and the two optimizations target entirely different memory categories. ZeRO reduces model state memory (parameters, gradients, optimizer states), while gradient checkpointing reduces activation memory by not storing intermediate activations and recomputing them during the backward pass at the cost of extra forward computation (roughly 33% more compute for transformers under full checkpointing).
Because they target different memory categories, they combine additively: using both together reduces both model state memory and activation memory simultaneously without any interference. For very large models with long sequences, both optimizations are typically needed at the same time. The combination allows training with batch sizes that would be impossible using either technique alone.
Interaction with Tensor Parallelism
Tensor parallelism, covered in the preceding chapter, splits individual layers across GPUs, so that each GPU holds a horizontal slice of each weight matrix. ZeRO Stage 3 instead splits parameters across the depth dimension (each GPU owns a contiguous set of parameters from all layers). Combining the two creates a more complex partitioning scheme where the same parameter may be split in two dimensions simultaneously.
In practice, tensor parallelism and ZeRO Stage 3 are rarely combined because they both target the same problem (parameter memory) and their combination creates awkward communication schedules. A more common pattern is to use tensor parallelism within a node (exploiting NVLink bandwidth) and ZeRO Stage 2 across nodes (low communication overhead), giving each level of the hierarchy a specialization. The Megatron-DeepSpeed integration supports this hierarchical combination.
Performance Benchmarks in Context
The original ZeRO paper reported training a 100B-parameter model on 400 V100 GPUs (32 GB each) with Stage 3, achieving over 38 teraflops per GPU, roughly 49% of the theoretical peak. This was the first demonstration of 100B-scale training outside of specialized supercomputing facilities. Throughput on these large-scale jobs was within 5% of the theoretical communication-compute overlap optimum, meaning ZeRO adds essentially no wall-clock overhead compared to baseline on well-provisioned hardware.
Subsequent work on LLaMA, Falcon, and Mistral training reported similar results: Stage 2 or Stage 3 with gradient checkpointing provides the right memory-compute tradeoff for models in the 7B-70B range on A100 clusters with NVLink within nodes and InfiniBand between nodes. The typical throughput efficiency is 45-55% of theoretical peak after accounting for communication, memory bandwidth constraints, and the overhead of mixed-precision operations.
Let's compare illustrative throughput scaling across stages by modeling compute and communication times:
def estimate_throughput(
params_billions,
num_gpus,
stage,
bandwidth_gbps=200,
exposed_communication_fraction=0.08,
):
"""
Simplified throughput model.
Models weak scaling: every GPU performs one unit of compute per step.
Communication time depends on volume, bandwidth, and overlap.
Returns relative throughput vs single-GPU baseline.
"""
compute_time_per_gpu = 1.0
# Ring collectives approach their asymptotic communication volume as N grows.
comm_gb = communication_volume(params_billions, stage, num_gpus)
ring_fraction = (num_gpus - 1) / num_gpus
bandwidth_effective = bandwidth_gbps / 8 # Gbps -> GB/s
raw_comm_time = ring_fraction * comm_gb / bandwidth_effective
# Most communication overlaps with layer computation. Only the exposed
# fraction extends the step time.
exposed_comm_time = raw_comm_time * exposed_communication_fraction
step_time = compute_time_per_gpu + exposed_comm_time
return num_gpus * compute_time_per_gpu / step_time
gpu_counts_perf = [8, 16, 32, 64, 128]
stages_perf = [1, 2, 3]
perf_data = {
stage: [estimate_throughput(7, nd, stage) for nd in gpu_counts_perf]
for stage in stages_perf
}
Stages 1 and 2 track near-ideal linear throughput scaling because their communication volume matches standard data parallelism, which practitioners already accept as overhead. Stage 3 falls below ideal due to 3x communication, but still provides substantial absolute speedup and enables training models that would otherwise require infeasible per-GPU memory. On faster interconnects, all three curves shift toward the ideal line, which is why newer clusters prioritize high network bandwidth alongside raw compute.