ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Samyam Rajbhandari∗, Jeff Rasley111Equal Contributors, Olatunji Ruwase, Yuxiong He {samyamr, jerasley, olruwase,
Abstract
Large deep learning models offer significant accuracy gains, but training billions to trillions of parameters is challenging. Existing solutions such as data and model parallelisms exhibit fundamental limitations to fit these models into limited device memory, while obtaining computation, communication and development efficiency. We develop a novel solution, Zero Redundancy Optimizer (ZeRO), to optimize memory, vastly improving training speed while increasing the model size that can be efficiently trained. ZeRO eliminates memory redundancies in data- and model-parallel training while retaining low communication volume and high computational granularity, allowing us to scale the model size proportional to the number of devices with sustained high efficiency. Our analysis on memory requirements and communication volume demonstrates: ZeRO has the potential to scale beyond 1 Trillion parameters using today’s hardware.
中文速览
训练拥有数百亿乃至万亿参数的超大深度学习模型时，单张GPU内存根本装不下，而现有的数据并行和模型并行方案要么内存浪费严重，要么通信开销巨大、难以扩展。微软研究团队提出了零冗余优化器（ZeRO，Zero Redundancy Optimizer），其核心思路是把原本在每张GPU上完整复制的优化器状态、梯度和模型参数分别切片分配到不同设备，让每张卡只存自己负责的那份，同时通过精心设计的通信调度保证整体通信量与传统数据并行相当。实验结果显示，ZeRO在400张V100 GPU上训练超过1000亿参数的模型时达到了15 Petaflops的聚合吞吐量，相比当时最优方案模型规模扩大8倍、训练速度提升10倍，还能让数据科学家无需改写模型代码就直接训练130亿参数的大模型。这项工作不仅将可训练模型的规模上限从百亿级推向万亿级，更催生了当时全球最大语言模型Turing-NLG（170亿参数），并以开源库DeepSpeed的形式向整个深度学习社区开放，对大模型训练的普及具有深远意义。
原文 arXiv:1910.02054;中英对照 + 大白话阅读 https://aha.fim.ai/paper/1910.02054v3