Nvidia挖矿专用芯片

现在的模型动辄数百、数千亿参数,普通人训不动怎么办?会昌网景

前不久,谷歌发布了参数量为 1.6 万亿的语言模型Swith Transformer,将 GPT-3 创下的参数量记录(1750 亿)推至新高。这些大模型的出现让普通研究者越发绝望:没有「钞能力」、没有一大堆 GPU 就做不了 AI 研究了吗?

在此背景下,部分研究者开始思考:如何让这些大模型的训练变得更加接地气?也就是说,怎么用更少的卡训练更大的模型?

为了解决这个问题,来自微软、加州大学默塞德分校的研究者提出了一种名为「 ZeRO-Offload 」的异构深度学习训练技术,可以在单个 GPU 上训练拥有 130 亿参数的深度学习模型,让普通研究者也能着手大模型的训练。与 Pytorch 等流行框架相比,ZeRO-Offload 将可训练的模型规模提升了 10 倍,而且不需要数据科学家对模型做出任何改变,也不会牺牲计算效率。

GPU机器学习

论文链接:https://arxiv.org/pdf/2101.06840.pdf

ZeRO-Offload 通过将数据和计算卸载(offload)至 CPU 来实现大规模模型训练。为了不降低计算效率,它被设计为最小化与 GPU 之间的数据往来,并在尽可能节省 GPU 内存的同时降低 CPU 的计算时间。因此,对于一个参数量为 100 亿的模型,ZeRO-Offload 可以在单个 NVIDIA V100 GPU 上实现 40 TFlops/GPU。相比之下,使用 PyTorch 训练一个参数量为 14 亿的模型仅能达到 30TFlops,这是在不耗尽内存的情况下所能训练的最大模型。ZeRO-Offload 还可以扩展至多 GPU 设置并实现线性加速,最多可在 128 个 GPU 上实现近似线性加速。

此外,ZeRO-Offload 还可以和模型并行一起使用,在一个 DGX-2 box AI 服务器上训练参数量超 700 亿的模型。与单独使用模型并行相比,这一参数量实现了 4.5 倍的规模提升。

在下文中,我们将结合 Medium 博主 LORENZ KUHN 的一篇博客来详细了解这篇论文。

ZeRO-Offload 是什么?

ZeRO-Offload 是一种通过将数据和计算从 GPU 卸载到 CPU,以此减少神经网络训练期间 GPU 内存占用的方法,该方法提供了更高的训练吞吐量,并避免了移动数据和在 CPU 上执行计算导致的减速问题。

借助 ZeRO-offload,使用相同的硬件能训练以往 10 倍大的模型,即使在单个 GPU 上也是如此。比如在一个 32GB RAM 的 V100 GPU 上训练百亿参数的 GPT-2。

此外,ZeRO-offload 还能实现在多 GPU 设置中的近似线性扩展。

对于研究者来说,ZeRO-offload 适用的情况包括:

想训练更大的模型,或者想更快地训练现在的模型,因为 ZeRO-offload 允许训练更大的 batch size;

你正在使用 PyTorch,并且愿意 / 能够使用微软的 DeepSpeed 库(ZeRO-offload 的其他实现形式暂未推出),你也可以尝试根据官方实现自行调整;

愿意接受一些建模时的限制,比如当前版本的 ZeRO-Offload 需要搭配使用 Adam 的混合精度训练。会昌网景

如何使用?

ZeRO-Offload 在微软的 DeepSpeed 库中实现,官方实现地址:https://github.com/microsoft/DeepSpeed/blob/6e65c2cc084ecfc393c67a2f64639e8d08d325f6/deepspeed/runtime/zero/stage2.py。

在 DeepSpeed 中设置完毕后,使用 ZeRO-Offload 就不需要太多额外的工作了,只需要修改一些标志和配置文件。

目前,Hugging Face 的 transformers 库与 DeepSpeed 进行了实验性集成,使用方法和基准测试结果参见:https://huggingface.co/blog/zero-deepspeed-fairscale。

GPU机器学习

Facebook 研究院的 fairscale 有 ZeRO 的部分实现,ZeRO-Offload 正是基于 ZeRO 这一多 GPU 内存优化方法构建的。目前还不支持 CPU 卸载。会昌网景