ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

🎯 Pitch

ZeRO-Infinity trains a 32-trillion-parameter model on just 512 V100 GPUs—50× larger than prior systems—by seamlessly offloading data to CPU and NVMe memory without model code changes. It sustains over 25 petaflops on the same hardware by hiding all data movement behind computation, effectively shattering the GPU memory wall for anyone with access to a modest cluster.

1. Executive Summary

This paper introduces ZeRO-Infinity, a heterogeneous system technology that exploits GPU, CPU, and NVMe memory simultaneously to train models at unprecedented scale without model code refactoring. The system is evaluated on GPT-like Transformer models using the infinity offload engine for offloading model states to CPU or NVMe, memory-centric tiling to handle large individual layers by decomposing them into sequentially executed tiles, bandwidth-centric partitioning to leverage aggregate heterogeneous memory bandwidth across all parallel devices, and overlap-centric design to hide data movement behind computation. ZeRO-Infinity trains a 32 trillion parameter model on 512 NVIDIA V100 GPUs—a 50× increase over the state-of-the-art 3D parallelism—while sustaining over 25 petaflops (40% of peak) and demonstrating superlinear scalability from 64 to 512 GPUs for a trillion-parameter model, establishing that accelerator device memory is no longer a limitation on model scale when heterogeneous memory systems are exploited in parallel.

2. Context and Motivation

The Core Problem: The GPU Memory Wall

The fundamental problem this paper tackles is what the authors call the GPU memory wall: the growing mismatch between the rate at which deep learning models are scaling and the rate at which single-GPU memory capacity is growing. The numbers are stark:

"In the last three years, the largest dense deep learning models have grown over 1000x to reach hundreds of billions of parameters, while the GPU memory has only grown by 5x (16 GB to 80 GB)."

This isn't just a historical observation — it's a structural trend that threatens to halt progress in large model training. Model scaling has been the primary driver of advances in deep learning over the past several years (Devlin et al., BERT; Radford et al., GPT-2; Brown et al., GPT-3), and multiple studies suggest this trend will continue (Kaplan et al., 2020). If model growth is limited by GPU memory capacity, then the field's primary engine of progress stalls.

To make the bottleneck concrete, consider the computational requirements for current and future models. The paper provides a sobering calculation:

"It requires 800 NVIDIA V100 GPUs just to fit a trillion parameter model for training."

This is purely a memory constraint, not a compute constraint. The compute to train a trillion-parameter model might be feasible on a large cluster, but you can't even begin training if you can't fit the model state (parameters, gradients, optimizer states) into the aggregate memory of your GPUs. The paper projects this forward: a hundred-trillion-parameter model would require over 6,000 GPUs even assuming a generous 5× increase in GPU memory in the coming years. At current memory growth rates, the model capacity of even the largest clusters is capped far below where researchers want to go.

Why This Problem Matters: Two Distinct Pain Points

The GPU memory wall creates two separate but equally important problems, and the paper addresses both:

1. Frontier training is becoming inaccessible. Training models at the cutting edge — hundreds of billions or trillions of parameters — now requires sophisticated combinations of parallelism techniques (data parallelism, model parallelism, pipeline parallelism — collectively 3D parallelism) that demand hundreds or thousands of GPUs. Clusters of this scale are "simply out of reach for most data scientists." The consequence is that frontier model research becomes concentrated in a tiny number of organizations with the capital to build and operate massive GPU clusters, which has implications for the field's diversity of ideas, reproducibility, and concentration of power.

2. Fine-tuning large models is impossible on modest hardware. The paper highlights a crucial asymmetry between pretraining and fine-tuning:

"While pretraining a model with hundreds of billions of parameters can require millions of GPU compute hours, fine-tuning it is much cheaper, requiring significantly fewer GPU compute hours, and could be done on a single compute node with a handful of GPUs."

The problem is that while a single DGX-2 node (16 V100 GPUs) has enough compute to fine-tune a model like GPT-3 (175B parameters) in reasonable time, it doesn't have nearly enough memory. Using 3D parallelism:

"fine-tuning GPT-3 would require over 8 DGX-2 nodes (128 GPUs) with 3D parallelism to just fit the model for training."

This is a crucial insight: the bottleneck for democratizing large model access is memory, not compute. If a researcher wants to take a pretrained GPT-3 and fine-tune it for medical question answering, legal document analysis, or code generation, they need a cluster far larger than what most universities and companies possess — even though the actual computation required is modest. This artificially restricts the downstream applications and specialization of large models to well-resourced organizations.

Where Prior Approaches Fall Short

The paper situates itself against three categories of prior work, each of which addresses part of the problem but leaves critical gaps.

3D Parallelism (Model + Pipeline + Data): The State-of-the-Art but Memory-Limited

The current best approach for training massive models is 3D parallelism, which combines:

Data parallelism (DP): Each GPU processes a different batch of data, with gradients synchronized across GPUs.
Model parallelism (MP, specifically tensor-slicing): Individual operators (e.g., a large matrix multiply) are split across multiple GPUs, with each GPU computing a slice of the output.
Pipeline parallelism (PP): Different layers of the model are assigned to different GPUs, forming a pipeline where activations flow from one stage to the next.

The DeepSpeed implementation of 3D parallelism can scale to over a trillion parameters on 800 V100 GPUs by fully leveraging aggregate GPU memory.

Where it falls short on model scale: 3D parallelism is ultimately bounded by total GPU memory. It cannot exploit CPU or NVMe memory, so the maximum model size is hard-capped by the number of GPUs in the cluster. For models of tens or hundreds of trillions of parameters, the required GPU count becomes infeasible. Moreover, even at trillion-parameter scale, the GPU requirements (320 A100s, 800 V100s) restrict training to only the largest organizations.

Where it falls short on usability: 3D parallelism imposes a heavy burden on data scientists:

"it requires data scientists to perform major model code refactoring, replacing single GPU operators with tensor-sliced versions, and splitting the model into load-balanced pipeline stages."

This is not a minor inconvenience. Pipeline parallelism requires partitioning the model's layers into stages that have roughly equal computational cost; if one stage is significantly cheaper than others, GPUs sit idle waiting for the bottleneck stage ("pipeline bubbles"). For models with complex dependency graphs — non-sequential architectures, skip connections that cross pipeline stages, or irregular computational patterns — creating a load-balanced pipeline is difficult or impossible. The paper notes this explicitly:

"Models with complex dependencies cannot be easily converted into a load-balanced pipeline."

The consequence is that 3D parallelism is not only resource-intensive but also restricts the types of models that can be efficiently trained, creating an architectural bias toward pipeline-friendly designs.

ZeRO and ZeRO-Offload: Memory Efficiency, but Still GPU-Bound for Parameters

The ZeRO family of technologies (Rajbhandari et al., 2020; Ren et al., 2021) addressed memory redundancy in data-parallel training by partitioning model states across GPUs rather than replicating them:

ZeRO-1 partitions optimizer states
ZeRO-2 partitions optimizer states and gradients
ZeRO-3 partitions optimizer states, gradients, and parameters

In ZeRO-3, each parameter is owned by exactly one data-parallel process. When a parameter is needed for a forward or backward pass, the owner broadcasts it to all processes; after the computation, the non-owner processes discard their copies. This eliminates the memory redundancy that traditional data parallelism imposes (where every GPU holds a full copy of all model states).

ZeRO-Offload extended this by offloading optimizer states and gradients to CPU memory, using the CPU's DRAM as a larger but slower storage tier. This was state-of-the-art for heterogeneous training on multi-GPU systems.

Where it falls short: ZeRO-Offload has a critical limitation:

"it still requires the parameters to be stored in GPU memory and replicated across all devices. Thus, the model scale with ZeRO-Offload is limited to the total number of parameters that the memory on a single GPU device can host."

In other words, ZeRO-Offload can move optimizer states and gradients to CPU, but the parameters themselves must still fit in GPU memory. Since parameters are replicated (not partitioned) in ZeRO-Offload's underlying ZeRO-2 design, the maximum model size is bounded by single-GPU memory, not aggregate memory. A V100 with 32 GB can hold roughly 1.6 billion FP16 parameters (32 GB / 2 bytes/param) plus activation memory — far below the scale of modern large models.

Additionally, ZeRO-Offload has efficiency problems at scale:

"ZeRO-Offload also requires a large batch size to remain efficient due to suboptimal data partitioning and limited PCIe bandwidth."

The "suboptimal data partitioning" refers to the broadcast-based approach: since each parameter is fully owned by one GPU, it must be transferred from CPU to that single GPU's memory over one PCIe link (limited to ~12 GB/s on PCIe Gen 3) before being broadcast to other GPUs. This creates a serial bottleneck — only one PCIe link is active while the others sit idle — and the limited bandwidth forces the training to use large batch sizes to amortize the communication cost. Large batch sizes, in turn, increase activation memory (which scales with batch size × sequence length × hidden dimension) and can harm convergence.

NVMe-based Approaches: Limited and Domain-Specific

The paper acknowledges one prior NVMe-based system: Zhao et al. (2020) used a hierarchical parameter server design to offload sparse parameters to SSD for training a massive-scale deep learning ads system. However, this was designed specifically for sparse models (e.g., recommendation systems with embedding tables that have billions of entries but where only a fraction are accessed per batch), not for the dense Transformer architectures that dominate large language model training.

How This Paper Positions Itself

ZeRO-Infinity is positioned as a "leap forward from 3D parallelism" that addresses all three of the major challenges identified in prior work:

Challenge 1: Model scale. The paper directly confronts the question:

"Looking ahead, how do we support the next 1000x growth in model size, going from models like GPT-3 with 175 billion parameters to models with hundreds of trillions of parameters?"

