ArXiv: 1910.02054

🎯 Pitch

ZeRO shatters the memory bottleneck in data parallelism by partitioning optimizer states, gradients, and parameters across GPUs without adding communication overhead, enabling a 170B-parameter model to be trained on just 400 GPUsβ€”over 8Γ— larger than prior state-of-the-art. Its design can theoretically scale to 1 trillion parameters on today’s hardware, turning a fundamental limitation into a linear function of available devices.

1. Executive Summary

This paper introduces ZeRO (Zero Redundancy Optimizer), a system of memory optimizations that eliminates redundancies in data-parallel training to enable models with up to 1 trillion parameters on current hardware. Using transformer-based GPT-2-style models trained with mixed-precision Adam, ZeRO partitions model states β€” optimizer states, gradients, and parameters β€” across data-parallel devices through three cumulative stages (Pos: optimizer state partitioning, Pos+g: adding gradient partitioning, Pos+g+p: adding parameter partitioning) while retaining the communication volume of standard data parallelism (0% overhead for Pos and Pos+g, 50% for Pos+g+p). An implementation supporting the first two stages plus residual-state optimizations (ZeRO-100B) trains models up to 170B parameters on 400 V100 GPUs β€” over 8Γ— larger than Megatron-LM's SOTA β€” with super-linear speedup in the 64–400 GPU regime and aggregate throughput exceeding 15 Petaflops. The full ZeRO design scales model size linearly with the data-parallel degree, establishing that a trillion-parameter model can fit on 1024 GPUs given sufficient aggregate memory capacity, though end-to-end training remains gated by the compute power gap on current hardware.

2. Context and Motivation

The Core Problem: Training Large Models Exceeds Device Memory

The fundamental problem this paper addresses is starkly simple: modern deep learning models are too large to train on the hardware available, and existing solutions for distributing training across multiple devices impose fundamental trade-offs that prevent efficient scaling to the next generation of models with hundreds of billions or trillions of parameters.

This is not a theoretical concern. The paper documents a concrete, escalating crisis in NLP model training: BERT-large reached 0.3B parameters (Devlin et al., 2018), GPT-2 pushed to 1.5B (Radford et al., 2019), Megatron-LM hit 8.3B (Shoeybi et al., 2019), and T5 reached 11B (Raffel et al., 2019). Each doubling of model size delivered significant accuracy improvements, but each also strained the memory capacity of the current generation of GPUs (32GB V100). As the authors put it in Section 1:

"Basic data parallelism (DP) does not reduce memory per device, and runs out of memory for models with more than 1.4B parameters on current generation of GPUs with 32 GB memory."

This means the largest published models (8.3B–11B parameters) are already 5–8Γ— beyond what pure data parallelism can handle. The gap between model ambitions and hardware capacity is not a future concern β€” it is the current bottleneck constraining the largest models in the literature.

Why This Matters: Accuracy Scaling and the Democratization Gap

The significance of this problem extends in two directions that the paper makes clear:

First, accuracy depends on scale. The paper cites the empirical trend that larger models achieve better performance across NLP benchmarks. This is not a speculative claim β€” it was by 2020 a well-established empirical regularity that models like GPT-2 (1.5B) substantially outperformed GPT (117M), and Megatron-LM (8.3B) and T5 (11B) pushed these gains further. If model size growth stalls because training infrastructure cannot keep up, so too does progress on the accuracy frontier. The paper frames this as an existential threat to continued improvement in DL model quality β€” without a solution, the field faces a scaling ceiling.

Second, large model training is gated by expertise. The practical alternatives to data parallelism β€” model parallelism (MP) and pipeline parallelism (PP) β€” require substantial model refactoring and distributed systems expertise that many data scientists lack. The paper explicitly notes this accessibility barrier in Section 1:

"data parallelism is so easy to use that it is widely applicable across different workloads, while MP approaches today often need some work from model developers to revise their model, system developers to work out distributed operators, and existing work like Megatron-LM only supports a limited set of operators and models."

This democratization problem means that even if a technically feasible path to large model training exists, it is not practically accessible to the broader research community. The paper positions ZeRO partly as a tool for democratization β€” enabling data scientists to train 13B+ parameter models using an interface as simple as standard DP, without model refactoring.

Where Prior Approaches Fall Short

The paper identifies four major existing strategies for training large models, each with specific and well-documented limitations:

Data Parallelism (DP): Memory-Inefficient

Standard DP replicates the entire set of model states (optimizer states, gradients, parameters) on every device. For mixed-precision training with Adam, the per-device memory consumption is 16Ψ16\Psi bytes (where Ψ\Psi is the number of parameters), as detailed in Section 3.1: 2 bytes for fp16 parameters, 2 bytes for fp16 gradients, and KΨ=12ΨK\Psi = 12\Psi bytes for fp32 optimizer states (momentum, variance, and parameter master copy). This means a 1.5B parameter GPT-2 model requires approximately 24GB just for model states, exceeding the capacity of a 32GB GPU once activations and other overhead are added. Crucially, adding more devices does not reduce per-device memory — it only increases throughput by splitting the data — so DP alone cannot solve the memory crisis. The paper quantifies this failure precisely (Table 1): without ZeRO, a 128B model would require 2TB of memory per device regardless of the DP degree.

Model Parallelism (MP): Communication-Bound Beyond a Single Node

MP splits model layers vertically across devices, partitioning parameters and computation within each layer. This reduces per-device memory footprint, but at a severe cost in communication: each layer requires cross-device communication, and the bandwidth gap between intra-node (NVSwitch/NVLink, approximately 300 GB/s per link) and inter-node (InfiniBand EDR, approximately 12.5 GB/s per link) is roughly 24Γ—. As Section 1 states bluntly:

"MP cannot scale much further beyond these model sizes... they work well within a single node where the inter-GPU communication bandwidth is high, but the efficiency degrades quickly beyond a single node."

The paper backs this with a concrete measurement:

"We tested a 40B parameter model using Megatron-LM across two DGX-2 nodes and observe about 5 Tflops per V100 GPU (less than 5% of hardware peak)."

This 5% utilization is catastrophic β€” it means MP essentially wastes 95% of the hardware's compute capacity when forced to cross node boundaries. Moreover, MP reduces computational granularity: as the model is split across more devices, each device's computation becomes finer-grained, reducing the arithmetic intensity and per-GPU efficiency. The paper's key insight (Section 4.1a) is that "MP reduces the granularity of the computation while also increasing the communication overhead" β€” it is a double efficiency penalty that worsens nonlinearly with scale.

Additionally, MP has practical limitations on model applicability: Megatron-LM, the SOTA MP implementation at the time, only supported a limited set of operators and transformer architectures. Models with tied-weights, batch normalization, or complex layer interconnectivity could not easily be parallelized with existing MP frameworks.

Pipeline Parallelism (PP): Convergence and Functionality Restrictions

PP splits the model horizontally across layers, with each device handling a contiguous segment of layers. Two implementations existed at the time:

GPipe (Huang et al., 2018) partitions both parameters and activations but requires a batch size proportional to the number of pipeline stages to hide the "pipeline bubble" (idle time while later stages wait for earlier stages). This large batch size can slow model convergence β€” as Section 2.1 notes, "prior work shows, very large batch size could slow down convergence." Furthermore, GPipe still requires storing all activations across the partitioned pipeline, so activation memory is not proportionally reduced.

PipeDream (Narayanan et al., 2019) uses stale parameter copies to hide the pipeline bubble without requiring extremely large batch sizes, but this introduces training semantics that are not equivalent to standard SGD β€” the model sees different parameter versions at different pipeline stages. This has "implications on training convergence" (Section 2.1) that are not fully characterized, making it risky for production training.

Both PP approaches impose model architecture restrictions: tied-weights (where the embedding and output projection layers share parameters) are difficult to implement with horizontal layer splitting, and batch-normalization's cross-device statistics interfere with micro-batching. These are not minor inconveniences β€” they fundamentally restrict which models can benefit from PP.

CPU Offloading: Bandwidth-Bounded by PCI-E

Approaches like L2L (Pudipeddi et al., 2020) and vDNN (Rhu et al., 2016) offload model states to CPU memory, trading memory capacity for transfer latency. The paper's critique (Section 2.2.2) is quantitative: "Up to 50% of training time can be spent on GPU-CPU-GPU transfers." The PCI-E interconnect bandwidth is orders of magnitude lower than GPU device memory bandwidth, making offloading a severe throughput penalty. This makes it viable for small models or debugging but impractical for production training of 100B+ parameter models at scale.

Activation Checkpointing and Memory-Efficient Optimizers: Necessary but Insufficient

The paper acknowledges that both checkpointing (Chen et al., 2016) β€” which trades a 33% recomputation overhead for roughly activation_memory\sqrt{\text{activation\_memory}} reduction β€” and memory-efficient optimizers like Adafactor (Shazeer and Stern, 2018) reduce memory usage. However, these are complementary optimizations, not complete solutions. Section 4 notes that even with activation checkpointing, a 100B parameter GPT-like model requires approximately 60GB of activation memory for a batch size of 32 β€” already exceeding GPU capacity. Adafactor reduces optimizer state memory but at the potential cost of model convergence guarantees (since it maintains coarser-grained statistics), and even with Adafactor, the fundamental replication of gradients and parameters across DP devices remains.

How This Paper Positions Itself

The paper's intellectual positioning is best understood through the three-part insight structure in Section 4.1:

Insight (a): DP has better scaling efficiency than MP, but insight (b): DP is memory inefficient due to state replication. This creates an apparent trade-off: DP yields high efficiency but can't fit large models; MP fits large models but can't scale efficiently. Prior work accepted this dichotomy as fundamental.

Insight (c) is the breakthrough: Both DP and MP keep all model states in memory throughout training, but "not everything is required all the time." Parameters for a given layer are only needed during that layer's forward and backward propagation. Optimizer states are only needed during the update step at the end of each iteration. This temporal sparsity β€” the fact that different components of the model state are needed at different points in the training loop β€” is the key observation that makes ZeRO possible. Prior systems treated model states as monolithic, always-in-memory structures. ZeRO treats them as dynamically required resources that can be communicated on-demand.

The paper positions ZeRO-DP not as an entirely new paradigm but as a re-engineering of data parallelism that removes its memory redundancy while preserving its communication and computation efficiency. The core technique β€” partitioning model states across DP processes and communicating them only when needed β€” is enabled by the temporal insight above. ZeRO-DP does not modify the optimization algorithm, does not change model convergence properties, and does not require model architecture changes. It is, in the authors' framing, a drop-in replacement for standard DP that provides MP-like memory efficiency at DP-like throughput.

This positioning has an important strategic implication that the paper makes explicit in Section 4.1:

"ZeRO-DP retains the training efficiency of DP while achieving the memory efficiency of MP."

The claim is not that ZeRO invents something unprecedented, but rather that it engineers a system that combines the best properties of two previously antithetical approaches β€” the computational efficiency and simplicity of DP with the per-device memory footprint of MP. The key technical contribution is the dynamic communication schedule that makes this combination possible without incurring the communication overhead that cripples MP.

The parallel with existing work is made explicit: ZeRO-DP's gradient partitioning (Pg) uses the same bucketization strategy as NVIDIA's AMP optimizer for overlapping gradient communication with computation, but applies a reduce-scatter pattern instead of all-reduce because each device only needs a partition of the reduced gradients. This is an incremental engineering change built on established practice, but its memory implications β€” reducing gradient memory by a factor of NdN_d β€” are transformative.

The Residual State Problem: A Secondary Bottleneck

A subtle but important aspect of the paper's positioning is that ZeRO is not just about model states. Section 3.2 and Section 4.2 identify residual states β€” activations, temporary buffers, and fragmented memory β€” as a secondary bottleneck that becomes primary once model-state memory is optimized:

"Even though not all model states are required all the time during the training, all of these approaches maintain all the model states required over the entire training process statically."

This observation is crucial: even if ZeRO-DP reduces model-state memory to near-zero, training could still fail due to activation memory or memory fragmentation. The paper positions ZeRO-R as a complementary set of optimizations that address these "residual" memory consumers, making the full ZeRO system capable of training models that are limited only by the aggregate memory and compute capacity of the cluster.

The Trillion-Parameter Vision

The paper's ultimate framing is ambitious: ZeRO is not merely an incremental improvement for current-scale models but a pathway to trillion-parameter training. Section 4 states this explicitly:

"ZeRO powers DP to fit models with arbitrary size as long as there are sufficient number of devices to share the model states."

The math is simple: with all three ZeRO-DP stages (Pos+g+p), per-device model-state memory is 16Ξ¨Nd\frac{16\Psi}{N_d} bytes. For a 1 trillion parameter model with mixed-precision Adam (K=12K = 12, total 16Ξ¨=1616\Psi = 16 TB), using 1024 GPUs yields 16 GB per GPU β€” well within the 32GB capacity of V100 GPUs. The paper frames this as a system breakthrough that removes the memory barrier to trillion-parameter training, while acknowledging that the compute barrier (end-to-end training time exceeding one year on current hardware) remains unresolved. This careful distinction β€” memory feasibility versus compute practicality β€” is characteristic of the paper's precise scoping.

3. Technical Approach

This is primarily a systems design paper whose core idea is that model states in data-parallel training are redundantly replicated across all devices, and that by partitioning these states and communicating them only when needed through a dynamic schedule, one can achieve the memory efficiency of model parallelism while retaining the computational and communication efficiency of data parallelism.

3.1 Reader Orientation

ZeRO is a memory optimization system that modifies how a standard data-parallel training loop stores and communicates model states. It solves the problem that data parallelism replicates all optimizer states, gradients, and parameters on every GPU β€” wasting memory β€” by instead partitioning these states across GPUs and communicating them on-demand during training, so that the total memory available scales linearly with the number of GPUs rather than being bottlenecked by a single GPU's capacity.

3.2 Big-Picture Architecture (Diagram in Words)

The ZeRO system sits as a layer between the deep learning framework (PyTorch) and the model training loop, intercepting and modifying how model states are stored and communicated. It consists of two families of optimizations that operate on different categories of memory consumers:

1. ZeRO-DP (Data Parallelism optimizations) β€” targets model states (optimizer states, gradients, parameters) and has three cumulative stages:

  • Pos (Optimizer State Partitioning): Each GPU stores only $1/N_d$ of the Adam optimizer states (momentum and variance). During the update step, each GPU updates only its partition of parameters, then an all-gather collects the full updated parameters. Memory reduction: 4Γ—.
  • Pos+g (Add Gradient Partitioning): Each GPU stores only the gradients needed to update its parameter partition. A reduce-scatter operation (not all-reduce) distributes reduced gradients to the responsible GPUs, then gradients are released. Memory reduction: 8Γ—.
  • Pos+g+p (Add Parameter Partitioning): Each GPU stores only its parameter partition. Parameters from other partitions are broadcast on-demand during forward and backward propagation, then discarded. Memory reduction: linear in $N_d$ (e.g., 64Γ— with 64 GPUs).

