{'fraction_non_alphanumeric': 0.027168659374205036, 'fraction_numerical': 0.012431103196811668, 'mean_word_length': 4.856192273314132, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

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). Therefore, the growth in model scale has been supported primarily though system innovations that allow large models to fit in the aggregate GPU memory of multiple GPUs. However, we are getting close to the GPU memory wall. It requires 800 NVIDIA V100 GPUs just to fit a trillion parameter model for training, and such clusters are simply out of reach for most data scientists. In addition, training models at that scale requires complex combinations of parallelism techniques that puts a big burden on the data scientists to refactor their model. In this paper we present ZeRO-Infinity, a novel heterogeneous system technology that leverages GPU, CPU, and NVMe memory to allow for unprecedented model scale on limited resources without requiring model code refactoring. At the same time it achieves excellent training throughput and scalability, unencumbered by the limited CPU or NVMe bandwidth. ZeRO-Infinity can fit models with tens and even hundreds of trillions of parameters for training on current generation GPU clusters. It can be used to fine-tune trillion parameter models on a single NVIDIA DGX-2 node, making large models more accessible. In terms of training throughput and scalability, it sustains over 25 petaflops on 512 NVIDIA V100 GPUs(40% of peak), while also demonstrating super linear scalability. An open source implementation of ZeRO-Infinity is available through DeepSpeed, a deep learning optimization library that makes distributed training easy, efficient, and effective.

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). Therefore, the growth in model scale has been supported primarily though system innovations that allow large models to fit in the aggregate GPU memory of multiple GPUs. However, we are getting close to the GPU memory wall. It requires 800 NVIDIA V100 GPUs just to fit a trillion parameter model for training, and such clusters are simply out of reach for most data scientists. In addition, training models at that scale requires complex combinations of parallelism techniques that puts a big burden on the data scientists to refactor their model.
In this paper we present ZeRO-Infinity, a novel heterogeneous system technology that leverages GPU, CPU, and NVMe memory to allow for unprecedented model scale on limited resources without requiring model code refactoring. At the same time it achieves excellent training throughput and scalability, unencumbered by the limited CPU or NVMe bandwidth. ZeRO-Infinity can fit models with tens and even hundreds of trillions of parameters for training on current generation GPU clusters. It can be used to fine-tune trillion parameter models on a single NVIDIA DGX-2 node, making large models more accessible. In terms of training throughput and scalability, it sustains over 25 petaflops on 512 NVIDIA V100 GPUs (40% of peak), while also demonstrating super linear scalability. An open source implementation of ZeRO-Infinity is available through DeepSpeed 1 .

