[论文解读] ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Samyam Rajbhandari, Olatunji Ruwase+3 more|arXiv (Cornell University)|Apr 16, 2021

Advanced Neural Network Applications参考文献 30被引用 57

一句话总结

ZeRO-Infinity 将模型状态和激活转移到 CPU/NVMe，使用以内存为中心的分块来在有限的 GPU 资源上训练万亿参数模型，实现高吞吐量和可访问性。

ABSTRACT

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.

研究动机与目标

• 解释极端规模 DL 的 GPU 内存墙挑战，并激发一个在不重构的前提下增加可达模型大小的系统动机。
• 介绍 ZeRO-Infinity 的异构内存卸载和分块技术，以让大规模模型适配现有集群。
• 展示在超大规模模型上的训练效率和可扩展性，包括万亿参数级别。
• 评估简化可用性的改进，降低或消除对模型重构或复杂并行性的需求。
• 提供开源实现细节并讨论对未来硬件设计的影响。

提出的方法

• 提出 infinity 卸载引擎，以同时利用 GPU、CPU 和 NVMe 内存。
• 引入以内存为中心的分块，将大算子分解为小块以降低工作内存。
• 定义带宽为中心的分区，以利用跨设备的总内存带宽。
• 开发以重叠为中心的设计，以在异构内存中隐藏通信与计算。
• 提供简易化实现，训练任意架构时避免模型重构。
• 描述基于 Transformer 的大模型训练的内存与带宽需求。

实验结果

研究问题

• RQ1 如何在实际的 GPU 集群上突破 GPU 内存墙，以训练具有数百万亿参数的模型？
• RQ2 是否可以通过利用 CPU 和 NVMe 内存，使缺乏大规模 GPU 集群的研究人员也能访问大型模型？
• RQ3 在不重构模型代码或采用复杂的 3D 并行的情况下，是否有可能训练万亿参数模型？
• RQ4 在异构内存下，维持高效训练所需的内存与带宽特性是什么？
• RQ5 在当前硬件上，ZeRO-Infinity 在吞吐量、可扩展性和可访问性方面的表现如何？

主要发现

• ZeRO-Infinity 可以在 32 NVIDIA DGX-2 节点（512 GPUs）上训练 32 万亿参数的模型。
• 在同一硬件上维持超过 25 petaflops 的吞吐量（峰值的 40%）。
• 它在万亿参数模型上展示了超线性可扩展性。
• 它使在单个 DGX-2 节点上对万亿参数进行微调成为可能，无需模型并行或重构。
• 以内存为中心的分块减少每个算子的工作内存，使极大层在不需要模型并行的情况下也能容纳。
• ZeRO-Infinity 的开源实现可在 DeepSpeed 中获取。