现在的模型动辄上千亿参数，普通人训练不出来怎么办？不久前，谷歌发布了参数量达1.6万亿的语言模型SwitchTransformer，将GPT-3创下的参数记录（1750亿）推向新高。这些大模型的出现，让普通研究人员更加绝望：没有“能力”，没有大量GPU，是不是做AI研究就不行了？在此背景下，一些研究者开始思考：如何让这些大模型的训练更加接地气？即如何用更少的卡片训练更大的模型？为了解决这个问题，来自微软和加州大学默塞德分校的研究人员提出了一种名为“ZeRO-Offload”的异构深度学习训练技术，可以在单个GPU上训练出130亿参数的深度学习模型，让普通的研究人员也可以开始训练大型模型。与Pytorch等流行框架相比，ZeRO-Offload将可训练模型的大小增加了10倍，而无需数据科学家对模型进行任何更改，并且不会牺牲计算效率。论文链接：https://arxiv.org/pdf/2101.06840.pdfZeRO-Offload通过将数据和计算卸载到CPU来实现大规模模型训练。为了不降低计算效率，设计上尽量减少与GPU的数据交换，减少CPU计算时间的同时尽可能节省GPU显存。因此，对于具有100亿个参数的模型，ZeRO-Offload可以在单个NVIDIAV100GPU上实现40TFlops/GPU。相比之下，使用PyTorch训练14亿参数的模型只能达到30TFlops，这是在不耗尽内存的情况下可以训练的最大模型。ZeRO-Offload还可以扩展到多GPU设置并实现线性加速，在多达128个GPU上实现近线性加速。此外，ZeRO-Offload还可以与模型并行使用，在DGX-2盒子AI服务器上训练出超过700亿参数的模型。与单独使用模型并行性相比，这一数量的参数实现了4.5倍的扩展。下面，我们将结合Medium博主LORENZKUHN的一篇博客，进一步了解这篇论文。什么是零卸载？ZeRO-Offload是一种在神经网络训练期间通过将数据和计算从GPU卸载到CPU来减少GPU内存占用的方法，它提供更高的训练吞吐量并避免移动数据和执行计算导致的减速问题。借助ZeRO-offload，即使在单个GPU上，也可以使用相同的硬件训练10倍大的模型。例如，在具有32GBRAM的V100GPU上训练具有数百亿参数的GPT-2。此外，ZeRO-offload可在多GPU设置中实现近似线性缩放。对于研究人员来说，ZeRO-offload适用于：你想训练更大的模型，或者你想更快地训练当前模型，因为ZeRO-offload允许以更大的batchsize进行训练；你正在使用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不需要太多额外的工作，只需修改几个标志和配置文件。目前，HuggingFace的变形金刚库正在实验性地与DeepSpeed集成。使用方法和基准测试结果参见：https://huggingface.co/blog/zero-deepspeed-fairscale。Facebook研究院的fairscale有ZeRO的部分实现，ZeRO-Offload是建立在ZeRO基础上的一种多GPU内存优化方法。当前不支持CPU卸载。ZeRO-Offload的工作原理ZeRO-Offload建立在零冗余优化器(ZeRO)之上。ZeRO是微软于2020年2月提出的一种万亿级模型参数训练方法，用于数据并行和模型并行训练中的内存优化，其中梯度、参数和优化器状态分布在多GPU内存中，没有任何冗余。这使GPU之间的通信开销保持在较低水平。与标准数据并行基准测试相比，三个阶段的内存和通信使用量为零。回顾一下ZeRO：为了解决数据并行和模型并行的问题，ZeRO提供了三阶段的优化方法，即优化器状态分割、梯度分割和参数分割，三个阶段依次实现。在优化器分裂状态下：ZeRO减少了3/4的内存，通信量与数据并行度相同；加入梯度分割：减少7/8内存，通信量与数据并行度相同；添加参数分段：内存减少与数据并行度呈线性关系。例如，在64个GPU上拆分时，内存可以减少到1/64。流量增加了50%。在去年9月的博客中，微软是这样介绍ZeRO-Offload的：ZeRO-Offload继承了ZeRO-2的优化器状态和梯度分割。但与ZeRO-2不同的是，ZeRO-Offload不会在每个GPU上保留优化器状态和梯度拆分，而是将两者卸载到主机CPU内存。在整个训练阶段，优化器状态都保存在CPU内存中；在反向传播过程中使用reduce-scatter计算梯度并在GPU上平均，然后每个数据并行线程将属于其split的梯度平均值卸载到CPU内存中（见下图中的goffload），剩下的就是丢弃。一旦梯度到达CPU，每个数据并行线程直接在CPU上并行更新优化器状态分区（参见下图中的pupdate）。之后，将参数拆分移回GPU，在GPU上执行all-gather操作以收集所有更新的参数（参见下面的gswap）。ZeRO-Offload还利用单独的CUDA流来耗尽通信和计算中的重叠，从而最大限度地提高训练效率。零卸载概述。值得注意的是，ZeRO-Offload是为与Adam进行混合精度训练而设计的。也就是说，当前版本的ZeRO-Offload使用的是Adam的优化版本DeepCPUAdam。这样做的主要原因是为了避免CPU计算成为整个过程的瓶颈。DeepCPUAdam比AdamPyTorch实现快6倍。实验结果最后，让我们看一下ZeRO-Offload论文中提供的一些实验结果。下面的图7显示了可以使用ZeRO-Offload技术在1、4或16个GPU（一个DGX-2）上训练的最大模型。下面的图11显示了随着GPU数量的增加，每个GPU的吞吐量。可以看出，随着GPU数量逐渐增加到128个，ZeRO-Offload可以实现几乎线性的吞吐量加速。下图8显示了PyTorch、L2L和ZeRO-Offload实现的每GPU吞吐量差异。可以看出，使用ZeRO-Offload实现的每GPU吞吐量比L2L平均高14%（最多高22%）。