EXTENDED INTRODUCTION
Deep learning (DL) has made tremendous advances in recent years, allowing it to become an integral part of our lives from powering our search engines to our smart home virtual assistants. Increased model size is at the center of these advancements [1][2][3], and multiple studies have shown that this trend will continue [4,5]. As a result, there has been significant investment in training huge models.
In the last three years, the largest trained dense model in deep learning has grown over 1000x, from one hundred million parameters (ELMo [6]) to over one hundred billion parameters (GPT-3 [4]). In comparison, the single GPU memory has increased by a meager 5x (16 GB to 80 GB). Therefore, the growth in model size has been made possible mainly through advances in system technology for training large DL models, with parallelism technologies such as Figure 1: ZeRO-Infinity can train a model with 32 trillion parameters on 32 NVIDIA V100 DGX-2 nodes (512 GPUs), 50x larger than 3D parallelism, the existing state-of-the-art.
The current state-of-the-art in large model training technology is three-dimensional parallelism (3D parallelism [13,14]), which combines model (tensor-slicing) and pipeline parallelism with data parallelism to efficiently scale DL training to trillions of parameters on hundreds or thousands of GPUs. For example, the DeepSpeed implementation of 3D parallelism can scale to over a trillion parameters on 800 NVIDIA V100 GPUs by fully leveraging the aggregate GPU memory of a cluster [15].
Despite the capabilities of 3D parallelism for large model training, we are now arriving at the GPU memory wall [16]. The aggregate GPU memory is simply not large enough to support the growth in model size. Even with the newest NVIDIA A100 GPUs with 80 GB of memory, 3D parallelism requires 320 GPUs just to fit a trillionparameter model for training, and scaling to a hundred trillion parameter model of the future would require over 6K GPUs even if we assume a 5x increase in GPU memory in the next few years. We can no longer sustain the continuous growth in the model scale with GPU memory as the bottleneck.
The GPU memory wall also limits data scientists from accessing even the large models of today, especially for fine tuning. Large models are first pretrained on large amounts of generic data, and through fine tuning the same model can be specialized for a wide variety of applications. 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. While such compute resources are accessible to many businesses and users, they are unfortunately restricted by the memory available on these compute nodes, which in turn limits the size of the model that can be fine tuned. It makes large model fine tuning inaccessible to most researchers and companies that do not have access to massive GPU clusters. For example fine-tuning GPT-3 would require over 8 DGX-2 nodes(128 GPUs) with 3D parallelism to just fit the model for training, even though a single DGX-2 node (16-GPUs) has enough compute to fine-tune it in a reasonable time.
In addition to the GPU memory wall, state-of-the-art for training massive models is also limited in terms of usability and flexibility. As discussed above, 3D parallelism requires combining data, model, and pipeline parallelism in sophisticated ways to get to hundreds of billions or trillions of parameters. While such a system can be very efficient, 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 also makes 3D parallelism inflexible in the types of models that it can support. Models with complex dependencies cannot be easily converted into a load-balanced pipeline.
Given the landscape of large model training, 3 questions arise: • 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? • How can we make large models of today accessible to more data scientists who don't have access to hundreds to GPUs? • Can we make large model training easier by eliminating the need for model refactoring and multiple forms of parallelism?
In this paper, we take a leap forward from 3D parallelism and present ZeRO-Infinity, a novel system capable of addressing all the aforementioned challenges of large model training.
Unprecedented Model Scale ZeRO-Infinity extends the ZeRO family of technology [11,12] with new innovations in heterogeneous memory access called the infinity offload engine. This allows ZeRO-Infinity to support massive model sizes on limited GPU resources by exploiting CPU and NVMe memory simultaneously. In addition, ZeRO-Infinity also introduces a novel GPU memory optimization technique called memory-centric tiling to support extremely large individual layers that would otherwise not fit in GPU memory even one layer at a time. With the infinity offload engine and memory-centric tiling, ZeRO-Infinity not only supports the next 1000x growth in model size, but also makes large models accessible to data scientists with limited GPU resources.
Excellent Training Efficiency ZeRO-Infinity introduces a novel data partitioning strategy for leveraging aggregate memory bandwidth across all devices, which we refer to as bandwidth-centric partitioning, and combines it with powerful communication overlapcentric design, as well as optimizations for high performance NVMe access in the infinity offload engine. Together, ZeRO-Infinity offers excellent training efficiency, despite offloading data to CPU or NVMe, unencumbered by their limited bandwidth.
Ease of Use With ZeRO-Infinity, 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 in ZeRO-Infinity discussed above aimed at reducing GPU memory requirements of large individual layers that would otherwise require model parallelism (tensor-slicing) to fit the layers in GPU memory. In addition, ZeRO-Infinity eliminates the need for manual model code refactoring, even when scaling to trillions of parameters via an ease inspired implementation that automates all of the communication and data partitioning required for training arbitrary model architectures.
The main contributions of this paper are as follows:  [7,17,18] and pipeline parallelism (PP) [7][8][9] can split the model among processes, vertically and horizontally, respectively. 3D parallelism [14,15] combines data, model, and pipeline parallelism to leverage the merits of each, allowing it to scale to trillions of parameters efficiently. While 3D parallelism can be highly efficient, it requires i) significant model code refactoring to split the model into model and pipeline parallel components, ii) models with complex dependency graphs are difficult to be expressed into load-balanced pipeline stages and iii) the model size is limited by the total available GPU memory. We refer the reader to Ben-Nun and Hoefler [19] for a thorough survey on parallelism in DL.
ZeRO: Zero Redundancy Optimizer ZeRO [11] removes the memory redundancies across data-parallel processes by partitioning the three model states (i.e., optimizer states, gradients, and parameters) across data-parallel processes instead of replicating them. By doing so, it boosts memory efficiency compared to classic data parallelism while retaining its computational granularity and communication efficiency. There are three stages in ZeRO corresponding to three model states: the first stage (ZeRO-1) partitions only the optimizer states, the second stage (ZeRO-2) partitions both the optimizer states and the gradients, and the final stage (ZeRO-3) partitions all three model states. In ZeRO-3, the parameters in each layer of the model are owned by a unique data parallel process. During the training, ZeRO-3 ensures that the parameters required for the forward or backward pass of an operator are available right before its execution by issuing broadcast communication collectives from the owner process. After the execution of the operator, ZeRO-3 also removes the parameters as they are no longer needed until the next forward or backward pass of the operator. Additionally, during the parameter update phase of training, ZeRO-3 ensures that each data-parallel process only updates the optimizer states corresponding to the parameters that it owns. Thus, ZeRO-3 can keep all the model states partitioned throughout the training except for the parameters that are required by the immediate computation.
Heterogeneous Training Approaches Out of several heterogeneous CPU memory based training approaches [20][21][22][23][24][25][26], ZeRO-Offload [12] is the state-of-the-art (SOTA) for large model training on multi-GPUs. ZeRO-Offload is built on top of ZeRO-2 and stores the gradients and the optimizer states in CPU memory. ZeRO-Offload leverages CPU memory in the absence of enough GPU devices to store the optimizer states and gradients. However, 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. ZeRO-Offload also requires a large batch size to remain efficient due to suboptimal data partitioning and limited PCIe bandwidth. We address these limitations of ZeRO-Offload with ZeRO-Infinity. In terms of NVMe based approaches, Zhao et al. [27] use a hierarchical parameter server-based design to offload sparse parameters to SSD for creating a massive scale DL Ads System. In contrast, ZeRO-Infinity is designed to be a generic DL system for training massive dense models.
Reducing Activation Memory Activations are the intermediate results produced during the forward propagation that need to be retained to compute the gradients during backward propagation. Multiple efforts have focused on reducing the memory required by activations through compression [28], activation checkpointing [29,30], or live analysis [31]. ZeRO-Infinity works together with activation checkpointing to reduce activation memory.
Adam Optimizer and Mixed Precision Training Adaptive optimization methods [32][33][34][35] are crucial to achieving SOTA performance and accuracy for effective model training of large models. Compared to SGD, by maintaining fine-grained first-order and second-order statistics for each model parameter and gradient at the cost of significant memory footprint. Adam [33] is the optimizer used most prominently in large model training.
Large model training is generally trained in mixed precision, where the forward and backward propagation are done in FP16 and the parameter updates in FP32 [36]. This leverages the performance acceleration of the tensor core units available on modern GPUs [37].

