ZeRO:面向万亿参数模型训练的内存优化
机器学习
2020-05-14 v3 分布式、并行与集群计算
机器学习
摘要
大型深度学习模型带来了显著的精度提升,但训练数十亿至数万亿参数颇具挑战。现有的数据与模型并行等方案存在根本局限,无法将这些模型放入有限的设备内存,同时兼顾计算、通信与开发效率。我们提出了一种新颖的方案——零冗余优化器(Zero Redundancy Optimizer, ZeRO),以优化内存,在大幅提升训练速度的同时增加可高效训练的模型规模。ZeRO消除了数据与模型并行训练中的内存冗余,同时保持低通信量与高计算粒度,使我们能够按设备数量比例扩展模型规模并维持高效率。我们对内存需求与通信量的分析表明:ZeRO有潜力在现今硬件上扩展至超过1万亿参数。我们实现并评估了ZeRO:它在400个GPU上以超线性加速训练超过100B参数的大型模型,达到15 Petaflops的吞吐量。这相比最先进水平实现了模型规模8倍增长与可达成性能10倍提升。在易用性方面,ZeRO可训练高达13B参数的大型模型(例如大于Megatron GPT 8.3B与T5 11B),且无需科学家较难应用的模型并行。最后但并非最不重要,研究人员已利用ZeRO的系统突破创建了全球最大语言模型(Turing-NLG, 17B参数),取得了破纪录的精度。
引用
@article{arxiv.1910.02054,
title = {ZeRO: Memory Optimizations Toward Training Trillion Parameter Models},
author = {Samyam Rajbhandari and Jeff Rasley and Olatunji Ruwase and Yuxiong He},
journal= {arXiv preprint arXiv:1910.02054},
year = {2020}
}