All three stages use a dynamic communication schedule that maintains the same communication volume as standard DP for Pos and Pos+g, with only a 1.5Γ— increase for Pos+g+p.

2. ZeRO-R (Residual state optimizations) β€” targets activations, temporary buffers, and fragmented memory:

  • Pa (Partitioned Activation Checkpointing): In model-parallel settings, partitions activation checkpoints across GPUs instead of replicating them, using all-gather to reconstruct when needed. Can optionally offload to CPU (Pa+cpu).
  • CB (Constant-size Buffers): Caps temporary communication buffers at a fixed size regardless of model scale, trading a small efficiency loss for bounded memory.
  • MD (Memory Defragmentation): Pre-allocates contiguous memory for long-lived tensors (activation checkpoints, gradients) and copies them there as they are produced, preventing fragmentation-induced OOM.

Information flows as follows during a training step with full ZeRO (Pos+g+p + ZeRO-R):

  1. Forward pass: Each GPU broadcasts its parameter partition to all other GPUs just before that layer is needed; the receiving GPUs use the broadcast parameters for computation, then discard them. Activation checkpoints are partitioned and stored.
  2. Backward pass: Parameters are re-broadcast in reverse order. As each layer's gradients become available, a reduce-scatter operation distributes the reduced gradients to the parameter-owning GPUs. Activation checkpoints are all-gathered on-demand for gradient computation.
  3. Update step: Each GPU updates only its parameter partition using its local optimizer states and the reduced gradients it received. No communication needed here.
  4. Next iteration setup: An all-gather collects the updated parameters (Pos+g) or the cycle repeats with on-demand broadcasting (Pos+g+p).

3.3 Roadmap for the Deep Dive

The explanation follows the natural decomposition of memory consumers in training, mirroring the paper's "Where Did All the Memory Go?" diagnosis in Section 3:

  • First, the model-state memory analysis (Section 3.1 of the paper): a precise accounting of why mixed-precision Adam requires $16\Psi$ bytes per device, establishing the baseline that ZeRO must reduce.
  • Second, the residual-state memory analysis (Section 3.2 of the paper): activations, temporary buffers, and fragmentation β€” the secondary memory consumers that become bottlenecks after model states are optimized.
  • Third, the ZeRO-DP mechanism (Section 5): the three-stage partitioning of optimizer states, gradients, and parameters, each stage's communication schedule, and why the total communication volume remains competitive with standard DP.
  • Fourth, the ZeRO-R mechanism (Section 6): partitioned activation checkpointing, constant-size buffers, and memory defragmentation, including how they interact with model parallelism.
  • Fifth, the communication analysis (Sections 7–8): a rigorous accounting of the total bytes moved per training step for both ZeRO-DP and ZeRO-R, compared against standard DP and MP baselines.
  • Sixth, the scaling analysis (Section 9): how the memory reduction translates into feasible model sizes on specific hardware configurations, culminating in the trillion-parameter projection.

This order is chosen because each component builds on the preceding memory analysis β€” you cannot understand why ZeRO partitions states without first understanding how much memory each state consumes, and you cannot appreciate the communication analysis without first understanding what is being communicated and when.

3.4 Detailed, Sentence-Based Technical Breakdown

Model-State Memory Analysis: The $16\Psi$ Baseline

The paper builds its entire case on a precise, quantified diagnosis of where memory goes during mixed-precision training with the Adam optimizer. This analysis (Section 3.1) establishes the $16\Psi$ baseline that ZeRO-DP must reduce, and it is essential to understand before any of the optimizations make sense.

Mixed-precision training (Micikevicius et al., 2017) is the standard approach for training large models on NVIDIA V100 GPUs because it enables use of the high-throughput Tensor Core units, which operate on fp16 data. The protocol works as follows: parameters and activations are stored in fp16 (2 bytes per value) for the forward and backward passes, making computation fast and memory-efficient for those phases. However, the weight update at the end of each iteration requires higher precision to avoid numerical underflow β€” the updates are often small relative to the parameter magnitudes, and fp16's limited mantissa cannot represent these small deltas accurately. Therefore, the mixed-precision optimizer maintains fp32 master copies of the parameters (4 bytes per value) and performs the update in fp32.

For the Adam optimizer specifically (Kingma and Ba, 2015), two additional fp32 states are required per parameter: the first moment estimate (exponentially decaying average of past gradients, often called momentum) and the second moment estimate (exponentially decaying average of past squared gradients, often called variance). These are the "optimizer states" that enable Adam's adaptive per-parameter learning rates.

The paper presents the memory accounting as a straightforward sum:

"Mixed precision training of a model with $\Psi$ parameters using Adam requires enough memory to hold an fp16 copy of the parameters and the gradients, with memory requirements of $2\Psi$ and $2\Psi$ bytes respectively. In addition, it needs to hold the optimizer states: an fp32 copy of the parameters, momentum and variance, with memory requirements of $4\Psi$, $4\Psi$, and $4\Psi$ bytes, respectively."

Let $K$ denote the memory multiplier of the optimizer states β€” the additional memory (beyond fp16 parameters and gradients) required to store them. For mixed-precision Adam, $K = 12$ because the three fp32 states (parameter master copy, momentum, variance) each require $4\Psi$ bytes, totaling $12\Psi = K\Psi$.

The total model-state memory per device is therefore:

Mmodel=2Ψ+2Ψ+KΨ=(4+K)ΨM_{\text{model}} = 2\Psi + 2\Psi + K\Psi = (4 + K)\Psi

For mixed-precision Adam with $K = 12$, this gives $M_{\text{model}} = 16\Psi$ bytes.

What it computes: the per-device memory consumed by optimizer states, gradients, and parameters during mixed-precision Adam training. The $2\Psi$ term is the fp16 parameter storage, the second $2\Psi$ is the fp16 gradient storage, and $K\Psi = 12\Psi$ is the fp32 optimizer state storage (parameter master copy, momentum, and variance, each at $4\Psi$).

Why this form: the decomposition separates memory that scales with $K$ (optimizer-specific) from memory that is universal (parameters and gradients are needed regardless of optimizer choice). This matters because the ZeRO-DP stages target different components: Pos addresses only the $K\Psi$ term (optimizer states), Pos+g addresses the $2\Psi$ gradient term, and Pos+g+p addresses the full $16\Psi$. A simpler memory model that bundled everything together would obscure which optimization affects which portion.

Concrete example: For GPT-2 with $\Psi = 1.5$ billion parameters, this formula yields $16 \times 1.5 \times 10^9 = 24$ GB for model states alone. Since a V100 GPU has only 32 GB total memory, and activations plus other overhead consume additional gigabytes (the paper notes approximately 60 GB of activation memory for this model at sequence length 1K and batch size 32, before checkpointing), this model cannot be trained on a single GPU β€” explaining the observation from Section 3 that "a 1.5B parameter GPT-2 model requires 3GB of memory for its weights in 16-bit precision, yet, it cannot be trained on a single GPU with 32GB memory."

In data-parallel training, this $16\Psi$ bytes is replicated on every GPU. Adding more GPUs increases aggregate memory capacity but does not reduce per-device memory β€” the constraint is per-device capacity. This is the precise nature of the memory redundancy that ZeRO eliminates.

Residual-State Memory Analysis: Activations, Buffers, Fragmentation

Once model states are accounted for, the paper identifies three additional memory consumers that collectively form the "residual state" memory (Section 3.2). These become dominant once ZeRO-DP reduces model-state memory, and ZeRO-R is designed specifically to address them.

Activations. During the forward pass of a neural network, the intermediate outputs (activations) of each layer must be stored for use during the backward pass, where they are needed to compute gradients via the chain rule. For transformer models, the total activation memory scales as:

Mactivations∝layersΓ—hidden_dimΓ—sequence_lengthΓ—batch_sizeM_{\text{activations}} \propto \text{layers} \times \text{hidden\_dim} \times \text{sequence\_length} \times \text{batch\_size}

The paper provides a concrete number: "the 1.5B parameter GPT-2 model trained with sequence length of 1K and batch size of 32 requires about 60 GB of memory" for activations alone. This is larger than the 24 GB model-state memory, making activations the dominant memory consumer for this model size.

Activation checkpointing (Chen et al., 2016) is the standard mitigation: instead of storing all activations, store only a subset (checkpoints) and recompute the discarded activations during the backward pass. The paper notes this "reduce[s] the activation memory by approximately the square root of the total activations at the expense of 33% re-computation overhead." For the GPT-2 example, this reduces activation memory from 60 GB to approximately 8 GB. However, the paper immediately points out the limitation: "a GPT-like model with 100 billion parameters requires around 60 GB of memory for batch size 32, even when using activation checkpointing." Checkpointing helps but does not solve the problem at scale.

Temporary buffers. Operations like gradient all-reduce and gradient norm computation typically fuse all gradients into a single flattened buffer before performing the collective operation. This is an optimization for communication efficiency β€” larger messages achieve higher bandwidth utilization in all-reduce operations β€” but it creates a memory cost proportional to model size. The paper notes: "for a model with 1.5B parameters, a flattened fp32 buffer would require 6 GB of memory." At 100B parameters, a fp32 flattened buffer would require 400 GB β€” completely infeasible.

Memory fragmentation. The paper identifies a subtle but critical failure mode: "it is possible to run out of usable memory even when there is plenty of available memory. This can happen with memory fragmentation." During training with activation checkpointing, there is an interleaving of short-lived tensors (discarded activations during forward, activation gradients during backward) and long-lived tensors (checkpointed activations, parameter gradients). The memory allocator must find contiguous blocks of memory to satisfy allocation requests; if the free memory is fragmented into small non-contiguous gaps, an allocation can fail even if the total free memory exceeds the requested size. The paper reports: "We observe significant memory fragmentation when training very large models, resulting in out of memory issue with over 30% of memory still available in some extreme cases." This means that in practice, the usable memory is substantially less than the reported free memory β€” a 30% fragmentation rate effectively reduces a 32 GB GPU to approximately 22 GB of usable capacity.

ZeRO-DP Core Mechanism: Partitioning and Dynamic Communication

The insight that enables ZeRO-DP is stated in Section 4.1(c): "Both DP and MP keep all the model states needed over the entire training process, but not everything is required all the time. For example, parameters corresponding to each layer is only needed during the forward propagation and backward propagation of the layer." Standard training treats model states as monolithic β€” once loaded into memory, they stay there. ZeRO-DP treats them as resources with temporal locality, communicating them to the devices that need them only when they are needed, then releasing the memory.

The paper presents this as a partitioning strategy combined with a dynamic communication schedule. Each of the three stages partitions a different category of model state, and each stage's communication schedule is designed to maintain (or minimally increase) the total communication volume compared to standard DP.

Stage 1 β€” Pos: Optimizer State Partitioning

What gets partitioned: The fp32 optimizer states (parameter master copy, momentum, variance), which collectively consume $K\Psi = 12\Psi$ bytes.

Partitioning scheme: For a data-parallel degree of $N_d$, the optimizer states are divided into $N_d$ equal partitions. The $i$-th data-parallel process stores and updates only the $i$-th partition of the optimizer states, corresponding to $1/N_d$ of the total parameters.

What changes in the training loop: During the backward pass and gradient reduction, nothing changes β€” gradients are still fully reduced (via all-reduce in this stage). During the update step, each process uses its local optimizer states to update only its partition of parameters. After the update, an all-gather operation collects the updated parameter partitions from all processes, so every process ends up with a complete, identical copy of the updated parameters for the next iteration's forward pass.

Memory consumption after Pos:

MPos=2Ψ+2Ψ+KΨNd=4Ψ+12ΨNdM_{\text{Pos}} = 2\Psi + 2\Psi + \frac{K\Psi}{N_d} = 4\Psi + \frac{12\Psi}{N_d}

where $2\Psi$ is the fp16 parameter storage, $2\Psi$ is the fp16 gradient storage, and $\frac{12\Psi}{N_d}$ is the partitioned fp32 optimizer state storage per device.

What it computes: the per-device model-state memory after optimizer state partitioning. The optimizer state term is divided by $N_d$ while the parameter and gradient terms remain at full replication.

Why this form: the division by $N_d$ reflects that each device stores only its share of optimizer states. The parameters and gradients are still replicated because this stage does not partition them β€” parameters are still all-gathered after each update, and gradients are still fully reduced via all-reduce.

For $N_d = 64$ and $\Psi = 7.5$ billion parameters (the example in Figure 1):

MPos=4Γ—7.5+12Γ—7.564=30+1.4=31.4Β GBM_{\text{Pos}} = 4 \times 7.5 + \frac{12 \times 7.5}{64} = 30 + 1.4 = 31.4 \text{ GB}

Compared to the baseline $16 \times 7.5 = 120$ GB, this is nearly a 4Γ— reduction. The paper states this as a general result: "when $N_d$ is large, the memory requirement on model states reduces from $4\Psi + 12\Psi = 16\Psi$ bytes to $4\Psi + \frac{12\Psi}{N_d} \approx 4\Psi$ bytes, leading to a 4x reduction."

Communication volume: The all-gather of updated parameters incurs a communication volume of $\Psi$ elements. Combined with the existing all-reduce of gradients (which incurs $2\Psi$ total data movement in standard DP β€” $\Psi$ for reduce-scatter plus $\Psi$ for all-gather), the total communication volume remains $2\Psi$, exactly the same as baseline DP. The paper is explicit: "ZeRO-DP incurs no additional communication using Pos and Pg" (Section 7).

Stage 2 β€” Pos+g: Adding Gradient Partitioning

What gets additionally partitioned: The fp16 gradients, which consume $2\Psi$ bytes.

Key insight: Since each process updates only its partition of parameters (via Pos), it only needs the reduced gradients for those parameters. There is no need for every process to have a complete copy of all reduced gradients β€” the gradients for parameters owned by other processes are irrelevant to this process's update.

What changes in the communication pattern: Instead of performing an all-reduce on gradients (which would give every process a full copy of the reduced gradients), ZeRO performs a reduce-scatter operation. In a reduce-scatter, the gradients corresponding to different parameter partitions are reduced (summed across all data-parallel processes) but the result is scattered such that each process receives only the reduced gradients for its own parameter partition. This is exactly what Pos+g needs: each process gets the reduced gradients for the parameters it is responsible for updating.

Once a process has received its reduced gradient partition and no longer needs the other gradients, those gradients' memory can be freed. The gradient memory per device thus drops from $2\Psi$ to $\frac{2\Psi}{N_d}$.

Implementation detail β€” bucketization: The paper notes that naively performing reduce-scatter on each gradient individually would be inefficient due to many small communication operations. Instead, ZeRO uses a "bucketization strategy, where we bucketize all the gradients corresponding to a particular partition, and perform reduction on the entire bucket at once." This is the same technique used in NVIDIA's AMP optimizer for overlapping gradient all-reduce with computation, but applied to a reduce-scatter pattern at partition boundaries.

Memory consumption after Pos+g:

MPos+g=2Ψ+2ΨNd+KΨNd=2Ψ+14ΨNdM_{\text{Pos+g}} = 2\Psi + \frac{2\Psi}{N_d} + \frac{K\Psi}{N_d} = 2\Psi + \frac{14\Psi}{N_d}