MEMORY REQUIREMENTS
This section characterizes the memory requirements for DL training. While our methodology is generic, we focus the concrete analysis on Transformer [38] based architectures since all of the SOTA models with over a billion parameters follow that. Our analysis assumes mixed precision training with the Adam optimizer since this recipe is the de facto standard for training Transformer based models.
The memory required for training can be categorized into two components: i) Model states including optimizer states, gradients, and model parameters, ii) Residual states primarily referring to activation memory. To study training on heterogeneous resources, we also characterize the GPU working memory, describing the minimum amount of memory that must be available on the GPU to support training, assuming the model and residual states can be successfully offloaded from GPU memory.
Memory for Model States: The model states are comprised of optimizer states, gradients, and parameters. For mixed precision training with Adam optimizer, the parameters and gradients are stored in FP16 while the optimizer states consist of FP32 momentum, variance, parameters, and gradients. In total, each parameter requires 20 bytes of memory. The total number of parameters in a Transformer based model primarily depends on the hidden dimension (ℎ ) and the number of Transformer layers ( ). Nearly all the parameters in a Transformer block come from four linear layers within each block with sizes: (ℎ , 3ℎ ), (ℎ , ℎ ), (ℎ , 4ℎ ) and (4ℎ , ℎ ), respectively. Thus, the total parameters in a Transformer based model and can be approximated as requiring a total memory 240 × × ℎ 2 (2) in bytes to store the model states. Figure 2a column 5 shows the memory required to store the model states of a GPT-3 like Transformer based model with 100 billion to a 100 trillion parameters created by varying hidden dimension and number of layers. To put the memory requirements in context, Figure 2b column 3 shows the aggregate GPU memory available on a single NVIDIA V100 DGX-2 box as well as a DGX-2 SuperPOD cluster. Note that it requires 64 GPUs to just fit the model states for a 100B parameter model. Fitting a trillion parameter model requires over 512 GPUs, while a 10 trillion parameter model is beyond the scope of even a massive 1536 GPU cluster.
Memory for Residual States: The residual states primarily consist of the activation memory, which depends on the model architecture, batch size ( ) and sequence length ( ), and it can be quite large. On the positive side, the memory required for activation can be significantly reduced via activation checkpointing [29], which trades off activation memory at the expense of 0.33x additional recomputation. Large models such as Turing-NLG 17.2B and GPT-3 175B were all trained using activation checkpointing. The memory required to store activation checkpoints is estimated as bytes where is the number of Transformer blocks between two activation checkpoints, and × × ℎ is the size of the input to each Transformer block. Figure 2a column 7 shows the memory required to store activation checkpoints for batch size of 32 and sequence length of 1024 assuming we store one activation per Transformer block. Many modern GPU clusters have 8-16 GPUs per node, and so we chose a batch size of 2-4 per GPU, resulting in a batch size of 32 as a conservative estimate of activation within each node. While the resulting activation checkpoints are orders of magnitude smaller than the full set of activations (column 6) , beyond a trillion parameters they still get too large to fit in GPU memory for the batch size and sequence length under consideration.
Model State Working Memory (MSWM) is the minimum amount of GPU memory required to perform forward or backward propagation on the largest single operator in the model after all the model states have been offload to CPU or NVMe. This is approximately given by the size of the parameters and gradients of that operator in the model, since there must be at least enough memory to hold the parameter and its gradient for backward propagation.
For a Transformer based model, the largest operator is a linear layer that transforms hidden states from ℎ to 4ℎ . The size of the parameter and gradients of this linear layer in bytes is Note that MSWM (Figure 2a Column 8) grows significantly beyond a 100 billion parameters, requiring multiple gigabytes in contiguous memory, which can result in running out of memory during training due to lack of enough contiguous memory to satisfy these requirements. State-of-art approaches like 3D Parallelism, addresses this issue via model parallelism, by splitting individual operator across multiple GPUs. In Sec. 5.1.3, we discuss a novel approach for addressing these massive model state working memory without requiring model parallelism. Activation Working Memory (AWM) is the memory required in the backward propagation for recomputing the activations before performing the actual backward propagation. This is the size of the activations between two consecutive activation checkpoints. For example, if we create one activation checkpoint per Transformer block, the memory is given by the size of the total activation per Transformer block. This is given in bytes by approximately Figure 2a column 8 shows that AWM gets large beyond 10 trillion parameters, even with = 1. Unlike MSWM that is only composed of a single parameter and gradient, AWM is composed of dozens of activations, and does not cause memory issues due to lack of contiguous memory as long as the total AWM can fit in GPU memory.

BANDWIDTH REQUIREMENTS
A critical question of offloading to CPU and NVMe memory is whether their limited bandwidth will hurt training efficiency. This section characterizes the impact of bandwidth on training efficiency.
We start from defining an efficiency metric. Assuming a workload execution without any compute and communication overlap, we can use the peak computational throughput ( ), data movement bandwidth ( ) and its arithmetic intensity ( ) to estimate the training efficiency.
The arithmetic intensity (AIT) of a workload is the ratio between the total computation and the data required by the computation. It describes the amount of computation per data movement. Higher AIT means a lower requirement on the data movement bandwidth, since for each data loaded the accelerator can do more computations. The efficiency metric can derived as follows: The efficiency can be written as a function of , and : We will use this simple efficiency equation to characterize the data movement bandwidth required for training massive models. But before that, we will first quantify for DL training workloads.

Quantifying AIT in DL training
Model states and activation checkpoints can have varying . We can quantify them by first identifying the total computation in each iteration of DL training, and then identifying the data movement volume for each of the model states and activations.
Total Computation per Iteration The total computation per iteration is dominated by the computation in the linear layers of the Transformer. For the forward propagation this can be approximated as a function of the number of parameters, sequence length, and batch size, given by 2 × × × . The cost of backward propagation is approximately twice that of forward propagation. Additionally, activation checkpointing requires an additional forward computation as part of recomputation during backward propagation. Therefore, the total computation per iteration is: AIT w.r.t. Parameters and Gradients During forward and back propagation, model parameters must be loaded from the source location to GPU registers at least twice, i) during forward, ii) during the actual backward, resulting in a data movement of 2× . In presence of activation checkpointing, the parameters may be loaded one additional time for re-computation during the backward pass, adding another 1 × . Furthermore, the gradients must be stored from the GPU registers to its final location at least once, adding a final 1 × in data movement. Therefore, assuming that parameters and gradients are stored at the same final location, the total data movement during the forward and backward pass would be 4 × , i.e. 2 × 4 × in bytes. The total computation per iteration is given by Sec. 4.1. Therefore the w.r.t parameter and gradients is × .
AIT w.r.t. Optimizer States During the optimizer step, the optimizer states must be read at least once, and the optimizer states must be written at least once. So the total data movement is 2 × _ , which is approximately 2×16× bytes. The total computation per iteration is given by Sec. 4.1. Therefore w.r.t optimizer states during a full training iteration is AIT w.r.t. Activation Checkpoints During the forward propagation activation checkpoints must be saved to their final location, and must be retrieved during the backward propagation. Therefore, the total data movement w.r.t activation checkpoints in bytes is given by 2× _ _ ℎ _ _ which is given by 4 × / × ℎ × × from Eq. (3). The total computation per iteration is given by Sec. 4.1. So the w.r.t activation checkpoints is given by

