微软研究人员推出了一个名为ZeRO++的新系统,该系统已经被开发用于优化大型AI模型的训练,解决了高数据传输开销和带宽有限的挑战。ZeRO++在现有的ZeRO优化基础上进行了扩展,提供了增强的通信策略,以提高训练效率,缩短训练时间和降低训练成本。
如Turing-NLG、ChatGPT和GPT-4等大型模型的训练需要跨多个GPU设备使用大量存储器和计算资源。DeepSpeed开发的ZeRO++引入了通信优化策略,以克服在每个GPU上使用小批量大小或在低带宽集群上训练时ZeRO的局限性。
ZeRO优化的家族,包括ZeRO-Inference,使模型状态在GPU之间进行分区,而不是复制,利用集体GPU内存和计算能力。然而,在训练过程中,ZeRO可能会产生高通信开销。ZeRO++通过融合三组通信优化解决了这个问题:量化权重通信(qwZ)、分层权重分区(hpZ)和量化梯度通信(qgZ)。
为了减少参数通信量,ZeRO++对权重进行了量化,利用基于块的量化来保留训练精度。这个优化的量化过程比基本的量化更快,更准确。为了在反向传播期间最小化通信开销,ZeRO++通过在每台机器内维护一个完整的模型副本来以GPU内存为代价进行通信。对于梯度通信,ZeRO++引入了一种新颖的量化梯度通信范式qgZ,以减少跨节点的流量和延迟。
这些通信优化导致通信量的大幅度减少。与ZeRO相比,ZeRO++实现了高达4倍的减少,提高了训练吞吐量和效率。当每个GPU使用小批量大小时,在高带宽集群中,ZeRO++比ZeRO-3提高了28%至36%的吞吐量。在低带宽集群中,ZeRO++与ZeRO-3相比实现了平均2倍的加速,使大型模型训练更加普及。
ZeRO++不仅限于训练场景,还扩展到使用人类反馈(RLHF)训练的对话模型中的强化学习。通过将ZeRO++与DeepSpeed-Chat集成,可以使RLHF训练受益于改进的生成和训练阶段,实现比ZeRO更高达2.25倍的更好的生成吞吐量和1.26倍的更好的训练吞吐量。
DeepSpeed发布了ZeRO++,以使大型模型训练更加高效和适用于AI社区。该系统旨在加速训练,减少通信开销并实现更大的批量大小,从而节省时间和资源。研究人员和实践者可以利用ZeRO++更有效地训练像ChatGPT这样的模型,并探索AI中的新可能性。