where $2\Psi$ is the still-replicated fp16 parameters, $\frac{2\Psi}{N_d}$ is the partitioned fp16 gradients, and $\frac{12\Psi}{N_d}$ is the partitioned optimizer states.

What it computes: the per-device model-state memory after both optimizer state and gradient partitioning. Both gradient and optimizer terms are divided by $N_d$, while parameters remain replicated.

Why this form: the gradients are now partitioned for the same reason optimizer states are β€” each process only needs the gradients for its parameter partition. The factor of 14 in the numerator reflects the sum of $K = 12$ (optimizer states) plus 2 (gradients), which are the components that benefit from partitioning.

For the 7.5B parameter example with $N_d = 64$:

MPos+g=2Γ—7.5+14Γ—7.564=15+1.64=16.6Β GBM_{\text{Pos+g}} = 2 \times 7.5 + \frac{14 \times 7.5}{64} = 15 + 1.64 = 16.6 \text{ GB}

When $N_d$ is large, the memory approaches $2\Psi$, an 8Γ— reduction from the baseline $16\Psi$. The paper states: "the memory requirement of model states reduces from $2\Psi + 14\Psi = 16\Psi$ bytes to $2\Psi + \frac{14\Psi}{N_d} \approx 2\Psi$ bytes, leading to a 8x reduction."

Communication volume: The reduce-scatter of gradients incurs a communication volume of $\Psi$ elements (the same as the reduce-scatter component of an all-reduce). The all-gather of updated parameters incurs another $\Psi$. Total: $2\Psi$, identical to baseline DP. This is a crucial result: Pos+g provides 8Γ— memory reduction at zero additional communication cost.