Bandwidth Requirements
Due to the variation in the AIT, model states and activation checkpoints have very different bandwidth requirements to achieve good efficiency. The former only depends on the batch size and sequence length, while the latter only depends on the frequency of activation checkpoints and hidden dimension size of the model. Besides AIT, the bandwidth requirement for efficiency also depends on , as shown in Eq. (6). Using , and we first show how efficiency varies with bandwidth w.r.t to different model and residual states, and then discuss the bandwidth requirements on these states for DL training to be efficient. Our methodology is generic and can be applied to understanding the bandwidth requirements on any current or future generation clusters. Here, we use NVIDIA V100 DGX-2 SuperPOD cluster as our example platform.
Using the expression from Sec. 4.1 and efficiency metric based on Eq. (6), Figure 3 shows the relationship between efficiency and available bandwidth w.r.t. parameter and gradients, optimizer states, and activation checkpoints. To produce these plots, we computed the based on expressions derived in Sec. 4.1, for varying batch sizes, sequence length and model configurations. More specifically, we use a sequence length of 1024, the same sequence length used for GPT-2 [2], Megatron-LM [7], and Turing-NLG [39]. We vary batch size range from 1 to 16 to capture large GPU and small GPU experiments, respectively. A small batch size per GPU is used when running on large number of GPUs, while a large batch size per GPU is used when training on relatively fewer GPUs to maintain a reasonable effective batch size for training. Our hidden size ranges from 8K-64K representing models with hundreds of billions of parameters, to tens of trillions of parameters as shown in Figure 2a.
To identify for this analysis, we use an empirical approach 4 . We ran models with aforementioned configurations on a single NVIDIA V100 DGX-2 box with all non-GPU communication turned off to simulate a virtually unlimited bandwidth scenario. The performance achieved ranged from 62-78 TFlops/GPU based on the hidden size of 8K-64K, respectively. We used the average of 70 TFlops/GPU to represent for the purpose of this analysis 5 .
Bandwidth w.r.t. Parameter and Gradients Figure 3a shows that with a bandwidth of over 70 GB/s for parameter and gradients, we can achieve over 50% efficiency for even the smallest batch size. At this bandwidth, the data movement in theory can be completely overlapped with the computation to achieve a 100% efficiency.
Bandwidth w.r.t. Optimizer States Figure 3b shows that optimizer states require nearly 4x higher bandwidth to achieve 50% efficiency compared to parameters and gradients. Furthermore, the optimizer states are updated at the end of the forward and backward propagation and cannot be overlapped with the computation. As a result they require significantly larger bandwidth to keep the overall DL workload efficient. For example achieving 90% efficiency with batch size of 2 per GPU requires nearly 1.5 TB/s of effective bandwidth, which is greater than even the GPU memory bandwidth.
Bandwidth w.r.t. activation memory Figure 3c also shows that with activation checkpointing enabled, a meager bandwidth of 2 GB/s is able to sustain over 50% efficiency even for a hidden size of 2 . The bandwidth requirement drops down to less than 1 GB/s once the hidden size grows over 8 .

ZERO-INFINITY DESIGN OVERVIEW
In this section we present an overview of the design choices in ZeRO-Infinity that enable it to achieve unprecedented model scale while offering excellent training efficiency and ease of use. A bird's eye view of ZeRO-Infinity is illustrated in Figure 4 and discussed below.

Design for Unprecedented Scale
Modern GPU clusters are highly heterogeneous in terms of memory storage. In addition to the GPU memory, they have CPU memory as well as massive NVMe storage that is over 50x larger than the GPU memory and nearly 20x larger than CPU memory (See Fig. 2b).
We developed ZeRO-Infinity, a parallel system for DL training that can transcend the GPU memory wall by exploiting these heterogeneous memory systems in modern GPU clusters. Figure 1 compares the maximum achieved model size of 3D parallelism and ZeRO-Infinity. ZeRO-Infinity supports one trillion parameters per NVIDIA V100 DGX-2 node, a 50x increase over 3D parallelism.

Infinity offload engine for model states.
ZeRO-Infinity is built on top of ZeRO-3 [11] which partitions all model states to remove memory redundancy as discussed in Sec. 2. Unlike any of the existing ZeRO family of technology, ZeRO-Infinity is designed with a powerful offload mechanism called the infinity offload engine which can offload all of the partitioned model states to CPU or NVMe memory, or keep them on the GPU based on the memory requirements. Note from Fig. 2a and Fig. 2b, even the model states required by a 100 trillion parameter model can fit in the aggregate NVMe memory of a DGX-2 cluster with 96 nodes (1536 GPUs). Therefore, the infinity offload engine allows ZeRO-Infinity to fit model states of models with hundreds of trillions of parameters. See Sec. 6 for more details.
5.1.2 CPU Offload for activations. In addition to model states, ZeRO-Infinity can offload activation memory to CPU memory, when necessary. Note that the activation checkpoints (0.76 TB) required by a 10 trillion parameter model can easily fit in the 1.5TB Figure 4: A snapshot of ZeRO-Infinity training a model with two layers on four data parallel (DP) ranks. Communication for the backward pass of the first layer is depicted. Partitioned parameters are moved from slow memory to GPU and then collected to form the full layer. After gradients are computed, they are aggregated, repartitoned, and then offloaded to slow memory. Layers are denoted with subscripts and DP ranks are denoted with superscripts. For example, (2) 0 is the portion of layer 0's parameters owned by (2) .
of CPU memory available on a DGX-2 system, while the 3 TBs of activation checkpoints required by a 100 trillion parameter is within reach of the CPU memory of the next generation hardware. Therefore, by offloading activation checkpoints to CPU memory, ZeRO-Infinity can fit the activation checkpoints of models with hundreds of trillions of parameters.

