多GPU通信效率提升4倍，RLHF生成提升2.25倍！DeepSpeed ZeRO++重磅升级
新智元报道
新智元报道
【新智元导读】多GPU集群训练大模型，ZeRO++通信量仅为ZeRO的四分之一。
过去半年，由ChatGPT引领的生成式大型语言模型技术，以其强大的「通用性」彻底颠覆了AI世界，普通人也可以很容易地使用AI工具来进行摘要、灵感创作、辅助编程、多语言翻译等任务。
不过，训练这种超大规模的模型往往需要数百个、甚至数千个GPU来存储和计算数据，比如训练5300亿参数的Megatron-Turning NLG就使用了超过4000块Nvidia A100 GPU
想要高效地利用硬件资源需要设计复杂的优化系统，将模型划分割成适合于单个设备内存的片段，然后跨设备进行高效的并行计算；同时，为了能够让深度学习社区更方便地对大型模型进行训练，这些优化操作必须易于使用。
2020年2月，微软开源了深度学习训练优化库DeepSpeed，并提供了内存优化技术ZeRO（零冗余优化器，Zero Redundancy Optimizer），提高了可训练模型的规模、训练速度、降低了成本、提升了可用性，极大地推动了大模型的训练，已应用于TNLG-17 B、Bloom-176 B、MPT-7 B、Jurrasic-1等模型的训练上。
仓库链接：https://github.com/microsoft/DeepSpeed
但在某些场景下，ZeRO系列技术会在GPU之间产生更高的数据传输开销，降低模型的训练效率，尤其是：
1. 总的GPU数量较多，导致分配到每个GPU上batch size较少，硬件间通信更频繁；
2. 在低端（low-end）集群上进行训练时，跨节点的网络带宽有限，导致GPU间的通信延迟较高。
为了解决这些问题，微软DeepSpeed团队发布了新一代基于ZeRO的通信优化策略系统ZeRO++，可以为大型模型训练提供更高的效率，不受batch size限制或跨设备带宽限制。
ZeRO++利用量化结合数据和通信重映射，将总通信量减少4倍，并且不会影响模型质量。
在数千个GPU上进行训练，或是在几百、几十个GPU上微调时，在每个GPU分配到的batch size都比较小的情况下，ZeRO++的吞吐量比ZeRO高出2.2倍，可以直接降低训练时间和成本。
在低带宽集群下，ZeRO++可以使其实现与4倍带宽集群相似的吞吐量，可以在更广泛的集群中实现高效的大型模型训练。
ZeRO++还可以加速使用RLHF技术的ChatGPT类模型的训练。
虽然ZeRO++主要是为训练而设计的，但通讯开销在训练和推理过程中都存在，其优化也适用于ZeRO-Inference，可以提高训练对话模型中使用的人工反馈强化学习（RLHF）等工作负载的效率。
与原始ZeRO相比，将ZeRO++与DeepSpeed-Chat集成后，可以将RLHF训练的生成速度提高2倍，强化学习训练阶段提高1.3倍。
ZeRO++
ZeRO++
ZeRO是数据并行的一种内存高效变体（memory efficient variation），其中模型states在所有GPU上进行分区而非复制，并在训练期间使用基于收集/广播（gather/broadcast）的通信集体进行重建，也使得ZeRO能够高效地利用聚合GPU内存和跨所有设备的计算，同时提供简单易用的数据并行训练。
假设模型大小为M，则在前向传递过程中，ZeRO执行所有收集/广播操作，以便在有需要之前收集完每个模型层的参数（通信量为M）；
在反向传递中，ZeRO对每一层的参数采用类似的通信模式以计算局部梯度（通信量为M）；
此外，ZeRO在使用reduce或reduce-scatter通信集合计算每个局部梯度之后立即对每个局部梯度进行平均和分区（通信量为M）。
所以ZeRO总共需要3M的通信量，均匀分布在两次全收集/广播和一次reduce/reduce-scatter操作中。
为了减少这些通信开销，ZeRO++设计了三组通信优化，分别针对上述三个通信集合。
ZeRO (qwZ)的量化权重通信
为了减少全收集（all-gather）过程中的参数通信量，研究人员选择采用权值量化，在通信之前将每个模型参数从FP16（两个字节）压缩到INT8（一个字节）数据类型，并在通信之后反量化权值。
不过只是简单地对权重进行量化可能会降低模型训练精度，为了在提速的同时保持训练精度，研究人员采用了基于块（block-based）的量化，即对模型参数的每个子集进行独立的量化。
目前尚不存在用于高性能、基于块量化的实现，所以研究人员从头开始实现高度优化的量化CUDA内核；与基本量化相比，准确性提高了3倍，速度提高了5倍。
ZeRO（hpZ）的分层权重划分
为了减少反向传递期间权重上的所有gather操作的通信开销，研究人员选择牺牲GPU内存换取通信效率。
具体来说，与在ZeRO中将整个模型权重分散在所有机器上不同，ZeRO++在每台机器内维护一个完整的模型副本，以更高的内存开销为代价，使得可以用机器内全收集/广播操作来替换昂贵的跨机器全收集/广播权重。
因为机器内的通信带宽更高，所以实际性能会得到提升。
用于ZeRO（qgZ）的量化梯度通信
使用reduce-scatter来降低梯度的通信成本更具挑战性，因为直接应用量化来减少通信量是不可行的，即使将基于块的权重量化为低精度，梯度降低也会累积并放大量化误差。
为了解决这个问题，可以在通信之前只量化梯度，但在任何reduce操作之前将梯度反量化到全精度。
为了高效地实现，研究人员发明了一个基于all-to-all，新的量化梯度通信范式qgZ，功能上等同于compressed reduce-scatter collective操作。
qgZ旨在解决两个问题：
1. 克服在INT 4/INT 8低精度中实现reduce-scatter导致的准确度下降；
2. 避免由长序列的量化和反量化（dequantization）步骤中导致的准确率下降和等待时间开销，这两个步骤是基于ring或tree的reduce scatter传统方法所必需的。
qgZ是一种全新的基于层次化all-to-all的方法，包括三个主要步骤：
1. 梯度切片重新排序
2. 节点内通信和reduction
3. 节点间通信和reduction
首先，在任何通信操作执行之前，对梯度进行切片并进行张量切片重新排序以保证最终的梯度放置，在通信结束时在每个GPU上是正确的。
然后量化重排序后的梯度切片，在每个节点内进行all-to-all通信，反量化接收到的梯度切片，并进行局部reductions
再次量化局部reduction的梯度，进行节点间all-to-all通信，反量化接收到的梯度，并计算最终的高精度梯度降低，得到最终结果。
使用层次化方法的原因是为了减少跨节点通信量，更确切地说，如果每个节点有N个GPU、模型大小为M和量化比为Z，单跳all-to-all会生成M*N/Z个跨节点流量。
相比之下，分层方法将每个GPU的跨节点流量从M/Z减少到M/（Z*N），总通信量可以降到M*N/（Z*N）= M/Z
通过重叠节点内和节点间通信以及融合CUDA内核（张量切片重新排序+节点内量化）和（节点内去量化+节点内归约+节点间量化）来进一步优化qgZ的端到端延迟。
总结
通过整合上述三个组件，ZeRO++将跨节点通信量从3M减少到0.75M。
使用qwZ将模型权重从M减少到0.5M；使用hpZ消除了反向传播过程中的跨节点all-gather，将通信从M减少到0；使用qgZ将反向传递期间的跨节点reduce-scatter通信从M降低到0.25M
ZeRO++加速LLM训练
ZeRO++加速LLM训练
在实验部分，研究人员展示了在384个Nvidia V100 GPU上使用真实LLM训练场景对ZeRO++的评估结果。
在高带宽集群中测试ZeRO++吞吐量的改进，对于不同的模型尺寸和micro batch size，使用4x Infiniband（IB）的400 Gbps跨节点互连，每个都以100 Gbps运行。
每个GPU使用1000 token，ZeRO++比ZeRO-3的吞吐量提高了28%到36%。
对于2000 micro batch size，ZeRO++比ZeRO-3实现了24%至29%的吞吐量增益。
在低带宽集群中，在100 Gbps网络等低网络环境中，ZeRO++的性能明显优于ZeRO-3
与ZeRO-3相比，ZeRO++在端到端吞吐量方面实现了高达2.2倍的加速，平均而言，ZeRO++比ZeRO-3基线实现了约2倍的加速。
ZeRO++结合DeepSpeed-Chat进行RLHF训练
ZeRO++结合DeepSpeed-Chat进行RLHF训练
RLHF训练背景
ChatGPT类模型的基础是大型语言模型，然后使用RLHF进行微调。
RLHF由生成（推理）阶段和训练阶段组成：在生成阶段，actor模型将部分对话作为输入，并使用一系列前向传递来生成回复。
在训练阶段，critic模型按质量对生成的回复进行排序，为actor模型提供强化信号；使用排名对actor模型进行微调，使其能够在后续迭代中生成更准确和适当的回复。
RLHF使用了四个模型（actor, reference, critic, reward），所以需要大量的内存，采用低秩自适应（LoRA）可以解决RLHF的内存压力。
LoRA冻结了预训练的模型权重，并将可训练的秩分解矩阵注入到Transformer架构的每一层中，从而显著减少了可训练参数的数量。
LoRA可以通过减少内存使用量、允许更大的batch size来加快RLHF的训练速度，从而大大提高吞吐量。
DeepSpeed-Chat结合ZeRO++
用LoRA进行RLHF训练是因为大多数模型权重都是冻结的，也意味着ZeRO++可以将这些冻结的权重保持在INT 4/8中量化，无需存储为FP 16并在每次通信操作之前重复量化。
以这种方式使用ZeRO++进行RLHF训练可以减少内存使用和通信量，通过减少通信以及由于减少内存使用而使用更大的batch size大小来提高训练吞吐量。
在生成阶段，ZeRO++使用hpZ保持每个节点内的所有权重通信，以利用更高的节点内通信带宽，减少通信量，进一步提高生成吞吐量。
研究人员对比了不同尺寸的actor模型的RLHF生成吞吐量，并比较了32个V100 GPU上的30 B和66 B actor模型的ZeRO和ZeRO++，结果显示，ZeRO++能够实现比ZeRO高2.25倍的RLHF生成吞吐量。
在16个V100 GPU上的训练阶段加速对比中，ZeRO++实现了比ZeRO高1.26倍的吞吐量，主要是由于ZeRO++实现了更低的通信和更大的batch size
微信扫码关注该文公众号作者