ZeRO-Offload：让十亿规模模型训练平民化
分布式、并行与集群计算
2021-01-19 v1 机器学习
摘要
大规模模型训练一直是少数人的竞技场，需要复杂的模型重构以及极其昂贵的GPU集群。ZeRO-Offload改变了大模型训练的格局，使大模型训练几乎对所有人开放。它能在单个GPU上训练超过130亿参数的模型，相比PyTorch等流行框架规模提升10倍，且无需数据科学家对模型做任何改动，也不牺牲计算效率。ZeRO-Offload通过将数据与计算卸载到CPU来实现大模型训练。为保持计算效率，其设计最小化了与GPU间的数据搬运，在最大化GPU内存节省的同时减少CPU计算时间。因此，对于100亿参数模型，ZeRO-Offload在单个NVIDIA V100 GPU上可达到40 TFlops/GPU，而PyTorch单独训练14亿参数（不溢出内存可训练的最大模型）仅达30 TFlops。ZeRO-Offload还设计为在有多GPU时可扩展，在最多128个GPU上提供近线性加速。此外，它可与模型并行协同，在单个DGX-2机上训练超过700亿参数的模型，较仅用模型并行模型规模提升4.5倍。通过结合计算与内存效率及易用性，ZeRO-Offload使大规模模型训练平民化，即便仅有一个GPU的数据科学家也能使用。
引用
@article{arxiv.2101.06840,
title = {ZeRO-Offload: Democratizing Billion-Scale Model Training},
author = {Jie Ren and Samyam Rajbhandari and Reza Yazdani Aminabadi and Olatunji Ruwase and Shuangyan Yang and Minjia Zhang and Dong Li and Yuxiong He},
journal= {arXiv preprint arXiv:2101.06840},
year = {2021}
}