Memory-centric tiling for working memory.
To reduce the working memory requirements of DL training for large models, ZeRO-Infinity introduces a novel technique called memory-centric tiling that exploits the data fetch and release pattern of ZeRO-3 to reduce the working memory requirements by breaking down a large operator into smaller tiles that can be executed sequentially.
For example, to reduce the working memory for a large linear operator, ZeRO-Infinity represents the operator as a mathematically equivalent sequence of smaller linear operators consisting of tiles of parameters from the original operator, and executes them sequentially. When combined with ZeRO-3, the parameter and gradients of each tile can be fetched and released one at a time, reducing the working memory proportional to the number of tiles. Therefore, ZeRO-Infinity can support operators of arbitrary sizes, without relying on model parallelism to fit them in limited GPU memory.

Design for Excellent Training Efficiency
Offloading all model states and activations to CPU or NVMe is only practical if ZeRO-Infinity can achieve high efficiency despite the offload. In reality this is extremely challenging since CPU memory is an order of magnitude slower than GPU memory bandwidth, while the NVMe bandwidth is yet another order of magnitude slower than the CPU memory bandwidth. Furthermore, reading and writing to these memory from GPU is even slower (see Fig. 2b).
On a system like the DGX-2, the bandwidth must be greater than 70GB/s, 1.5TB/s, and 1-4 GB/s w.r.t. parameter and gradients, optimizer states, and activation checkpoints, respectively for DL training to be efficient, based on our analysis in Sec. 4. Here we discuss how ZeRO-Infinity achieves the necessary bandwidths to achieve excellent efficiency.

Efficiency w.r.t Parameter and
Gradients. The data movement bandwidth for parameters and gradients must be greater than 70GB/s, close to the GPU-GPU bandwidth available on DGX-2 clusters [40]. Therefore, a DL parallel training solution like ZeRO-3 [11] where parameters are broadcasted from the owner GPU to the rest before using them in forward or backward propagation can run efficiently as long as the communication is overlapped.
On the contrary, a meager 12 GB/s PCIe bandwidth from a single GPU to CPU memory or NVMe (see Fig. 2b) or vice-versa is simply not sufficient to support heterogeneous training at scale 6 . Therefore, existing heterogeneous solutions like ZeRO-Offload where the parameters must be first moved from CPU to owner GPU before broadcasting requires significantly large batch sizes per GPU to achieve enough necessary to be efficient under the limited bandwidth. This poses two problems: i) for massive models the activation memory will get too large to fit even in CPU memory, and ii) the effective batch size becomes too large when scaling to hundreds or thousands of GPUs for effective convergence.
ZeRO-Infinity addresses these challenges in two ways: i) bandwidthcentric partitioning: a novel data mapping and parallel data retrieval strategy for offloaded parameters and gradients that allows ZeRO-Infinity to achieve virtually unlimited heterogeneous memory bandwidth (details in Sec. 6.1), and ii) an overlap centric design that allows ZeRO-Infinity to overlap not only GPU-GPU communication with computation but also NVMe-CPU and CPU-GPU communications over the PCIe (details in Sec. 5.1.3).

Efficiency w.r.t Optimizer States.
Unlike parameters and gradients that are consumed and produced sequentially during the forward and backward propagation, optimizer states can be updated in parallel, all at once. This property is leveraged by both ZeRO-3 and ZeRO-Offload, that store and update the optimizer states in GPU and CPU memory, respectively, in parallel across all available GPUs and CPUs. As a result the aggregate GPU or CPU memory bandwidth can get much higher than the required 1.5TB/s with increase in GPU or CPU count.
Since ZeRO-Infinity is built upon ZeRO-3, it can also leverage the aggregate GPU and CPU memory bandwidth as well as the aggregate CPU compute for optimizer step, when offloading optimizer states to CPU memory. However, with NVMe offload, it is necessary to bring the data from NVMe to CPU memory and back in chunks that can fit in the CPU memory to perform the optimizer step, one chunk at a time. The optimizer step is therefore limited by the NVMe-CPU memory bandwidth: while ZeRO-Infinity can achieve aggregate NVMe bandwidth across multiple nodes, it is crucial to achieve near peak NVMe bandwidth per node, to allow supporting the necessary bandwidth of over 1.5 TB/s with as few nodes, and as small batch size as possible. Furthermore, the process of bringing data in and out of NVMe to CPU memory, or from CPU memory to GPU memory can cause CPU memory fragmentation 6 CPU and NVMe bandwidth are in the order of 100 GB/s and 25 GB/s, respectively, but reading data from CPU or NVMe to a single GPU is limited by the achievable PCIe bandwidth which is around 10-12 GB/s in both GPU and CPU that can result in out of memory even with plenty of memory still available.
The infinity offload engine can not only achieve near peak NVMe bandwidth, it can also allows ZeRO-Infinity to overlap NVMe to CPU reads with CPU to NVMe writes, as well as the CPU computation for the optimizer step at the same time to allow ZeRO-Infinity to remain efficient with a modest batch size on small number of GPUs and with small batch sizes on large numbers of GPUs. At the same time, it minimizes memory fragmentation by carefully reusing temporary buffers for data movement. We discuss the optimizations in infinity offload engine and in detail in Sec. 6.

Efficiency w.r.t Activations.
On a DGX-2 node, each GPU can read and write data at about 3 GB/s to CPU memory in parallel over the PCIe allowing activation checkpoints to be offloaded to CPU memory while retaining over 80% efficiency for hidden size larger 8 or larger. To also allow for high efficiency at smaller hidden sizes, ZeRO-Infinity can decrease the frequency of activation checkpoints as well as effectively overlap the communication of activation checkpoints both to and from CPU memory with the forward and backward computation on the GPU.