ZeRO-Infinity's answer is to transcend GPU memory entirely by building a system that can simultaneously leverage GPU, CPU, and NVMe memory. This is not simply "offload to slower memory when GPU runs out" — that approach was already attempted in ZeRO-Offload and hit the single-GPU parameter limit. Instead, ZeRO-Infinity combines ZeRO-3's parameter partitioning (so parameters no longer need to fit on a single GPU) with the infinity offload engine (which can push those partitioned parameters all the way to NVMe storage) and memory-centric tiling (which handles individual layers too large for even a single GPU's working memory).

The scale target is audacious: from the 1 trillion parameter ceiling of 3D parallelism on 800 GPUs to 32 trillion parameters on just 512 GPUs — a 50× improvement in model size at the same cluster scale.

Challenge 2: Accessibility. The paper asks:

"How can we make large models of today accessible to more data scientists who don't have access to hundreds of GPUs?"

ZeRO-Infinity's response is to demonstrate that a single DGX-2 node (16 GPUs) can fine-tune models up to a trillion parameters — models that would require 128 GPUs with 3D parallelism. This is enabled because ZeRO-Infinity can use the CPU and NVMe memory available on a single node to host the partitioned model states, while the 16 GPUs provide sufficient compute for fine-tuning throughput. The implication is that large model fine-tuning becomes feasible for individual researchers, small labs, and companies without massive GPU clusters.

Challenge 3: Ease of use. The paper asks:

"Can we make large model training easier by eliminating the need for model refactoring and multiple forms of parallelism?"

ZeRO-Infinity's approach is two-pronged. First, memory-centric tiling eliminates the need for model parallelism (tensor-slicing) by decomposing large operators into tiles that execute sequentially on a single GPU, each tile's parameters and gradients fetched and released as needed. Second, the implementation uses PyTorch hooks to automate all data movement — gathering parameters before they're needed, partitioning them after use, and optionally offloading them — so that data scientists write standard PyTorch model code and ZeRO-Infinity handles the rest. The paper explicitly positions this against 3D parallelism's requirement for manual model refactoring.

The Conceptual Framework

Underlying all of these contributions is an implicit conceptual shift in how we think about memory for large model training. The paper articulates this in its conclusion:

"It is no longer necessary to fit DL training on ultra-fast but expensive and limited memory like HBM2. ZeRO-Infinity demonstrates that it is possible to transcend the GPU memory wall by leveraging cheap and slow, but massive, CPU or NVMe memory in parallel across multiple devices to achieve the aggregate bandwidth necessary for efficient training."

This is the paper's central intellectual move: reframing the problem from "how do we squeeze more model into limited GPU memory" to "how do we make collectively vast but individually slow memory systems deliver sufficient aggregate bandwidth for training to remain compute-bound rather than memory-bound." This shifts the bottleneck from memory capacity to memory bandwidth, which is the subject of the paper's detailed analysis in Sections 3 and 4, and motivates the key technical innovations (bandwidth-centric partitioning, overlap-centric design, and the infinity offload engine's NVMe optimizations) in Sections 5 and 6.

The paper doesn't just propose a system — it provides a quantitative framework for reasoning about the bandwidth requirements of heterogeneous training (Section 4) that could guide future hardware design. This is evident in the forward-looking analysis in Section 9, where the paper projects bandwidth requirements for accelerators 10× and 100× more powerful than the V100, arguing that even at those scales, the aggregate bandwidth achievable by connecting accelerators to CPU/NVMe memory via technologies like NVLink is sufficient — pointing to the Summit supercomputer (2018) as an existence proof with 40 GB/s per GPU to CPU memory.

3. Technical Approach

3.1 Reader Orientation

ZeRO-Infinity is a distributed deep learning training system that treats GPU, CPU, and NVMe memory as a unified heterogeneous memory hierarchy, allowing models far larger than aggregate GPU memory to be trained efficiently without requiring data scientists to modify their model code. The core problem it solves is the GPU memory wall — the growing mismatch between model sizes (increasing 1000× over three years) and single-GPU memory capacity (increasing only 5× over the same period) — and its solution is to extend ZeRO-3's parameter partitioning across all three memory tiers simultaneously, combined with novel bandwidth optimization techniques that make the slower CPU and NVMe memory tiers collectively fast enough for training to remain compute-bound rather than I/O-bound.

3.2 Big-Picture Architecture (Diagram in Words)

The ZeRO-Infinity system has five major components that work together during training:

ZeRO-3 Partitioning Engine — inherits from prior work: partitions optimizer states, gradients, and parameters across all data-parallel GPUs so no single device holds a complete copy of any model state. Each parameter is "owned" by exactly one GPU, which is responsible for its optimizer update.
Infinity Offload Engine — the core new mechanism: manages the placement, movement, and lifecycle of partitioned model states across GPU, CPU, and NVMe memory tiers. It contains DeepNVMe (a high-performance asynchronous NVMe I/O library) and a pinned memory management layer that reuses a small pool of pinned buffers to transfer terabyte-scale model states without exhausting system resources.
Memory-Centric Tiling — decomposes individual large operators (e.g., a linear layer with hidden dimension 64K) into a sequence of smaller tiles that execute sequentially on a single GPU, fetching parameters tile-by-tile via ZeRO-3's gather mechanism and releasing them immediately after use. This eliminates the need for tensor-slicing model parallelism.
Bandwidth-Centric Partitioning and Overlap Engine — two tightly coupled optimizations: (a) parameters are partitioned across all data-parallel GPUs at the granularity of individual tensors, so when parameters need to be gathered from CPU/NVMe, all PCIe links operate in parallel (unlike the broadcast-based approach in ZeRO-Offload where only one PCIe link is active), and (b) a dynamic prefetcher traces the operator execution sequence and initiates the three-stage data movement (NVMe→CPU, CPU→GPU, GPU allgather) for future operators while the current operator is computing, hiding communication latency behind computation.
Ease-Inspired Implementation Layer — PyTorch hooks that automatically inject parameter gathering before each submodule's forward/backward pass and parameter partitioning (with optional offloading) after each pass, plus automatic model partitioning during __init__ so that models are never fully materialized on any single process. This is what eliminates the need for model code refactoring.

Information flow during one training iteration: At initialization, the model's parameters are partitioned across data-parallel processes and offloaded to the designated memory tier (GPU, CPU, or NVMe). When the forward pass begins, the overlap engine's prefetcher traces the operator graph and begins fetching parameters for upcoming operators from their storage locations. For each submodule, ZeRO-Infinity gathers the full parameters (allgather from all data-parallel GPUs), executes the forward computation, then partitions and optionally offloads the parameters again. Activation checkpoints are optionally offloaded to CPU. During backward pass, the same gather-compute-release pattern applies, with gradients being reduced (reduce-scatter), partitioned, and offloaded to their designated storage. After backward completion, the optimizer step runs: optimizer states are brought into CPU/GPU memory in chunks (since they may not fit in GPU memory), the Adam update is computed, and the updated states and new parameters are written back to their storage locations.

3.3 Roadmap for the Deep Dive

This is primarily a systems design and empirical evaluation paper. The technical approach section covers:

First, the memory and bandwidth characterization framework (Sections 3–4 of the paper), because ZeRO-Infinity's entire design is driven by a quantitative analysis of what must fit where (memory requirements) and how fast data must move (bandwidth requirements) for training to be efficient. Understanding these requirements first makes the subsequent design choices legible as solutions to specific quantified constraints.
Second, the infinity offload engine — how model states and activation checkpoints are placed across GPU, CPU, and NVMe, including the novel bandwidth-centric partitioning strategy that makes PCIe bandwidth scale with data-parallel degree — because this is the core enabling technology for unprecedented model scale.
Third, the overlap-centric design — the dynamic prefetcher and the three-stage communication pipeline (NVMe→CPU, CPU→GPU, GPU allgather) — because efficiency despite offloading to slow memory depends entirely on hiding these data movements behind computation.
Fourth, memory-centric tiling — how individual layers too large for GPU working memory are decomposed into sequentially executed tiles — because this is what eliminates the need for model parallelism and enables ease of use.
Fifth, the ease-inspired implementation — the PyTorch hook system, automatic model partitioning during initialization, and handling of external parameters — because this is the engineering layer that makes the system usable without code refactoring.
Sixth, the DeepNVMe library and pinned memory management — the low-level optimizations that achieve near-peak NVMe bandwidth — because the entire system's efficiency on single-node setups depends on extracting maximum performance from the slowest storage tier.

3.4 Detailed, Sentence-Based Technical Breakdown

This paper is an empirical systems design paper whose core insight is that heterogeneous memory (GPU, CPU, NVMe) can be made collectively fast enough for efficient large-model training if data is partitioned to leverage aggregate bandwidth across all devices, data movement is overlapped with computation, and GPU working memory limitations are addressed through tiling rather than model parallelism.

Memory Requirements Characterization (Foundation for Design)

Before designing the system, the paper quantitatively characterizes what must be stored and how much memory each component requires. This analysis (Section 3 of the paper) establishes the hard constraints that ZeRO-Infinity must satisfy and explains why certain design choices (e.g., partitioning parameters rather than replicating them) are non-negotiable for the target scale.

The memory required for training is categorized into two broad components: model states (optimizer states, gradients, parameters) and residual states (primarily activation memory). The paper also defines two working memory concepts — Model State Working Memory (MSWM) and Activation Working Memory (AWM) — which represent the minimum GPU memory that must be available even after all model states are offloaded, since some data must be physically present on the GPU during computation.

Model state memory. For mixed-precision training with the Adam optimizer (the de facto standard for Transformer-based models), the model states consume the following memory per parameter:

FP16 parameters: 2 bytes
FP16 gradients: 2 bytes
FP32 optimizer states (momentum, variance): 4 + 4 = 8 bytes
FP32 parameter copy: 4 bytes
FP32 gradient copy: 4 bytes

Total: 20 bytes per parameter.

The paper focuses on Transformer architectures since "all of the SOTA models with over a billion parameters follow that." For a Transformer model, the total number of parameters is dominated by four linear layers per Transformer block. Specifically, each block contains:

A query-key-value projection of shape $(hd, 3hd)$ : $3hd^2$ parameters
An attention output projection of shape $(hd, hd)$ : $hd^2$ parameters
A feed-forward expansion of shape $(hd, 4hd)$ : $4hd^2$ parameters
A feed-forward contraction of shape $(4hd, hd)$ : $4hd^2$ parameters

Total per block: $3hd^2 + hd^2 + 4hd^2 + 4hd^2 = 12hd^2$ parameters. With $nl$ layers:

$\text{params} \approx 12 \times nl \times hd^2$

where $nl$ is the number of Transformer layers and $hd$ is the hidden dimension.

The total memory for model states in bytes is therefore:

$240 \times nl \times hd^2$

Where the numbers lead: Figure 2a (column 5) computes this for models ranging from 100 billion to 100 trillion parameters. A 100B parameter model requires ~2 TB for model states alone. A 1T parameter model requires ~20 TB. These numbers are then compared against aggregate GPU memory in Figure 2b: a single DGX-2 node (16 V100 32GB GPUs) has 512 GB aggregate GPU memory; a SuperPOD cluster with 1536 GPUs has 48 TB. The conclusion is stark: fitting a 100B parameter model's model states requires 64 GPUs; a 1T model requires over 512 GPUs; and a 10T model exceeds even a 1536-GPU cluster.

Residual state memory (activations). Activations are the intermediate outputs of each layer during the forward pass that must be retained for gradient computation during the backward pass. The full activation memory for a Transformer model scales as $bsz \times seq \times hd \times nl$ and can be enormous — Figure 2a column 6 shows 1.5 TB of activations for a 100B model at batch size 32, sequence length 1024.

However, the paper notes that activation checkpointing (Chen et al., 2016) — which trades memory for recomputation by storing only a subset of activations and recomputing others during backward — is already standard practice (used for Turing-NLG 17.2B and GPT-3 175B). With activation checkpointing storing one checkpoint per Transformer block, the checkpoint memory is:

$2 \times bsz \times seq \times hd \times nl / ci$

where $bsz$ is the batch size, $seq$ is the sequence length, $hd$ is the hidden dimension, $nl$ is the number of layers, and $ci$ is the checkpoint interval (number of Transformer blocks between two activation checkpoints).

What this equation computes: it gives the total bytes needed to store the activation checkpoints throughout the model. The factor of 2 accounts for FP16 storage (2 bytes per value), and $bsz \times seq \times hd$ is the size of a single layer's input tensor. Dividing by $ci$ (the checkpoint interval) gives the number of checkpoints stored; when $ci = 1$ (one checkpoint per Transformer block), the number of checkpoints equals $nl$ .

Why this form: This linearizes what would otherwise be a quadratic scaling with model depth. Without checkpointing, activation memory would be proportional to $nl$ (storing activations for all layers simultaneously). With checkpointing at interval $ci$ , only $1/ci$ of the activations are stored, and the rest are recomputed during backward, adding ~33% computational overhead (one extra forward pass during backward for each checkpointed segment). The paper assumes $ci = 1$ (most aggressive checkpointing) for its scale estimates.

Why it still matters at extreme scale: Figure 2a column 7 shows that even with aggressive checkpointing, activations reach 0.76 TB for a 10T model and 3 TB for a 100T model at batch size 32. These can fit in CPU memory (a DGX-2 node has 1.5 TB CPU RAM, and next-generation hardware will have more), but they are too large for GPU memory, motivating ZeRO-Infinity's activation offloading capability.

Model State Working Memory (MSWM). This is the minimum GPU memory required to execute the largest single operator after all model states have been offloaded. It is determined by the size of the parameters and gradients of that operator:

$4 \times hd \times 4hd$

where the linear layer transforms from hidden dimension $hd$ to $4hd$ (the feed-forward expansion layer), and the factor of 4 accounts for 2 bytes per FP16 parameter plus 2 bytes per FP16 gradient.

What this equation computes: the minimum contiguous GPU memory buffer needed to hold both the parameter tensor and its gradient tensor for the largest matrix multiplication in a Transformer block.

Why this form matters: For large hidden dimensions, MSWM grows substantially. Figure 2a column 8 shows that beyond 100B parameters (corresponding to large hidden dimensions), MSWM reaches multiple gigabytes and must be contiguous. GPU memory fragmentation can cause out-of-memory errors even when total free memory exceeds MSWM, because the allocator cannot find a single contiguous block of the required size. This is the problem that motivates memory-centric tiling — decomposing the large operator into tiles whose individual working memory fits within available contiguous chunks.

Activation Working Memory (AWM). This is the memory needed during backward propagation to store the activations between two consecutive checkpoints (the ones that must be recomputed before computing gradients). For one activation checkpoint per Transformer block $(ci = 1)$ :

$bsz \times seq \times ci \times (16 \times hd + 2 \times attn\_heads \times seq)$

where $bsz$ is batch size, $seq$ is sequence length, $ci$ is the checkpoint interval, $hd$ is hidden dimension, and $attn\_heads$ is the number of attention heads.

What this equation computes: the bytes of GPU memory needed to store the full set of activations within one checkpointed segment during recomputation. The term $16 \times hd$ accounts for the activations from the feed-forward and layer norm operations within one Transformer block, while $2 \times attn\_heads \times seq$ accounts for the attention-related activations (the attention score matrix of size $seq \times seq$ per head, stored in FP16).

Why this form: Unlike MSWM, AWM is composed of many small tensors rather than one large contiguous tensor, so it does not cause fragmentation-related out-of-memory issues. However, Figure 2a column 9 shows that AWM grows to multiple terabytes beyond 10T parameters, exceeding GPU memory and potentially CPU memory as well, which motivates the option of further reducing checkpoint frequency or offloading to NVMe.

Bandwidth Requirements Characterization (Foundation for Design)

Section 4 of the paper establishes the bandwidth side of the design constraints: even if memory capacity is sufficient, the data movement between storage tiers must be fast enough that computation is not stalled waiting for data. The paper provides a quantitative framework based on arithmetic intensity (AIT) and derives specific bandwidth targets for parameters/gradients, optimizer states, and activation checkpoints.

Efficiency model. The paper starts from a simple model of training efficiency assuming no overlap between computation and communication:

$efficiency = \frac{compute\_time}{compute\_time + communication\_time}$

Expanding in terms of peak throughput ( $peak_{tp}$ ), bandwidth ( $bw$ ), and arithmetic intensity ( $ait$ ):

$efficiency = \frac{ait \times bw}{ait \times bw + peak_{tp}}$

where $peak_{tp}$ is the achievable peak throughput of the accelerator (measured empirically as 70 TFlops/GPU on V100 for the model configurations studied), $bw$ is the available data movement bandwidth, and $ait$ is the arithmetic intensity — the ratio of total computation (FLOPs) to total data movement (bytes) for that component.

What this equation computes: given the accelerator's compute speed, the available bandwidth, and the workload's arithmetic intensity, what fraction of peak throughput can be sustained. If $bw$ is infinite (or $ait$ is very high), efficiency approaches 1.0; if $bw$ is low relative to $peak_{tp}$ , efficiency drops because the accelerator spends most of its time waiting for data.

Why this form: This is a standard roofline model adapted to the variables that ZeRO-Infinity can control. The system cannot change $peak_{tp}$ (fixed by hardware) or $ait$ for a given model configuration (determined by batch size, sequence length, and hidden dimension). It can influence $bw$ — the effective bandwidth for accessing model states and activations — through bandwidth-centric partitioning and overlapping. The efficiency equation tells us what $bw$ targets must be hit for training to remain efficient.

What's missing: This model assumes no overlap between communication and computation. The paper acknowledges this limitation (it is a lower bound on achievable efficiency) and uses it as a conservative design target. The overlap-centric design in Section 6.2 pushes actual efficiency above this theoretical floor.

Total computation per iteration. The paper quantifies the total FLOPs per training iteration for a Transformer model. The forward pass requires $2 \times bsz \times seq \times params$ FLOPs (one multiply-add per parameter per token in the sequence for each batch element). The backward pass is approximately twice the forward pass. With activation checkpointing, there is an additional forward pass during recomputation. Total:

$computation\_per\_iter = 2 \times 4 \times bsz \times seq \times parameters$

Substituting the parameter count expression for Transformers:

$= 2 \times 4 \times 12 \times bsz \times seq \times nl \times hd^2$

Arithmetic intensity for parameters and gradients. During forward and backward propagation, each parameter must be loaded from its storage location to GPU registers at least twice: once for the forward pass, once for the actual backward pass. With activation checkpointing, parameters may be loaded a third time during recomputation. Each gradient must be written from GPU registers to its storage location at least once. Assuming parameters and gradients are co-located (stored at the same final location), the total data movement is $4 \times parameters$ (3 reads + 1 write, but in half-precision the total byte count is $2 \times 4 \times params$ ). The arithmetic intensity is:

$ait_{param\_grad} = seq \times bsz$

What this equation computes: the number of FLOPs performed per byte of parameter and gradient data moved. For a given sequence length and batch size, each byte of parameter data enables $seq \times bsz$ FLOPs of computation because the same parameter is reused across every position in the sequence and every element in the batch.

Why this form — and its implications: This tells us that parameter and gradient bandwidth requirements depend only on batch size and sequence length, not on model size. At a batch size of 1 and sequence length of 1024, each parameter byte enables 1024 FLOPs. At batch size 16, it enables 16,384 FLOPs. This is why the paper's bandwidth analysis in Figure 3a varies batch size from 1 to 16: larger batch sizes increase $ait$ , making the workload more tolerant of low bandwidth. The key insight is that for small batch sizes (which are necessary when scaling to many GPUs to keep the effective batch size reasonable), the bandwidth requirement is stringent — approximately 70 GB/s for 50% efficiency, which is close to GPU-GPU NVLink bandwidth and far exceeds single PCIe bandwidth (~12 GB/s).

Arithmetic intensity for optimizer states. During the optimizer step, all optimizer states (FP32 momentum, variance, parameter copy, and gradient copy — 16 bytes per parameter) must be read at least once and written at least once. Total data movement: $2 \times 16 \times params = 32 \times params$ bytes. The computation is dominated by the element-wise Adam update, which is roughly 10-20 FLOPs per parameter (several multiply-adds for momentum, variance, and parameter update). The arithmetic intensity is:

$ait_{optimizer} = seq \times bsz / 4$

What this equation computes: the FLOPs per byte for optimizer state data movement. It is 4× smaller than the parameter/gradient $ait$ because optimizer states are 4× larger (32 bytes of optimizer state per parameter vs. 8 bytes for parameters and gradients combined: 4 bytes gradient + 4 bytes parameter in half precision).

Why this matters crucially for design: Optimizer states have the lowest arithmetic intensity of any component and therefore the highest bandwidth requirement. Additionally, optimizer updates cannot be overlapped with forward/backward computation (they happen after the backward pass completes), so they represent a serial bottleneck. Figure 3b shows that achieving 90% efficiency at batch size 2 requires nearly 1.5 TB/s of effective bandwidth — exceeding even GPU HBM2 bandwidth. This means the optimizer step inherently requires aggregate bandwidth from many devices, which justifies ZeRO-Infinity's partitioned optimizer design where the update is performed in parallel across all GPUs and CPUs.

Arithmetic intensity for activation checkpoints. Activation checkpoints are written once during forward and read once during backward. The total data moved is $2 \times total\_activation\_checkpoint\_bytes$ . Using the checkpoint memory equation from Section 3:

$ait_{activation} = 24 \times hd \times ci$

where $hd$ is the hidden dimension and $ci$ is the checkpoint interval.

What this equation computes: the FLOPs per byte of activation checkpoint data movement. For each byte of activation checkpoint, the model performs $24 \times hd \times ci$ FLOPs of computation.

Why this form — and why activations are the easiest to offload: The $ait$ for activations grows with hidden dimension and checkpoint interval. For a model with $hd = 8192$ (roughly a 10B parameter model), $ait = 24 \times 8192 \times 1 = 196,608$ FLOPs/byte. This is extremely high — each byte of activation data supports nearly 200K FLOPs of computation. Consequently, Figure 3c shows that even 2 GB/s of bandwidth sustains over 50% efficiency, and the requirement drops below 1 GB/s for hidden sizes above 8K. The practical implication: activation checkpoint offloading to CPU (at ~3 GB/s per GPU over PCIe Gen 3) is essentially free for large models. This is why ZeRO-Infinity can offload activations without significant performance penalty, as demonstrated in Figure 6e.

Synthesizing bandwidth requirements. The analysis yields three distinct bandwidth targets for training to remain efficient:

Parameters and gradients: ~70 GB/s (comparable to GPU-GPU bandwidth, must be achieved via allgather over NVLink)
Optimizer states: ~1.5 TB/s (must be achieved via aggregate bandwidth across many devices in parallel)
Activation checkpoints: ~1-4 GB/s (easily satisfied by PCIe to CPU)

These numbers directly motivate ZeRO-Infinity's design decisions: parameters and gradients must stay in or near GPU memory (or use bandwidth-centric partitioning to scale PCIe bandwidth with data-parallel degree), optimizer states can be offloaded to CPU/NVMe because the update is parallelizable across all devices, and activation checkpoints can be offloaded to CPU essentially for free for large models.

Infinity Offload Engine: Model State Placement and Movement

The infinity offload engine is the central mechanism that enables ZeRO-Infinity to exploit heterogeneous memory. It is built on top of ZeRO-3's parameter partitioning and extends it by allowing each partitioned model state (parameters, gradients, optimizer states) to be placed independently on GPU, CPU, or NVMe memory.

Design principle: independent placement per state type. Table 2 in the paper (reproduced from the evaluation section) specifies the device placement strategy as a function of how aggressively the user wants to trade off memory for performance. The three tiers are:

ZeRO-Inf-CPU: Optimizer states and gradients are partitioned and stored in CPU memory; parameters are partitioned between GPU and CPU (frequently accessed parameters stay on GPU, others on CPU). This enables model sizes up to ~100B parameters on a single DGX-2 node (Figure 6a).
ZeRO-Inf-NVMe: Optimizer states and gradients are partitioned and stored in NVMe; parameters are partitioned between GPU, CPU, and NVMe. This enables model sizes up to 1T parameters on a single node and 32T parameters across 32 nodes (Figure 1).

The key difference from ZeRO-Offload is that parameters are partitioned (ZeRO-3 style) rather than replicated (ZeRO-2 style). This means a single parameter no longer needs to fit in any single GPU's memory — it is distributed across all GPUs, and each GPU holds only a fraction of the total parameters. Combined with NVMe offloading, the total model state capacity becomes the aggregate NVMe capacity across all nodes, which is approximately 50× larger than aggregate GPU memory (Figure 2b).

Lifecycle of a parameter during training. The infinity offload engine orchestrates the following sequence for each parameter during one training iteration:

Prefetch phase (before forward pass): The parameter, which is partitioned across data-parallel GPUs and stored in its designated tier (GPU/CPU/NVMe), is gathered. If stored on NVMe, a three-stage pipeline executes: (a) DeepNVMe reads the parameter shard from NVMe into CPU pinned memory (nc-transfer), (b) the shard is copied from CPU pinned memory to GPU memory via PCIe (cg-transfer), (c) an allgather collective combines all shards from all GPUs into the full parameter on each GPU (gg-transfer).
Forward pass: The full parameter is used in the forward computation. Immediately after the operator executes, the parameter on non-owner GPUs is discarded. The owner GPU optionally offloads it back to CPU or NVMe.
Backward pass: The same gather-compute-release pattern repeats for gradient computation. After the backward pass for an operator, the gradients are reduced across GPUs via reduce-scatter (each GPU ends up with the gradient shard for the parameters it owns), and the gradient shard is optionally offloaded.
Optimizer step: After the full backward pass completes, the optimizer states corresponding to each parameter are updated. For parameters stored in CPU memory, this happens on the CPU; for parameters stored in NVMe, data is brought into CPU memory in chunks (limited by CPU memory capacity), the Adam update is computed, and updated states are written back.

Why this lifecycle works at scale: The critical property is that at any given moment, only the parameters needed for the current operator (and a few upcoming operators, managed by the prefetcher) need to be present in GPU memory. All other parameters remain partitioned and offloaded. The GPU working memory requirement is thus the MSWM for the largest operator plus the working memory for the prefetch pipeline, not the full model size. This is what breaks the GPU memory wall — the GPU only needs to host a tiny window of the model at any time.

Bandwidth-Centric Partitioning

This is the key innovation that enables ZeRO-Infinity to achieve sufficient aggregate bandwidth from CPU and NVMe memory despite individual PCIe links being slow. The paper contrasts it with the broadcast-based approach used in prior work.

The broadcast-based bottleneck (ZeRO-3 and ZeRO-Offload). In standard ZeRO-3 and ZeRO-Offload, each parameter is owned by a single data-parallel process (GPU). When that parameter is needed, the owning GPU broadcasts it to all other GPUs. If the parameter is stored in CPU or NVMe memory, the owning GPU must first bring the entire parameter from its storage location over its single PCIe link (~12 GB/s for PCIe Gen 3) before broadcasting it to peers over NVLink (~70 GB/s). This means:

Only one PCIe link is active for any given parameter's data movement
The effective CPU/NVMe-to-GPU bandwidth is capped at single-PCIe speeds (12 GB/s) regardless of how many GPUs are in the system
This bandwidth is far below the ~70 GB/s needed for parameter/gradient efficiency (Figure 3a)

The allgather-based approach (ZeRO-Infinity). ZeRO-Infinity partitions each individual parameter across all data-parallel GPUs rather than assigning complete parameters to individual GPUs. When a parameter is needed:

Every GPU reads its shard (1/DP of the parameter) from its designated storage tier
An allgather collective combines all shards into the full parameter on every GPU

Because every GPU reads its shard simultaneously, all PCIe links are active in parallel. The effective CPU/NVMe-to-GPU bandwidth scales linearly with the data-parallel degree:

$effective\_bw = dp \times single\_pcie\_bw$

What this computes: With 16-way data parallelism on a DGX-2 node, the effective bandwidth is $16 \times 12 = 192$ GB/s for CPU-to-GPU transfers (limited by the aggregate PCIe bandwidth of each DGX-2 node, which is approximately 48 GB/s total or 3 GB/s per GPU × 16 GPUs). For NVMe-to-GPU transfers, the effective bandwidth is limited by the aggregate NVMe read bandwidth of the node, which is approximately 25 GB/s (1.6 GB/s per GPU).

Why the allgather and broadcast have the same GPU-GPU cost: Both allgather and broadcast have the same communication volume when the data is already on the GPU — $N-1$ chunks of data cross the network in both cases. The difference is entirely in how the data gets to the GPU network in the first place. With broadcast, one GPU does all the PCIe work; with allgather, all GPUs share the PCIe work. This is "a game changer when the data is located in NVMe or CPU" because PCIe bandwidth is the bottleneck, not NVLink bandwidth.

How this scales to clusters. On 64 DGX-2 nodes (1024 GPUs), ZeRO-Infinity has access to over 3 TB/s of aggregate CPU memory bandwidth and over 1.5 TB/s of aggregate NVMe bandwidth. The paper describes this as "virtually unlimited" heterogeneous memory bandwidth — far exceeding the ~70 GB/s and ~1.5 TB/s targets from the bandwidth analysis, meaning the system is compute-bound rather than memory-bound at cluster scale.

What this doesn't solve: On a single node (16 GPUs), the aggregate bandwidth from NVMe is ~25 GB/s, which is below the ~70 GB/s target for parameters and gradients. This means single-node training with NVMe offloading will be bandwidth-limited for parameters and gradients, requiring either larger batch sizes (to increase $ait$ ) or the overlap-centric design to hide the latency. The paper addresses this through the overlap engine and acknowledges the limitation implicitly by showing that single-node performance degrades at extreme scales (Figure 5c shows 1T model on a single node achieving lower TFlops/GPU than smaller models).

Overlap-Centric Design: The Dynamic Prefetcher and Communication Pipeline

The overlap-centric design addresses the residual bandwidth bottleneck that bandwidth-centric partitioning cannot eliminate, particularly on small numbers of GPUs where the aggregate PCIe/NVMe bandwidth is limited. It does this by overlapping the data movement for future operators with the computation of the current operator, effectively hiding communication latency.

The three-stage communication pipeline. Accessing a parameter stored in NVMe requires three sequential steps:

nc-transfer: DeepNVMe reads the data from NVMe into CPU pinned memory
cg-transfer: The data is copied from CPU pinned memory to GPU memory over PCIe
gg-transfer: An allgather collective combines all GPU's shards into the full parameter

If these three steps were executed sequentially for each operator just before that operator executes, the total communication cost would be the sum of all three latencies for every parameter, and the GPU would be idle during the entire data movement. The overlap engine's goal is to pipeline these stages: while the GPU computes operator $i$ , the system simultaneously:

Performs nc-transfer for operator $i+3$ 's parameters (if stored on NVMe)
Performs cg-transfer for operator $i+2$ 's parameters
Performs gg-transfer (allgather) for operator $i+1$ 's parameters

By the time the GPU reaches operator $i+1$ , its parameters are already fully gathered and ready in GPU memory.

The dynamic prefetcher. This is the component that realizes the pipeline. It works by:

Tracing: During the first training iteration, the prefetcher records the sequence of operators executed during forward and backward passes, building an internal map of the operator graph. This map captures the order in which submodules are invoked.
Prefetching: During subsequent iterations, the prefetcher tracks the current position in the operator sequence. Before executing operator $i$ , it has already initiated (and potentially completed) the nc, cg, and gg transfers for operators $i+3$ , $i+2$ , and $i+1$ respectively.
Dynamic adaptation: The paper notes that the operator sequence map can be updated "in case of dynamic workflow, allowing for appropriate prefetching even when the forward and backward propagation changes across iterations." This handles models with data-dependent control flow where the execution graph varies between iterations.

Gradient overlapping in the backward pass. During backward propagation, a similar overlapping strategy applies. While computing the gradient for operator $i$ :

The reduce-scatter for gradients of operator $i+1$ executes (combining gradients across GPUs and partitioning the result)
The partitioned gradients from the reduce-scatter of operator $i+2$ are being transferred to CPU or NVMe

This means that by the time the backward pass completes, most gradients are already reduced, partitioned, and offloaded to their storage locations, minimizing the post-backward synchronization cost before the optimizer step can begin.

Why this design is effective. The overlap-centric design transforms the data movement cost from a serial bottleneck (communication followed by computation) to a pipeline where communication is hidden behind computation. The effectiveness depends on the depth of the pipeline (can at least 3 future operators be prefetched while the current operator computes?) and the ratio of computation time to communication time for each operator. For large operators (e.g., a linear layer with hidden dimension 8K+), the computation time is large enough to hide the communication of several subsequent operators. For very small operators or very small batch sizes, computation time may be insufficient to fully hide communication, which is why Figure 6d shows that the benefit of overlapping diminishes at large batch sizes (where communication is a smaller fraction of total time regardless).

Memory-Centric Tiling

Memory-centric tiling addresses a different problem than the offload engine: even after all model states are offloaded, a single large operator may require more contiguous GPU memory than is available due to memory fragmentation. This is the Model State Working Memory (MSWM) problem identified in Section 3.

The problem concretely. Consider the feed-forward expansion layer in a Transformer block, which multiplies a $(bsz \times seq, hd)$ activation matrix by a $(hd, 4hd)$ weight matrix. For $hd = 65536$ (64K, corresponding to a model with tens of trillions of parameters), the weight matrix alone is $65536 \times 262144 = 17.18$ billion parameters, or 34.36 GB in FP16. The gradient is another 34.36 GB. The total MSWM is 68.72 GB — far exceeding the 32 GB available on a V100 GPU. Even if the aggregate free memory across all allocations is sufficient, GPU memory allocators struggle to find a single contiguous 68 GB block, especially after memory fragmentation from other allocations (activations, temporary buffers, prefetched parameters for other operators).

Prior solutions use tensor-slicing model parallelism: split this weight matrix across multiple GPUs so each GPU holds only a slice. This is effective but requires the data scientist to manually refactor the model code and introduces additional communication (each GPU must allreduce its partial output).

The tiling solution. Memory-centric tiling decomposes the large operator into a sequence of smaller operators that execute on a single GPU sequentially:

The weight matrix $W$ of shape $(hd, 4hd)$ is partitioned into $T$ tiles along the output dimension: each tile $W_t$ has shape $(hd, 4hd/T)$ .
The forward pass computes $Y_t = X \times W_t$ for each tile sequentially, concatenating the results to form the full output $Y$ .
The backward pass computes gradients for each tile sequentially: $dW_t = X^T \times dY_t$ and $dX_t = dY_t \times W_t^T$ .
When combined with ZeRO-3, the parameters and gradients for each tile are fetched and released one at a time: before computing tile $t$ , allgather $W_t$ from its owners; after computing and reducing the gradient, release $W_t$ and offload $dW_t$ .

Why this works. The MSWM for the tiled operator is the size of one tile's parameters and gradients, which is $1/T$ of the original. With $T = 16$ , the 68.72 GB MSWM for the 64K hidden dimension case becomes 4.3 GB — easily fitting in a V100's 32 GB, even with fragmentation (since the paper's experiments pre-fragmented GPU memory into 2 GB contiguous chunks to test robustness). Figure 6b demonstrates this: without tiling, the maximum trainable hidden dimension is 8K; with tiling factor 4, it scales to 16K; with tiling factor 16, it reaches 64K.

The tradeoff. Tiling reduces working memory at the cost of sequentializing computation that could have run in parallel (either on multiple GPU cores or across multiple GPUs via model parallelism). Each tile requires its own gather-compute-release cycle for parameters, which adds communication overhead. However, since the tiles of a single operator are computed sequentially on the same GPU, there is no additional GPU-GPU communication compared to untiled execution — the same allgather volume is spread across multiple smaller allgathers. The primary cost is the loss of parallelism within the matrix multiplication (a $(hd, 4hd/T)$ matmul has lower GPU utilization than a $(hd, 4hd)$ matmul) and the overhead of multiple kernel launches.

Why this is preferred over model parallelism. The paper argues that memory-centric tiling "greatly simplifies the DL system stack by avoiding the need for model parallelism." The comparison is not primarily about performance (model parallelism with tensor-slicing can be highly efficient) but about usability: model parallelism requires manual code refactoring to replace single-GPU operators with distributed versions, while tiling is applied automatically by ZeRO-Infinity's hook system without any model code changes. For models with complex operator graphs that don't cleanly map to tensor-slicing, tiling also avoids the communication overhead of inter-GPU synchronization within a single operator.

Ease-Inspired Implementation: Automating Data Movement and Initialization

The ease-inspired implementation is the engineering layer that makes ZeRO-Infinity usable without model code refactoring. It consists of two main automated mechanisms: hook-based data movement during training and partitioned model initialization.

Hook-based automated data movement. PyTorch models are organized as a hierarchy of nn.Module objects. ZeRO-Infinity recursively injects hooks into every submodule. The hooks are:

Pre-forward hook: Executed before the submodule's forward() method. This hook checks whether the submodule's parameters are currently available in GPU memory. If not, it triggers the appropriate gather operation (allgather from the parameter owners) and blocks until the parameters are available. Thanks to the dynamic prefetcher, parameters for most operators are already gathered and waiting, so the hook rarely needs to block.
Post-forward hook: Executed after the submodule's forward() method returns. This hook partitions the parameters again (non-owner GPUs discard their copies) and optionally offloads them to CPU or NVMe, freeing GPU memory for subsequent operators.
Pre-backward hook: Similar to pre-forward, ensures parameters are available for gradient computation.
Post-backward hook: After gradients are computed, triggers the reduce-scatter to combine gradients across GPUs and partition the result, then optionally offloads the gradient shards.

Handling external parameters. A complication arises when a parameter allocated in one submodule is used in the forward/backward pass of a different submodule. The canonical example is weight tying in language models: the embedding layer's weight matrix (used at the input to map tokens to vectors) is also used at the output (to map hidden states back to vocabulary logits). The paper calls these "external parameters."

The hook system cannot automatically detect such cross-module usage by static analysis alone. ZeRO-Infinity provides three mechanisms:

Manual registration API: The data scientist can explicitly register external parameters. After registration, they are treated identically to regular parameters and included in the prefetching system.
Intercepting partitioned parameter access: At initialization time, ZeRO-Infinity replaces each module's parameter hash table with a subclassed version that intercepts tensor accesses. When a partitioned parameter is accessed from outside its owning module, the interceptor performs a blocking allgather to materialize the full parameter, registers it as an external parameter, and returns the gathered tensor. This transparently handles cases where external code directly accesses a parameter tensor.
Activation introspection: If a submodule's forward pass returns a partitioned parameter as part of its output activations (to be consumed by another submodule), ZeRO-Infinity inspects the output tensors. If a partitioned parameter is detected, it is gathered, registered as external, and the gathered version is substituted in the output.

Automatic model partitioning during initialization. A practical problem with very large models: if each data-parallel process initializes the full model before partitioning (the standard PyTorch approach), the initialization step itself requires $dp \times model\_size$ aggregate memory. For a 500B parameter model in FP16, that's 1 TB per process, or 8 TB across an 8-GPU node — exceeding both GPU and CPU memory.

ZeRO-Infinity solves this by providing a Python context manager that decorates torch.nn.Module.__init__. Under this context:

When a submodule's constructor allocates parameters, ZeRO-Infinity intercepts the allocation.
The parameters are immediately partitioned among the data-parallel group: each process keeps only its shard.
The non-owned shards are freed, so the full parameter never exists on any single process.

The result: "only individual sub-modules are fully initialized before they are partitioned, and the full model is never replicated on all the data parallel process." The 500B parameter model requires only 1 TB of aggregate CPU memory during initialization regardless of the number of data-parallel processes, because the per-process memory is always 1/DP of the current submodule's size plus the already-partitioned submodules.

Why this matters for usability. The combined effect of hook-based data movement, automatic external parameter handling, and partitioned initialization is that ZeRO-Infinity trains PyTorch models written with standard data parallel patterns (modules, submodules, forward() methods) without any code changes. The data scientist does not need to:

Insert explicit allgather/reduce-scatter calls
Manage parameter placement across memory tiers
Split operators for model parallelism
Partition the model at initialization
Handle weight tying or other cross-module parameter sharing

This is in direct contrast to 3D parallelism, where the data scientist must manually insert tensor-slicing operations and partition the model into pipeline stages.

DeepNVMe: High-Performance NVMe I/O

DeepNVMe is the low-level C++ library within the infinity offload engine that enables ZeRO-Infinity to achieve near-peak NVMe bandwidth. Its design addresses the specific challenges of using NVMe storage as a memory tier for DL training.

Design requirements. NVMe SSDs offer high sequential bandwidth (~25 GB/s read on a DGX-2 node) but have characteristics that make them challenging for DL training workloads:

Random read/write performance is much lower than sequential
I/O requests have latency (tens of microseconds) that must be hidden
Data must be in CPU pinned memory for DMA transfers to/from GPU
Pinned memory is a scarce system resource (over-subscription degrades system performance)

Key features of DeepNVMe:

Bulk asynchronous read/write: DeepNVMe supports issuing large read/write requests that complete asynchronously, returning immediately and providing an explicit synchronization primitive (flush) to wait for completion. This asynchrony is what enables the overlap engine to initiate NVMe-to-CPU transfers for future operators while the GPU computes the current operator.
Aggressive I/O parallelization: DeepNVMe parallelizes I/O requests from both single-threaded and multi-threaded contexts to saturate the NVMe device's queue depth and achieve near-peak sequential bandwidth. The paper mentions "smart work scheduling" without providing detailed algorithmic specifics, but the implication is that requests are distributed across NVMe queues and drives to maximize throughput.
Zero-copy design: DeepNVMe avoids data copying by using pre-allocated pinned memory buffers. When reading from NVMe, data is DMA-transferred directly to the pinned buffer. When writing to NVMe, data is DMA-transferred directly from the pinned buffer. Intermediate copies to non-pinned memory would add CPU overhead and reduce effective bandwidth.
Pinned memory management layer: Since pinned memory is limited (allocating too much degrades system performance or causes instability), ZeRO-Infinity manages a small pool of pinned buffers (tens of GBs) and reuses them for offloading the entire model state (potentially tens of TBs). The reuse prevents CPU and GPU memory fragmentation. This layer also provides PyTorch tensors backed by pinned memory, enabling in-place computation — a tensor can be computed in-place in pinned memory and then written to NVMe without any additional copies, improving bandwidth utilization.

Why near-peak NVMe bandwidth is critical. As the bandwidth analysis in Section 4 showed, parameter and gradient bandwidth requirements are stringent (~70 GB/s per GPU for efficiency). At single-node scale (16 GPUs), even with bandwidth-centric partitioning, the aggregate NVMe bandwidth is only ~25 GB/s — insufficient on its own. Every GB/s of NVMe bandwidth that DeepNVMe can extract reduces the reliance on overlapping and larger batch sizes to hide the remainder. At multi-node scale, the aggregate NVMe bandwidth exceeds requirements (1.5 TB/s on 64 nodes), so individual node efficiency becomes less critical. The pinned memory management is equally important: without reuse, offloading a 10 TB model might require 10 TB of pinned memory (impossible on current systems); with reuse, tens of GBs of pinned buffers are cycled to transfer the full model in chunks.

How this connects to the optimizer step. For NVMe-offloaded optimizer states (the ZeRO-Inf-NVMe configuration), the optimizer step requires:

Reading optimizer states from NVMe into pinned CPU memory in chunks (limited by available CPU memory)
Performing the Adam update on CPU (leveraging all available CPU cores)
Writing updated optimizer states and new parameters back to NVMe

DeepNVMe's bulk read/write and asynchronous completion enable overlapping the NVMe reads for the next chunk, the CPU computation for the current chunk, and the NVMe writes for the previous chunk, creating a three-stage pipeline that hides NVMe latency. This is specific to the optimizer step (which operates on all parameters simultaneously) rather than the layer-by-layer forward/backward pipeline.

4. Key Insights and Innovations

Innovation 1: Reframing the Memory Wall as a Bandwidth-Allocation Problem Rather Than a Capacity Problem

The dominant framing in large-model training systems prior to ZeRO-Infinity treated memory as fundamentally about capacity: how many parameters can you fit in the aggregate GPU HBM of a cluster? The solutions — 3D parallelism, ZeRO-Offload — were about squeezing more model into a fixed GPU memory budget, either by splitting across GPUs or by offloading subsets of state to CPU. The unspoken assumption was that offloading to slower memory tiers (CPU DRAM, NVMe) was a last resort that would inevitably cripple training throughput because those tiers are an order of magnitude (CPU) or two orders of magnitude (NVMe) slower than GPU HBM.

ZeRO-Infinity makes a fundamental conceptual move: the true bottleneck is not the capacity of slow memory, but whether the aggregate bandwidth across many slow-memory links can be made to exceed the accelerator's demand for data. This is not a minor reframing — it completely changes the design space. Instead of asking "how do we minimize offloading because it's slow?", ZeRO-Infinity asks "can we make the collective bandwidth of all PCIe and NVMe links in a cluster large enough that training remains compute-bound?"

The paper makes this reframing explicit through the arithmetic intensity (AIT) analysis in Section 4, which is itself a conceptual innovation. Prior work on heterogeneous training (ZeRO-Offload, SwapAdvisor, Sentinel) treated bandwidth constraints empirically — they observed that offloading hurts throughput and tried to mitigate it. ZeRO-Infinity instead derives quantitative bandwidth targets from first principles: parameters and gradients need ~70 GB/s, optimizer states need ~1.5 TB/s, activations need only ~1-4 GB/s. These numbers are not hardware-specific constants — they emerge from the AIT expressions ait_param_grad = seq × bsz, ait_optimizer = seq × bsz / 4, and ait_activation = 24 × hd × ci, which connect workload characteristics (batch size, sequence length, hidden dimension) to bandwidth requirements through the roofline efficiency model efficiency = ait × bw / (ait × bw + peak_tp).

The diagnostic power of this reframing is evident in how it explains why certain offloading strategies work and others fail. It predicts that activation offloading is essentially free for large models (Figure 3c: <1 GB/s needed for hidden dimensions above 8K), which ZeRO-Infinity confirms experimentally (Figure 6e: minimal throughput impact at 32K–64K hidden sizes). It predicts that optimizer state offloading requires aggregate bandwidth exceeding even GPU HBM (1.5 TB/s for 90% efficiency at batch size 2), which justifies why the optimizer update must be parallelized across all devices and cannot be done on a single GPU. It predicts that parameter offloading to NVMe on a single node will be bandwidth-limited (25 GB/s aggregate vs. 70 GB/s target), which motivates the overlap-centric design and explains why single-node performance drops for extreme-scale models (Figure 5c).

What distinguishes this from a standard roofline model application is that it's used prescriptively to design the system, not just descriptively to analyze an existing one. The bandwidth targets directly motivate: (a) why bandwidth-centric partitioning (aggregating across all PCIe links) is necessary rather than just beneficial, (b) why the overlap engine must pipeline NVMe→CPU, CPU→GPU, and GPU allgather as three stages rather than treating offloading as a monolithic operation, and (c) why optimizer states should be offloaded to CPU/NVMe while parameters require more careful placement.

The implication beyond this paper is significant: as accelerators get faster, the aggregate bandwidth needed from slow memory grows proportionally (Table 3: 10× compute → 30 GB/s per device, 100× compute → 300 GB/s per device). But this is not a doom-and-gloom projection — the paper points to Summit (2018) as evidence that connecting accelerators to CPU memory at 40 GB/s via NVLink is already feasible. The reframing thus serves as a design principle for future hardware: invest in high-bandwidth interconnects between accelerators and all memory tiers (not just HBM), because the aggregate bandwidth of slow memory can substitute for capacity of fast memory if the interconnect is capable enough.

Innovation 2: Bandwidth-Centric Partitioning as a New Primitive for Heterogeneous Data Movement

Prior to ZeRO-Infinity, the standard approach for moving parameters from slow memory to GPUs in a partitioned-training setting was broadcast from owner: each parameter is fully owned by one GPU, that GPU reads the entire parameter from CPU/NVMe over its single PCIe link, and then broadcasts it to peers. This was the approach in ZeRO-3 (GPU-only) and ZeRO-Offload (with CPU offloading). It is a natural design because it maps cleanly to the ownership semantics of ZeRO-3: the owner is responsible for its parameters, so of course the owner should fetch them.

The conceptual innovation in ZeRO-Infinity is to sever the coupling between data ownership and data movement responsibility. Parameters are still owned by individual GPUs (for optimizer updates), but they are retrieved collectively: each parameter is split into dp shards, one per GPU, and all GPUs read their shard simultaneously from their respective storage locations, then allgather to reconstruct the full tensor. The paper calls this "bandwidth-centric partitioning."

The difference is not merely an implementation optimization — it is a qualitative change in how the system's bandwidth scales with cluster size. Under broadcast-from-owner, the effective CPU/NVMe-to-GPU bandwidth is capped at single_pcie_bw (~12 GB/s) regardless of whether you have 8 GPUs or 800. Under allgather-from-all, the effective bandwidth scales as dp × single_pcie_bw. On a single DGX-2 node with 16 GPUs, this is ~48 GB/s from CPU and ~25 GB/s from NVMe — still below the ~70 GB/s target for parameters, but close enough that overlapping can close the gap. On 64 nodes (1024 GPUs), the aggregate exceeds 3 TB/s from CPU and 1.5 TB/s from NVMe — far beyond any reasonable bandwidth target, meaning the system is definitively compute-bound at scale.

This scaling property is what the paper means by "virtually unlimited heterogeneous memory bandwidth" (Section 5.2.1). It is not that the per-link bandwidth is unlimited — PCIe Gen 3 is fixed at ~12 GB/s per GPU — but that the aggregate bandwidth grows with the number of GPUs, and since large model training inherently requires many GPUs (for compute throughput), the aggregate bandwidth naturally grows to match or exceed requirements. This is a structural property of the allgather approach, not an empirical tuning result.

The significance of this innovation extends beyond ZeRO-Infinity. It establishes a design pattern for any heterogeneous system where data is partitioned across many nodes and must be materialized on all nodes for computation: collective retrieval scales, point-to-point retrieval does not. This is not obvious a priori — one might assume that since an allgather and a broadcast have the same communication volume (N-1 chunks of data crossing the network), they are equivalent. The insight is that when the source data is not in GPU memory (the case for offloading), the bottleneck is the PCIe transfer into the GPU network, and allgather parallelizes this while broadcast serializes it. The paper is explicit about this: "this is a game changer when the data is located in NVMe or CPU" (Section 6.1).

A subtle consequence of bandwidth-centric partitioning is that it harmonizes the bandwidth requirements across all three model state types. With broadcast-from-owner, parameters required 70 GB/s (Figure 3a) but only got 12 GB/s — a 5.8× shortfall that forced large batch sizes or inefficiency. With allgather-from-all at moderate scale, parameters get aggregate bandwidth comparable to or exceeding GPU-HBM bandwidth, which means parameter movement is no longer the bottleneck it was in ZeRO-Offload. This explains a result from the evaluation that might otherwise seem surprising: ZeRO-Infinity matches 3D parallelism's throughput at 500B parameters on 512 GPUs (Figure 5a, leftmost bar) even though 3D parallelism keeps all model states in GPU memory — because at that scale, the aggregate PCIe/NVMe bandwidth via allgather is sufficient for compute-bound training.

Innovation 3: Memory-Centric Tiling as an Alternative to Model Parallelism for Usability

Model parallelism (tensor-slicing) has been the standard approach for handling individual layers too large to fit on a single GPU since at least Megatron-LM (Shoeybi et al., 2019). It is effective — by splitting matrix multiplications across GPUs, the per-GPU memory for parameters, activations, and gradients of that layer is reduced proportionally. But it comes with a fundamental usability cost: it requires the data scientist to refactor their model code, replacing single-GPU operators with distributed versions that include inter-GPU communication (allreduce of partial outputs). This is not a one-time cost — it must be done for every new model architecture, and not all architectures map cleanly to tensor-slicing (e.g., operators with irregular spatial dimensions or complex dependency patterns).

Memory-centric tiling offers a qualitatively different solution: decompose the large operator into a temporal sequence of smaller operators on the same GPU, rather than a spatial partition across multiple GPUs. The tile's parameters and gradients are fetched via ZeRO-3's gather mechanism before each tile executes and released after. The working memory per tile is 1/T of the original operator (where T is the number of tiles), so even a 68 GB operator becomes manageable (4.3 GB with T = 16).

The conceptual innovation here is not the idea of tiling per se — tiled matrix multiplication is a standard technique in high-performance computing. The innovation is recognizing that ZeRO-3's parameter fetch-and-release lifecycle already provides the infrastructure to make tiling work transparently for model training. In ZeRO-3, parameters are gathered before each operator and released after. With tiling, the operator is simply decomposed into T sub-operators, each of which triggers its own gather of the relevant parameter tile and release after use. The interaction between tiling and ZeRO-3 is seamless: the hook system that automates parameter movement (Section 7.1) doesn't need to know about tiling — it just sees T sequential sub-operators, each with its own parameter set, and applies the standard gather-compute-release pattern.

What makes this a fundamental rather than incremental contribution is that it eliminates an entire axis of parallelism from the usability burden. The paper explicitly frames this: "data scientists no longer have to adapt their model to multiple forms of parallelism like in 3D parallelism. This is possible due to memory-centric tiling" (Section 5.3). The claim is not that tiling is more performant than model parallelism (it likely isn't — sequentializing a large matmul across tiles on one GPU has lower utilization than splitting it across multiple GPUs that compute in parallel). The claim is that tiling makes model parallelism unnecessary for memory reasons, which removes one of the two major model-code-refactoring requirements of 3D parallelism (the other being pipeline partitioning, which ZeRO-Infinity avoids by not using pipeline parallelism at all).

The evidence in Figure 6b supports the claim concretely: without tiling, the maximum trainable hidden dimension on a V100 with 2 GB memory fragmentation is 8K (corresponding to a model of tens of billions of parameters). With tiling factor 16, it scales to 64K (corresponding to tens of trillions of parameters). This is a ~8× increase in the largest trainable hidden dimension, achieved without any model parallelism and without any model code changes.

A subtle but important point: tiling for working memory reduction is possible because of a property that is specific to ZeRO-3's parameter lifecycle, not general to all partitioned training systems. In pipeline-parallel or model-parallel systems, parameters are persistently resident on their assigned GPUs — there is no mechanism to "release and re-fetch" a parameter tile because parameters aren't moved during training. ZeRO-3's design (gather before use, release after use) is what creates the opportunity to interleave the fetches of individual tiles. This is an example of a system design creating emergent capabilities that weren't part of the original motivation — ZeRO-3 was designed to eliminate memory redundancy, not to enable tiling, but the gather-release pattern it introduced turns out to be the exact mechanism needed to make operator tiling transparent.

Innovation 4: The Infinity Offload Engine's Unified Heterogeneous Memory Abstraction

Prior heterogeneous training systems treated GPU, CPU, and (in rare cases) NVMe memory as separate tiers with distinct APIs and usage patterns. ZeRO-Offload (Ren et al., 2021) could offload optimizer states and gradients to CPU, but parameters had to stay in GPU memory. Zhao et al. (2020) could offload sparse embedding parameters to SSD, but only for embedding tables with sparse access patterns in recommendation models. No prior system offered a unified abstraction where any model state (parameter, gradient, optimizer state) could reside on any memory tier (GPU, CPU, NVMe) and be moved between tiers automatically during training.

The infinity offload engine provides exactly this abstraction. It builds on ZeRO-3's parameter partitioning — which already means no single GPU holds a complete copy of any model state — and extends it with a placement layer that assigns each partitioned tensor to GPU, CPU, or NVMe storage based on a user-specified policy (Table 2). The three policies form a continuum: GPU-only (ZeRO-3 baseline, limited to aggregate GPU memory), CPU-offloaded (ZeRO-Inf-CPU, model size limited by aggregate CPU memory), and NVMe-offloaded (ZeRO-Inf-NVMe, model size limited by aggregate NVMe capacity).

The conceptual contribution is not any single mechanism (DeepNVMe, pinned memory management, bandwidth-centric partitioning) but the demonstration that these mechanisms compose into a system where the storage hierarchy is transparent to both the training loop and the user. The paper's ease-inspired implementation (Section 7) is crucial here: the data scientist writes standard PyTorch model code, ZeRO-Infinity intercepts parameter allocations and accesses via hooks, and whether a parameter is fetched from GPU HBM, CPU DRAM, or NVMe flash is determined by the offload policy, not by model code. The transition from ZeRO-Inf-CPU to ZeRO-Inf-NVMe requires only a configuration change — no model refactoring, no manual data movement, no changes to the training loop.

This unified abstraction is what enables the paper's headline scalability results. Figure 6a shows the progression: data parallelism alone maxes out at 1.4B parameters (limited by GPU memory with full replication), ZeRO-2 and ZeRO-Offload reach 13B (optimizer/gradient partitioning removes redundancy), ZeRO-Inf-CPU reaches ~100B (parameter partitioning removes the single-GPU bottleneck), and ZeRO-Inf-NVMe reaches 1T on a single node and 32T across 512 GPUs (NVMe offloading transcends GPU and CPU memory limits). The jump from 13B to 1T on a single node — a ~77× increase — is only possible because the infinity offload engine treats NVMe as just another tier in a unified memory hierarchy, not as a special-case storage for sparse parameters.

What distinguishes this from simple "swap to disk" approaches (common in operating systems and some prior ML work like vDNN, Rhu et al., 2016) is that the offloading is semantics-aware rather than page-based. A page-based swap system moves fixed-size memory pages between tiers based on access patterns, with no knowledge of the training loop's structure. The infinity offload engine knows which parameters belong to which layer, when each layer will be needed (via the dynamic prefetcher's operator sequence trace), and how to orchestrate the three-stage pipeline (nc-transfer, cg-transfer, gg-transfer) to hide latency. This knowledge allows it to achieve near-peak NVMe bandwidth (via DeepNVMe's bulk asynchronous I/O) and to overlap data movement with computation (via the overlap engine), which a generic page-based system cannot match.

The significance of this innovation extends to how it changes the economics of large model training. NVMe storage is approximately 50× cheaper per GB than GPU HBM and 20× cheaper than CPU DRAM, while being 50× larger in capacity than GPU memory on a typical cluster (Figure 2b). By making NVMe a first-class memory tier, ZeRO-Infinity dramatically reduces the capital cost of training a model of a given size: a 1T parameter model that would require 320 A100 GPUs (at ~ $10K each, ~$ 3.2M in GPU cost alone) with 3D parallelism can be trained on 16 V100 GPUs with NVMe offloading. The tradeoff is throughput, not feasibility — and for fine-tuning workloads where total compute is modest (as the paper argues in Section 1), this tradeoff is overwhelmingly favorable.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper evaluates on GPT-like Transformer-based language models with a fixed sequence length of 1024. Model configurations vary hidden dimension and number of layers to produce models with different parameter counts, ranging from hundreds of millions to tens of trillions of parameters. The specific configurations (hidden dimension, number of layers, batch size per GPU, model parallelism degree, FP16 parameter storage device, and optimizer state storage device) for each experiment are detailed in Table 1 and Appendix A (Tables 4–8).
Base model. All experiments use GPT-like Transformer architectures, chosen because "all of the SOTA models with over a billion parameters follow that" architecture. The evaluation does not train to convergence on a specific dataset like C4 or the Pile; instead, it measures training throughput and maximum achievable model scale, which are properties of the model architecture, hardware configuration, and system software rather than the training data distribution. This is standard practice for systems papers focused on memory and throughput.
Metrics. The primary metrics are training throughput (measured in TFlops/GPU) and maximum trainable model size (measured in number of parameters). For throughput, the paper uses the achievable peak of 70 TFlops/GPU on NVIDIA V100 as a reference point, derived empirically by running models on a single DGX-2 node with all non-GPU communication disabled to simulate unlimited bandwidth. The efficiency metric defined in Section 4 (efficiency = ait × bw / (ait × bw + peak_tp)) provides the theoretical framework for understanding when training should be compute-bound vs. bandwidth-bound. For scalability experiments, superlinear scaling is assessed by comparing aggregate throughput at 64 GPUs vs. 512 GPUs while keeping batch size per node constant (weak scaling).
Baselines. The paper compares against several baselines depending on the experiment:
- Data parallel (DDP): PyTorch's distributed data parallel (Li et al., 2020), where all model states are replicated on each GPU. This is the baseline for experiments without model parallelism.
- 3D parallelism: The state-of-the-art combining data, model (tensor-slicing), and pipeline parallelism, implemented via Megatron-LM (Shoeybi et al., 2019) and DeepSpeed (Microsoft, 2020). This is the primary baseline for model scale and throughput at the cluster level.
- ZeRO-2: Partitions optimizer states and gradients but replicates parameters (Rajbhandari et al., 2020).
- ZeRO-Offload: Offloads optimizer states and gradients to CPU memory while keeping parameters on GPU (Ren et al., 2021). This is the baseline for heterogeneous training comparisons (Figure 6c).
- ZeRO-3: Partitions all model states (optimizer states, gradients, and parameters) across GPUs but keeps everything in GPU memory (Rajbhandari et al., 2020). This is the baseline against which ZeRO-Infinity's offloading extensions are compared.
Generation budget / compute accounting. In this systems paper, "compute" is measured in terms of GPU hardware resources (number of V100 GPUs and DGX-2 nodes) rather than a per-sample generation budget. The key accounting is: (a) memory capacity — how many parameters can fit given aggregate GPU, CPU, and NVMe memory across all nodes (Figure 2b summarizes these capacities); (b) bandwidth — the achievable data movement rates between memory tiers (PCIe, NVLink, NVMe) that determine whether training is compute-bound; (c) computational throughput — measured as TFlops/GPU sustained during training. For the efficiency analysis in Section 4, peak_tp is measured empirically at 70 TFlops/GPU on V100 for the model configurations under study. Communication volume is measured in bytes of data moved per training iteration, with the arithmetic intensity (FLOPs/byte) determining the bandwidth necessary to keep the accelerator fed.
Cross-validation / statistical protocol. No cross-validation or statistical protocol is used. This is standard for systems papers where the outputs are deterministic given the hardware and software configuration — measured throughput and maximum model size are not subject to sampling variability in the way that ML benchmark accuracy is. Reproducibility is addressed through open-source implementation in DeepSpeed rather than statistical methodology.

Main Quantitative Results

Model Scale: 50× Increase Over 3D Parallelism

The paper's headline result is that ZeRO-Infinity trains a model with 32 trillion parameters on 512 NVIDIA V100 GPUs (32 DGX-2 nodes), compared to approximately 650 billion parameters with 3D parallelism on the same hardware. This represents a 50× increase in maximum trainable model size (Figure 1). The claim is not that 3D parallelism cannot scale further with more GPUs — it's that on a fixed 512-GPU cluster, ZeRO-Infinity can train models 50× larger than what 3D parallelism can fit.

Figure 6a breaks down how this scale is achieved incrementally on a single DGX-2 node (16 GPUs). Each successive strategy removes a memory bottleneck:

Data parallelism alone (DDP): 1.4 billion parameters. Every GPU stores a full copy of all model states (parameters, gradients, optimizer states), so model size is limited by single-GPU memory (~32 GB on V100).
ZeRO-2 (optimizer/gradient partitioning): 13 billion parameters (~9× increase). By partitioning optimizer states and gradients across GPUs, the per-GPU memory for these components is reduced by the data-parallel degree (16×), but parameters are still fully replicated, so the parameter memory (2 bytes per parameter in FP16, plus activation memory) remains the bottleneck.
ZeRO-Offload (CPU offload of optimizer/gradients): Also ~13 billion parameters. Offloading optimizer states and gradients to CPU frees GPU memory, but parameters must still fit on a single GPU because ZeRO-Offload is built on ZeRO-2 (parameter replication). This demonstrates that CPU offloading alone without parameter partitioning does not increase maximum model size beyond what GPU parameter memory allows.
ZeRO-Inf-CPU (parameter partitioning + CPU offload): ~100 billion parameters (~7.7× over ZeRO-2). By partitioning parameters across GPUs (ZeRO-3 style), the per-GPU parameter memory requirement drops to 1/DP of the total. Combined with CPU offloading of partitioned optimizer states and gradients, the capacity ceiling becomes aggregate CPU memory rather than single-GPU memory.
ZeRO-Inf-NVMe (NVMe offload of all model states): 1 trillion parameters on a single DGX-2 node (~700× over data parallel). By extending the infinity offload engine to NVMe, the capacity ceiling becomes the aggregate NVMe storage across the node's GPUs, which is ~50× larger than aggregate GPU memory (Figure 2b).

Why the 50× figure matters beyond raw scale: The 32T parameter result on 512 GPUs demonstrates that ZeRO-Infinity's scaling is not asymptotic — it actually achieves the scale that the bandwidth and memory analysis predicts is feasible. The paper emphasizes that even a 100 trillion parameter model's model states (requiring ~2 PB of storage for 20 bytes/param) would fit in the aggregate NVMe of 96 DGX-2 nodes (1536 GPUs per Figure 2b), establishing headroom for at least another 3× growth beyond the demonstrated result.

Training Throughput: Matching 3D Parallelism at Equal Scale, Then Scaling Beyond

Figure 5a compares ZeRO-Infinity's training throughput against 3D parallelism across model sizes from 500 billion to 20 trillion parameters on 512 GPUs:

At 500 billion parameters (the largest model 3D parallelism can fit on 512 GPUs): ZeRO-Infinity and 3D parallelism achieve nearly identical throughput. The paper states this indicates ZeRO-Infinity is "on par with the training efficiency of the state-of-the-art." This is a critical validation: it shows that the overhead of offloading to CPU/NVMe does not degrade performance when the model could fit in GPU memory — the infinity offload engine's overlapping and bandwidth-centric partitioning successfully hide the offloading cost.
At 1 trillion parameters: ZeRO-Infinity achieves approximately 44 TFlops/GPU. 3D parallelism cannot train a model of this size on 512 GPUs (it runs out of memory). Figure 5b provides a weak scaling analysis for the 1T model: from 4 nodes (64 GPUs) to 32 nodes (512 GPUs), aggregate throughput increases from 2.8 petaflops to over 25 petaflops, exceeding perfect linear scaling. This superlinear behavior is attributed to the linear increase in aggregate PCIe and NVMe bandwidth with additional nodes — more nodes means more parallel PCIe links for parameter and gradient movement, and more CPU cores for parallel optimizer updates, both of which reduce the communication bottleneck relative to per-GPU computation.
At 5 trillion parameters: 43 TFlops/GPU.
At 10 trillion parameters: 43 TFlops/GPU.
At 20 trillion parameters: 34 TFlops/GPU — a performance drop that the paper attributes not to NVMe bandwidth saturation but to an "extremely small batch size per GPU" (Table 1: batch size 1.25 at 20T with 8-way model parallelism on 512 GPUs). The small batch size is forced by limited CPU memory for activation checkpoints (3 TB for a 100T model per Figure 2a) — when activation memory cannot fit in CPU, the only recourse is to reduce batch size or increase activation checkpoint frequency, both of which reduce GPU utilization.

The sustained 25+ petaflops on 512 GPUs corresponds to approximately 40% of peak (512 × 125 TFlops theoretical peak for V100 = 64 petaflops, though the paper uses the 70 TFlops achievable peak per GPU as the reference, giving 512 × 70 = 35.8 petaflops theoretical and 25/35.8 ≈ 70% of achievable peak, or 25/64 ≈ 39% of hardware peak). The paper claims "over 25 petaflops" and "40% of peak" in the abstract — the precise baseline for "peak" is ambiguous (hardware theoretical vs. achievable) but the sustained throughput is documented.

Superlinear Scalability from 64 to 512 GPUs

Figure 5b demonstrates that ZeRO-Infinity achieves superlinear scaling when training a 1 trillion parameter model from 4 nodes (64 GPUs) to 32 nodes (512 GPUs) under weak scaling (constant batch size per node). At 4 nodes, aggregate throughput is approximately 2.8 petaflops (~44 TFlops/GPU). Perfect linear scaling to 32 nodes would predict 2.8 × 8 = 22.4 petaflops. The measured throughput "over 25 petaflops" exceeds this, indicating superlinearity.

The paper attributes this to two mechanisms that improve per-GPU efficiency as the node count increases:

Aggregate PCIe and NVMe bandwidth increases linearly with node count (bandwidth-centric partitioning means each GPU reads/writes over its own PCIe link in parallel). At 4 nodes, the aggregate NVMe bandwidth is ~100 GB/s (4 × 25 GB/s per node); at 32 nodes, it's ~800 GB/s. The parameter/gradient bandwidth target of ~70 GB/s per GPU translates to ~4.5 TB/s aggregate at 64 GPUs and ~35.8 TB/s at 512 GPUs — the system moves from being bandwidth-constrained to compute-bound as nodes are added, which means per-GPU throughput improves rather than staying constant.
CPU compute for optimizer updates scales with node count. The Adam optimizer step is an element-wise operation on all parameters that can be parallelized across all available CPU cores. More nodes provide more CPU cores, reducing the optimizer step time per node even though the per-node optimizer state volume stays constant (since model parameters are partitioned across all nodes).

A subtle point in this result: the paper notes that ZeRO-Infinity "already achieves over 2.8 petaflops (44 Tflops/GPU) with just 4 nodes," which demonstrates that "the aggregated NVMe bandwidth is sufficient to achieve good efficiency even at a modest scale." This is important because it addresses the concern that NVMe offloading might only be practical at massive cluster sizes — even 4 nodes (64 GPUs) provide enough aggregate bandwidth for reasonable efficiency.

Democratizing Large Model Training: Single-Node Results

Figure 5c shows the throughput of training models from 10 billion to 1 trillion parameters on a single DGX-2 node (16 GPUs) without any model parallelism:

10 billion parameters: Over 40 TFlops/GPU.
50–100 billion parameters: Over 40 TFlops/GPU.
500 billion parameters: Measurable throughput (exact value not quoted in text, but the bar chart in Figure 5c indicates it's in a comparable range to the 50–100B results, likely above 35 TFlops/GPU).
1 trillion parameters: Measurable throughput on a single node — the paper emphasizes this is sufficient for fine-tuning workloads on models like GPT-3 (175B parameters), which would require 128 GPUs with 3D parallelism.

The paper explicitly contrasts this with 3D parallelism, which "is unable to scale to models with over 20 billion parameters" on the same single DGX-2 node. The 20B limit for 3D parallelism on a single node matches expectations: with 16 GPUs, pipeline parallelism has limited depth, model parallelism splits operators but doesn't reduce optimizer memory, and data parallelism replicates model states — so even with aggressive partitioning, fitting a 20B+ parameter model (400 GB in FP16 for parameters alone, plus 320 GB for optimizer states) into 512 GB of aggregate GPU memory is infeasible.

These single-node results support the paper's claim about "democratizing" large model access. A researcher with access to a single DGX-2 node (a common configuration in university and industry labs, costing on the order of $100K rather than the millions required for multi-node clusters) can fine-tune models up to a trillion parameters. The throughput (over 40 TFlops/GPU for models up to 100B) is sufficient that fine-tuning — which typically requires orders of magnitude less compute than pretraining — is practical.

Ablation Studies and Robustness Checks

ZeRO-Infinity vs. ZeRO-Offload for gradient offloading: Figure 6c compares the backward propagation time for an 8 billion parameter model when offloading gradients to CPU memory with ZeRO-Infinity vs. ZeRO-Offload. At 64 GPUs, ZeRO-Infinity achieves "nearly 2×" speedup compared to ZeRO-Offload. The measured difference is attributed to bandwidth-centric partitioning: ZeRO-Offload uses broadcast-from-owner (one PCIe link active), while ZeRO-Infinity uses allgather-from-all (all PCIe links active in parallel). At 64 GPUs (4 DGX-2 nodes), the aggregate CPU-to-GPU bandwidth with ZeRO-Infinity is ~192 GB/s (64 × 3 GB/s per GPU) vs. ~12 GB/s for ZeRO-Offload, matching the ~16× difference in active PCIe links. The actual 2× speedup (rather than 16×) suggests the backward pass is not purely communication-bound — computation and GPU-GPU allgather also contribute to execution time.

Prefetching and overlap impact by batch size: Figure 6d shows the relative throughput difference with communication overlapping and prefetching enabled vs. disabled for an 8 billion parameter model on 64 GPUs, across batch sizes ranging from 2 to 16 per GPU. The key finding is that prefetching and overlapping are crucial at small batch sizes but their impact diminishes as batch size increases. This is consistent with the arithmetic intensity analysis in Section 4: small batch sizes have low ait = seq × bsz, meaning each byte of parameter data supports fewer FLOPs, so communication time is a larger fraction of total iteration time and overlapping provides more benefit. At batch size 16, the ait is 16× higher than at batch size 1, communication is a smaller fraction of iteration time, and the marginal benefit of overlapping is correspondingly smaller. The paper does not quote specific percentages, but the trend in Figure 6d confirms the theoretical prediction.

Activation checkpoint offloading overhead by hidden size: Figure 6e measures the training throughput impact of offloading activation checkpoints to CPU memory, for hidden dimensions ranging from 2K to 64K. The finding: CPU offloading reduces throughput by up to 1.2× for small hidden sizes (2K), but the impact becomes "minimal" for hidden sizes 32K and 64K. This matches the bandwidth analysis from Figure 3c and Equation 11: ait_activation = 24 × hd × ci. At hd = 2048, the AIT is ~49K FLOPs/byte, and the required bandwidth for high efficiency is roughly 2 GB/s (Figure 3c), which is comparable to the per-GPU PCIe bandwidth to CPU (~3 GB/s) — so PCIe contention can cause a throughput drop. At hd = 32768, the AIT is ~786K FLOPs/byte, and the required bandwidth drops below 1 GB/s, making PCIe bandwidth more than sufficient even under contention. This ablation validates a key theoretical claim: activation offloading is essentially free for large models.

Maximum hidden size with and without memory-centric tiling: Figure 6b measures the largest trainable hidden dimension on a single DGX-2 node with GPU memory pre-fragmented into 2 GB contiguous chunks (to simulate realistic memory fragmentation conditions). Without memory-centric tiling, the maximum trainable hidden dimension is 8192 (8K). With tiling factors of 2, 4, 8, and 16, the maximum trainable hidden dimensions are 16K, 32K, 64K, and 64K respectively. The 64K hidden dimension corresponds to models with tens of trillions of parameters (per Equation 1, hd = 64K enables 12 × nl × (65536)^2 ≈ 12 × 200 × 4.3B ≈ 10.3T parameters). This ablation demonstrates that memory-centric tiling addresses the MSWM fragmentation problem identified in Section 3 (Equation 4): without tiling, the contiguous memory requirement for the largest linear layer (4 × hd × 4hd bytes) exceeds what the fragmented GPU allocator can provide beyond hd = 8K, even with sufficient aggregate free memory. The paper notes that this "greatly simplifies DL system stack by avoiding the need for model parallelism."

Maximum model size with different device placement strategies: Figure 6a (discussed in the scale results above) also serves as an ablation showing the marginal contribution of each successive ZeRO stage to maximum model size on a single DGX-2 node. The progression (1.4B → 13B → 100B → 1T) isolates the effect of optimizer/gradient partitioning (ZeRO-2), parameter partitioning (ZeRO-3), and NVMe offloading (ZeRO-Inf-NVMe). The key takeaway is that NVMe offloading provides the largest single jump (100B → 1T, a 10× increase), but it is only possible because parameter partitioning (ZeRO-3) removed the single-GPU parameter bottleneck first. This demonstrates the composability of the ZeRO family: each stage addresses a specific memory redundancy or capacity limit, and they can be combined incrementally.

Negative result: Performance degradation at 20T parameters. Figure 5a shows a drop from 43 TFlops/GPU at 10T to 34 TFlops/GPU at 20T (a ~21% decrease). The paper attributes this to "an extremely small batch size per GPU at 20T scale as a result of limited CPU memory to store activation checkpoints" (Table 1 shows batch size 1.25 per GPU at 20T). This is a genuine bottleneck: as model size grows, activation memory (which scales with bsz × seq × hd × nl) grows, and even with CPU offloading, the per-node CPU memory (1.5 TB on a DGX-2) eventually limits how many activation checkpoints can be stored, forcing a batch size reduction that leaves GPU compute units underutilized. The paper acknowledges this can be addressed by offloading activation checkpoints to NVMe in a future implementation.

Critical Assessment

Claim 1: ZeRO-Infinity trains 32 trillion parameter models on 512 GPUs — a 50× increase over 3D parallelism.

What was tested: The paper demonstrates that ZeRO-Infinity can allocate memory for and execute training iterations on a model with 32 trillion parameters on 512 V100 GPUs. Figure 1 and the abstract report this as the maximum achieved model size.

What this actually demonstrates: The 32T result is a memory capacity and throughput demonstration, not evidence that such a model can be successfully trained to convergence. The paper measures that training iterations execute without out-of-memory errors and at a reported throughput (~25 petaflops aggregate), but it does not report validation loss curves, convergence behavior, or final model quality for the 32T configuration. This is standard for systems papers (the claim is about infrastructure capability, not model quality), but it means the 32T figure should be understood as "ZeRO-Infinity can execute training iterations" rather than "ZeRO-Infinity can produce a useful 32T-parameter model." The distinction matters because very large models trained with extreme offloading might have subtle numerical issues (e.g., stale optimizer states due to NVMe latency, or convergence problems from the large effective batch sizes needed to maintain throughput) that are not captured by throughput measurements.

Additionally, Table 1 reveals that the 20T and 32T configurations use model parallelism (mp=4, mp=8 respectively) — the 32T model is not trained with pure data parallelism. This means the "without model parallelism" ease-of-use claim (Section 5.3) applies to models up to approximately 1T parameters (as shown in Figure 5c for single-node), but the extreme-scale results still require some model parallelism for the largest hidden dimensions where memory-centric tiling alone may be insufficient or where the efficiency of tiled matrix multiplications is too low. The paper is transparent about this (mp values are listed in Table 1), but the headline "without requiring model code refactoring" requires qualification: for models beyond a certain hidden dimension threshold, some model parallelism may still be needed.

The 50× claim compared to what baseline: The comparison is ZeRO-Infinity on 512 GPUs vs. 3D parallelism on the same 512 GPUs. This is a fair comparison on fixed hardware. However, 3D parallelism can scale beyond 650B parameters by adding more GPUs — the paper notes it scales to over a trillion parameters on 800 GPUs. So the 50× is not a claim about absolute capability (ZeRO-Infinity can train models 50× larger than 3D parallelism ever can) but about resource efficiency (ZeRO-Infinity can train 50× larger models on the same hardware). The distinction is important and the paper makes it clear, but the shorthand "50× increase over 3D parallelism" could be misinterpreted.

Claim 2: ZeRO-Infinity sustains over 25 petaflops on 512 V100 GPUs (40% of peak).

What was tested: Figure 5b and the accompanying text report over 25 petaflops aggregate throughput for a 1T parameter model on 512 GPUs, and Figure 5a shows per-GPU TFlops for models ranging from 500B to 20T parameters.

Assessment: The throughput numbers are credible for a systems paper and compare well to 3D parallelism at the same scale. However, several important details affect interpretation:

The "40% of peak" figure uses an undefined baseline. V100 theoretical peak is 125 TFlops/GPU (tensor core FP16). The paper uses 70 TFlops/GPU as "achievable peak" for their model configurations, which is 56% of theoretical. The 40% figure in the abstract (~25 petaflops / 64 petaflops theoretical) differs from the 70% figure implied by the achievable peak baseline (25 / 35.8 ≈ 70%). The abstract's "40% of peak" appears to reference hardware theoretical peak, while Figure 5's TFlops/GPU numbers (which range from 34–49) reference the 70 TFlops achievable baseline. This inconsistency in which "peak" is being used makes it difficult to assess how close the system is to the roofline.
The throughput numbers are for specific model configurations (Table 1). The 49 TFlops/GPU figure is for a 5T parameter model with batch size 3, while the 34 TFlops/GPU figure is for a 20T model with batch size 1.25. The throughput variation is driven primarily by batch size (smaller batch size → lower GPU utilization → fewer TFlops/GPU) rather than by offloading overhead. This is consistent with the paper's own analysis (Section 4: ait = seq × bsz, so smaller batch size means more bandwidth-sensitive workload), but it means the throughput claims are configuration-dependent.
The superlinear scaling result (Figure 5b) is striking but the paper does not provide detailed breakdown of where the superlinear gains come from (e.g., separate measurements of optimizer step time, communication time, and computation time at different scales). Without this breakdown, the reader cannot assess whether the superlinearity is primarily from better overlapping, reduced optimizer step time, or reduced communication overhead — or whether it would persist at even larger scales.

Claim 3: Superlinear scalability for a trillion-parameter model from 64 to 512 GPUs.

What was tested: Weak scaling from 4 nodes to 32 nodes with constant batch size per node.

Assessment: The superlinear result is plausible given the bandwidth-centric partitioning design: more nodes means proportionally more aggregate PCIe bandwidth, and if the 4-node configuration was partially bandwidth-limited, adding nodes reduces the bandwidth bottleneck faster than it adds computational capacity, improving per-GPU efficiency. This is a genuine strength of the architecture.

However, the paper does not show where the superlinearity saturates. If bandwidth-centric partitioning is the mechanism, then the superlinearity should saturate when aggregate PCIe/NVMe bandwidth exceeds the ~70 GB/s per GPU target — at that point, the system becomes compute-bound and further nodes should show linear (not superlinear) scaling. The 32-node experiment (512 GPUs) represents aggregate bandwidth of ~1.5 TB/s from NVMe, which significantly exceeds requirements. It would strengthen the claim to show the scaling curve leveling off to linear at some intermediate node count, confirming the mechanistic explanation.

Additionally, weak scaling (constant batch size per node) means the effective batch size increases with node count (from batch size ~28 per node × 4 nodes = 112 to batch size ~28 × 32 = 896). This is a valid weak scaling design, but it means the per-GPU batch size stays constant — so the per-GPU ait stays constant and the improved efficiency comes entirely from reduced communication overhead at larger scale, not from changes in the workload's bandwidth sensitivity. Showing that the superlinearity holds even at fixed per-GPU batch size is a strong validation of bandwidth-centric partitioning.

Claim 4: Single-node fine-tuning of up to a trillion parameter model, democratizing access.

What was tested: Figure 5c demonstrates throughput for models from 10B to 1T parameters on a single DGX-2 node (16 V100 GPUs).

What was not tested: The paper does not report time-to-convergence for fine-tuning a specific model (e.g., GPT-3 175B) on a specific downstream task. The claim "makes large model fine-tuning accessible" is an inference from the throughput numbers — if you can execute training iterations at over 40 TFlops/GPU, then fine-tuning (which might require, say, 10^18 FLOPs for a full fine-tune of a 175B model on a moderate dataset) would complete in a reasonable wall-clock time. But this depends on assumptions about the fine-tuning compute budget that are not validated in the paper. A demonstration of actual fine-tuning (e.g., fine-tuning GPT-3 on SQuAD or a similar task and reporting wall-clock time and final accuracy) would substantially strengthen the accessibility claim.

Missing experiments that would strengthen the paper:

Convergence results for NVMe-offloaded training. The paper measures training throughput but never reports whether models trained with NVMe offloading converge to the same loss as GPU-only or CPU-offloaded training. Numerical issues from the optimizer step being performed on CPU (potentially with different floating-point reduction order or precision) or from parameters being stored on NVMe (with potential for silent data corruption at extreme scale) could affect model quality.
Comparison against a version of 3D parallelism that also uses CPU memory. The paper compares against GPU-only 3D parallelism. A hybrid 3D parallelism that offloads optimizer states to CPU (similar to combining Megatron-LM with ZeRO-Offload) would be a stronger baseline for the "state of the art" claim, since it would show that ZeRO-Infinity's gains are not achievable by simply adding CPU offloading to existing 3D parallelism.
Detailed breakdown of time spent in nc-transfer, cg-transfer, gg-transfer, computation, and optimizer step. The throughput numbers are aggregate; without a breakdown, it's difficult to verify the paper's claims about which component is the bottleneck at different scales. For instance, the claim that the 20T model's lower throughput is due to small batch size rather than NVMe bandwidth could be verified by showing that GPU utilization (fraction of time spent computing vs. waiting for data) drops at small batch sizes.
Evaluation on a model with non-Transformer architecture. The paper focuses exclusively on GPT-like Transformer models because "all of the SOTA models with over a billion parameters follow that." This is fair for the target audience, but the memory-centric tiling and ease-of-use claims (no model code refactoring) would be strengthened by demonstrating them on a model with a more complex dependency graph — e.g., a Transformer with cross-layer attention, a mixture-of-experts model, or a graph neural network — where 3D parallelism's pipeline and model parallelism would struggle with load balancing.
Sensitivity to NVMe device characteristics. DeepNVMe achieves near-peak sequential NVMe bandwidth through aggressive I/O parallelization. Consumer-grade NVMe devices, cloud-attached NVMe (e.g., AWS instance store), and different filesystems may have different performance characteristics. The paper's experiments use a specific hardware configuration (DGX-2 nodes with unspecified NVMe drives); evaluating on a range of NVMe hardware would establish the robustness of the near-peak bandwidth claim.

Conditions where claims hold and where they weaken:

Model scale claims hold when: (a) the model architecture is a standard Transformer with linear layers that can be tiled, (b) the NVMe storage is configured for high sequential I/O (not shared with other workloads), (c) the cluster has sufficient CPU memory for activation checkpoints at the target batch size and sequence length. The 20T throughput drop shows that CPU activation memory can become a bottleneck at extreme scale.
Throughput claims hold when: (a) the batch size per GPU is sufficient for reasonable GPU utilization (Figure 6d shows overlapping benefits diminish at small batch sizes, and the 20T result shows throughput drops at batch size 1.25), (b) the model's hidden dimension is large enough that activation offloading overhead is negligible (Figure 6e). For models with small hidden dimensions (<8K), activation offloading to CPU imposes a measurable throughput penalty.
Superlinear scaling holds when: the system is partially bandwidth-limited at small scale and becomes compute-bound at large scale. If the system is already compute-bound at small scale (e.g., training a small model that fits entirely in GPU memory with no offloading), scaling would show linear (not superlinear) behavior. The paper does not evaluate this regime, which is reasonable since ZeRO-Infinity is designed for models that cannot fit in GPU memory.
Ease-of-use claims hold when: (a) the model can be expressed as a PyTorch nn.Module hierarchy (supported by the hook injection mechanism), (b) external parameters are either detected automatically by the activation introspection or manually registered. For models with particularly exotic parameter sharing patterns that the automated detection misses, manual registration is required, which is a minor but nonzero code change.
Democratization claims weaken when: the target workload requires very long sequence lengths (which increase activation memory and may exceed single-node CPU capacity), or when the fine-tuning compute budget is so large that single-node throughput is insufficient (e.g., full fine-tuning of a 1T model on a very large dataset). The paper acknowledges this implicitly: single-node training is for fine-tuning, not pretraining.

6. Limitations and Trade-offs

6.1 No Convergence or Model Quality Results at Extreme Scale

The assumption or constraint. ZeRO-Infinity is evaluated exclusively on training throughput (TFlops/GPU) and maximum feasible model size. The paper demonstrates that training iterations execute without out-of-memory errors for models up to 32 trillion parameters, but it never reports validation loss, convergence behavior, or final model quality for any model trained with NVMe offloading. The paper does not acknowledge this gap as a limitation; it treats training throughput as the sole metric of interest, which is standard for systems papers focused on memory and communication but leaves the end-to-end training viability unvalidated.

The consequence. Practitioners cannot determine from this paper whether models trained with aggressive NVMe offloading actually converge to useful solutions. Several mechanisms unique to ZeRO-Infinity's heterogeneous memory design could plausibly affect convergence: the Adam optimizer step is computed on CPU rather than GPU for NVMe-offloaded states (potentially with different floating-point reduction order, leading to subtle numerical divergence after many iterations); parameter updates may experience variable latency between gradient computation and application (since optimizer states and parameters are written back to NVMe asynchronously through the pinned memory pipeline); and the small batch sizes forced by CPU activation memory limits (batch size 1.25 at 20T, Table 1) may require learning rate adjustments or affect optimization dynamics in ways not reflected in throughput measurements. For a practitioner deciding whether to use ZeRO-Infinity for a real training run, the absence of even a single convergence curve means the entire framework is validated only for memory capacity and throughput, not for producing working models.

What evidence exists in the paper. None. The evaluation section (Section 8) contains no loss curves, no accuracy metrics, and no comparison of final model quality between ZeRO-Infinity configurations and GPU-only baselines. Figure 5a and 5b report only TFlops/GPU. Table 1 lists model configurations used for throughput measurements but provides no training duration or convergence status. This is a gap in the evidence, not a contested finding.

Mitigation status. Not addressed. The paper does not mention convergence as a concern and does not suggest it as future work. This is a defensible scope choice for a systems paper — the contribution is about memory and communication infrastructure, not optimization dynamics — but it means the central claim ("ZeRO-Infinity can train 32T parameter models") is unvalidated in the sense that matters to ML practitioners.

6.2 Difficulty Estimation Is Never Applied to the Models Studied

EDIT: This limitation was mistakenly included. Removing to stay within 4-6 limitation count.

6.2 Activation Memory Becomes a Bottleneck at Extreme Scale That NVMe Offloading Does Not Currently Solve

The assumption or constraint. ZeRO-Infinity offloads activation checkpoints to CPU memory but not to NVMe (Section 5.1.2). The paper notes that activation checkpoints for a 10T parameter model at batch size 32 and sequence length 1024 require 0.76 TB (Figure 2a, Column 7), which fits in a DGX-2 node's 1.5 TB CPU memory, but a 100T model would require 3 TB, exceeding even next-generation hardware projections. For models beyond 10T parameters, per-GPU batch size must be reduced to keep activation memory within CPU limits. This is directly observable in Table 1: the 5T model uses batch size 3, the 10T model uses batch size 2, and the 20T model uses batch size 1.25.

The consequence. Activation memory becomes the de facto scaling bottleneck at extreme model sizes, even though the infinity offload engine has removed model state memory as a constraint. The paper's headline result (32T parameters on 512 GPUs) is achieved without specifying the activation checkpoint configuration or batch size used for that specific run, but the trend from Table 1 shows that batch size per GPU drops as model size increases. At batch size 1.25 for the 20T model, GPU utilization suffers significantly — Figure 5a shows a 21% throughput drop from 10T to 20T (43 TFlops/GPU → 34 TFlops/GPU), which the paper explicitly attributes to "an extremely small batch size per GPU at 20T scale as a result of limited CPU memory to store activation checkpoints." This creates a tension with the paper's central narrative: model states are no longer the memory bottleneck, but activations still are, and they cannot be offloaded to NVMe in the current implementation. A 100T parameter model — which the paper claims would fit in the aggregate NVMe of 96 DGX-2 nodes (Section 5.1.1) — might be untrainable in practice because activation memory at any reasonable batch size would exceed CPU capacity, and reducing batch size further would make throughput prohibitive.

What evidence exists in the paper. The throughput drop at 20T (Figure 5a) is the clearest evidence. Table 1 documents the declining batch sizes. Figure 2a Column 7 quantifies the activation checkpoint memory for models from 100B to 100T parameters, showing that it grows to 3 TB at 100T — beyond the 1.5 TB CPU memory on a DGX-2 node. The paper acknowledges this implicit bottleneck in Section 8.2, noting that the 10T→20T performance drop "can be improved by increasing the CPU memory or offloading activation checkpoints to NVMe in a future implementation."

Mitigation status. Partially acknowledged, not solved. The paper mentions NVMe offloading of activation checkpoints as future work but provides no analysis of whether activation AIT (Equation 11: ait_activation = 24 × hd × ci) would make this efficient — activations have very high AIT for large hidden dimensions (Figure 3c), so NVMe offloading might be viable, but it would add a fourth data movement stage to the already three-stage parameter pipeline, potentially overwhelming the overlap engine on limited PCIe bandwidth.

6.3 The "Without Model Code Refactoring" Claim Has Boundary Conditions Not Fully Characterized

The assumption or constraint. The paper claims that data scientists "no longer have to adapt their model to multiple forms of parallelism" (Section 5.3) and that ZeRO-Infinity "eliminates the need for manual model code refactoring" (Section 1). This claim is grounded in memory-centric tiling (which removes the need for tensor-slicing model parallelism) and the ease-inspired implementation (which automates data movement via hooks). However, the extreme-scale results in Table 1 reveal that some model configurations at the largest scales do use model parallelism: the 5T and 10T models use mp=4, and the 20T model uses mp=8.

The consequence. The "no model code refactoring" claim applies to models up to approximately 1T parameters (as demonstrated in the single-node experiments, Figure 5c, where models up to 1T are trained with mp=1). For larger models — particularly those with hidden dimensions large enough that memory-centric tiling alone is insufficient or inefficient — model parallelism is still required. The paper does not characterize the threshold at which model parallelism becomes necessary, what hidden dimensions or layer sizes trigger it, or whether the model parallelism used at the largest scales requires the same kind of manual code refactoring that 3D parallelism demands. A data scientist planning to train a 10T parameter model cannot determine from this paper whether they will need to refactor their code.

The tiling approach itself has an implicit architectural constraint: it decomposes operators along the output dimension, which works cleanly for linear layers (matrix multiplications can be tiled along either dimension) but may not apply to operators with irregular spatial dimensions, custom CUDA kernels, or complex data-dependent computation patterns. The paper only evaluates on standard Transformer linear layers. For models with custom operators (e.g., efficient attention variants, Mixture-of-Experts routing, sparse operations), the memory-centric tiling approach may not be transparently applicable, and the ease-of-use claim does not extend to such architectures without additional engineering.

What evidence exists in the paper. Table 1 lists mp=4 for 5T and 10T models and mp=8 for 20T models. The paper does not discuss why model parallelism is used for these configurations, whether tiling could have substituted, or what the tradeoff was. The single-node experiments (Figure 5c) all use mp=1, demonstrating the claim holds at that scale. The ease-inspired implementation (Section 7) describes hooks and automatic initialization for data-parallel training with ZeRO-3 partitioning but does not describe how model parallelism interacts with this automation — it is unclear whether the mp>1 configurations required additional code changes beyond what ZeRO-Infinity automates.

Mitigation status. Not addressed. The paper does not discuss the interaction between ZeRO-Infinity's automation and model parallelism, does not specify when model parallelism is necessary, and does not characterize what fraction of the 32T model's training code required manual intervention. The claim "without requiring model code refactoring" should be qualified with the scale at which it holds, but the paper never provides this qualification.

6.4 DeepNVMe Performance Is Hardware-Specific and Not Characterized Across Storage Configurations

The assumption or constraint. The infinity offload engine's efficiency depends critically on DeepNVMe achieving near-peak sequential NVMe bandwidth (Section 6.3). The paper's experiments use a specific hardware configuration: NVIDIA DGX-2 nodes with unspecified NVMe drives, running in a cluster with 800 Gbps internode communication. DeepNVMe's optimizations — aggressive I/O parallelization, pinned memory management, zero-copy design — are described at a high level, but the paper provides no sensitivity analysis showing how throughput degrades with different NVMe hardware (consumer-grade SSDs, cloud-attached instance storage, different filesystems, or degraded drives).

The consequence. A practitioner attempting to reproduce ZeRO-Infinity on different hardware — an academic cluster with slower NVMe, a cloud instance with network-attached SSDs, or a shared cluster where NVMe bandwidth is contended by other jobs — cannot predict their training throughput from this paper. The aggregate NVMe bandwidth numbers in Figure 2b (25 GB/s read per DGX-2 node) are hardware-specific. If a target system achieves only 10 GB/s aggregate NVMe read bandwidth, the parameter/gradient bandwidth per GPU drops proportionally (from ~1.6 GB/s to ~0.6 GB/s in a 16-GPU node), potentially pushing the system from compute-bound to severely bandwidth-bound, especially at small batch sizes. The paper's bandwidth analysis (Section 4) provides a theoretical framework for predicting this (Equation 6, Figure 3a), but the framework is never applied to characterize minimum NVMe bandwidth requirements for different model configurations and batch sizes.

Additionally, NVMe endurance (write cycles before failure) is not discussed. Training a trillion-parameter model involves writing terabytes of optimizer states and parameters to NVMe per iteration. Over thousands of iterations (typical for large model training), this could approach or exceed consumer NVMe write endurance limits, potentially causing silent data corruption or drive failure. Enterprise NVMe drives have higher endurance but are more expensive, partially offsetting the cost advantage of NVMe over GPU HBM that the paper emphasizes.

What evidence exists in the paper. None. The only bandwidth measurements are for the DGX-2 cluster (Figure 2b). There are no experiments varying NVMe configuration, no measurements of NVMe latency under contention, and no discussion of endurance. DeepNVMe is described as achieving "near peak sequential read and write bandwidths" but no numbers are given for achieved vs. theoretical peak on the specific drives used.

Mitigation status. Not addressed. The paper treats NVMe performance as a solved problem via DeepNVMe's optimizations but does not characterize the robustness of that solution. Future work on "offloading activation checkpoints to NVMe" (Section 8.2) would compound this issue, as it adds additional NVMe traffic to an already bandwidth-constrained link.

6.5 Latency and Wall-Clock Time Are Not Analyzed — Throughput Alone Overstates Practical Efficiency for Some Use Cases

The assumption or constraint. The paper's efficiency metric is training throughput (TFlops/GPU), which measures the rate of floating-point operations executed, not the wall-clock time to complete a training run or the latency of individual training iterations. The overlap-centric design (Section 6.2) hides data movement latency behind computation, which improves throughput but does not eliminate the latency of the serial portions of the training loop — particularly the optimizer step, which "cannot be overlapped with the computation" (Section 4.2) and which, for NVMe-offloaded optimizer states, involves reading and writing terabytes of data through the CPU-NVMe interface.

The consequence. For throughput-sensitive workloads (large-scale pretraining where total FLOPs dominate), the throughput metric is appropriate and ZeRO-Infinity's achieved TFlops/GPU are impressive. However, for latency-sensitive workloads — interactive fine-tuning, rapid experimentation cycles, or any scenario where the user cares about time-to-result rather than FLOPs-per-dollar — the unanalyzed latency of the optimizer step and the pipeline depth of the overlapping prefetcher could make ZeRO-Infinity slower in wall-clock time than a system that uses more GPUs to keep all model states in HBM.

Consider the single-node fine-tuning use case: ZeRO-Infinity can fine-tune a 1T parameter model on 16 GPUs (Figure 5c), while 3D parallelism requires 128 GPUs for the same model (Section 1). ZeRO-Infinity is clearly more hardware-efficient. But if the ZeRO-Infinity iteration takes 4× longer in wall-clock time than an iteration on the 128-GPU 3D parallelism setup (due to serial optimizer step and NVMe latency), the total time-to-fine-tune might be comparable or even longer, despite using 8× fewer GPUs. The paper provides no data to evaluate this tradeoff.

A related concern: the dynamic prefetcher requires one warm-up iteration to trace the operator graph before it can begin prefetching (Section 6.2). For fine-tuning workloads with a small number of total iterations, this warm-up overhead is proportionally larger. The paper does not measure the prefetcher's warm-up cost or characterize its impact on short training runs.

What evidence exists in the paper. The paper provides only throughput numbers (TFlops/GPU), not iteration times, optimizer step times, or end-to-end training durations. The bandwidth analysis (Section 4.2) notes that optimizer states "cannot be overlapped with the computation" and require "significantly larger bandwidth to keep the overall DL workload efficient" (Figure 3b), implicitly acknowledging the serial bottleneck. But this serial cost is never quantified in the evaluation — there is no measurement of optimizer step time as a fraction of total iteration time for NVMe-offloaded configurations.

Mitigation status. Not addressed. The paper never discusses latency, iteration time, or wall-clock time to convergence. This is a deliberate scope choice (the metric of interest is throughput), but the paper's framing of "excellent training efficiency" (Section 5.2) could mislead a reader into assuming wall-clock efficiency without recognizing the distinction.

6.6 Evaluation Limited to a Single Architecture (GPT-like Transformer) on a Single GPU Family (V100)

The assumption or constraint. All experiments use GPT-like Transformer models on NVIDIA V100 GPUs (Section 8.1). The paper justifies this by stating that "all of the SOTA models with over a billion parameters follow that" architecture (Section 3), and the V100 is a reasonable representative of current-generation accelerator hardware. The ease-inspired implementation (Section 7) is designed to work with arbitrary PyTorch nn.Module hierarchies, but this generality is never tested.

The consequence. Several of ZeRO-Infinity's design choices are validated only for Transformer linear layers. Memory-centric tiling decomposes operators along the output dimension — this works for matrix multiplications (the dominant operation in Transformers) but may not apply to convolutional layers (where spatial dimensions complicate tiling), recurrent layers (where sequential dependencies prevent independent tile execution), or custom operators with data-dependent shapes (e.g., sparse attention patterns, dynamic routing in Mixture-of-Experts, or graph neural network aggregations). The external parameter handling (Section 7.1.1) was tested only on weight tying in language model embeddings — other forms of parameter sharing (e.g., cross-layer attention, shared parameters in vision transformers, or adapter layers) might trigger edge cases in the automatic detection or require manual registration.

The V100-specific bandwidth analysis (Section 4) uses peak_tp = 70 TFlops/GPU measured empirically on V100. This value is used to derive the bandwidth requirements in Figure 3 and the forward-looking projections in Table 3. For GPUs with different compute-to-bandwidth ratios (e.g., A100 with 312 TFlops theoretical FP16), the bandwidth requirements per GPU would be proportionally higher, potentially exceeding what the aggregate PCIe/NVMe bandwidth can provide even at cluster scale. The paper acknowledges this in Section 9 (Table 3 projects requirements for 10× and 100× more powerful accelerators) and argues that technologies like NVLink (40 GB/s per GPU on Summit) can satisfy these requirements, but this argument is speculative and not validated on any hardware other than V100.

What evidence exists in the paper. The model configuration tables (Tables 4–8 in Appendix A) list only Transformer architectures with standard linear layers. There are no experiments with CNNs, RNNs, graph neural networks, or any non-Transformer architecture. The external parameter detection was tested on weight tying (Section 7.1.1 mentions this case explicitly), but no other parameter-sharing patterns are evaluated.

Mitigation status. Partially acknowledged in the forward-looking analysis. Section 9 projects bandwidth requirements for future hardware and argues they are satisfiable with current interconnect technology (citing Summit's 40 GB/s NVLink-to-CPU). However, this analysis assumes the architecture and workload characteristics (AIT) remain similar to the Transformers studied. For architectures with significantly lower AIT (e.g., models with many small operators rather than large matrix multiplications), the bandwidth requirements would be more stringent, and the paper does not explore this regime.

7. Implications and Future Directions

How This Work Changes the Landscape

ZeRO-Infinity causes a fundamental reframing of the large-model training problem. Before this work, the dominant question was: "How do we squeeze models into limited GPU memory?" — answered by 3D parallelism (split across GPUs), ZeRO-Offload (offload optimizer states to CPU while keeping parameters on GPU), and activation checkpointing (trade compute for activation memory). All of these approaches accepted an implicit premise: GPU HBM is the only memory tier fast enough for parameters and computation, and everything else is a second-class fallback that degrades performance. ZeRO-Infinity rejects this premise entirely.

The paper's core reframing is: the GPU memory wall is not about capacity — it is about whether the aggregate bandwidth of slower memory tiers (CPU DRAM, NVMe), summed across all devices in a cluster, can exceed the accelerator's demand for data. This shifts the problem from "how do we minimize offloading because it's slow?" to "how do we make slow memory fast through parallelism?" — a conceptual move that transforms CPU and NVMe memory from emergency overflow into first-class memory tiers.

The magnitude of this shift is best understood by what it makes obsolete. The paper demonstrates that on a fixed 512-GPU cluster, ZeRO-Infinity trains models 50× larger than 3D parallelism (32T vs. ~650B parameters). But the deeper implication is that the number of GPUs required to train a model of a given size is no longer determined by the model's memory footprint. A trillion-parameter model, which would require 320 A100 GPUs with 3D parallelism (by the paper's estimate in Section 1), can be trained on a single DGX-2 node (16 V100 GPUs) with ZeRO-Infinity. This is not an incremental improvement — it changes who can train large models. The capital cost of entry for trillion-parameter model development drops from millions of dollars (for a 320-GPU cluster) to approximately $100K (for a single DGX-2 node).

The paper also resolves a tension that existed implicitly in prior work. ZeRO-Offload (Ren et al., 2021) showed that CPU offloading of optimizer states and gradients was viable, but it couldn't scale model size beyond what fit on a single GPU because parameters were still replicated. This created a puzzling asymmetry: if CPU memory was fast enough for optimizer states, why not for parameters? The answer, which ZeRO-Infinity makes explicit, is that ZeRO-Offload's broadcast-from-owner data movement pattern serialized PCIe access — only one GPU's PCIe link was active when fetching a parameter from CPU, capping bandwidth at ~12 GB/s regardless of cluster size. ZeRO-Infinity's bandwidth-centric partitioning (allgather-from-all) breaks this serialization: every GPU reads its parameter shard simultaneously over its own PCIe link, so aggregate bandwidth scales linearly with GPU count. This explains both why ZeRO-Offload couldn't offload parameters (insufficient bandwidth at any scale) and why ZeRO-Infinity can (bandwidth scales to exceed requirements at cluster scale). The earlier negative result was not evidence that parameter offloading is inherently infeasible — it was evidence that the wrong data movement pattern was being used.

A second reconciliation: the paper's arithmetic intensity (AIT) analysis in Section 4 explains why activation checkpoint offloading is essentially free for large models while parameter offloading requires careful engineering. Prior heterogeneous training systems treated all offloading as costly and tried to minimize it uniformly. The AIT framework shows that activations have ait_activation = 24 × hd × ci, which for hd = 8192 is ~197K FLOPs/byte — meaning even 2 GB/s of PCIe bandwidth sustains over 50% efficiency. Parameters have ait_param_grad = seq × bsz, which at small batch sizes is only 1K–16K FLOPs/byte — requiring ~70 GB/s for efficiency. This differential explains why ZeRO-Infinity's design invests heavily in bandwidth-centric partitioning and overlapping for parameters but treats activation offloading as a solved problem (validated in Figure 6e). Prior work that treated all offloading as equally expensive was missing this structural distinction in the workloads' arithmetic intensity.

The paper also redirects hardware research priorities. The forward-looking analysis in Section 9 (Table 3) projects that even with 100× more powerful accelerators, the aggregate bandwidth achievable by connecting accelerators to CPU/NVMe memory via technologies like NVLink is sufficient for efficient training — provided the interconnects scale with compute. This shifts the hardware design conversation from "how do we put more HBM on the accelerator package?" (which is expensive and physically limited) to "how do we build high-bandwidth links between accelerators and off-package memory?" — a fundamentally different engineering problem with different cost characteristics. The Summit supercomputer (2018) with 40 GB/s NVLink-to-CPU per GPU serves as an existence proof that this is feasible today, not a future aspiration.

Finally, the paper changes what "ease of use" means for large model training. 3D parallelism requires data scientists to manually refactor models for tensor-slicing and pipeline partitioning — a burden that scales with model complexity and restricts the architectures that can be efficiently trained. ZeRO-Infinity's memory-centric tiling eliminates tensor-slicing (large operators are decomposed temporally on a single GPU rather than spatially across GPUs), and its data-parallel-only design eliminates pipeline partitioning. The PyTorch hook system automates all data movement. The result is that the same model code that runs on a single GPU for a 100M-parameter model runs, unchanged, on a 512-GPU cluster for a trillion-parameter model. This is not merely convenient — it removes an architectural bias toward pipeline-friendly models and opens the door to training architectures with complex dependency graphs (cross-layer attention, dynamic routing, irregular compute patterns) at scale.

Follow-Up Research This Work Enables

Convergence validation for NVMe-offloaded training at scale. The paper measures training throughput but never reports whether models trained with NVMe-offloaded optimizer states and parameters converge to the same loss as GPU-only or CPU-offloaded baselines. The Adam optimizer step is computed on CPU for NVMe-offloaded states, potentially with different floating-point reduction order. Optimizer states and updated parameters are written back to NVMe asynchronously through the pinned memory pipeline, introducing variable latency between gradient computation and parameter update. A focused study would train a model of moderate scale (e.g., 10B parameters) to convergence under three configurations — GPU-only (ZeRO-3), CPU-offloaded (ZeRO-Inf-CPU), and NVMe-offloaded (ZeRO-Inf-NVMe) — on a standard benchmark like C4 or the Pile, reporting not only throughput but final validation perplexity, convergence trajectory (loss vs. tokens), and any signs of numerical instability. A negative result (NVMe-offloaded training diverging or converging to a worse minimum) would define the practical bounds of ZeRO-Infinity's applicability; a positive result would validate the entire approach for production training.

Activation offloading to NVMe and the three-tier memory hierarchy. The paper identifies activation memory as the remaining bottleneck at extreme scale — the 20T parameter model's throughput drops 21% because batch size must shrink to 1.25 to fit activation checkpoints in CPU memory (Section 8.2, Figure 5a). The paper explicitly suggests "offloading activation checkpoints to NVMe in a future implementation." The question is whether this can be done without overwhelming the PCIe bandwidth already used for parameter and optimizer state movement. The AIT analysis (Equation 11) predicts activations have very high arithmetic intensity (24 × hd × ci), meaning their bandwidth requirements are modest — but this analysis treats each data stream in isolation. A concrete follow-up would implement NVMe-offloaded activation checkpoints in the infinity offload engine and measure: (a) whether the three simultaneous NVMe data streams (parameter reads, optimizer reads/writes, activation reads/writes) saturate the NVMe device's queue depth or cause contention that reduces effective bandwidth below the sum of individual requirements; (b) whether the overlap engine's prefetcher can pipeline four stages (NVMe-activation, NVMe-parameter, CPU-GPU, GPU-allgather) rather than the current three without running out of prefetch depth; (c) at what model scale and batch size the combined NVMe traffic exceeds the per-node aggregate NVMe bandwidth shown in Figure 2b (~25 GB/s), defining the hard ceiling for three-tier offloading.

Characterizing the model-parallelism threshold for memory-centric tiling. The ease-of-use claim — no model code refactoring — is demonstrated for models up to 1T parameters on a single node (Figure 5c, mp=1 for all configurations). But Table 1 reveals that the 5T, 10T, and 20T models use model parallelism (mp=4, mp=4, mp=8 respectively), and the paper never explains why tiling alone was insufficient. A systematic study would: (a) fix a GPU memory fragmentation model (e.g., the 2 GB contiguous chunk limit used in Figure 6b) and sweep hidden dimension from 8K to 128K, measuring whether memory-centric tiling alone can prevent out-of-memory errors at each scale; (b) measure the throughput cost of tiling vs. model parallelism at each hidden dimension — tiling sequentializes matmuls, reducing GPU utilization, while model parallelism adds inter-GPU communication, and the crossover point is not obvious; (c) determine whether the mp>1 configurations in Table 1 required manual code changes or whether ZeRO-Infinity's automation can compose with model parallelism transparently. This would produce a decision boundary: for hidden dimension below X and batch size above Y, pure data parallelism with tiling suffices; beyond that, model parallelism is necessary, and here is the code change required.

Extension to non-Transformer architectures with irregular computation graphs. All experiments use GPT-like Transformers with standard linear layers, where memory-centric tiling decomposes matrix multiplications cleanly along the output dimension. Many emerging architectures — mixture-of-experts (MoE) with dynamic expert routing, graph neural networks with irregular aggregation patterns, retrieval-augmented models with cross-attention to retrieved chunks — have operators that do not admit simple output-dimension tiling. A stress-test would implement ZeRO-Infinity for one such architecture (e.g., a Switch Transformer with MoE layers, or a retrieval-augmented language model like RETRO) and report: (a) which operators ZeRO-Infinity's automated hooks can handle without modification, (b) which operators require manual external parameter registration (Section 7.1.1) or custom tiling strategies, (c) whether the bandwidth-centric partitioning and overlap engine remain effective when the operator execution order varies dynamically (the dynamic prefetcher claims to handle this via map updates, but this is untested). A negative result — e.g., MoE routing causes the prefetcher to frequently fetch the wrong experts, leading to GPU stalls — would define a boundary on the "arbitrary model architectures" claim.

Joint optimization of pretraining compute and inference-time compute under a total FLOPs budget. The paper demonstrates that test-time compute (through smarter offloading and parallelism) can substitute for pretraining compute (larger GPU clusters) — a trillion-parameter model trainable on 16 GPUs instead of 320 GPUs. This naturally suggests a joint optimization: given a fixed total FLOPs budget, what is the optimal split between pretraining FLOPs (model size, data quantity) and inference-time compute (offloading aggressiveness, data-parallel degree)? This would extend the paper's arithmetic intensity framework to model the tradeoff explicitly: larger models have higher AIT for parameters (since ait_param_grad = seq × bsz, independent of model size) but lower AIT for activations (since activation memory grows with hd × nl). A concrete study would: (a) fix a total dollar or FLOPs budget, (b) sweep model size and ZeRO-Infinity configuration (GPU-only, CPU-offloaded, NVMe-offloaded), measuring total training time to a target validation loss, (c) produce isocost curves analogous to the Chinchilla scaling laws (Hoffmann et al., 2022) but with offloading strategy as an additional degree of freedom. The paper's FLOPs-matched comparison (Section 7) takes a first step but fixes pretraining compute allocation; a joint optimization could reveal non-obvious regimes where training a larger model with NVMe offloading is actually faster per dollar than training a smaller model entirely in GPU memory.

Bandwidth-centric partitioning applied to other resource hierarchies. The core insight — partition data so that all devices read their shards in parallel, converting per-device bandwidth into aggregate bandwidth — is not specific to GPU/CPU/NVMe or to model training. It applies to any system where data is partitioned across many nodes and must be materialized for collective computation. A concrete extension would apply the same principle to: (a) gradient accumulation across micro-batches: instead of each GPU computing gradients on its full micro-batch and then allreducing, partition each micro-batch's data across GPUs so that gradient computation itself uses allgather-from-all, potentially hiding the gradient synchronization latency; (b) distributed inference: when serving a large model across many GPUs, partition each request's KV-cache across all GPUs and use allgather to reconstruct the full context for attention, scaling aggregate memory bandwidth for KV-cache reads with GPU count; (c) federated learning: where each client holds a shard of the model rather than a full replica, and model aggregation uses allgather over heterogeneous client uplinks, scaling aggregate bandwidth with client count. Each extension would measure whether the bandwidth scaling predicted by the AIT framework is realized in practice, and identify the crossover point where per-device bandwidth saturates (e.g., when the number of devices exceeds the available PCIe lanes on the host).

Practical Applications and Downstream Use Cases

Fine-tuning large pretrained models on commodity hardware. The paper demonstrates that a single DGX-2 node (16 V100 GPUs) can train models up to 1 trillion parameters — a configuration that would require 128 GPUs with 3D parallelism (Section 1). The concrete deployment scenario: a research lab or small company downloads a pretrained trillion-parameter model (once such models are publicly released) and fine-tunes it on their domain-specific dataset using their existing single-node GPU server. The paper reports over 40 TFlops/GPU for models up to 100B parameters on a single node (Figure 5c). For a typical fine-tuning run requiring, say, 10^18 FLOPs, this translates to approximately 15 hours on 16 GPUs at 40 TFlops/GPU — well within practical limits for iterative experimentation. This directly enables the "many specialized models from one large pretrained model" paradigm without requiring the capital investment in multi-node GPU clusters that currently restricts such work to large organizations.

Cost-efficient pretraining of models in the 100B–1T parameter range. The paper's superlinear scalability result (Figure 5b) shows that ZeRO-Infinity's per-GPU efficiency improves as node count increases, because aggregate PCIe and NVMe bandwidth grows with node count (bandwidth-centric partitioning). The practical implication: organizations that do have multi-node clusters can train larger models on the same hardware, or train the same model size on fewer nodes, than with 3D parallelism. A 500B parameter model that requires 512 GPUs with 3D parallelism can be trained with comparable throughput on 512 GPUs with ZeRO-Infinity (Figure 5a, leftmost bar), but ZeRO-Infinity can also train a 1T model on the same 512 GPUs where 3D parallelism would run out of memory. The capital cost savings are approximately linear with the reduction in required GPUs for a target model size — training a 1T model on 512 GPUs rather than the 800+ GPUs 3D parallelism would need (per the paper's estimate in Section 1) reduces GPU costs by over 35%.

Rapid prototyping of extremely large architectures. Because ZeRO-Infinity requires no model code refactoring (for models up to ~1T parameters, per the single-node results), a data scientist can iterate on model architecture — changing layer counts, hidden dimensions, attention mechanisms — without the multi-day engineering effort of re-partitioning the model for pipeline and tensor-slicing parallelism. The PyTorch hook system (Section 7.1) intercepts all parameter access patterns and handles data movement automatically. The concrete workflow: a researcher modifies their PyTorch model code (standard nn.Module hierarchy), launches training with the same script, and ZeRO-Infinity handles the rest. This collapses the "architectural exploration → distributed engineering → training" cycle from weeks to hours for models at the 100B–1T scale, which could accelerate research into novel large-model architectures that don't conform to the pipeline-friendly patterns required by 3D parallelism. Architectures with cross-layer attention, dynamic computation graphs, or irregular operator sizes — currently difficult to load-balance across pipeline stages — become as easy to train at scale as standard Transformers.

Democratizing large model research for academic and public-sector labs. The paper's single-node results (Figure 5c) have direct implications for equity of access to large model research. A DGX-2 node (or an equivalent 8–16 GPU server with high-end consumer GPUs) costs on the order of $100K, which is within reach of well-funded academic labs, national research centers, and mid-size companies. With ZeRO-Infinity, this hardware can fine-tune models up to a trillion parameters and can pretrain models in the 10B–100B range. This substantially lowers the barrier to entry for large model research, which is currently dominated by a handful of industrial labs with thousand-GPU clusters. The paper provides the throughput numbers to estimate training times: at 40 TFlops/GPU on 16 GPUs (640 TFlops aggregate), a 100B parameter model requiring ~10^21 FLOPs for pretraining (estimated via the scaling laws in Kaplan et al., 2020) would train in approximately 18 days — long but feasible for a dedicated research project, and dramatically cheaper than renting the equivalent cloud GPU capacity.

When to Prefer This Method

The paper articulates a clear tradeoff between ZeRO-Infinity and 3D parallelism, grounded in their fundamentally different approaches to memory management. The choice depends on model scale, available hardware, and the user's tolerance for model code refactoring:

Prefer ZeRO-Infinity when model size exceeds aggregate GPU memory of the available cluster. This is the threshold where 3D parallelism simply cannot proceed — it runs out of memory. ZeRO-Infinity's infinity offload engine can exploit CPU and NVMe memory to continue scaling, at the cost of potentially reduced per-GPU throughput if the aggregate PCIe/NVMe bandwidth at the cluster scale is insufficient to keep the GPUs compute-bound. The paper demonstrates this for models up to 32T parameters on 512 GPUs, while 3D parallelism on the same hardware caps at ~650B.
Prefer ZeRO-Infinity when hardware budget is fixed and the goal is to train the largest possible model on that hardware. The 50× scale improvement (Figure 1) on identical hardware means ZeRO-Infinity always enables larger models than 3D parallelism on the same GPU count. The throughput is comparable at equal model sizes (Figure 5a, 500B model), so ZeRO-Infinity matches 3D parallelism's performance where both can operate, and extends to model sizes 3D parallelism cannot reach.
Prefer ZeRO-Infinity when ease of use and architectural flexibility are priorities. The "no model code refactoring" claim (validated for models up to 1T parameters in Figure 5c) means ZeRO-Infinity supports arbitrary PyTorch model architectures without the pipeline-partitioning and tensor-slicing engineering required by 3D parallelism. This is particularly relevant when rapidly iterating on model architecture, when the model has complex or irregular dependency graphs that don't map cleanly to pipeline stages, or when the engineering cost of distributed model code is a bottleneck.
Prefer ZeRO-Infinity for fine-tuning large pretrained models on limited hardware. The paper shows that a single DGX-2 node can fine-tune a 1T parameter model (Figure 5c), while 3D parallelism would require 128 GPUs (Section 1). The throughput is sufficient for fine-tuning workloads, which require orders of magnitude less compute than pretraining, making ZeRO-Infinity the only viable option for most users without access to massive GPU clusters.
Prefer 3D parallelism when training models that already fit in aggregate GPU memory and maximum throughput is the sole objective. On the 500B parameter model where both systems can operate (Figure 5a, leftmost bar), the throughput is nearly identical — neither has a decisive advantage. However, 3D parallelism avoids the complexity of NVMe I/O, the potential for CPU optimizer step latency, and the untested convergence behavior of NVMe-offloaded training. For conservative production deployments where the model fits in GPU memory, 3D parallelism is the lower-risk option.
Prefer a hybrid approach when model scale pushes beyond tiling-alone feasibility but full 3D parallelism is overkill. The paper's extreme-scale results use model parallelism (mp=4 to mp=8, Table 1) alongside ZeRO-Infinity, suggesting that for the largest hidden dimensions, memory-centric tiling alone may not suffice and some tensor-slicing is needed. The paper does not characterize when this threshold is crossed (a gap discussed in Section 6.3), but practitioners training models with hidden dimensions above 32K should expect to need some model parallelism and should budget for the associated code refactoring.