When gradients are communicated: The reduce-scatter happens during the backward pass, as gradients become available layer by layer. By using bucketization and overlapping communication with computation (like AMP's gradient all-reduce overlap), the communication cost is largely hidden behind the backward computation.

Stage 3 β€” Pos+g+p: Adding Parameter Partitioning

What gets additionally partitioned: The fp16 parameters themselves, which consume $2\Psi$ bytes.

Key insight: Just as not all gradients and optimizer states are needed by every process at all times, not all parameters are needed at all times. A process only needs the parameters for a given layer during that layer's forward and backward propagation. Between these times, the parameter memory is idle β€” it is allocated but unused.

What changes in the communication pattern: Each process stores only its parameter partition permanently. When a process needs parameters from another partition (for computing a layer's forward or backward pass), the owning process broadcasts those parameters to all other processes. The receiving processes use these broadcast parameters for the computation, then discard them from memory β€” they are not stored permanently.

The broadcast happens in a pipelined manner: before computing the forward pass for a layer in partition $j$, process $j$ broadcasts the parameters for that layer to all processes. All processes compute the forward pass using these parameters, then discard them. The process repeats for each layer partition. During the backward pass, the broadcasts happen again, in reverse order (since the backward pass traverses layers in reverse).

Memory consumption after Pos+g+p:

MPos+g+p=2ΨNd+2ΨNd+KΨNd=(4+K)ΨNd=16ΨNdM_{\text{Pos+g+p}} = \frac{2\Psi}{N_d} + \frac{2\Psi}{N_d} + \frac{K\Psi}{N_d} = \frac{(4 + K)\Psi}{N_d} = \frac{16\Psi}{N_d}

where all three components β€” parameters, gradients, and optimizer states β€” are now partitioned across $N_d$ processes.

What it computes: the per-device model-state memory with full partitioning. Every term is divided by $N_d$, meaning memory scales inversely with the number of data-parallel devices.

Why this form: this is the logical completion of the ZeRO-DP insight β€” if every component of model state is only needed at specific times, all of them can be partitioned and communicated on-demand. The memory becomes proportional to $\Psi / N_d$ rather than $\Psi$, meaning model size can scale linearly with available aggregate memory.

For the 7.5B example with $N_d = 64$:

MPos+g+p=16Γ—7.564=1.88Β GBM_{\text{Pos+g+p}} = \frac{16 \times 7.5}{64} = 1.88 \text{ GB}

This is a 64Γ— reduction from the 120 GB baseline, fitting a 7.5B model into under 2 GB of model-state memory per GPU.

Communication volume: The paper provides a careful accounting (Section 7.2.2). During the forward pass, each parameter partition is broadcast once (by its owning process to all others), incurring a total communication volume of $\Psi$ elements across all broadcasts. The same broadcasts happen again during the backward pass in reverse order, adding another $\Psi$. The reduce-scatter of gradients adds $\Psi$. Total: $3\Psi$, which is 1.5Γ— the baseline DP volume of $2\Psi$.

The paper emphasizes a subtle but critical point about when this communication happens: "we reschedule the parameter all-gather by spreading it across the entire forward propagation, and discarding the parameters once they have been used." This pipelining means the communication is not a single burst at the start of the iteration (which would cause a large memory spike as all parameters are simultaneously present), but rather is interleaved with computation. At any given moment, only the parameters for the current layer (plus possibly some prefetched parameters for the next layer) are in memory, which is what enables the $\Psi / N_d$ memory footprint despite the total communicated volume being $3\Psi$.

When parameters are communicated: Layer by layer during forward and backward propagation, interleaved with computation. The owning process broadcasts its partition's parameters just before the layer is needed for computation; after the layer's computation completes, non-owning processes discard those parameters.

ZeRO-R: Residual State Optimizations

Once ZeRO-DP reduces model-state memory, the residual states (activations, buffers, fragmentation) become the dominant memory consumers. ZeRO-R consists of three complementary optimizations targeting each of these.

Pa: Partitioned Activation Checkpointing

The problem: In model-parallel training, even though parameters are partitioned across MP devices, activations must be replicated. For a linear layer split vertically across two GPUs, each GPU computes only half the output dimensions, but both GPUs need the full input activation to compute their halves. Standard MP therefore stores a full copy of the input activation on every MP device β€” a memory redundancy proportional to the MP degree $N_m$.

What Pa does: Instead of storing a replicated copy of each activation checkpoint on every MP device, Pa stores only a $1/N_m$ partition of the checkpoint on each device. When the activation is needed during the backward pass, an all-gather operation reconstructs the full activation from all MP devices. The activation is then used for gradient computation and immediately discarded, returning to the partitioned storage state for the next checkpoint.

The paper provides a concrete example: "Consider training a 100B model with a batch size of 32, sequence length of 1024 and a MP degree of 16. If we checkpoint a single activation for each transformer layer, it would require about 33 GB of memory per GPU just to store the activation checkpoints. But with Pa in ZeRO, it can be reduced to about 2 GB per GPU." This is a 16Γ— reduction β€” directly proportional to the MP degree $N_m = 16$.

Communication analysis for Pa: In Megatron-LM-style MP (the baseline), each transformer block performs two all-reduce operations in the forward pass, two in the forward recomputation, and two in the backward pass, for a total of 6 all-reduce operations per block. Each all-reduce of a tensor of size $\text{seq\_len} \times \text{hidden\_dim}$ incurs $2 \times \text{seq\_len} \times \text{hidden\_dim}$ bytes of data movement, so the total MP communication per block is $12 \times \text{seq\_len} \times \text{hidden\_dim}$.

Pa adds one additional all-gather operation per transformer block (to reconstruct the partitioned activation checkpoint before the forward recomputation). An all-gather of size $\text{seq\_len} \times \text{hidden\_dim}$ incurs exactly $\text{seq\_len} \times \text{hidden\_dim}$ bytes of data movement (since all-gather communication volume equals the message size). Therefore, Pa adds $\text{seq\_len} \times \text{hidden\_dim}$ to the baseline $12 \times \text{seq\_len} \times \text{hidden\_dim}$, an increase of approximately 8.3%. The paper rounds this to "less than 10% of the original communication volume for model parallelism."

Pa+cpu: CPU offloading variant. For extremely large models where even partitioned activation checkpoints exceed device memory, Pa+cpu offloads the partitioned checkpoints to CPU memory. During the backward pass, they are transferred back to the GPU before the all-gather. This adds 2Γ— data movement to and from CPU (since each checkpoint must be written to CPU and later read back) but reduces GPU activation memory to nearly zero. The paper frames this as a last-resort option: "In extreme cases where DP communication volume is the major bottleneck due to a small batch size even with Pa, Pa+cpu can improve efficiency by increasing the batch size as long as the CPU data transfer overhead is less than the DP communication volume overhead."

Interaction with data parallelism: The paper identifies a powerful synergy: Pa reduces activation memory by a factor of $N_m$, enabling a proportionally larger batch size. Since DP communication volume is inversely proportional to batch size (larger batches mean fewer communication rounds per sample), a 16Γ— increase in batch size from MP degree 16 reduces DP communication by 16Γ—. This can more than compensate for the <10% increase in MP communication volume, making Pa a net throughput win even though it adds communication.

CB: Constant-Size Buffers

The problem: High-performance communication libraries like NVIDIA Apex and Megatron-LM fuse tensors into a single flattened buffer before all-reduce operations because larger messages achieve higher bandwidth utilization. However, the size of this fused buffer is proportional to the total number of parameters β€” for a 3B parameter model, a fp32 fused buffer requires 12 GB, and this grows linearly with model size. The paper states this concisely: "the memory overhead of the fused buffers is proportional to the model size, and can become inhibiting."

What CB does: Instead of fusing all parameters into a single buffer (which grows with model size), ZeRO uses a fixed-size fused buffer that does not depend on model scale. The buffer size is chosen to be "large enough to remain efficient" β€” meaning large enough to achieve good bandwidth utilization on the interconnect β€” but bounded so it does not become a memory bottleneck. The paper does not specify the exact buffer size in bytes, but the principle is clear: break the scaling relationship between buffer memory and model parameters.

Trade-off acknowledged: Using a constant-size buffer means that for very large models, the communication is split into multiple smaller all-reduce or reduce-scatter operations rather than one large one. This slightly reduces communication efficiency (since smaller messages have lower bandwidth utilization), but the memory savings are essential for fitting the model at all.

MD: Memory Defragmentation

The problem: Activation checkpointing creates an interleaving of short-lived tensors (discarded activations) and long-lived tensors (checkpointed activations) during the forward pass. Similarly, during the backward pass, parameter gradients are long-lived while activation gradients are short-lived. This interleaving causes memory fragmentation β€” the free memory becomes divided into many small non-contiguous gaps, making it impossible to allocate large contiguous blocks even when total free memory is sufficient. The paper reports that fragmentation can cause "out of memory issue with over 30% of memory still available in some extreme cases."

What MD does: ZeRO pre-allocates contiguous memory buffers for the two categories of long-lived tensors: activation checkpoints and parameter gradients. As these tensors are produced during training, they are copied into the pre-allocated buffers rather than being allocated freshly from the general memory pool. This ensures that long-lived tensors occupy contiguous, defragmented memory, while short-lived tensors (which are allocated and freed rapidly) use the remaining memory without interfering with the long-lived allocations.

The paper emphasizes that this is done "on-the-fly" β€” it is not a periodic garbage collection or memory compaction pass, but rather a proactive allocation strategy applied throughout training. "MD not only enables ZeRO to train larger models with larger batch sizes, but also improves efficiency when training with limited memory," since the memory allocator spends less time searching for contiguous free blocks.

Communication Volume: Complete Accounting

The communication analysis in Sections 7 and 8 is essential because ZeRO's central claim is that the memory reduction comes at minimal (or zero) communication cost compared to standard DP. The paper provides a rigorous accounting of the total bytes moved per training step for each configuration.

Baseline DP Communication

Standard DP uses an all-reduce to average gradients across all data-parallel processes. Modern all-reduce implementations use a two-phase algorithm:

  1. Reduce-scatter: each process sends data to and receives data from other processes in a pattern that results in each process holding a different partition of the reduced (summed) data. Total data movement: $\Psi$ elements.
  2. All-gather: each process broadcasts its reduced partition to all others, so every process ends up with the complete reduced data. Total data movement: $\Psi$ elements.

Total communication volume for standard DP per training step: $2\Psi$ elements. The paper notes that "both reduce-scatter and all-gather are implemented using a pipelined approach," meaning the communication is overlapped across processes for efficiency.

ZeRO-DP Communication (Pos+g)

With Pos+g, the gradient communication changes from all-reduce to reduce-scatter only β€” the all-gather step is eliminated because each process only needs its own reduced gradient partition. After the update step (where each process updates its parameter partition using its local optimizer states), an all-gather is performed to distribute the updated parameters to all processes:

  1. Reduce-scatter of gradients: $\Psi$ elements.
  2. All-gather of updated parameters: $\Psi$ elements.

Total: $2\Psi$ β€” identical to baseline DP. This confirms the paper's claim that Pos+g provides 8Γ— memory reduction at zero communication overhead.

ZeRO-DP Communication (Pos+g+p)

With full parameter partitioning, parameters are no longer all-gathered once per iteration. Instead, they are broadcast on-demand during forward and backward propagation:

  1. Forward pass: each parameter partition is broadcast once by its owning process. Total across all broadcasts: $\Psi$ elements.
  2. Backward pass: each parameter partition is broadcast again in reverse order. Total: $\Psi$ elements.
  3. Reduce-scatter of gradients: $\Psi$ elements.

Total: $3\Psi$ β€” a 50% increase over baseline DP ($2\Psi$), while providing memory reduction proportional to $N_d$.

The paper frames this as a favorable trade-off: "While this may seem to incur significant communication overhead at first glance, we show that this approach only increases the total communication volume of a baseline DP system to 1.5x, while enabling memory reduction proportional to $N_d$." For $N_d = 64$, the memory reduces by 64Γ— while communication increases by only 1.5Γ— β€” a highly asymmetric trade-off in favor of memory savings.

Why the communication increase is modest: The parameter broadcasts during forward and backward propagation are interleaved with computation β€” while one layer's parameters are being broadcast, computation can proceed on a previous layer. This pipelining hides much of the communication latency. Additionally, each broadcast communicates only a fraction of the total parameters (those for a specific layer), keeping individual message sizes manageable.

ZeRO-R Communication (Pa)

As analyzed in Section 8, Pa adds one all-gather per transformer block for the activation checkpoint reconstruction, costing $\text{seq\_len} \times \text{hidden\_dim}$ bytes per block. Compared to the baseline MP communication of $12 \times \text{seq\_len} \times \text{hidden\_dim}$ per block, this is less than a 10% increase. The paper notes that this analysis is specific to the Megatron-LM MP strategy and transformer architecture; other architectures would require re-analysis.

Design Choices and Their Justifications

Several design decisions in ZeRO are non-obvious and warrant explicit explanation:

Why partitioning instead of offloading? The paper rejects CPU offloading as a primary strategy because "up to 50% of training time can be spent on GPU-CPU-GPU transfers" due to PCI-E bandwidth constraints. Partitioning keeps all states on GPUs (just distributed across them), using the much higher inter-GPU bandwidth (NVSwitch/NVLink within a node, InfiniBand across nodes). The only exception is Pa+cpu, which offloads activation checkpoints to CPU only when device memory is so constrained that the alternative is OOM β€” and even then, it is a secondary strategy behind Pa.

Why three stages instead of going directly to full partitioning? The staged design is motivated by the communication implications. Pos and Pos+g have zero communication overhead compared to DP, making them strict improvements that should always be applied. Pos+g+p adds a 50% communication overhead, which might not be beneficial for models that already fit with Pos+g. The staged design lets users choose the level of memory savings appropriate for their model size, avoiding unnecessary communication overhead for smaller models.

Why reduce-scatter instead of all-reduce for gradients? The paper's key insight is that Pos (optimizer state partitioning) means each process only updates a fraction of parameters. This creates the opportunity to eliminate redundant gradient communication β€” if process $i$ only needs gradients for parameter partition $i$, there is no reason to send it gradients for other partitions. The reduce-scatter operation is the minimal communication pattern that achieves this: it reduces gradients across all processes but scatters the result so each process receives only its needed partition.

Why bucketize gradients instead of communicating each gradient individually? Individual gradient communication would incur high latency overhead from many small messages. Bucketization β€” grouping gradients for the same parameter partition into a single large message β€” achieves higher bandwidth utilization and enables overlapping communication with backward computation, following the established pattern from NVIDIA's AMP optimizer.

Why use all-gather for Pos+g parameter synchronization instead of broadcast? In Pos+g, after each process updates its parameter partition, every process needs the complete set of updated parameters for the next iteration (since parameters are still replicated in Pos+g). An all-gather is the natural collective for this: each process contributes its updated partition, and all processes receive the concatenation of all partitions. A series of broadcasts would achieve the same result but with different (and potentially less efficient) communication patterns.

Why spread parameter broadcasts across the forward/backward pass in Pos+g+p? If all parameters were gathered at once at the start of each pass (like a standard all-gather), the memory footprint would temporarily spike to $2\Psi$ as all parameters are simultaneously present. By interleaving broadcasts with layer computation and discarding parameters after each layer, the peak parameter memory remains at approximately $\Psi / N_d$ plus the size of one layer's parameters β€” achieving the memory reduction that makes Pos+g+p valuable. The cost is more complex scheduling, but the memory benefit is the entire point of this stage.

The Trillion-Parameter Scaling Projection

The paper's scaling analysis (Section 9 and Table 1) translates the memory formulas into concrete model size limits for specific hardware configurations. This is not a theoretical exercise β€” it is a direct application of the per-device memory equations to answer the question: "Given $G$ GPUs each with $M$ GB of memory, what is the largest model that can be trained?"

The calculation: For mixed-precision Adam, total model-state memory is $16\Psi$ bytes distributed across $N_d$ data-parallel devices with ZeRO-DP stage $s$:

  • Pos: $16\Psi / 4 = 4\Psi$ per device (asymptotically for large $N_d$)
  • Pos+g: $16\Psi / 8 = 2\Psi$ per device (asymptotically)
  • Pos+g+p: $16\Psi / N_d$ per device

Setting this equal to the available per-GPU memory (minus residual state overhead) and solving for $\Psi$ gives the maximum trainable model size.

Concrete examples from Table 1:

For $N_d = 64$ (64 GPUs, each with 32 GB):

  • Pos: $4\Psi = 32$ GB β†’ $\Psi = 8$ GB / 4 bytes per fp32 parameter β‰ˆ 2B parameters in fp32, but for fp16 models with fp32 optimizer states, the actual limit depends on the precise memory accounting. The paper's Table 1 shows 7.5B for Pos with $N_d = 64$ (31.4 GB out of 32 GB).
  • Pos+g: the memory equation gives $2\Psi \approx 32$ GB β†’ $\Psi \approx 16$ GB / ... β‰ˆ higher capacity. Table 1 shows Pos+g can train models up to 14.4B with $N_d = 64$.
  • Pos+g+p: $16\Psi / 64 = 0.25\Psi$ bytes per parameter β†’ $32$ GB / $0.25$ bytes per parameter β†’ 128B parameters. Table 1 confirms: with Pos+g+p and $N_d = 64$, ZeRO can train up to 128B parameter models.

For $N_d = 1024$: $16\Psi / 1024 = 0.015625\Psi$ bytes per parameter β†’ for a 32 GB GPU, about 2 trillion parameters in theory. The paper's Table 1 shows 1 trillion (1T) as a conservative estimate, accounting for residual state overhead.

Combined with model parallelism: The paper shows (Table 2) that ZeRO-DP can be combined with MP to further reduce per-device memory. With MP degree $N_m$, each device's memory is reduced by an additional factor of $N_m$, giving a total reduction of $N_d \times N_m$ with full Pos+g+p. For 1024 GPUs with $N_d = 64$ and $N_m = 16$, the total memory reduction factor is $64 \times 16 = 1024$, making it theoretically possible to fit a model consuming $1024 \times 32 \text{ GB} = 32 \text{ TB}$ of model states. The paper notes this could accommodate a 2 trillion parameter model (2T in Table 2).

The compute gap caveat: The paper is careful to distinguish memory feasibility from training practicality. The compute analysis (Section 9) estimates that a 1 trillion parameter model would require on the order of 3000Γ— more computation per sample than BERT-Large (330M parameters), and with realistic assumptions about sequence length and data requirements, "training a 1T model would take 140 days, assuming the same hardware and similar computational efficiency," and "it would require an exa-flop system to train a 1T parameter model in a reasonable time." The ZeRO system solves the memory problem but acknowledges that the compute problem for trillion-parameter training remains open.

4. Key Insights and Innovations

Innovation 1: Temporal Sparsity as a Memory Optimization Strategy β€” Rejecting the Static-Allocation Assumption

The most fundamental conceptual move in this paper is the recognition that model states in training are not static, monolithic entities that must live in memory for the entire iteration, and that exploiting their temporal locality enables a fundamentally different memory-scaling regime.

Prior work β€” whether data parallelism (which replicates all states), model parallelism (which partitions states vertically but keeps each partition resident), or pipeline parallelism (which partitions horizontally but similarly keeps each stage's states resident) β€” all operated under an implicit assumption: once a model state is loaded into device memory, it stays there for the duration of the training step. This was not an explicit design choice that anyone defended; it was simply the default behavior of deep learning frameworks. Tensors are allocated, used, and typically retained until the end of the backward pass when the computation graph is freed. The optimizer states, in particular, are only needed during the update step at the very end of each iteration, yet they sat in memory throughout the forward and backward passes consuming capacity that could be used for activations or larger batch sizes.

The paper's Section 4.1(c) articulates what should have been obvious in retrospect: "not everything is required all the time." Parameters for layer 5 of a 100-layer transformer are completely irrelevant during the computation of layers 1–4 and 6–100. Optimizer states are irrelevant during the entire forward and backward computation. Gradients, once reduced for a particular parameter partition, are irrelevant to all processes except the one that owns that partition.

This insight is not merely an optimization trick β€” it is a reframing of what model states are. Before ZeRO, the field treated model states as storage costs β€” things you pay a fixed memory price for per device, no matter what. After ZeRO, model states become communication resources β€” things you can move between devices on-demand, paying a bandwidth cost instead of a capacity cost. This is the same conceptual shift that virtual memory brought to operating systems: not every byte of a program's address space needs to be in physical RAM at all times. The paper does not make this analogy explicitly, but it is the precise intellectual move: ZeRO creates a virtualized model-state space where the aggregate memory of the cluster acts as a pool, and each device keeps only the working set it currently needs.

The significance of this reframing extends beyond the specific mechanisms of ZeRO-DP. It opens a design space where any training state with temporal locality β€” not just optimizer states, gradients, and parameters β€” becomes a candidate for on-demand communication. The paper itself applies this logic to activations via ZeRO-R (Pa), but the principle generalizes: embedding tables, attention caches, even temporary computation results could be managed this way. The temporal-sparsity lens is the intellectual contribution; ZeRO-DP is the first systematic application of that lens to the most memory-hungry components of training.

Evidence that this is the paper's core conceptual contribution, not just an implementation detail: the communication analysis in Sections 7–8 is structured entirely around when communication happens and how it interleaves with computation. The paper could have presented ZeRO-DP as "we partition states and use collective operations to reassemble them," but that would miss the point. The contribution is the dynamic schedule β€” the temporal choreography that exploits the fact that different model states are needed at different points in the iteration. Without the temporal insight, partitioning would either require all states to be gathered at iteration start (defeating the memory savings, since peak memory would spike) or would be so communication-inefficient as to be useless. The temporal insight is what makes the difference between "we can partition model states" (obvious, and done by MP already) and "we can partition model states while keeping communication volume close to DP and peak memory close to the per-device partition size" (non-obvious, and the paper's contribution).

The paper's three-stage design (Pos, Pos+g, Pos+g+p) can be read as progressively exploiting deeper levels of temporal sparsity: optimizer states are needed only at the very end of the iteration (weakest temporal constraint, but largest memory savings β€” 4Γ—); gradients are needed only until they are reduced (intermediate constraint, additional 2Γ—); parameters are needed only during specific layers (strictest temporal constraint, but enables linear scaling). This staged structure is not arbitrary β€” it reflects increasing exploitation of temporal locality, with each stage's communication overhead corresponding to how tightly the state's usage is constrained in time.

Innovation 2: Data Parallelism Can Achieve Model-Parallel Memory Efficiency β€” Resolving a Fundamental Trade-off

Before ZeRO, the systems community accepted a seemingly inescapable trade-off: data parallelism gives high compute efficiency but poor memory scaling; model parallelism gives good memory scaling but poor compute efficiency. This was treated as a law of nature, not an engineering choice. The paper's Section 4.1 explicitly lays out this dichotomy and then systematically dismantles it.

The trade-off arose from a specific architectural decision: DP replicates all model states to achieve independent per-device computation with minimal communication (only gradient averaging, 2Ξ¨ volume per step). MP partitions model states to reduce per-device memory but at the cost of per-layer communication (much higher volume proportional to activation sizes and number of layers) and reduced computational granularity (each device computes only a fraction of each layer, reducing arithmetic intensity). The field's dominant assumption was that these two properties β€” memory efficiency and compute/communication efficiency β€” were fundamentally coupled. You could have one or the other, not both.

ZeRO's demonstration that DP can achieve MP-like memory efficiency at DP-like communication cost is therefore not just an engineering improvement; it is a proof of concept that this trade-off was an artifact of implementation, not a fundamental constraint. The mechanism is the partitioning-plus-dynamic-communication strategy described in Section 3, but the conceptual contribution is the decoupling of memory scaling from computational granularity. In ZeRO-DP, each device still computes the full forward and backward pass for its data sample β€” the computational granularity is identical to standard DP. The only difference is where the parameters come from: local memory (standard DP) or a broadcast from a peer (ZeRO-DP Pos+g+p). Since the broadcast is pipelined with computation, it adds latency but does not fundamentally alter the per-device compute pattern.

This decoupling has profound architectural implications that the paper gestures toward but does not fully explore. It means that data parallelism is not inherently memory-inefficient β€” it was only memory-inefficient because we were replicating states that didn't need to be replicated. Once that redundancy is eliminated, DP becomes the strictly dominant paradigm: it has better computational granularity than MP, better communication scaling (the 1.5Γ— overhead of Pos+g+p is modest compared to MP's per-layer communication), and far better usability (no model refactoring). The question shifts from "should we use DP or MP?" to "under what circumstances, if any, does MP provide advantages that DP with ZeRO cannot?"

The paper addresses this directly in Section 1's FAQ-style discussion "ZeRO and MP," where it identifies two residual use cases for MP: (1) reducing activation memory in conjunction with ZeRO-R (Pa), since Pa's memory reduction is proportional to the MP degree, and (2) controlling the aggregate batch size when DP alone would require batches too large for good convergence. Both are secondary considerations β€” MP becomes a knob for fine-tuning memory and batch size, not the primary mechanism for fitting large models. This is a complete inversion of the pre-ZeRO landscape, where MP was the only viable approach for models beyond ~1.4B parameters.

The evidence for this decoupling is spread across multiple results: Figure 2 shows ZeRO-100B maintaining high throughput (38+ TFlops/GPU) on models up to 100B parameters where the Megatron baseline collapses to ~5 TFlops/GPU; Table 1 shows memory per device dropping linearly with DP degree under Pos+g+p; Section 7.2.2 shows the communication overhead is only 1.5Γ—. Together, these results demonstrate that the DP-MP trade-off was not fundamental β€” it was solvable through a re-engineering of how DP manages model states. This is not merely an incremental improvement over standard DP; it is a category shift in what DP can achieve, transforming it from a memory-bound approach to one limited only by aggregate cluster capacity.

Innovation 3: Verifier Over-Optimization as a First-Class Phenomenon in Test-Time Scaling

[Note: The paper provided is the ZeRO paper, which does not involve verifiers, reward models, or test-time compute strategies. This innovation from the example analysis does not apply to the ZeRO paper and should not be included. I am omitting it and will instead identify innovations that are actually present in the ZeRO paper.]

Innovation 3: Memory Fragmentation as a Diagnosed, Quantified Bottleneck β€” Not Just "Out of Memory"

A distinctive contribution that separates this paper from prior systems work on large-model training is its explicit diagnosis and quantification of memory fragmentation as a first-class failure mode, with a targeted solution (MD: Memory Defragmentation) rather than treating it as an opaque "ran out of memory" error.

Prior work on training large models typically reported memory consumption in terms of the sum of tensor sizes β€” parameters, gradients, optimizer states, activations β€” and concluded that if this sum exceeded device capacity, the model could not be trained. The implicit assumption was that memory was a fungible resource: if you have X GB of tensors and a Y GB GPU, you need X < Y. The paper's Section 3.2 shatters this assumption with a concrete, quantified observation:

"We observe significant memory fragmentation when training very large models, resulting in out of memory issue with over 30% of memory still available in some extreme cases."

This is a diagnostic contribution, not just an optimization. The paper identifies why fragmentation occurs: the interleaving of short-lived and long-lived tensors during activation checkpointing (where checkpointed activations are long-lived and non-checkpointed activations are short-lived) and during the backward pass (where parameter gradients are long-lived and activation gradients are short-lived). This interleaving pattern is specific to the activation checkpointing training regime β€” it is not a general property of all PyTorch programs β€” which means it is a systematic artifact of how large models are necessarily trained, not a random occurrence.

The significance of this diagnosis is that it explains a class of failures that would otherwise be mysterious. A practitioner running a large model, seeing "CUDA out of memory" with 30% free memory reported, would reasonably conclude that the memory estimation was wrong or that there was a memory leak. The paper shows that neither is the case β€” the memory is genuinely free, but it is not contiguous, and large allocations (for fused gradient buffers or activation tensors) require contiguous blocks. Without this diagnosis, the practitioner might waste time hunting for memory leaks or reducing batch size unnecessarily; with it, they can apply targeted defragmentation (MD) or increase the DP/MP degree to reduce per-device tensor sizes below the fragmentation threshold.

The defragmentation solution (MD) is elegant in its simplicity: pre-allocate contiguous buffers for the two categories of long-lived tensors (activation checkpoints and gradients) and copy them there as they are produced. This is not a novel memory management technique β€” it is essentially a slab allocator, a concept from systems programming dating back decades β€” but its application to the specific fragmentation pattern of DL training, and its integration into a holistic memory optimization system, is novel. The fact that MD "not only enables ZeRO to train larger models with larger batch sizes, but also improves efficiency when training with limited memory" (Section 6.3) suggests that fragmentation-induced allocation overhead was previously degrading performance even when it did not cause outright OOM β€” a subtle point that most prior work overlooked entirely.

Evidence: Figures 7–8 show the impact of different ZeRO configurations on max cached memory and throughput, but the qualitative evidence β€” the 30% fragmentation figure, the description of OOM-with-free-memory β€” is itself the contribution. By naming and quantifying the fragmentation problem, the paper gives the community a vocabulary and a diagnostic framework for understanding a failure mode that was previously just "it doesn't fit."

Innovation 4: The FLOPs-Matched Comparison Re-frames the Pretraining vs. Inference Compute Trade-off

[Note: This innovation from the example analysis also does not apply to the ZeRO paper, which does not perform FLOPs-matched comparisons between pretraining and inference. I am instead identifying innovations present in the ZeRO paper.]

Innovation 4: Super-Linear Speedup as Evidence That Memory Efficiency Unlocks Latent Compute Capacity

The paper reports a result that is, on its face, paradoxical: super-linear speedup when scaling from 64 to 400 GPUs for a 60B parameter model (Figure 3). In classical parallel computing, super-linear speedup is rare and usually indicates that the baseline was operating in a degraded regime β€” for example, when adding more processors allows the working set to fit in cache, reducing memory latency. ZeRO's super-linear speedup has a precise and novel explanation that constitutes a conceptual contribution: ZeRO-DP's memory reduction enables larger per-GPU batch sizes, which increases arithmetic intensity and thus GPU utilization.

The mechanism is straightforward once explained but was not obvious before this work. As the data-parallel degree Nd increases, Pos+g reduces per-device memory consumption (the gradient and optimizer state terms shrink as 14Ξ¨/Nd14\Psi/N_d), freeing device memory that can be used to increase the per-GPU batch size. Larger batch sizes mean larger matrix multiplies, which better utilize the GPU's tensor cores and memory bandwidth. In the 64–400 GPU range, the memory freed is sufficient to increase batch sizes significantly, and the resulting compute utilization gains more than compensate for the added communication overhead of more data-parallel processes.

This is not merely a performance engineering result; it is a reframing of what limits training throughput. Before ZeRO, the conventional wisdom was that scaling DP hits diminishing returns because adding more GPUs increases communication overhead without reducing per-device computation. ZeRO's super-linear speedup shows the opposite: in the memory-constrained regime where large models operate, adding GPUs can increase per-device efficiency because the primary bottleneck is not communication but memory pressure β€” the inability to fit large enough batch sizes to keep the GPU compute units fed.

The paper explicitly expects this behavior to continue: "We expect this trend to continue further for more GPUs." This implies that for very large models, the optimal DP degree is not determined by communication scaling (as in classical DP) but by the point where per-GPU batch sizes become large enough to saturate compute β€” or by convergence constraints that limit total batch size. This is a fundamentally different scaling regime than the one the field was accustomed to, and it emerges directly from ZeRO's memory optimization.

Evidence: Figure 3 shows throughput per GPU increasing from approximately 32 TFlops at 64 GPUs to over 38 TFlops at 400 GPUs, a roughly 19% increase in per-GPU throughput from a 6.25Γ— increase in GPU count. The aggregate throughput therefore scales by approximately 7.4Γ— β€” well above the 6.25Γ— proportional scaling. The paper attributes this explicitly to "Pos+g reduces per GPU memory consumption of ZeRO-100B with increase in DP degree, allowing ZeRO-100B to fit larger batch sizes per GPU, which in turn improves throughput as a result of increasing arithmetic intensity" (Section 10.3). This is not a hand-wavy claim; it is a mechanistic explanation grounded in the memory formulas from Section 5.

The super-linear speedup also has a subtle implication for the trillion-parameter vision. If memory pressure is the primary constraint, and ZeRO alleviates it linearly with device count, then adding GPUs not only enables larger models (via the 16Ξ¨/Nd16\Psi/N_d scaling) but also makes each GPU more efficient by enabling larger batch sizes. This creates a virtuous cycle: more GPUs β†’ more aggregate memory β†’ larger feasible models β†’ larger per-GPU batch sizes β†’ higher utilization. The paper does not fully explore this feedback loop, but the super-linear speedup data strongly suggests it exists.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The evaluation uses GPT-2-style transformer-based language models trained with mixed-precision Adam. The paper does not train on a standard benchmark dataset to completion; rather, it measures training throughput (TFlops/GPU), model size capacity (maximum trainable parameters), and scalability (throughput scaling with GPU count). Model configurations vary in hidden dimension, number of layers, and attention heads, ranging from 1.5B to 170B parameters (Table 4, Appendix Tables 5–10). Perplexity is reported only for the Turing-NLG result (Figure 5), evaluated on Webtext-103.

  • Base model(s). All experiments use GPT-2-like transformer architectures with varying configurations of layers and hidden dimensions. The specific parameter counts tested include: 1.5B, 8B, 40B, 60B, 80B, 100B, 120B, 140B, and 170B for the throughput comparison (Figure 2); 60B for the scalability experiment (Figure 3); and 1.16B–13B for the DP-only democratization experiment (Figure 4). The Turing-NLG result uses a 17B-parameter model. These are not pretrained public checkpoints β€” they are architectural instantiations used to benchmark training system performance.

  • Metrics. The primary metric is per-GPU training throughput measured in Teraflops (TFlops), computed from iteration time and known FLOP counts for the model architecture. The paper also reports aggregate throughput in Petaflops, maximum trainable model size (parameters), and maximum cached memory per GPU. For super-linear scalability (Figure 3), throughput per GPU is plotted against GPU count. For Turing-NLG (Figure 5), validation perplexity on Webtext-103 is reported over training iterations. Memory consumption is measured via PyTorch's reported maximum cached memory during each training iteration (Figure 7).

  • Baselines. Two baselines are used:

    1. PyTorch Distributed Data Parallel (DDP) β€” for experiments without model parallelism. This is standard data parallelism where all model states are replicated on every GPU.
    2. Megatron-LM (Shoeybi et al., 2019) β€” for experiments with model parallelism. The paper uses the open-source NVIDIA implementation dated September 2019. The most recent Megatron-LM results at the time reported scaling to 16B parameters on 512 V100 GPUs (32 DGX-2 nodes).

    No other baselines (e.g., GPipe, PipeDream, or CPU-offloading approaches) are directly benchmarked β€” comparisons to these are analytical rather than empirical (Section 2).

  • Generation budget / compute accounting. Compute is measured as wall-clock throughput per GPU in TFlops, calculated from measured iteration time for the specific model configuration. Since these are systems benchmarking experiments rather than task-evaluation experiments, there is no "generation budget" in the sense of number of sampled outputs. Instead, the paper controls for total GPU count, model parallelism (MP) degree, and batch size per GPU. The key fairness consideration is that ZeRO-100B and the baseline Megatron-LM are given the same hardware resources (same GPU count and type, same interconnect), and the comparison is throughput achieved for training the same model architecture. The paper notes that the baseline sometimes uses fewer GPUs than ZeRO (e.g., 256 or 384 vs. 400) because the MP degree must divide the model dimensions, which actually gives the baseline a communication advantage (fewer GPUs = less communication). Per-GPU throughput is reported rather than aggregate throughput to maintain comparability.

  • Cross-validation / statistical protocol. There is no cross-validation or statistical significance reporting because these are systems throughput benchmarks, not model accuracy evaluations. Results are engineering measurements of iteration time and memory consumption on specific hardware (400 V100 GPUs across 25 DGX-2 nodes, 800 Gbps internode bandwidth). The paper provides detailed configuration tables (Appendix Tables 5–10) specifying exact layer counts, hidden dimensions, attention heads, and batch sizes for each experiment, enabling exact reproduction on equivalent hardware.

Main Quantitative Results

Model Size Capacity: ZeRO-100B Trains Models 8Γ— Larger than Megatron-LM

Figure 2 (and its supporting Table 5) reports the headline finding: ZeRO-100B (Pos+g + ZeRO-R) combined with model parallelism can efficiently train models up to 170B parameters on 400 V100 GPUs, compared to Megatron-LM's practical ceiling of approximately 40B parameters. The paper states this as an "over 8x increase in model size compared to SOTA."

The degradation of the Megatron-LM baseline is stark. At 40B parameters, Megatron-LM achieves some usable throughput (the exact number is not stated, but the paper describes it as the largest model Megatron can train "with acceptable throughput," approximately 16–20B in a DGX-2 system). When the authors tested a 40B parameter model across two DGX-2 nodes (requiring inter-node MP communication), they observed "about 5 Tflops per V100 GPU (less than 5% of hardware peak)." This near-total throughput collapse is the empirical justification for ZeRO's design: MP cannot scale beyond a single node because inter-node bandwidth (InfiniBand EDR at ~12.5 GB/s per link) is roughly 24Γ— lower than intra-node bandwidth (NVSwitch at ~300 GB/s per link).

For ZeRO-100B, throughput remains high across the full range. The paper reports:

"ZeRO-100B achieves a sustained throughput of 15 PetaFlops (over 30% of the peak) on average for models with 8B to 100B parameters."

At 100B parameters, ZeRO-100B achieves approximately 38 TFlops/GPU (visible in Figure 2), compared to Megatron-LM's roughly 5 TFlops/GPU for the 40B model that required cross-node MP. The paper states this as a "10x improvement in training speed compared to SOTA for the same model size."

For models beyond 100B, ZeRO-100B's throughput decreases slightly (Figure 2 shows a modest dip for 120B, 140B, and 170B), which the paper attributes to "lack of enough memory to run larger batch sizes" β€” as the model consumes more of the available memory, the batch size must be reduced, decreasing arithmetic intensity and thus GPU utilization.

The model configurations used (Table 4, Appendix Table 5) reveal that the largest 170B model uses 212 layers with hidden dimension 8192 and 64 attention heads, with a per-GPU batch size of 12 and total batch size of 300 across 400 GPUs with MP degree 16. The baseline Megatron-LM, attempting to run the same 170B model, required 256 GPUs (the closest power-of-2 configuration that could fit) with MP degree 256 and could only achieve a batch size of 2 (total batch size 2, since DP=1 with 256-way MP leaves no data parallelism) β€” a configuration that is essentially non-functional.

Super-Linear Scalability: Per-GPU Throughput Increases with GPU Count

Figure 3 (and Appendix Table 6) demonstrates that for a fixed 60B parameter model, per-GPU throughput increases from approximately 32 TFlops at 64 GPUs to over 38 TFlops at 400 GPUs, a roughly 19% improvement in per-device performance even as the GPU count increases 6.25Γ—. This is super-linear scaling: the aggregate throughput increases by a factor of approximately 7.4Γ— when GPU count increases 6.25Γ—.

The mechanism, as explained in Sections 5.2 and 10.3, is that Pos+g reduces per-GPU memory consumption as the DP degree increases (since the gradient and optimizer state terms shrink as 14Ξ¨/N_d). This freed memory enables larger per-GPU batch sizes: the table shows batch size per GPU increasing from 16 at 64 GPUs to 64 at 400 GPUs, while the total batch size scales from 64 to 1600. Larger batch sizes improve arithmetic intensity (ratio of computation to memory access), better utilizing the GPU's tensor cores and memory bandwidth.

The paper explicitly expects this trend to continue: "We expect this trend to continue further for more GPUs." This is not pure extrapolation β€” it follows from the memory formula: as N_d increases, the 14Ξ¨/N_d term continues to shrink, freeing additional memory for larger batch sizes, until either the model-state memory becomes negligible (limited by residual states) or the batch size reaches a convergence-imposed ceiling.

Democratization: 13B Parameters Without Model Parallelism

Figure 4 (and Appendix Table 10) shows that ZeRO-100B can train models up to 13B parameters using data parallelism alone (no MP), on 128 GPUs, achieving throughput over 40 TFlops/GPU. This is significant because it eliminates the need for model parallelism β€” and its associated model refactoring complexity β€” for models up to this scale.

In comparison, standard DDP (PyTorch Distributed Data Parallel) without ZeRO runs out of memory at 1.4B parameters, with throughput below 20 TFlops/GPU. ZeRO thus enables a 9.3Γ— increase in trainable model size with pure DP (13B vs. 1.4B), while simultaneously doubling throughput.

The model configurations in Appendix Table 10 show the scaling: the 13B model uses 62 layers, hidden dimension 4096, 32 attention heads, and a batch size of 2 per GPU (256 total batch size across 128 GPUs). The baseline can only reach 1.38B parameters (40 layers, hidden dimension 1536) with a batch size of 1 per GPU.

The paper notes a practical implication: "in the absence of the communication overhead from MP, these models can be trained with lower-end compute nodes without very fast intra-node interconnect such as NVLINK or NVSwitch." This is because DP communication is limited to gradient reduction (and parameter all-gather in Pos+g), which is less sensitive to interconnect bandwidth than the per-layer communication of MP.

Memory and Performance Analysis: ZeRO Optimizations Have Cumulative Benefit

Figures 6–8 and Table 3 present an ablation of five ZeRO configurations applied to 40B, 60B, 100B, and 170B models:

  • C1: Pos + CB + MD
  • C2: Pos + CB + MD + Pa
  • C3: Pos+g + CB + MD
  • C4: Pos+g + CB + MD + Pa
  • C5: Pos+g + CB + MD + Pa+cpu

Maximum model size (Figure 6): Under a fixed MP of 16, the largest trainable model increases from 40B (C1) to 60B (C2, from Pa reducing activation memory by 16Γ—) to 140B (C4, from Pos+g halving model-state memory vs. Pos) to 150B (C5, from Pa+cpu offloading partitioned activations to CPU). The jumps correspond directly to the memory reduction mechanisms: Pa targets activations, Pos+g targets model states, and Pa+cpu provides an additional activation memory reduction at the cost of CPU transfer overhead.

Maximum cached memory (Figure 7): For a 40B model, C2 reduces cached memory compared to C1 (Pa partitions activations). The difference between C2 and C3 depends on relative model-state vs. activation sizes β€” "can increase when activation memory is larger, or decrease when the model states are larger." For the 100B model, C5 shows a noticeable reduction from C4 (Pa+cpu offloads activations), while for the 40B model, the reduction is not significant because activations are proportionally smaller. This directly validates the design choice to make Pa+cpu optional β€” it matters most for the largest models.

Maximum achievable performance (Figure 8): Performance improvements correspond to memory reduction β€” lower memory β†’ larger batch sizes β†’ higher throughput. For the 60B model, C5 drops performance below C4 despite lower memory, because "C5 incurs activation movement to and from the CPU, this will result in worse performance in most cases." The exception is the 170B model, which cannot run without C5 β€” the CPU offloading cost is the price of fitting the model at all. The paper states: "During training, Pa+cpu is turned on only when it is beneficial."

A specific throughput figure from Figure 8: for the 170B model, C5 with Pa+cpu is the only configuration that executes without OOM, achieving a throughput that the paper does not quote exactly but which Figure 8 shows as the sole data point for that model size.

Turing-NLG: SOTA Language Model at 17B Parameters

Figure 5 reports results for Turing-NLG, a 17B-parameter model trained end-to-end using ZeRO-100B. Achievements:

  • Size: 17B parameters, the largest published model as of May 12, 2020.
  • Accuracy: Webtext-103 perplexity of 10.21, establishing a new SOTA.
  • Training throughput: 41.4 TFlops/GPU sustained.

The validation perplexity curve (Figure 5) shows Turing-NLG over 300K training iterations progressively outperforming the previous SOTA Megatron-LM 8.3B model. The exact perplexity trajectory and the point of crossover cannot be precisely read from the small figure, but the trend is consistently downward and below the Megatron-LM baseline.

This result serves as an existence proof: ZeRO is not just a laboratory benchmark but has been used to train a production-scale model with published SOTA results. The paper characterizes this as "the system breakthroughs of ZeRO" powering "the world's largest language model."

Ablation Studies and Robustness Checks

The paper's ablation strategy differs from that of a machine learning paper β€” since ZeRO is a systems design, the "ablations" are the five configurations (C1–C5) that incrementally enable different combinations of ZeRO-DP stages and ZeRO-R optimizations. Each configuration represents a point in the design space, and the experiments measure the marginal contribution of each optimization.

  • Pos vs. Pos+g (C1/C2 vs. C3/C4): Enabling gradient partitioning halves model-state memory compared to Pos alone. In Figure 6, this translates to a jump from 60B (C2) to 140B (C4) maximum trainable model size. The communication overhead is zero (Section 7.2.1). This is a strict improvement β€” there is no scenario where Pos+g is worse than Pos alone β€” which justifies its inclusion as the default for ZeRO-100B.

  • Pa: Partitioned Activation Checkpointing (C1 vs. C2, C3 vs. C4): Figure 6 shows Pa increases the maximum trainable model from 40B (C1) to 60B (C2) under Pos, and from some intermediate size to 140B under Pos+g. The paper quantifies the activation memory reduction: for a 100B model with MP=16, batch size 32, sequence length 1024, activation checkpoints drop from ~33 GB per GPU to ~2 GB β€” a 16Γ— reduction proportional to the MP degree (Section 6.1). The communication overhead is <10% of the baseline MP volume (Section 8). Pa is particularly important because it specifically targets the activation memory bottleneck that becomes dominant once ZeRO-DP reduces model-state memory.

  • Pa+cpu: CPU Offloading of Partitioned Activations (C4 vs. C5): This is the only ablation that shows a negative performance impact in some configurations. Figure 8 shows C5 has lower throughput than C4 for the 60B model because the CPU-GPU data movement overhead dominates any benefit from the additional batch size enabled by lower memory. However, for the 170B model, C5 is the only configuration that can run β€” without it, the model hits OOM. The paper's policy is to enable Pa+cpu selectively: "During training, Pa+cpu is turned on only when it is beneficial." This is a nuanced design choice β€” the system does not blindly apply all optimizations, but adapts based on whether the memory savings justify the communication cost.

  • CB: Constant-Size Buffers: The impact of CB cannot be isolated from Figures 6–8 because all configurations (C1–C5) include CB. The paper's argument for CB is qualitative: "the memory overhead of the fused buffers is proportional to the model size, and can become inhibiting" (Section 6.2). The design trades a small communication efficiency loss (from smaller messages) for bounded memory overhead β€” essential for the largest models but the exact throughput impact is not quantified independently.

  • MD: Memory Defragmentation: Like CB, MD is included in all configurations and not independently ablated. The paper's evidence is the qualitative observation that fragmentation causes OOM with "over 30% of memory still available in some extreme cases" (Section 3.2). The throughput improvement from MD is not isolated, but the paper claims it "also improves efficiency when training with limited memory" (Section 6.3) because the memory allocator spends less time searching for contiguous blocks.

Non-trivial finding on activation memory: One of the more subtle results is that Pa's benefit depends on the relative sizes of model states and activations. The paper notes that for the 40B model, C5 does not show a memory reduction from C4 (Figure 7), but it does for the 100B model. This is because "the activation memory for 100B is much larger for the decrease to be noticeable." This implies a design principle: Pa+cpu is not universally beneficial but becomes valuable only when activation memory dominates, which occurs for very large models or very deep architectures. The paper's adaptive policy for enabling Pa+cpu operationalizes this insight.

Negative result with ReST-like training: While not part of the main ZeRO-100B evaluation, the paper's mention that "ZeRO does not change the model optimization method or affect model convergence" (Section 2.2.3) is a claim about robustness. The Turing-NLG result (Figure 5) provides empirical support β€” the model trains successfully to SOTA perplexity β€” but this is a single model, and the paper does not systematically evaluate convergence across a range of model sizes or architectures under ZeRO. The claim that ZeRO does not affect convergence is plausible (the optimizer sees identical gradient information; only the storage location changes), but it is not rigorously tested.

Critical Assessment

Claim 1: ZeRO-100B trains models 8Γ— larger than SOTA (170B vs. ~20B)

Supported, with a qualification about the baseline. Figure 2 clearly shows ZeRO-100B running 170B parameters while the Megatron-LM baseline β€” which the paper identifies as the SOTA for large-model training infrastructure β€” collapses to non-functional throughput levels beyond 40B. The 8Γ— figure (170B / ~20B) is conservative given that Megatron-LM's practical ceiling in a single DGX-2 is 16–20B, and the paper measured 5 TFlops/GPU for a 40B model across two nodes (which is already 95% below peak).

However, the "SOTA" designation deserves scrutiny. The paper benchmarks only against Megatron-LM, which was the dominant MP framework at the time. Other approaches β€” GPipe (Huang et al., 2018), PipeDream (Narayanan et al., 2019), and various CPU-offloading strategies β€” are discussed analytically in Section 2 but not empirically compared. If GPipe with sufficient hardware could train a 100B+ model (at some throughput), the 8Γ— claim becomes a comparison against one specific system rather than against all alternatives. The paper's analytical arguments against PP (batch size constraints, tied-weights difficulties, convergence implications) are reasonable, but without empirical throughput numbers, we cannot quantify how much larger ZeRO's model size advantage actually is compared to the full space of alternatives. The 8Γ— figure is best understood as "8Γ— larger than the best published system (Megatron-LM) that was widely used and fully open-sourced at the time."

Claim 2: Super-linear speedup in the 64–400 GPU regime

Supported, with the mechanism clearly explained. Figure 3 shows per-GPU throughput increasing from ~32 to ~38 TFlops as GPU count scales from 64 to 400. The explanation β€” ZeRO-DP memory reduction enables larger batch sizes, improving utilization β€” is mechanistically grounded in the memory formulas and consistent with the batch size data in Appendix Table 6. The super-linear behavior is not magic; it is exactly what one would expect when a memory-constrained workload is given more aggregate memory and can trade it for increased arithmetic intensity.

The limitation is that this result is demonstrated for a single model size (60B) and a single hardware configuration (V100 GPUs on DGX-2 nodes). The paper expects the trend to continue for more GPUs, but this extrapolation depends on the continued availability of memory to convert into batch size β€” if the model-state memory term 14Ξ¨/N_d becomes negligible compared to activation memory at very high DP degrees, the batch-size scaling would saturate. The paper does not specify at what GPU count this saturation would occur, making the extrapolation plausible but unverified.

Claim 3: ZeRO democratizes large model training by supporting 13B parameters without MP

Supported. Figure 4 clearly shows ZeRO-100B training 13B-parameter models with DP alone at over 40 TFlops/GPU, compared to the baseline DDP's ceiling of 1.4B parameters at under 20 TFlops/GPU. The 9.3Γ— increase in model size and 2Γ— increase in throughput with a simpler programming model (no model refactoring) substantiates the democratization claim.

A weakness: the largest model trained with DP alone (13B) is demonstrated on 128 GPUs. The paper does not explore the lower bound β€” what is the minimum GPU count needed to train a 13B model with ZeRO-powered DP? This would matter for practitioners with smaller clusters. Based on the memory formulas, a 13B model with Pos+g requires approximately 2Ξ¨ + 14Ξ¨/N_d bytes per GPU. At N_d = 16 (16 GPUs), this is 2Γ—13 + 14Γ—13/16 = 26 + 11.4 = 37.4 GB, which would not fit in a 32 GB V100, so more than 16 GPUs are needed. At N_d = 32, it's 26 + 5.7 = 31.7 GB β€” just fitting. This lower-bound analysis would strengthen the democratization claim by showing the minimum cluster size needed.

Claim 4: ZeRO has the potential to scale to 1 trillion parameters

Supported as a memory feasibility claim, but critically gated by compute constraints. Table 1 shows that with Pos+g+p and N_d = 1024, per-device model-state memory for a 1T parameter model is approximately 15.6 GB β€” well within a 32 GB V100 GPU. Table 2 shows that with a combined 16-way MP and 64-way DP on 1024 GPUs, the theoretical maximum model size reaches 2 trillion parameters. These calculations follow directly from the memory formulas in Section 5 and are algebraically correct.

However, the paper is admirably honest about the limitations of this projection. Section 9 explicitly states that training a 1T parameter model end-to-end would take "over a year" on current hardware, and "it would require an exa-flop system to train a 1T parameter model in a reasonable time." The ZeRO system solves the memory problem but not the compute problem. The trillion-parameter claim is about fit, not practical training time. The paper does not empirically demonstrate even a partial training run at this scale; it is an extrapolation from the memory formulas validated at much smaller scales (up to 170B).

This is not a weakness of the experimental work β€” building and running a 1T parameter model was infeasible in 2020 β€” but it means the claim should be understood as "ZeRO's memory scaling, if the formulas continue to hold and sufficient hardware is available, would allow a 1T model to fit on 1024 GPUs." The 170B result provides strong evidence that the formulas hold at that scale, but the extrapolation to 1T (roughly 6Γ— larger) assumes no unexpected bottlenecks emerge, such as activation memory scaling nonlinearly, communication overhead growing faster than the 1.5Γ— factor, or memory fragmentation patterns changing at extreme scales.

Claim 5: ZeRO has zero communication overhead for Pos+g and only 1.5Γ— for Pos+g+p

Supported analytically but not directly measured empirically. Sections 7.2.1 and 7.2.2 provide a careful accounting of the communication volume per training step for Pos+g (2Ξ¨, identical to baseline DP) and Pos+g+p (3Ξ¨, 1.5Γ— baseline). The analysis is convincing because it follows directly from the definitions of the collective operations (reduce-scatter, all-gather, broadcast) and does not depend on empirical measurement.

However, the paper does not provide a direct empirical comparison of communication time for ZeRO vs. baseline DP at identical model sizes and batch sizes. The throughput results (Figures 2–4, 8) represent the combined effect of memory reduction and communication changes, making it difficult to isolate whether the communication overhead for Pos+g+p is actually 1.5Γ— in wall-clock time (as opposed to volume). Communication time depends on bandwidth utilization, which can vary with message size, and the Pos+g+p parameter broadcasts are spread across the forward/backward passes in a way that may overlap with computation differently than the baseline all-reduce. The analytical claim is well-supported; the empirical communication latency impact is confounded with batch size effects in the throughput numbers.

Missing Experiments That Would Strengthen the Paper

  • Direct communication time benchmarks: Isolate the communication component of iteration time for ZeRO-DP (Pos, Pos+g, Pos+g+p) vs. baseline DP at a fixed model size and batch size where all configurations fit in memory. This would validate the analytical communication volume claims with empirical latency data.
  • Convergence comparison: Train a model to completion with ZeRO and without ZeRO (at a scale where both fit) and confirm identical loss curves. While the paper argues ZeRO does not change optimization semantics, subtle differences in floating-point reduction order during reduce-scatter vs. all-reduce could theoretically affect convergence. The Turing-NLG result provides some evidence but is a single model.
  • Comparison to non-MP alternatives at scale: Benchmark against GPipe or activation-offloading approaches at a model size where both can run (e.g., 10–20B parameters, where MP + CPU offloading might compete). The paper's analytical arguments against these approaches are reasonable, but empirical throughput numbers would be more convincing.
  • Minimum GPU count for different model sizes: A sweep showing the smallest cluster needed to train a given model size with ZeRO-DP at each stage. This would be directly useful for practitioners deciding whether ZeRO enables their hardware to handle a target model size.
  • Sensitivity to interconnect bandwidth: Evaluate ZeRO-DP throughput at varying interconnect speeds (e.g., within a node vs. across nodes vs. with degraded links). This would characterize how robust the 1.5Γ— communication overhead is to bandwidth constraints.

Bottom Line

The experiments convincingly establish that ZeRO-100B (Pos+g + ZeRO-R) enables training models up to 170B parameters β€” 8Γ— larger than what Megatron-LM could practically support β€” and that the super-linear speedup from 64 to 400 GPUs results from memory reduction enabling larger batch sizes. The memory formulas that project to trillion-parameter feasibility are validated up to the 170B scale, and the analytical communication volume claims are rigorous. The main experimental gaps are the lack of direct communication latency measurements, convergence comparisons, and competitive benchmarks against non-MP alternatives β€” though these gaps are understandable given the systems benchmarking focus and the practical constraints of running experiments at this scale (requiring up to 400 V100 GPUs). The democratization claim (13B with DP alone) is the most directly validated, with a clear head-to-head comparison against standard DDP.

6. Limitations and Trade-offs

The Compute-Power Gap: Memory Feasibility Does Not Imply Practical Training

The most fundamental limitation of ZeRO is that it solves the memory bottleneck for training trillion-parameter models but does nothing to address the compute bottleneck. The paper's headline projection β€” that ZeRO with Pos+g+p can fit a 1 trillion parameter model on 1024 GPUs β€” is a statement about memory capacity, not about practical end-to-end training time.

The paper is explicit and transparent about this distinction in Section 9:

"Running a model with a trillion parameters efficiently is no longer impossible! ... training a trillion parameter model end-to-end within an acceptable time range, however, could still require significant amount of compute power, which is lacking in today's AI clusters."

The consequence is stark. The paper's own analysis estimates that a 1T parameter model trained on 1024 V100 GPUs would require over 140 days at minimum, assuming the same sequence length and data efficiency as BERT-Large. In practice, "both data samples and sequence length are likely to increase with the increased model size requiring over a year to train." The paper further estimates that "it would require an exa-flop system to train a 1T parameter model in a reasonable time." At the time of writing (2020), no such system was available β€” the largest AI clusters operated in the hundreds of petaflops range, roughly an order of magnitude short.

A practitioner reading this paper and hoping to train a trillion-parameter model would therefore face a harsh reality: ZeRO allows the model to fit in GPU memory, but they would still need either dramatically more GPUs or a fundamentally different approach to computation (sparsity, mixture-of-experts, etc.) to finish training within months. The paper's memory formulas tell you how many GPUs you need to fit the model; the compute analysis tells you that fitting and training in a reasonable time are decoupled problems, and ZeRO addresses only the first.

What evidence exists? Section 9 provides the rough calculation: BERT-Large (330M parameters) trains in 67 minutes on 1024 V100 GPUs. A 1T parameter model would have approximately 3000Γ— more computation per sample (1 trillion / 330 million β‰ˆ 3000), leading to the ~140 day estimate under optimistically identical conditions. Section 9 also references this as a "Compute Power Gap" β€” the title of the relevant discussion. The empirical results (Figures 2–4) stop at 170B parameters and do not attempt to measure training throughput above that scale, so the extrapolation to 1T is purely based on the FLOP-count scaling assumption and has not been empirically tested at scale.

Mitigation status: The paper does not attempt to address the compute gap β€” it explicitly scopes ZeRO as a memory solution and acknowledges the compute limitation as future work. The framing is that "when such compute capacity becomes available, we hope ZeRO will provide the system technology to run the 1T models efficiently." This is honest but leaves the trillion-parameter vision as a conditional promise: ZeRO removes the memory barrier, but the compute barrier remains entirely unresolved.

Difficulty Estimation Cost: Amortized in the Headline Numbers

[Note: This limitation is erroneously carried over from the paper analysis example β€” it refers to the LLM test-time compute paper, not the ZeRO paper. The ZeRO paper does not involve difficulty estimation or oracle difficulty bins. This limitation should be replaced with one that is actually present in the ZeRO paper.]

ZeRO-DP Stage 3 (Pos+g+p) Adds 50% Communication Overhead β€” Quantified but Not Empirically Isolated

The paper's analytical communication analysis (Section 7.2.2) establishes that Pos+g+p increases total communication volume from 2Ξ¨ (baseline DP and Pos+g) to 3Ξ¨ β€” a 1.5Γ— increase. This is analytically rigorous and relatively modest compared to the N_dΓ— memory reduction it enables. However, the paper does not provide direct empirical measurement of the communication overhead in isolation, separate from the throughput effects of increased batch sizes.

The consequence is that a practitioner cannot easily determine whether the 1.5Γ— volume increase translates to a 1.5Γ— communication time increase in practice, or something different. Several factors could cause the empirical overhead to diverge from the analytical overhead:

  • Message size effects: Pos+g+p communicates parameters layer-by-layer via broadcasts interleaved with computation. Each broadcast carries only the parameters for one layer, making individual messages much smaller than the all-reduce of all gradients in standard DP. Smaller messages achieve lower bandwidth utilization on interconnects, potentially making the effective communication time higher than the volume ratio suggests.
  • Overlap efficiency: The paper claims the parameter broadcasts can be "pipelined to avoid the memory overhead" and "interleaved with computation" (Section 7.2.2). Whether this pipelining actually hides the communication latency depends on the computation-to-communication ratio per layer. For shallow-but-wide models (few layers, large hidden dimension), the computation per layer is substantial and communication may be fully hidden. For deep-but-narrow models (many layers, small hidden dimension), the per-layer computation may be too brief to hide the broadcast latency, making the communication overhead fully exposed.
  • Interconnect heterogeneity: Within a DGX-2 node, NVSwitch provides ~300 GB/s per link, making parameter broadcasts cheap. Across nodes, InfiniBand EDR provides only ~12.5 GB/s per link β€” 24Γ— less. A practitioner with a multi-node cluster would experience the 1.5Γ— volume increase very differently depending on whether the DP processes are within the same node or spread across nodes. The paper does not benchmark Pos+g+p separately for intra-node vs. inter-node configurations.

What evidence exists? None directly. The ZeRO-100B implementation used for experimental evaluation only includes Pos+g (not Pos+g+p), so all measured throughput figures in Figures 2–8 reflect the zero-overhead regime of Pos+g. The communication volume analysis for Pos+g+p (Section 7.2.2) is purely analytical. Table 1 projects memory sizes with Pos+g+p for a future implementation that the paper plans to "release ... by end of May 2020." The 1.5Γ— volume increase has not been validated with empirical timing data.

Mitigation status: The paper does not measure the communication overhead of Pos+g+p empirically and does not discuss how it might vary with model architecture or hardware topology. The analytical argument strongly suggests the overhead is modest and manageable, but a practitioner deploying Pos+g+p would be taking this on analytical faith without empirical evidence at the time of publication. The paper's plan to release the full implementation (Section 1) suggests the authors intended to close this gap, but the published paper does not include those measurements.

Activation Memory Remains a Hard Ceiling for Extremely Large Models β€” Even with Pa and Pa+cpu

While ZeRO-DP eliminates model-state memory as the primary bottleneck, activation memory becomes the dominant remaining constraint for the largest models, and the paper's mitigations (Pa and Pa+cpu) have fundamental limits.

The paper quantifies the scale of this problem in Section 6.1:

"Consider training a 100B model shown in Table 4 with a batch size of 32, sequence length of 1024 and a MP degree of 16. If we checkpoint a single activation for each transformer layer, it would require about 33 GB of memory per GPU just to store the activation checkpoints."

This is on a 32 GB V100 GPU β€” meaning activations alone, even with checkpointing (one checkpoint per layer), would overflow device memory before accounting for any model states, temporary buffers, or fragmentation. The activation memory for a 100B model at batch size 32 is approximately 33 GB, which already exceeds the 32 GB capacity. This means that even with ZeRO-DP reducing model-state memory to near-zero, the model cannot be trained at this batch size without Pa.

Pa reduces activation memory by the MP degree N_m (to ~2 GB in the 100B example with MP=16), which solves the problem for the given configuration. However, this reduction is capped by the feasible MP degree. If N_m = 16 reduces activations to 2 GB, then an even larger model β€” say 500B parameters with proportionally larger activations β€” might require N_m = 64 or N_m = 128 to fit, which may exceed the number of GPUs available within a single node. Since MP across nodes incurs the 24Γ— bandwidth penalty that cripples Megatron-LM (as demonstrated by the 5 TFlops/GPU measurement for a cross-node 40B model in Section 1), Pa's effectiveness is bounded by the intra-node GPU count (typically 8 or 16).

Pa+cpu provides an escape hatch by offloading partitioned activations to CPU, but at a cost the paper quantifies in Section 8 as "2x added data movement to and from CPU memory compared to Pa." The empirical results (Figure 8) show that Pa+cpu worsens throughput for the 60B model compared to Pa alone, and is only beneficial when the model literally cannot run without it (the 170B case). This is not a fundamental solution β€” it is a last-resort trade-off that trades throughput for feasibility.

The consequence for the trillion-parameter vision is significant: even if ZeRO-DP reduces model-state memory to 16 GB on 1024 GPUs (as Table 1 shows), the activation memory for a 1T model could still be the gating factor. The paper does not project activation memory requirements for a 1T model, but if activations scale proportionally with model parameters (roughly proportional to number of layers Γ— hidden dimension), a 1T model would have approximately 10Γ— the activation memory of a 100B model β€” roughly 330 GB per GPU in the batch-size-32 example, or approximately 20 GB even after Pa with MP=16. That 20 GB, plus the 16 GB of model states, plus temporary buffers and fragmentation overhead, would push against or exceed the 32 GB capacity of a V100 GPU.

What evidence exists? The 33 GB activation figure for a 100B model (Section 6.1) is the key data point. Figure 8 shows Pa+cpu being necessary for the 170B model but degrading performance for the 60B model. These together demonstrate that activation memory is not fully solved by ZeRO β€” it is managed, not eliminated, and the management bounds become tight at extreme scale.

Mitigation status: The paper acknowledges this implicitly through the existence of Pa and Pa+cpu, but does not provide a projection of activation memory scaling to 1T parameters or a discussion of whether the MP-degree cap makes activation memory a hard ceiling. The suggestion that "Pa is turned on only when it is beneficial" (Section 10.5) is a practical mitigation within the tested range but does not address whether the approach fundamentally scales.

Single Hardware Platform: V100 GPUs on DGX-2 with NVSwitch β€” No Characterization of Hardware Sensitivity

All experiments in the paper are conducted on a single hardware configuration: 400 NVIDIA V100 GPUs across 25 DGX-2 nodes, with NVSwitch intra-node interconnect (~300 GB/s per link) and InfiniBand EDR inter-node interconnect (800 Gbps aggregate, or approximately 12.5 GB/s per link). The paper does not evaluate ZeRO on any other GPU architecture (e.g., T4, A100, TPU), any other interconnect topology, or any other cluster scale.

The consequence is that several of the paper's key numerical claims may not transfer to different hardware environments:

  • The super-linear speedup (Figure 3) depends on the ability to convert freed memory into larger batch sizes, which in turn depends on the GPU's compute capacity and memory bandwidth. On a GPU with different memory-to-compute ratios (e.g., an A100 with 40 GB or 80 GB of HBM2e and higher peak throughput), the batch-size sweet spot would shift, potentially changing the scaling curve.
  • The 24Γ— bandwidth gap between intra-node and inter-node communication is specific to the V100 + DGX-2 + InfiniBand EDR configuration. Newer hardware (DGX A100 with NVSwitch 3.0 and InfiniBand HDR/NDR) would narrow this gap, potentially making cross-node MP more viable and altering the trade-off between ZeRO-DP and MP. The paper's strong claim that "MP cannot scale much further beyond these model sizes ... the efficiency degrades quickly beyond a single node" (Section 1) is tied to the specific bandwidth ratios of the tested hardware.
  • The trillion-parameter projection assumes V100-class memory capacity (32 GB/GPU). On GPUs with different memory capacities, the required DP degree for a given model size would change proportionally. The paper's Table 1 is parameterized by N_d, so it can be recomputed, but the per-GPU residual state overhead (activations, buffers, fragmentation) does not necessarily scale linearly with GPU memory β€” a 40 GB A100 might have proportionally more or less headroom than a 32 GB V100 for the same model-state memory partition.
  • The communication analysis assumes InfiniBand-class interconnects. On clusters with Ethernet or lower-bandwidth interconnects, the 1.5Γ— communication overhead of Pos+g+p would be more painful, and the paper's claim that it is modest might not hold.

What evidence exists? The hardware specification is stated in Section 10.1 without qualification. No sensitivity analysis is provided β€” there are no experiments varying GPU type, interconnect speed, or node count independently of GPU count. The paper treats its hardware platform as the universal reference, and all numerical claims about throughput (38 TFlops/GPU, 15 Petaflops aggregate, super-linear speedup regime) are implicitly conditioned on this specific hardware.

Mitigation status: The paper does not discuss hardware sensitivity as a limitation. The open-source release of DeepSpeed (announced in Section 1) would allow the community to evaluate on other platforms, but the published results provide no guidance on what a practitioner with different hardware should expect. A practitioner with an older cluster (pre-NVSwitch, or with PCI-E-only interconnects) or a newer one (A100 with more memory and higher compute) cannot directly extrapolate from the reported throughput numbers.

Evaluation Limited to GPT-2-Style Transformer Architectures β€” No Evidence for Other Model Families

All experiments in the paper use GPT-2-like transformer-based language models with varying hyperparameters (layers, hidden dimensions, attention heads) as documented in Appendix Tables 5–10. While transformers were (and remain) the dominant architecture for large-scale NLP, the paper makes no attempt to evaluate ZeRO on other model architectures that were prominent at the time or that have different memory and communication characteristics:

  • Encoder-decoder transformers (T5, BART): These have different activation memory patterns (cross-attention activations in the decoder) and different parameter distributions (encoder and decoder parameters may be asymmetrically sized). The activation analysis in Section 3.2 assumes a decoder-only architecture's linear scaling of activations with layers Γ— hidden_dim Γ— seq_len Γ— batch_size; encoder-decoder models would have a different formula.
  • Models with tied weights (input/output embedding sharing): The paper mentions that tied-weights are "difficult to implement" with PP (Section 2.1) and that MP has issues with models where "MP cannot divide the model evenly" (Section 1). ZeRO-DP partitions parameters evenly across DP processes, which would also interact non-trivially with tied weights β€” if the embedding matrix is shared, which partition does it belong to? The paper does not discuss this case.
  • Vision models (ResNet, EfficientNet, ViT): The memory consumption patterns of convolutional networks (where activations scale with spatial dimensions rather than sequence length) differ substantially from transformers. The paper's analyses of activation memory (Section 3.2) and communication (Sections 7–8) are specific to transformer architectures.
  • Mixture-of-experts (MoE) models: While not widespread in 2020, MoE models were beginning to emerge (Shazeer et al., 2017) and have fundamentally different parameter-to-computation ratios (many parameters, sparse activation). ZeRO's model-state partitioning would need to account for the expert imbalance β€” some parameter partitions might be much larger than others if experts are unevenly distributed.
  • Models with batch normalization: The paper mentions that "batch-normalization [is] difficult to implement due to ... micro-batching" in PP (Section 2.1) but does not discuss how ZeRO-DP handles batch-norm statistics, which require cross-device synchronization beyond simple gradient averaging.

The consequence is that a practitioner with a non-GPT architecture cannot be certain that ZeRO's memory scaling formulas apply as stated. The 16Ξ¨ model-state memory formula depends only on the optimizer (mixed-precision Adam) and is architecture-agnostic, so that portion likely transfers. But the activation memory scaling, the communication analysis for specific layer types, and the interaction of parameter partitioning with non-standard layer configurations (weight tying, layer-specific parameter sizes) are architecture-dependent and untested.

What evidence exists? The paper specifies that "the models presented in this section are GPT-2 like transformer based models" (Section 10.1). The Turing-NLG result (Figure 5) is also a GPT-style language model. No other architectures are mentioned as tested. The analytical memory formulas for model states (Section 3.1) are architecture-independent (they depend only on Ξ¨), but the activation analysis (Section 3.2) uses the transformer-specific formula: "The activation memory of a transformer-based model is proportional to the number of transformer layers Γ— hidden dimensions Γ— sequence length Γ— batch size."

Mitigation status: The paper does not acknowledge the architecture-specificity of its evaluation as a limitation. The claims about democratization ("data scientists can thus experiment freely with large models without worrying about parallelism," Section 1) implicitly assume that the benefits transfer to whatever architecture the data scientist is using, but the evaluation provides evidence only for GPT-2-style transformers. The open-source release of DeepSpeed would allow testing on other architectures, but the paper provides no guidance or caveats for architecture-specific considerations.

7. Implications and Future Directions

How This Work Changes the Landscape

ZeRO fundamentally reframes how the field should think about the relationship between parallelism strategy and memory capacity in large-model training. Before this work, the systems community operated under an implicit dichotomy: data parallelism gave you high compute efficiency but replicated model states wastefully, while model parallelism gave you memory efficiency at the cost of communication overhead and reduced computational granularity. The two properties β€” memory scaling and compute/communication efficiency β€” appeared to be inescapably coupled. You chose one at the expense of the other, and scaling to larger models meant accepting the communication penalty of MP or the batch-size constraints of PP.

ZeRO decouples these properties by recognizing that model states are not monolithic, always-resident entities but rather resources with temporal locality β€” optimizer states are only needed during the update step, gradients are only needed until they are reduced, and parameters are only needed during the specific layers they govern. By partitioning states across data-parallel devices and communicating them on-demand through a dynamic schedule, ZeRO achieves the memory efficiency of model parallelism (per-device memory scales as ~1/N_d with the data-parallel degree) while retaining the computational granularity and communication volume of data parallelism (zero overhead for Pos and Pos+g, 1.5Γ— for Pos+g+p). This is not an incremental optimization β€” it is a category shift in what data parallelism can achieve, transforming it from a memory-bound approach limited to ~1.4B parameters on 32 GB GPUs into a strategy that scales model size linearly with aggregate cluster memory.

The magnitude of this shift is evident in the numbers: ZeRO-100B (implementing only the first two ZeRO-DP stages) trains models 8Γ— larger than Megatron-LM's practical ceiling (170B vs. ~20B parameters) while delivering 10Γ— higher throughput. The full ZeRO design projects to trillion-parameter feasibility on 1024 GPUs β€” a scale that was, prior to this work, considered fundamentally unreachable without exotic hardware or algorithmic breakthroughs. The trillion-parameter projection is particularly significant not because it promises imminent trillion-parameter training (the paper is honest about the compute gap β€” Section 9 estimates over a year of training time on current hardware), but because it reframes the bottleneck from memory capacity to compute capacity. Before ZeRO, the conversation was "we cannot fit large models in memory." After ZeRO, the conversation becomes "we can fit them, but can we afford to train them?" This shifts research attention from memory optimization (where ZeRO provides a solution) to compute optimization (sparsity, mixture-of-experts, model compression) β€” a more productive direction because compute capacity improves with hardware generations, while memory-per-GPU has grown more slowly.

The paper also reconciles a practical tension in the field between the accessibility of data parallelism and the necessity of model parallelism for large models. Data parallelism was the default for most practitioners because it required no model refactoring β€” the same code that runs on a single GPU runs on many. Model parallelism, by contrast, demanded careful partitioning of layers, custom communication operators, and architecture-specific implementations (Megatron-LM only supported a limited set of transformer operators). This created a democratization gap: only well-resourced teams with systems expertise could train models beyond ~1.4B parameters. ZeRO's demonstration that 13B-parameter models can be trained with pure data parallelism (Figure 4, achieving over 40 TFlops/GPU on 128 GPUs without any MP) means that the barrier to entry for large-model training drops dramatically. A data scientist who understands standard PyTorch Distributed Data Parallel can use ZeRO-powered DP to train models an order of magnitude larger than what standard DP allows, without learning model parallelism. This is not just a convenience β€” it expands the pool of researchers who can experiment with large models, potentially accelerating progress in the field.

The paper also redirects research attention around model parallelism. Before ZeRO, MP was the primary strategy for training models beyond a few billion parameters, and significant research effort went into optimizing MP implementations (Megatron-LM, Mesh-TensorFlow, FlexFlow). ZeRO's demonstration that DP with state partitioning can achieve comparable or better memory efficiency at lower communication cost makes MP a secondary tool rather than the primary scaling mechanism. The paper identifies two residual use cases for MP (reducing activation memory via Pa, and controlling aggregate batch size for convergence), but these are fine-tuning knobs, not foundational necessities. This implies that future research on distributed training infrastructure should prioritize improving ZeRO-style state partitioning (better communication scheduling, more aggressive overlapping, support for heterogeneous hardware) over further optimizing fine-grained model parallelism β€” a substantial reallocation of research effort.

Finally, the paper's diagnosis of memory fragmentation as a first-class failure mode (Section 3.2, quantifying OOM-with-30%-free-memory) establishes a new lens for thinking about training memory. Prior work treated memory as a fungible resource where total tensor size was the only metric. ZeRO shows that fragmentation β€” the interleaving of short-lived and long-lived tensors during activation checkpointing β€” can silently reduce usable capacity by over 30%. The defragmentation solution (MD, Section 6.3) demonstrates that proactive memory management (pre-allocating contiguous buffers for long-lived tensors) can recover this capacity. This diagnosis generalizes beyond ZeRO: any training system that uses activation checkpointing on large models will encounter fragmentation, and the paper provides both the vocabulary to describe it and a proven mitigation strategy.

Follow-Up Research This Work Enables

Empirical measurement of Pos+g+p communication overhead across hardware topologies. The paper provides a rigorous analytical accounting of Pos+g+p communication volume (3Ξ¨, or 1.5Γ— baseline DP) but does not empirically measure the wall-clock communication overhead in isolation from batch-size effects. A critical follow-up would benchmark Pos+g+p on a fixed model size and batch size (where all configurations fit in memory) across three hardware settings: intra-node (NVSwitch), inter-node (InfiniBand EDR), and mixed (some DP processes per node, some across nodes). The key measurement is the fraction of iteration time spent in communication for Pos+g+p vs. baseline DP vs. Pos+g at the same batch size, isolating the overhead of the per-layer parameter broadcasts. A negative result β€” where Pos+g+p communication overhead significantly exceeds the 1.5Γ— volume ratio due to small-message inefficiency in the per-layer broadcasts β€” would bound the practical DP degree for which Pos+g+p is beneficial and would motivate research on more efficient broadcast scheduling (e.g., fusing parameter broadcasts across multiple layers into larger messages).

Convergence equivalence under ZeRO-DP reduce-scatter vs. all-reduce. The paper asserts that ZeRO "does not change the model optimization method or affect model convergence" (Section 2.2.3) because the optimizer sees identical gradient values regardless of whether they are aggregated via all-reduce or reduce-scatter. However, floating-point reduction is non-associative β€” the order in which gradients are summed across devices differs between reduce-scatter (which reduces different gradient partitions on different processes) and all-reduce (which produces identical full reduced gradients everywhere). A controlled experiment would train two identical models β€” one with standard DP (all-reduce), one with ZeRO-DP Pos+g (reduce-scatter) β€” at a scale where both fit (e.g., a 1B-parameter model on 8 GPUs), running multiple seeds, and compare loss curves, final perplexity, and checkpoint-level parameter differences. If differences emerge, they would constrain the claim that ZeRO is optimization-neutral and motivate research into bit-exact reduction schemes for partitioned training. The Turing-NLG result (Figure 5) provides anecdotal evidence of convergence but is a single model and does not include a DP baseline comparison.

Activation memory scaling laws for trillion-parameter models. The paper identifies activation memory as the residual bottleneck after ZeRO-DP reduces model-state memory (Section 6.1: a 100B model requires ~33 GB of activation memory per GPU even with checkpointing at batch size 32). A systematic study would project activation memory requirements for 200B, 500B, and 1T parameter transformer models under various configurations: with and without Pa (partitioned activation checkpointing), at different MP degrees, with different checkpointing strategies (one checkpoint per layer vs. fewer), and across decoder-only vs. encoder-decoder architectures. The output would be a set of scaling curves showing, for a given GPU memory capacity, the maximum feasible batch size as a function of model size and available MP degree. This would identify the true memory ceiling of ZeRO-based training β€” the point at which activation memory, not model-state memory, prevents scaling β€” and would reveal whether Pa's reduction (proportional to MP degree, capped by intra-node GPU count) is sufficient for trillion-parameter models or whether more aggressive activation compression or offloading is required.

ZeRO-DP for mixture-of-experts (MoE) models. MoE architectures (Shazeer et al., 2017) have a fundamentally different parameter-to-computation ratio than dense models: they have many parameters (due to multiple expert sub-networks) but only a fraction are activated per token, meaning the optimizer states, gradients, and parameters are larger relative to the computation and activations. ZeRO-DP partitions model states evenly across DP processes, but MoE models introduce load imbalance β€” different experts may be accessed at different frequencies, causing uneven gradient and optimizer state sizes across partitions. A study would adapt ZeRO-DP's partitioning to account for expert popularity, perhaps using dynamic repartitioning or expert-aware sharding that sizes partitions by expected access frequency rather than parameter count. The evaluation would measure per-GPU memory balance, communication volume (which may increase if popular experts require more frequent parameter broadcasts), and throughput compared to a dense model of equivalent FLOPs. The paper's democratization of large-model training would extend naturally to MoE if the partitioning can handle expert imbalance gracefully.

ZeRO-DP on heterogeneous hardware (different GPU memory capacities per node). The paper's experiments use uniform 32 GB V100 GPUs, but production clusters increasingly contain mixed hardware (e.g., some nodes with 40 GB A100s, others with 80 GB A100s, or older V100 nodes). ZeRO-DP's equal partitioning assumes uniform per-device memory, which breaks down when some GPUs have more capacity than others. A study would implement and evaluate proportional partitioning, where each device's partition size is proportional to its available memory, allowing devices with more memory to hold larger model-state partitions and process larger batch sizes. The key metrics would be: (1) whether the communication schedule can handle non-uniform partition sizes without introducing stragglers (processes with smaller partitions may finish computation faster and wait for larger-partition processes to complete communication), (2) the aggregate throughput compared to uniform partitioning (which would be capped by the smallest-memory device), and (3) the impact on model convergence when different devices process different batch sizes. This would make ZeRO practical for the heterogeneous clusters that many research labs and companies actually operate.

ZeRO-powered self-improvement loops for language model training. The paper's democratization claim β€” that ZeRO enables 13B-parameter training with pure DP β€” opens the possibility of iterative self-improvement at scales previously restricted to large industrial teams. In a self-improvement loop (similar to STaR or ReST), a model generates training data, the data is filtered for quality, and the model is fine-tuned on the filtered data. The memory bottleneck has historically limited self-improvement to smaller models or required complex MP setups. A concrete experiment would use ZeRO-powered DP to train an 8B-parameter model, use it to generate completions on a reasoning benchmark (e.g., GSM8K or MATH), filter for correctness, fine-tune on the correct completions, and repeat for multiple iterations β€” all on a modest cluster (e.g., 32–64 GPUs) without MP. The key measurement is whether the self-improvement loop yields meaningful accuracy gains given the compute budget, and whether ZeRO's memory efficiency allows larger per-GPU batch sizes during fine-tuning that improve training stability. This would test whether ZeRO genuinely democratizes the kind of iterative large-model training that was previously the domain of well-resourced teams.

Practical Applications and Downstream Use Cases

Training 10B–100B parameter models on modest clusters without model parallelism expertise. The most immediate practical application of ZeRO-100B is enabling teams with access to 64–400 V100 or A100 GPUs to train models in the 10B–100B parameter range using pure data parallelism, without hiring MP specialists or refactoring model code for Megatron-style parallelism. Figure 4 demonstrates this concretely: on 128 GPUs, ZeRO-100B trains a 13B-parameter model at over 40 TFlops/GPU with DP alone, while standard DDP fails beyond 1.4B parameters. For a team with a 128-GPU cluster, this means they can train a model 9Γ— larger than what their existing PyTorch DDP code supports, simply by replacing DDP with ZeRO-powered DP β€” no architecture changes, no custom operators, no pipeline stage tuning. The throughput numbers (38+ TFlops/GPU for models up to 100B parameters, Figure 2) mean the training time for a given model size and data budget is comparable to what a well-tuned MP implementation would achieve, but with dramatically lower engineering overhead. For academic labs, small companies, or teams within larger organizations that lack dedicated distributed systems expertise, this is transformative: it makes large-model research accessible without the systems engineering barrier that previously restricted it to a handful of industrial labs.

Cost-efficient batch inference and fine-tuning for large language models. While the paper focuses on training, the memory reduction from ZeRO-DP applies equally to fine-tuning and (with some adaptation) to inference. For fine-tuning, ZeRO-DP Pos+g allows a 13B-parameter model to be fine-tuned on as few as 8–16 GPUs (depending on sequence length and batch size) without MP β€” the memory formulas in Table 1 show that at N_d = 16, Pos+g reduces model-state memory for a 13B model to approximately 2Γ—13 + 14Γ—13/16 = 26 + 11.4 = 37.4 GB, which fits within a 40 GB A100. This means organizations can fine-tune large open-source models (e.g., LLaMA-13B) on domain-specific data using a single DGX node, without the complexity of sharding the model across GPUs. For batch inference (scoring thousands of examples), ZeRO-DP's memory reduction allows larger per-GPU batch sizes, improving throughput by increasing arithmetic intensity β€” the same mechanism that produces super-linear speedup during training (Figure 3). A deployment running batch inference on a 60B-parameter model could use ZeRO-DP to fit batch sizes 2–4Γ— larger than standard DP, reducing total inference time proportionally, at the cost of a one-time parameter broadcast overhead that is amortized over the larger batch.

Scaling diffusion models and vision transformers to higher resolutions. Although the paper evaluates only on language models, the memory scaling formulas are architecture-agnostic for model states. The primary architecture-specific constraint is activation memory, which for vision transformers (ViT) and diffusion models scales with image resolution (analogous to sequence length in language models). A lab training a 5B-parameter diffusion model at high resolution (e.g., 1024Γ—1024 images) would encounter activation memory as the dominant bottleneck, just as the paper identifies for large language models. Applying ZeRO-R Pa (partitioned activation checkpointing) with a modest MP degree (e.g., MP=4 or MP=8 within a node) would reduce activation memory proportionally β€” the paper's Section 6.1 example shows a 16Γ— reduction for MP=16 on a language model, and the same scaling applies to any architecture where activations are replicated across MP devices. The communication overhead of Pa (under 10% of baseline MP volume, Section 8) is architecture-independent, depending only on the activation tensor sizes. ZeRO therefore generalizes naturally to the vision domain, where high-resolution training faces the same memory-capacity constraints that ZeRO addresses for large language models.

Enabling self-hosted large-model training for privacy-sensitive applications. In regulated industries (healthcare, finance, legal), training or fine-tuning large models on proprietary data requires on-premise hardware β€” data cannot be sent to cloud APIs. On-premise GPU clusters are typically much smaller than cloud-scale deployments (often 8–64 GPUs). ZeRO-100B's democratization result (13B parameters on 128 GPUs with DP alone, Figure 4) means that even a modest on-premise cluster can train models that were previously only feasible at cloud scale. For example, a hospital with a single DGX A100 node (8Γ— 80 GB GPUs) could use ZeRO-DP Pos+g to fine-tune a 70B-parameter LLaMA-style model on patient records β€” the per-GPU model-state memory would be approximately 2Γ—70 + 14Γ—70/8 = 140 + 122.5 = 262.5 GB across 8 GPUs, or ~33 GB each, fitting comfortably within 80 GB. This would enable privacy-preserving large-model adaptation in domains where sending data to a cloud provider is legally or ethically prohibited, expanding the reach of large-model technology into sensitive applications.