Design for Ease of Use
With ZeRO-Infinity, 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 in ZeRO-Infinity discussed in Sec. 5.1.3 aimed at reducing GPU memory requirements of large individual layers that would otherwise require model parallelism (tensor-slicing) to fit the layers in GPU memory.
In addition, ZeRO-Infinity is implemented in PyTorch in a way that eliminates the need for manual model code refactoring even when scaling to trillions of parameters. This is made possible through an ease-inspired implementation with two automated features: i) automated data movement to gather and partition parameters right before and after they are required during the training. ZeRO-Infinity does this by injecting i) pre forward/backward hooks into PyTorch submodules that trigger allgather collectives to collect the parameters required before its forward/backward pass and ii) post forward/backward hooks that trigger parameter/gradient partiitoning and optionally offloading them to CPU or NVMe (see Sec. 7.1 for details). ii) automated model partitioning during initialization such that models that can not fit within single GPU or CPU memory can still be initialized without requiring manual partitioning of the model across data parallel processes. ZeRO-Infinity achieves this by wrapping the constructor of all module classes so that parameters of each submodule are partitioned and offloaded immediately after they are created during initialization. The entire model is never fully instantiated on a single data parallel process (see Sec. 7.2 for details).

EFFICIENCY OPTIMIZATIONS
In this section, we deep dive into the optimizations introduced in Sec. 5 that allow ZeRO-Infinity to achieve excellent efficiency.

Bandwidth-Centric Partitioning
ZeRO-Infinity implements a novel data mapping and retrieval strategy to address the NVMe and CPU memory bandwidth limitations. Unlike ZeRO [11] and ZeRO-Offload [12], where parameters of each layer are owned by a single data parallel process, which broadcasts them to the rest when needed, ZeRO-Infinity partitions individual parameters across all the data parallel process, and uses an allgather instead of a broadcast when a parameter needs to be accessed. Note that both broadcast and allgather communication collectives have the same communication cost when it comes to data movement volume if the data is located on the GPU. Therefore, this makes no difference for a GPU-only training. However, this is a game changer when the data is located in NVMe or CPU.
In the broadcast-based approach, since each parameter is fully owned by one of the data parallel processes, the parameter must be first communicated from its source location to the GPU memory via the PCIe before the broadcast can happen. Note that only a single PCIe can be active for this process, while all the PCIe links connected to all the other GPUs are idle. On the contrary, with the partitioned parameter and allgather based approach in ZeRO-Infinity, all PCIe links are active in parallel, each bringing in 1/ ℎ portion of the parameter where is the data parallel degree. As a result, the effective communication bandwidth between NVMe or CPU to the GPU, increases linearly with the degree. For example, with broadcast-based approach, the CPU/NVMe to GPU bandwidth stays constant at about 12 GB/s with PCIe Gen 3, even with 16-way data parallelism on the DGX-2 box. However, with the all-gather-based approach, the effective achievable bandwidth increases to about 48/25 GB/s (3.0/1.6 GB/s per GPU), respectively (see Fig. 2b), limited only by the max aggregate PCIe bandwidth and max NVMe bandwidth per DGX-2 node. From here, the bandwidth grows linearly with more nodes. When training a massive model at massive scale, ZeRO-Infinity can therefore offer significantly more heterogeneous memory bandwidth than necessary (virtually unlimited) for the training to remain efficient. For example, on 64 DGX-2 nodes, ZeRO-Infinity has access to over 3TB/s of CPU memory bandwidth and over 1.5TB/s of NVMe bandwidth.

Overlap Centric Design
While ZeRO-Infinity can leverage sufficient heterogeneous memory bandwidth on a multi-node setup, the bandwidth can still be a bottleneck on a single GPU or single node setup. Even the GPU-GPU allgather communication has a big impact on efficiency when running with a small batch size (Fig. 3). Furthermore, accessing NVMe memory requires a three step process: i) read data from NVMe to CPU memory (nc-transfer), ii) copy the data from CPU memory to GPU memory (cg-transfer), iii) execute allgather to construct the full parameter on all GPUs (gg-transfer). The sequential nature of these data movements means that if done naively, the total communication time would be the sum of each of these three data movement cost, resulting in poor efficiency even if the bandwidth for data movement at each of these stages is individually sufficient.
To address these issues, ZeRO-Infinity has an overlap engine that not only overlaps GPU-GPU communication with GPU computation, but also overlaps the NVMe to CPU, and CPU to GPU communication, all at the same time. The overlap engine has two components: i) A dynamic prefetcher for overlapping the data movement required to reconstruct parameters before they are consumed in the forward or backward pass, and ii) a communication and offload overlapping mechanism for executing the data movement required by gradients in parallel with the backward computation.
The dynamic prefetcher in ZeRO-Infinity traces the forward and backward computation on that fly, constructing an internal map of the operator sequence for each iteration. During each iteration, the prefetcher keeps track of where it is in the operator sequence and prefetches the parameter requires by the future operators. The prefetcher is aware of the three step communication process, and therefore can overlap the nc-transfer for one parameter, with cgtransfer and gg-transfer of other parameters. For instance, before executing the ℎ operator, the prefetcher can invoke nc, cg, and gg-transfer for parameters required by + 3, + 2, and + 1 operators, respectively. Note that all of these data movement can happen in parallel with the execution of the ℎ operator. Furthermore, ZeRO-Infinity can update the operator sequence map in case of dynamic workflow, allowing for appropriate prefetching even when the forward and backward propagation changes across iterations.
Similarly, in the backward pass, ZeRO-Infinity can overlap the reduce-scatter for gradients of the parameters in ( + 1) ℎ operator with the computation of the ℎ operator, while simultaneous transferring the partitioned gradients from the reduce-scatter of the gradients of the ( + 2) ℎ operator to the CPU or NVMe.
With this powerful overlap centric design, ZeRO-Infinity hides significant portions of data movement even when training with a small number of GPUs and small batch size per GPU.

