ZeRO: Memory Optimization Towards Training A Trillion Parameter Models

By Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He. 2019

616 citations recorded by the source corpus.

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, background): “Many other works aim to increase of the scale of models, while limiting communication overheads (Rajbhandari et al., 2020; Lepikhin et al., 2020; Li et al., 2020; Rasley et al., 2020; Rajbhandari et al., 2021; Ren et al., 2021; Narayanan et al., 2021a).”
Unified Scaling Laws for Routed Language Models (2022, methodology): “All were optimized with AdamW [Loshchilov and Hutter, 2018] and ZeRO Stage 1 was used to shard the optimizer state [Rajbhandari et al., 2020].”
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model (2022, background): “The Zero Redundancy Optimizer (ZeRO) [55] is a collection of optimizations that improve the memory efﬁciency of data parallelism by partitioning the replicated data among data-parallel workers.”
Efficient Large Scale Language Modeling with Mixtures of Experts (2021, methodology): “Recent work has explored sharding model parameters, optimizer state and gradients across workers (Xu et al., 2020; Rajbhandari et al., 2020), enabling training of models with more than one trillion parameters using only data parallelism without the added complexity introduced by model parallel…”
Improving language models by retrieving from trillions of tokens (2021, methodology): “All models use ZeRO to shard the optimiser state (Rajbhandari et al., 2020).”
Scaling Up Models and Data with t5x and seqio (2022, methodology): “Some major diﬀerentiators of t5x are its use of JAX and Flax for model expression, its support for TPU (including TPU v4), and its Gin-based conﬁguration system that allows uses to modify nearly everything about the model and training procedure. t5x also doesn’t support pipeline parallelism, a major”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “We present GSPMD, a system that uses simple tensor sharding annotations to achieve different parallelism paradigms in a unified way, including data parallelism, in-layer model parallelism, and novel strategies like image spatial partitioning [11] and weight-update/optimizer-state sharding [30, 40].”
PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation (2021, neutral): “The third regime is optimizer model parallelism [29] (Figure 4(d)), which aims to reduce the redundant optimizer memory and computation consumption resulted from data parallelism.”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, methodology): “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 ”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “We note that we have only considered ZeRO-3 without tensor parallelism.”
ZeRO-Offload: Democratizing Billion-Scale Model Training (2021, methodology): “To be more specific, distributed parallel DL training technologies such as pipeline parallelism [10], model parallelism [28], and ZeRO [21] (Zero Redundancy Optimizer) allow transcending the memory boundaries of single GPU/accelerator device by splitting the model states (parameters, gradients and o”
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020, background): “Zero (Rajbhandari et al., 2019) presents a set of optimizations to reduce memory redundancy in parallel training devices, by partitioning weights, activations, and optimizer state separately, and it is able to scale models to 170 billion parameters; in comparison, GShard is more general in the…”
Memory-Efficient Pipeline-Parallel DNN Training (2020, methodology): “DeepSpeed (Rajbhandari et al., 2019) partitions optimizer state over data-parallel replicas instead of replicating it, us-ing a technique called ZeRO.”
GPT-NeoX-20B: An Open-Source Autoregressive Language Model (2022, methodology): “We extend AdamW with the ZeRO optimizer (Rajbhandari et al., 2020) to reduce memory consumption by distributing optimizer states across ranks.”
Sustainable AI: Environmental Implications, Challenges and Opportunities (2021, background): “In addition, significant research has gone into algorithmic approaches to efficiently scale training [59], [60] by reducing communication cost via compression [61], [62], pipelining [63], and sharding [64], [65].”