[论文解读] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Samyam Rajbhandari, Jeff Rasley+2 more|arXiv (Cornell University)|Oct 4, 2019

Advanced Neural Network Applications参考文献 21被引用 72

一句话总结

ZeRO 引入对优化器状态、梯度和参数的内存高效分区（以及残留内存优化），以通过数据并行实现对万亿参数模型的训练，获得大幅加速并使大模型训练民主化。

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. We implement and evaluate ZeRO: it trains large models of over 100B parameter with super-linear speedup on 400 GPUs, achieving throughput of 15 Petaflops. This represents an 8x increase in model size and 10x increase in achievable performance over state-of-the-art. In terms of usability, ZeRO can train large models of up to 13B parameters (e.g., larger than Megatron GPT 8.3B and T5 11B) without requiring model parallelism which is harder for scientists to apply. Last but not the least, researchers have used the system breakthroughs of ZeRO to create the world's largest language model (Turing-NLG, 17B parameters) with record breaking accuracy.

研究动机与目标

• 激发对训练万亿参数模型的需求，并识别现有数据、模型和流水线并行中的内存瓶颈。
• 提出 ZeRO，在不牺牲效率的前提下消除模型状态和残差中的内存冗余。
• 开发并分析三种分区阶段（优化器状态、梯度、参数）及残留内存优化，以在可用硬件尺度下扩展模型规模。
• 展示实际训练改进，展示与模型并行（MP）的集成，并提供开源工具以民主化大模型训练。

提出的方法

• 将优化器状态在数据并行进程中进行分区（P_os），使每个 rank 持有状态的 1/N_d。
• 将梯度按每个参数分区（P_g）进行分区，并在分区边界执行分桶化的 reduce-scatter。
• 将参数在数据并行进程之间分区（P_p），对非本地参数进行按需广播以用于前向/反向传播。
• 将 ZeRO-DP 与 ZeRO-R 结合，以对激活进行分区并管理残留内存（检查点、缓冲区和碎片化），包括可选的 CPU 卸载（P_a、P_a+cpu）。
• 使用动态通信计划，在降低通信的同时实现内存收益与计算解耦；分析通信开销（在 P_p 下最多 1.5 倍）。
• 在有益时支持与模型并行（MP）结合，但展示在许多情况下 ZeRO-DP 可以匹配或超过 MP，并实现无需大规模模型重构的训练。

实验结果

研究问题

• RQ1 ZeRO 是否能对模型状态进行分区，在不牺牲数据并行效率的前提下显著降低每设备的内存？
• RQ2 优化器状态、梯度和参数的内存可以随着数据并行度 N_d 的增加节省多少？
• RQ3 如何管理残留内存（激活、缓冲区、碎片化）以进一步支持极大规模模型？
• RQ4 应用 ZeRO-DP 和 ZeRO-R 后，对通信量和可扩展性有何影响？
• RQ5 ZeRO 如何实现对数千亿参数模型的训练，在实际硬件上的实用加速是多少？

主要发现

数据并行度	7.5B 模型 (GB) P_os	7.5B 模型 (GB) P_os+g	7.5B 模型 (GB) P_os+g+p	128B 模型 (GB) P_os	128B 模型 (GB) P_os+g	128B 模型 (GB) P_os+g+p	1T 模型 (GB) P_os	1T 模型 (GB) P_os+g	1T 模型 (GB) P_os+g+p
1	120	120	120	2048	2048	2048	16000	16000	16000
4	52.5	41.3	30	896	704	512	7000	5500	4000
16	35.6	21.6	7.5	608	368	128	4750	2875	1000
64	31.4	16.6	1.88	536	284	32	4187	2218	250
256	30.4	15.4	0.47	518	263	8	4046	2054	62.5
1024	30.1	15.1	0.12	513	257	2	4011	2013	15.6

• ZeRO 使数据并行训练能够拟合远大于传统 DP/MP 限制的模型，内存降低大致与 N_d 成正比地减小参数和梯度的内存，取决于阶段，优化器状态的降低幅度为 4x–8x。
• 通过三个阶段（P_os、P_g、P_p），每个设备的内存占用可以显著减少，在大型 GPU 集群上实现万亿参数模型。
• ZeRO-DP 在 64 路 DP 下可以用更低的内存训练 7.5B 模型，且在 P_os+g+p 下，7.5B 模型的内存降至大约 1.88 GB，从而实现更大规模。
• ZeRO-R 进一步减少残留激活内存，使用恒定大小的缓冲区提高效率，并进行按需内存碎片整理以缓解碎片导致的内存不足。
• 实现结果显示在 400 颗 Nvidia V100 GPU 上训练 100B 参数模型，单颗 GPU 超过 38 TFlops，总性能超过 15 PetaFlops；所称数据表明最大模型规模可达原有的 8 倍，且在当时可达到的性能约提高 10 倍。
• ZeRO 作为 DeepSpeed 库的一部分发布，旨在民主化大模型训练，并通过进一步分区（P_os+g+p）向 1 万亿参数规模扩展。