Infinity Offload Engine
The infinity offload engine is composed of two main components: DeepNVMe, a powerful C++ NVMe read/write library in the infinity offload engine that supports bulk read/write requests for asynchronous completion, and explicit synchronization requests to flush ongoing read/writes. The support for asynchrony allows ZeRO-Infinity to overlap these requests with GPU/GPU or GPU/CPU communication or computation.
Most importantly, DeepNVMe is capable of achieving near peak sequential read and write bandwidths on the NVMe storage device. It achieves this high performance through a number of optimizations, including aggressive parallelization of I/O requests (whether from a single user thread or across multiple user threads), smart work scheduling, avoiding data copying, and memory pinning.
Pinned memory management layer To ensure high performance tensor reads (or writes) from (to) NVMe/CPU storage, the source (or destination) tensors must reside in pinned memory buffers. However, pinned memory buffers are scarce system resources, and their oversubscription by a single process can degrade overall system performance or cause system instability. This layer manages the limited supply of pinned memory by reusing a small amount (tens of GBs) for offloading the entire model states (up to tens of TBs) to CPU or NVMe. The reuse of memory buffer prevents memory fragmentation in CPU and GPU memory. This layer also provides PyTorch tensors with pinned memory data, allowing inplace computation of the tensors so that they can then be written to NVMe without any further copies to improve bandwidth.
ZeRO-Infinity is implemented on top of PyTorch, and is designed to be used without any model code refactoring similar to standard data-parallel training in PyTorch. This section details some of the challenges faced in implementing such a system.

Automating Data Movement
ZeRO-Infinity must coordinate the movement of tensors comprising the model parameters, gradients, and optimizer states. When a tensor is not in active use, it remains partitioned among workers and potentially offloaded to CPU or NVMe memory. The system must ensure that the tensors are resident in GPU memory in time for use and then later re-partitioned.
PyTorch models are expressed as a hierarchy of modules that represent the layers of a neural network. For example, Transformer architectures [38] contain submodules such as self-attention and feedforward networks. The self-attention submodules are further comprised of linear transformations and other submodules.
ZeRO-Infinity recursively injects hooks into the submodules of a model to automate the required data movement. At the start of a submodule's forward pass, these hooks ensure that the submodule's parameters are available for computations, otherwise it will execute the appropriate allgather collectives and block until the parmaeters become available. The overlap-centric design detailed in Sec. 6.2 is critical to minimizing stalls due to parameter communication. At the end of the submodule's forward pass, we partition the parameters again and optionally offload them. The backward pass is handled in a similar fashion.

Auto Registration of External Parameters.
In the ideal case, a submodule's parameters and gradients are only accessed within its own forward and backward passes, making it straight forward to identify and automate the data movement as discussed in section above. However, some model architectures are exceptions, where parameters defined and allocated in a submodule is used in forward and backward propagation of a different submodule. For example, language models such as GPT [41] share the weights of the embedding layer at both the beginning and the end of the network to map words to vectors and vice versa. We refer to the parameters that are used across module boundaries as external parameters. In presence of external parameters, it is difficult to know which parameters to gather at the beginning of a submodule's forward and backward pass.
One way to address this is to register external parameters with ZeRO-Infinity so that they are collected for the forward and backward passes of the submodule that access them. After registration, an external parameter is treated like all others and will be included in the prefetching system as described in Sec. 6.2. We provide APIs for manual registration of external parameters.
In order to improve user experience, we also provide mechanisms to detect these scenarios and automatically register external parameters so the the user does not have to make any code change: Intercepting partitioned parameter accesses PyTorch modules store their tensor parameters in a hash table. At the initialization time, we replace the hash table with a subclassed type that overrides the tensor accesses. When a partitioned parameter is accessed, we do a blocking allgather on the parameter, register it as an external parameter, and then return the gathered parameter.
Activation introspection A submodule may return a parameter from its forward pass to be consumed by another submodule's forward and backward passes. For example, Megatron-LM returns bias vectors from the forward pass of linear layers and they are consumed by the parent Transformer layer modules. We inspect the activation outputs returned from each submodule's forward pass for partitioned parameters. If one is discovered, we collect and register it as an external parameter.

Automatic Model Partitioning during Initialization
If the model is large, then it may not be possible to fully initialize the model with traditional data parallel approach, replicating it on each data parallel process before it can be partitioned for ZeRO-Infinity. For example, a 500 billion parameter model will occupy 1 TB of memory in half precision, and thus a system with 8 GPUs per node requires 8 TB of aggregate CPU or GPU memory just for the initial data parallel allocation step. This is beyond the GPU or CPU memory available on a node. To address this limitation, the parameters corresponding to each layer of the model must be partitioned at the time of initialization, and not after the entire model is initialized. To do this, we provide a Python ZeRO-Infinity context which decorates the __init__ method of torch.nn.Module, so that parameters allocated under each module/sub-module are partitioned immediately after its initialization among the group of data parallel processes.
As a 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. In the example above, the 500 billion parameter model can therefore be fully partitioned during its initialization requiring only 1 TB of aggregate CPU memory regardless of the total number of data parallel process.

EVALUATION
This section evaluates ZeRO-Infinity, demonstrating that it achieves excellent training efficiency and scalability for models with tens of trillion parameters. We also show the impact of various technologies within ZeRO-Infinity on model scale and performance.

Methodology
Hardware. We conducted our experiments on a cluster of up to 512 V100 SXM3 32 GB GPUs (32 DGX-2 nodes) with 800 Gbps internode communication bandwidth.
Baseline. For experiments without model parallelism (mp), we use torch's distributed data parallel (DDP [42]) as a baseline. For experiments with model parallelism, we use Megatron-LM [7]. As a baseline for each experiment we use the relevant state-of-the-art method among 3D Parallelism [13], ZeRO [11], or ZeRO-Offload [12].
Model Configurations. We use GPT-like Transformer based models. We fix the sequence length to 1024 and vary the hidden dimension and number of layers to obtain models with different number of parameters. Table 1   used throughout our evaluation, (see Appendix A) for additional configurations.

