AB
AiBoss站
百科

什么是零冗余优化器(Zero Redundancy Optimizer)?

零冗余优化器(Zero Redundancy Optimizer,ZeRO)是 arXiv 论文《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》提出的一种分布式训练内存优化方案。它通过消除数据并行与模型并行训练中的内存冗余,在保持较低通信量和较高计算粒度的前提下,让可训练模型规模随设备数量成比例扩展。

零冗余优化器(Zero Redundancy Optimizer,ZeRO)是一种面向大规模深度学习训练的分布式内存优化方案。按提出该方法的论文所述,它的核心目标是消除数据并行(data parallelism)与模型并行(model parallelism)训练中存在的内存冗余,从而在有限的设备内存中容纳更大的模型,同时保持较低的通信量与较高的计算粒度。它要解决的问题是:当模型参数从数十亿增长到数万亿量级时,单张加速卡的内存装不下模型状态,而传统并行策略在内存、通信与开发效率之间存在难以调和的取舍。

为什么重要

论文指出,大型深度学习模型能带来显著的精度收益,但训练数十亿到数万亿参数的模型十分困难。已有的数据并行与模型并行方案存在根本性限制:它们难以把这样的模型放进有限的设备内存,同时又难以兼顾计算效率、通信效率和开发效率。

具体来说,数据并行会在每张设备上复制一份完整的模型状态,设备越多,重复存储的部分就越多;模型并行则需要把模型切分到不同设备上,切分方式与模型结构强耦合,论文认为这类方法对科研人员而言更难应用。两种路线都让「模型能有多大」受制于单设备内存,而不是受制于集群整体的算力与显存总量。ZeRO 的出发点正是把这部分被重复占用的内存释放出来,用于存放更大的模型。

工作机制

按论文的表述,ZeRO 的做法可以概括为以下几点:

  • 消除内存冗余:识别数据并行与模型并行训练中重复保存的模型状态,把这些冗余部分去掉,而不是简单地增加设备或缩小模型。
  • 保留低通信量与高计算粒度:在削减内存占用的同时,论文强调方案维持了较低的通信数据量和较高的计算粒度,避免为了省内存而牺牲扩展效率。
  • 规模随设备数成比例扩展:论文称,借助上述设计,可高效训练的模型规模能够与设备数量成比例地增长,并保持较高的效率。
  • 以内存与通信分析为依据:论文对内存需求和通信量做了分析,并据此论证该方案在当时的硬件条件下有扩展到一万亿(1 Trillion)参数以上的潜力。

需要说明的是,上述机制描述来自论文自身的表述;论文并未在摘要中给出逐阶段的分片细节,因此这里不展开论文未明确写出的实现步骤。

典型例子

论文报告了若干实现与评估结果,这些数字均出自该论文,属于作者自述的实验结论:

  • 在 400 张 GPU 上训练超过 1000 亿(100B)参数的模型,取得了超线性加速(super-linear speedup),吞吐达到 15 Petaflops。
  • 论文称,相对于当时的最先进水平,这相当于模型规模提升 8 倍、可达性能提升 10 倍。
  • 在可用性方面,ZeRO 无需模型并行即可训练最多 130 亿(13B)参数的模型,论文举例称这一规模大于 Megatron GPT 的 8.3B 与 T5 的 11B。
  • 论文提到,研究人员利用 ZeRO 的系统突破构建了当时世界上最大的语言模型 Turing-NLG(170 亿参数),并取得创纪录的准确率。

这些结果说明 ZeRO 在论文的实验设定下同时覆盖了「极大模型 + 多设备」和「中等模型 + 免模型并行」两类场景,但它们是对特定硬件与模型配置的测量,不应外推为任意集群上的通用性能承诺。

边界与常见误解

首先,ZeRO 常被误解为一种新的并行范式。按论文的定位,它是在数据并行与模型并行训练的基础上消除内存冗余的优化方案,而不是取代并行本身;论文强调的是在保留低通信量与高计算粒度的前提下提升可训练规模。

其次,论文中「有潜力扩展到一万亿参数以上」是基于内存需求与通信量分析得出的判断,而非已经完成的训练结果;实际可达到的规模取决于当时的硬件条件。论文实际报告的最大训练规模是 100B 量级,以及无需模型并行的 13B 量级,两者不应混为一谈。

第三,消除冗余并非没有代价。论文在摘要中把低通信量作为设计目标之一,说明内存节省与通信开销之间存在需要平衡的关系;具体开销随并行配置和集群互连条件变化,论文摘要未给出统一的代价结论。

最后,论文报告的加速比、吞吐与模型规模数字都来自作者自己的实验环境,属于该论文的主张,不能视为行业共识或对其他系统的性能保证。

参考资料