Model Size and Speed
Model Size ZeRO-Infinity trains models over 32 trillion parameters compared to about 650B parameters with 3D parallelism, state of the art, offering a leap of 50x in model scale ( Figure 1). Model Speed Figure 5a shows the performance of ZeRO-Infinity on up to 20 trillion parameter models on 512 GPUs. For the 500B model (close to the largest that 3D parallelism can run on these resources), ZeRO-Infinity and 3D parallelism achieve nearly identical throughput, indicating ZeRO-Infinity is on par with the training efficiency of the state-of-the-art. When increasing the model size further, 3D parallelism simply runs out of memory, while ZeRO-Infinity trains up to 20 trillion parameter models (40x larger) with excellent throughput of up to 49 TFlops/GPU. At the extreme-scale, Figure 5a shows a performance drop from 10T (43 TFlops/GPU), and 20T (34 TFlops/GPU). This drop is not due to NVMe bandwidth as both model sizes use NVMe offload, but instead due to an extremely small batch size per GPU (Table 1) at 20T scale as a result of limited CPU memory to store activation checkpoints. This can be improved by increasing the CPU memory or offloading activation checkpoints to NVMe in a future implementation.  Figure 5c shows performance of training 10B to 1T models on a single node (16 GPUs) with ZeRO-Infinity without any model parallelism. With models up to 100 billion parameters, ZeRO-Infinity achieves excellent performance of over 40 TFlops/GPU, making it possible to fine-tuning models such as GPT-3 with just a single DGX-2 box. In contrast, 3D parallelism is unable to scale to models with over 20 billion parameters. These results demonstrate two aspects of ZeRO-Infinity: i) Accessibility to fine-tuning large models with up to a trillion parameters on a single NVIDIA DGX-2 node, empowering users without access to large GPU clusters. ii) Ease-of-Use: Models at this scale can be trained using ZeRO-Infinity without combining model or pipeline parallelism, or requiring model code refactoring, making it easy for data scientists to scale up their models.

Impact of System Features on Model Scale
We show the impact of different device placement strategies on model scale and impact of memory-centric tiling (Sec. 5.1.3) on maximum hidden size using a single DGX-2 system (16 GPUs).
Maximum model size Figure 6a shows the effect of different device placement and partitioning strategies (see Table 2) on maximum model size. By using data parallelism alone we're limited to only 1.4B parameters, due to limited GPU memory and significant    model state redundancies. As we introduce optimizer/gradient partitioning and offloading to CPU with ZeRO-2 and ZeRO-Offload, we are able to scale up 9x to 13B parameters on a single node. Partitioning and offloading parameter states to CPU in ZeRO-Infinity allows us to almost reach 100B parameters. However, the final major jump in scale comes from offloading model states to NVMe which finally gets us to 1T parameters, resulting in a 700x increase in model size relative to data parallelism alone.
Maximum Hidden Size We evaluate the impact of memorycentric tiling in enabling large hidden sizes in the presence of memory fragmentation. We train a single layer transformer model with different hidden sizes and tiling factors to identify the largest hidden size that can be trained with and without tiling. To keep memory fragmentation consistent across all the experiments, we pre fragment the total GPU memory into 2 GB contiguous chunks so that all memory allocation requests larger than 2GB will fail. Figure 6b shows the largest hidden size that can be trained without memory-centric tiling is 8K, while we can even train a massive hidden size of 64K using memory-centric tiling factor of 16. With memory-centric tiling, ZeRO-Infinity greatly simplifies DL system stack by avoiding the need for model parallelism, making it easy for data scientists to train with large hidden sizes.

Impact of System Features on Performance
We evaluate the effects of the infinity offload engine (Sec. 5), bandwidthcentric partitioning (Sec. 6.1), overlap-centric design (Sec. 6.2), and activation checkpoint offloading (Sec. 4.1) on training speed.
ZeRO-Infinity vs ZeRO-Offload Figure 6c shows the impact of offloading gradients to CPU memory with ZeRO-Infinity vs ZeRO-Offload on the back propagation time of an 8B parameter model. ZeRO-Infinity leverages the aggregate PCIe bandwidth across GPUs to offload the gradients, resulting in a speedup of nearly 2x at 64 GPUs compared to ZeRO-Offload which is limited by single PCIe bandwidth.
Prefetching and Overlapping Figure 6d shows the relative throughput difference with communication overlapping and prefetching turned on and off for an 8B parameter model with 64 GPUs. The figure shows that prefetching and overlapping are crucial to achieving good performance at small batch sizes per GPU, while its impact diminishes at large batch sizes.
Activation checkpoint offload Figure 6e shows that CPU offloading of activation checkpoints in ZeRO-Infinity reduces the training throughput by up to 1.2x for small hidden sizes, but for hidden sizes 32K and 64K, the impact is minimal, demonstrating that it is possible to offload activation checkpoints to CPU memory without impacting efficiency for large hidden sizes.  Table 3: Bandwidth (bw) requirements for ZeRO-Infinity to remain efficient on a cluster of 512 accelerator devices with 10x and 100x more achievable compute than NVIDIA V100 GPUs.

CONCLUSION & FUTURE IMPLICATIONS
In this paper, we presented ZeRO-Infinity, a novel heterogeneous system technology that leverages GPU, CPU, and NVMe memory to allow for unprecedented model scale that is accessible and easy to use while achieving excellent efficiency. It offers a paradigm shift in how we think about memory for large model training. 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 on current generation of GPU clusters.
As we look into the future, the GPUs and other accelerators will become more powerful, and this aggregate bandwidth required for efficient training will also increase. Table 3 shows that even when the compute of the accelerators increases by 10x compared to the NVIDIA V100 GPUs, on a cluster with 512 of them, ZeRO-Infinity only requires a bandwidth of 30 GB/s between each accelerator and the slow memory to remain efficient. In fact, this is already possible with today's technology by connecting accelerators to slow memory via NVLink [43]. For example, the Summit Supercomputer launched in 2018 [44] connects NVIDIA V100 GPUs with the CPU memory at 40GB/s per GPU.
It is clear that with ZeRO-Infinity, accelerator device memory is no longer a limitation on model scale or training efficiency. However, training models with tens or hundreds of trillions of parameters in a reasonable time still requires massive leaps in compute, and running efficiently on these future devices requires a proportional leap in device-to-device bandwidth ( Table 3).
We hope that, with device memory no longer a limitation, ZeRO-Infinity will inspire a more compute and device-device bandwidth focused innovations of ultra-powerful accelerator devices and supercomputing clusters in the future to support the next 1000x growth in model scale and the advancements that they can offer.         Figure 6(e)