外贸网站整站程序,网站如何做微信推广方案,智慧团建网页手机版官网,上海做seo引言
今天带来ZeRO: Memory Optimizations Toward Training Trillion Parameter Models的论文笔记。
大型深度模型提供了显著的准确性提升#xff0c;但训练数十亿到数万亿个参数是具有挑战性的。现有的解决方案#xff0c;如数据并行和模型并行#xff0c;存在基本的局限…引言
今天带来ZeRO: Memory Optimizations Toward Training Trillion Parameter Models的论文笔记。
大型深度模型提供了显著的准确性提升但训练数十亿到数万亿个参数是具有挑战性的。现有的解决方案如数据并行和模型并行存在基本的局限性难以将这些模型适配到有限的设备内存中同时又能达到计算、通信和开发效率。
作者开发了一种新颖的解决方案零冗余优化器(Zero Redundancy Optimizer, ZeRO)以优化内存大幅提升训练速度同时增加可以高效训练的模型大小。ZeRO在数据和模型并行训练中消除了内存冗余同时保持低通信量和高计算粒度可以根据设备数量按比例扩展模型大小保持高效的可持续性。
1. 总体介绍
深度学习模型正变得越来越大模型规模的增加带来了显著的准确性提升。在自然语言处理领域Transformer技术使得大模型得以问世。为了使模型规模从数十亿参数继续增长到数万亿参数我们面临着训练它们的挑战——这些模型显然无法适应单一设备的内存例如GPU或TPU单纯增加设备数量也无法有效扩展训练。
基本的数据并行(data parallelism, DP)并不能减少每个设备的内存在32GB内存的GPU上参数超过14亿的模型很容易就会出现内存不足的情况。其他现有的解决方案如流水线并行(Pipeline Parallelism,PP)、模型并行(Model Parallelism,MP)、**CPU卸载(CPU-Offloading)**等虽然在功能、可用性以及内存和计算/通信效率之间进行了权衡但这些都是实现快速和大规模训练的关键因素。 DP 将模型复制到多个GPU上每个GPU设备可以并行接受不同的数据批次计算的梯度结果也不一致因此需要在模型间交换梯度计算平均值。由于模型在每个GPU上复制了一份产生了很多GPU内存消耗。 MP 当模型本身的参数量很大无法在单卡上加载时无法进行数据并行。模型并行选择将模型按层拆分每个GPU只加载模型参数的一部分在前向传播过程中汇总数据将依次通过各个GPU带来的时间消耗很大。 PP 流水线并行是MP的一种变体将每批输入数据拆分成很多小的微批次以减少设备的空闲等待时间。只有在整个模型处理完每个微批次后才会更新模型参数因此当其他GPU设备仍在处理上一个微批次时空闲的GPU就可以开始处理下一个微批次。 Megatron-LM提出了1F1B流水线策略即一个前向通道和一个后向通道。1F1B流水线策略引入了任务调度机制使得下游设备能够在等待上游计算的同时执行其他可并行的任务从而提高设备的利用率。 在现存的训练大模型的解决方案中模型并行或许是最有前景的。然而MP无法进一步扩展到更大的模型。模型并行是通过将模型垂直切分将每层的计算和参数分配到多个设备上从而需要在每一层之间进行大量通信。因此它们在单个节点内工作良好因为节点间的GPU通信带宽很高但在超出单节点时效率迅速下降。
作者在两个DGX-2节点上使用Megatron-LM测试了一个40亿参数的模型观察到每个V100 GPU的性能约为5T flops(不到硬件峰值的5%)。
那么如何克服现有解决方案的局限性更高效地训练大模型呢为了解答这个问题作者首先分析了现有系统在模型训练中的内存消耗全景并将其分类为两个部分1)对于大模型大部分内存被模型状态(model states)占用包括优化器状态(optimizer states, 如Adam中的动量和方差)、梯度(gradients)和参数(parameters)2剩余的内存被激活值(activation)、临时缓冲区( temporary buffers)和不可用的碎片内存(unusable fragmented memory)消耗将其统称为剩余状态(residual states)。
作者开发了ZeRO——零冗余优化器(Zero Redundancy Optimizer)旨在优化这两部分的内存效率同时实现高计算和通信效率。由于这两部分面临不同的挑战分别来讨论其解决方案。
优化模型状态内存
模型状态在训练过程中通常消耗最多的内存但现有的方法如数据并行(DP)和模型并行(MP)并未提供令人满意的解决方案。DP在计算/通信效率方面表现良好但内存效率较差而MP则可能导致计算/通信效率不理想。具体而言DP在所有数据并行进程中复制所有模型状态导致冗余内存消耗而MP通过将这些状态进行划分以获得高内存效率但通常会导致过于细粒度的计算和昂贵的通信从而降低扩展效率。此外所有这些方法在整个训练过程中都静态维护所需的所有模型状态尽管在训练期间并不是所有的模型状态都是随时需要的。
因此作者开发了ZeRO-DP(ZeRO支持的数据并行)它在实现DP的计算/通信效率的同时也达成了MP的内存效率。ZeRO-DP通过对模型状态进行划分而不是复制消除了数据并行进程之间的内存状态冗余并通过在训练期间使用动态通信调度保持DP的计算粒度和通信量从而保持计算/通信效率。 图1比较采用ZeRO-DP优化的三个阶段下每个设备的模型状态内存消耗。 Ψ Ψ Ψ表示模型大小(参数量 K K K表示优化器状态的内存乘子 N d N_d Nd表示数据并行度。在此示例中我们假设模型大小为 Ψ 7.5 B Ψ 7.5B Ψ7.5B数据并行度 N d 64 N_d 64 Nd64基于使用Adam优化器的模型并行训练优化器状态的内存乘子 K 12 K 12 K12Adam优化器会以FP32保存模型参数备份一阶动量和二阶动量也都采用FP32格式存储因此是12。
ZeRO-DP主要有三个优化阶段(图1)分别对应优化器状态、梯度和参数的划分。当累积启用时
优化器状态划分( P o s P_{os} Pos)内存减少4倍与DP的通信量相同添加梯度划分( P o s g P_{osg} Posg)内存减少8倍与DP的通信量相同添加参数划分 ( P o s g p (P_{osgp} (Posgp)内存减少与数据并行程度 N d N_d Nd成线性关系。
例如在64个GPU上划分( N d 64 N_d 64 Nd64)将实现64倍的内存减少。通信量有适度增加约50%。
ZeRO-DP消除了内存冗余使集群的全部聚合内存容量可用。在启用所有三个阶段的情况下ZeRO能够在仅有1024个NVIDIA GPU上训练一个万亿参数模型。一个采用16位精度的优化器(如Adam)的大规模万亿参数模型大约需要16TB的内存来存储优化器状态、梯度和参数。将16TB除以1024得到16GB这在GPU(具有32GB设备内存)可接受的范围内。
优化剩余状态内存
在ZeRO-DP提升了模型状态的内存效率之后激活、临时缓冲区和不可用内存碎片所消耗的剩余内存可能成为次要的内存瓶颈。作者开发了ZeRO-R以分别优化这三种因素所消耗的剩余内存。 对于激活(在前向传播中存储以便执行反向传播)检查点(checkpointing)1技术有所帮助但对于大模型而言仍不够。因此ZeRO-R通过激活分区(activation partitioning)优化激活内存识别并消除现有模型并行方法中的激活复制同时在适当时将激活卸载到CPU。 ZeRO-R为临时缓冲区定义了适当的大小以期在内存和计算效率之间达到平衡。 在训练过程中由于不同张量的生命周期差异导致内存碎片化。由于碎片化而缺乏连续内存可能导致内存分配失败尽管可用内存充足。ZeRO-R基于不同张量的生命周期主动管理内存防止内存碎片化。
ZeRO-DP和ZeRO-R结合在一起形成了一个强大的深度学习训练内存优化系统统称为ZeRO。
ZeRO与模型并行
由于ZeRO消除了数据并行中的内存低效性自然会问我们是否还需要模型并行何时需要ZeRO如何与模型并行协同工作在使用ZeRO的情况下模型并行作为单纯为了适应大模型而采用的选项变得不那么吸引人。ZeRO-DP在减少每个设备的内存占用方面至少与模型并行同样有效有时在模型并行无法均匀划分模型时甚至更有效。它的扩展效率也相当或更好。此外数据并行性易于使用广泛适用于不同的工作负载。
尽管如此仍然存在一些情况下我们希望利用模型并行1)在与ZeRO-R结合使用时模型并行可以减少非常大型模型的激活内存占用。2) 对于较小的模型当仅使用数据并行时聚合批量大小过大而导致收敛效果不佳时模型并行也可以带来好处。在这些情况下可以将ZeRO与模型并行结合使用以适应可接受的聚合批量大小。
ZeRO可以与模型并行结合从而在每个设备上实现理论上的最大内存减少减少倍数为 N d × N m N_d × N_m Nd×Nm其中 N d N_d Nd为数据并行度 N m N_m Nm为模型并行度。
实现与评估 ZeRO中的完整优化集使我们能够在今天的高端硬件集群上运行万亿参数的模型然而硬件计算能力仍然太有限训练时间可能过长(超过1年)。因此作者实现的重点是高效支持参数量是SOTA的10倍(约1000亿参数)的模型同时仍然在当前硬件的计算能力范围内。作者实现并评估了ZeRO的部分优化称为ZeRO-100B——这是ZeRO-DP的 P o s g P_{osg} Posg与ZeRO-R的结合。结果显示
模型规模 结合模型并行ZeRO-100B能高效运行170亿参数的模型如图2所示。这相比于SOTA模型规模增加了超过8倍。
速度 改进的内存效率提升了吞吐量和训练速度。如图2所示ZeRO能在一个400个Nvidia V100 GPU集群上运行100亿参数的模型每个GPU超过38 TFlops整体性能超过15 Petaflops。这在相同模型规模下相比于SOTA训练速度提升超过10倍。
可扩展性 在64-400个GPU的范围内观察到超线性加速其性能在双倍增加GPU数量时超过两倍。这是ZeRO-DP的特性在增加数据并行度时减少模型状态的内存占用使我们能在每个GPU上适配更大的批量大小从而提高性能。
大型模型训练的民主化 ZeRO-100B使数据科学家能够在没有任何需要重构模型的模型并行或流水线并行的情况下训练最多130亿参数的模型。数据科学家因此可以自由实验大型模型而不用担心并行性。相比之下现有系统在1.4亿参数的模型上会出现内存不足的情况。
新一代最先进模型 ZeRO支持了最大的语言模型参数为170亿并且取得了创纪录的准确性即Turing-NLG。
2. 相关工作
2.1 数据、模型和流水线并行
并行化是大规模训练大型模型的关键策略。对于能够适应设备内存进行训练的模型可以使用数据并行DP将训练扩展到多个设备。在数据并行中模型参数在每个设备上被复制。在每一步中mini-batch被均匀地分配到所有数据并行进程使每个进程在不同的数据样本子集上执行前向和后向传播并使用跨进程的平均梯度在本地更新模型。
当模型无法适应设备内存时模型并行MP和流水线并行PP会在进程之间分割模型分别从垂直和水平的方式进行。
流水线并行水平分割模型通过在不同设备上运行每个部分并使用micro-batch处理来隐藏流水线气泡。由于水平分割和micro-batch某些模型功能如绑定权重和批归一化的实现变得困难。
流行的流水线并行实现如G-pipe在分割模型参数和总激活值时需要一个与流水线分区数量成比例的批量大小来隐藏流水线气泡。大的批量大小会影响收敛速度同时还需要大量内存来存储激活值。
PipeDream的另一种流水线并行实现保留多个过时参数副本以隐藏流水线气泡而不会显著增加批量大小这使得其内存效率较低。此外该实现与标准的深度学习训练并不等同并对训练收敛有影响。相比之下ZeRO在不产生流水线并行的功能、性能和收敛相关限制的情况下获得了与PP相同或更好的内存效率。
2.2 非并行方法减少内存
2.2.1 减少激活内存
多项研究致力于通过压缩、激活检查点或实时分析(live analysis)来降低激活的内存占用。这些方法是互补的可以与ZeRO协同工作。实际上ZeRO-R中的激活内存减少与激活检查点一起工作。
2.2.2 CPU卸载
一些研究利用当前计算节点的异构特性通过算法设计或虚拟化内存将模型状态卸载到CPU内存中使得训练时间的多达50%可能用于GPU-CPU-GPU的传输。与这些方法不同ZeRO显著减少内存消耗而不是将模型状态存储到由于PCI-E带宽限制而严重受限的CPU内存中。在罕见情况下ZeRO-R可能会将激活检查点卸载到CPU中以提升性能。
2.2.3 内存高效优化器
另一些研究关注通过维护粗粒度的模型参数和梯度统计来减少自适应优化方法的内存消耗这可能对模型收敛性保证产生影响。ZeRO与这些努力是正交的其优化不会改变模型优化方法或影响模型收敛而是有效减少每个设备的优化器状态和梯度的内存占用。自适应优化方法对于大模型的有效训练达到SOTA性能和准确性至关重要。与随机梯度下降相比由于需要维护每个模型参数和梯度的细粒度一阶和二阶统计信息这会产生显著的内存占用。ZeRO能够将这些优化器的内存占用降低几个数量级使得这些复杂的优化方法在内存资源有限的硬件上变得可行。这也使得开发和使用更复杂、内存需求更高的优化器成为可能可能会带来更好的收敛性。
3 内存都去哪儿了
我们看看当前训练系统的内存消耗。例如一个有15亿参数的GPT-2模型在16位精度下需要3GB的内存来存储其权重(或参数)然而它无法在32GB内存的单个GPU上使用Tensorflow或PyTorch进行训练。这让人不禁要问所有内存都去哪儿了。在模型训练过程中大部分内存被模型状态(model states)所占用即由优化器状态、梯度和参数组成的张量。除了这些模型状态剩余的内存则被激活、临时缓冲区和我们称之为剩余状态的碎片内存所消耗。
3.1 模型状态优化器状态、梯度和参数
在训练过程中设备内存的大部分被模型状态所占用。以Adam为例。Adam需要存储两个优化器状态1)时间平均动量2)梯度的方差以计算更新。因此使用Adam训练模型时必须有足够的内存来存储梯度的动量和方差的副本。此外还需要有足够的内存来存储梯度和权重本身。在这三种类型的参数相关张量中优化器状态通常占用最多的内存尤其是在进行混合精度训练时。
混合精度训练 在当前一代NVIDIA GPU上训练大型模型的最先进方法是采用混合精度(fp16/32)训练其中参数和激活以fp16存储从而能够使用这些GPU高吞吐量的张量核心单元。在混合精度训练中前向和反向传播都使用fp16权重和激活进行。然而为了在反向传播结束时有效计算和应用更新混合精度优化器还需要保持参数的fp32副本以及所有其他优化器状态的fp32副本。
以Adam为具体例子使用Adam训练一个具有 Ψ Ψ Ψ个参数的模型的混合精度版本需要足够的内存来存储参数和梯度的fp16副本需要的内存为 2 Ψ 2Ψ 2Ψ和 2 Ψ 2Ψ 2Ψ字节。此外还需要存储优化器状态参数、动量和方差的fp32副本需要的内存为 4 Ψ 4Ψ 4Ψ、 4 Ψ 4Ψ 4Ψ和 4 Ψ 4Ψ 4Ψ字节。我们用 K K K表示优化器状态的内存倍增因子即存储它们所需的额外内存为 K Ψ KΨ KΨ字节。在混合精度Adam中K 12。总之这导致内存需求为 2 Ψ 2 Ψ K Ψ 16 Ψ 2Ψ 2Ψ KΨ 16Ψ 2Ψ2ΨKΨ16Ψ字节。对于一个像GPT-2这样的模型具有15亿个参数这意味着至少需要24GB的内存这显著高于仅存储fp16参数所需的微薄的3GB内存。
3.2 剩余内存消耗
在训练过程中激活值会占用大量内存。以具体例子来说参数为15亿的GPT-2模型在序列长度为1K和批量大小为32时需要约60 GB的内存。激活检查点(或激活重计算)是一种常见的方法可以在重计算开销为33%的情况下将激活内存降低约为总激活的平方根。这将使该模型的激活内存消耗减少到约8 GB。
尽管显著减少了内存但对于更大的模型即使使用激活检查点激活内存仍然可能变得相当庞大。例如一个具有1000亿参数的类GPT的模型在批量大小为32时仍需要约60 GB的内存。
用于存储中间结果的临时缓冲区在较大模型中也会消耗相当可观的内存。诸如梯度all-reduce或梯度norm计算等操作通常会将所有梯度融合到一个单一的扁平化缓冲区中以提高吞吐量。例如跨设备的全规约带宽会随着大消息大小的增加而提高。虽然梯度本身通常存储为fp16张量但根据操作不同融合的缓冲区可能是fp32张量。当模型规模较大时这些临时缓冲区的大小是相当不容小觑的。例如对于一个具有15亿参数的模型扁平化的fp32缓冲区将需要6 GB的内存。
内存碎片化 到目前为止我们讨论了训练过程中实际的内存消耗。除此之外即使有大量可用内存也可能会出现可用内存不足的情况。这种情况可能是由于内存碎片化造成的。如果没有足够的连续内存来满足内存请求即使总可用内存大于请求的内存该请求也会失败。在训练非常大的模型时我们观察到显著的内存碎片化导致出现内存不足的问题某些极端情况下仍然有超过30%的内存可用。
4 ZeRO洞察与概述
ZeRO有两组优化措施1) ZeRO-DP旨在减少模型状态的内存占用2) ZeRO-R则针对减少剩余内存消耗。我们将介绍这些优化的概述和背后的洞察这使得ZeRO能够在保持高效的同时减少内存占用。这里效率很关键如果没有这个约束像将所有参数状态移动到CPU内存或任意增加MP度数等简单解决方案都可以减少内存占用。 洞察独一无二的观点 4.1 ZeRO-DP
ZeRO驱动的DP基于三个关键洞察
DP的扩展效率优于MP因为MP降低了计算的粒度同时增加了通信开销。在某个临界点之后较低的计算粒度会降低每个GPU的效率而增加的通信开销则会妨碍跨GPU的可扩展性尤其是在跨节点边界时。相反DP具有更高的计算粒度和更低的通信量从而实现更高的效率。DP在内存使用上效率较低因为模型状态在所有数据并行进程中冗余存储。相反MP通过对模型状态进行分区来获得内存效率。DP和MP在整个训练过程中都保持所有所需的模型状态但并不是所有状态在所有时间都需要。例如层参数仅在该层的前向传播和反向传播期间需要。
基于这些思考ZeRO-DP在保持DP的训练效率的同时实现了MP的内存效率。ZeRO-DP通过分区模型状态而不是复制它们并使用动态通信调度利用模型状态的内在时间特性同时最小化通信量。
通过这样做ZeRO-DP在增加DP度数的同时线性减少每个设备的模型内存占用同时保持通信量接近默认DP的水平从而保持效率。
4.2 ZeRO-R
4.2.1 减少激活内存
两个关键洞察是
MP对模型状态进行分区但通常需要复制激活内存。例如如果我们将线性层的参数垂直拆分在两个GPU上并行计算则每个GPU都需要完整的激活来计算其分区。对于像GPT-2或更大规模的模型算术强度(每次迭代的计算量与每次迭代的激活检查点量之比)非常大并且随着隐藏层维度的增加而线性增加使得即使带宽较低也能够隐藏激活检查点的数据移动开销。
ZeRO通过在GPU之间分区激活检查点减少了MP中的内存冗余并使用allgather根据需要重构它们。激活内存的占用量根据MP度数成比例减少。对于非常大的模型ZeRO甚至可以选择将激活部分移动到CPU内存中同时由于这些模型的算术强度大仍然能够实现良好的效率。
4.2.2 管理临时缓冲区
ZeRO使用恒定大小的缓冲区以避免随着模型规模的增加而导致临时缓冲区膨胀同时确保缓冲区的大小足够以保持效率。
4.2.3 管理内存碎片
内存碎片是短期内存对象与长期内存对象之间交错的结果。在前向传播期间激活检查点是长期存在的而重新计算的激活则是短期存在的。类似地在反向计算中激活梯度是短期存在的而参数梯度是长期存在的。基于这一洞察ZeRO通过将激活检查点和梯度移动到预分配的连续内存缓冲区进行即时内存碎片整理。这不仅增加了内存的可用性还通过减少内存分配器寻找空闲连续内存所需的时间提高了效率。
5. 深入ZeRO-DP
虽然现有的数据并行(DP)方法在每个设备上复制模型状态并引入显著的内存开销但 ZeRO-DP 通过在数据并行进程之间对优化器状态、梯度和参数进行分区消除了这种内存冗余。图1定量并可视化了使用和不使用 ZeRO-DP 的内存需求。该图展示了在分区后(1)优化器状态、(2)梯度和(3)参数冗余的累计内存占用。将它们称为 ZeRO-DP 的三个优化阶段 P o s P_{os} Pos、 P g P_g Pg 和 P p P_p Pp以下将详细说明。
5.1 Pos优化器状态分区
对于 DP 程度 N d N_d Nd将优化器状态分为 N d N_d Nd 个相等的分区使得第i个数据并行进程只更新与第i个分区对应的优化器状态。因此每个数据并行进程只需存储和更新总优化器状态的 1 N d \frac{1}{N_d} Nd1部分同时只更新 1 N d \frac{1}{N_d} Nd1的参数。在每个训练步骤结束时在所有数据并行进程之间执行一次 all-gather以获取所有数据并行进程的完全更新参数。
内存节省如图1所示经过优化状态分区后的内存消耗从 4 Ψ K Ψ 4Ψ KΨ 4ΨKΨ 降低到 4 Ψ K Ψ N d 4Ψ \frac{KΨ}{N_d} 4ΨNdKΨ。以图1中的具体例子为例使用 64 路数据并行的 P o s P_{os} Pos 处理( N d 64 N_d64 Nd64)一个 75 亿参数的模型需要 31.4GB 的内存而使用标准 DP 则需要 120 GB。此外当 N d N_d Nd 较大时模型状态的内存需求从 4 Ψ 12 Ψ 16 Ψ 4Ψ 12Ψ 16Ψ 4Ψ12Ψ16Ψ字节降至 4 Ψ 12 Ψ N d ≈ 4 Ψ 4Ψ \frac{12Ψ}{N_d} ≈ 4Ψ 4ΨNd12Ψ≈4Ψ字节得到 4 倍的减少。
5.2 Pg梯度分区
由于每个数据并行进程只更新其对应的参数分区因此它只需要对应参数聚合(reduce)的梯度。因此在反向传播过程中当每层的梯度可用时仅在负责更新相应参数的数据并行进程中进行reduce。之后不再需要这些梯度可以释放其内存。这将需要持有梯度的内存占用从 2 Ψ 2Ψ 2Ψ 字节减少到 2 Ψ N d \frac{2Ψ}{N_d} Nd2Ψ。
实际上这是一种 Reduce-Scatter 操作其中对应于不同参数的梯度被reduce到不同的进程。为了在实践中提高效率使用了一种分桶策略将所有与特定分区对应的梯度归入一个桶并对整个桶进行一次性reduce。执行在分区边界的reduce而不是all-reduce以减少内存占用并重叠计算和通信。
内存节省通过消除梯度和优化器状态冗余将内存占用进一步降低到 2 Ψ 14 Ψ N d ≈ 2 Ψ 2Ψ \frac{14Ψ}{N_d} ≈ 2Ψ 2ΨNd14Ψ≈2Ψ。如图1中的例子所示使用 P o s g P_{osg} Posg 处理的 75 亿参数模型在 64 路数据并行情况下仅需 16.6 GB 的内存而使用标准 DP 则需要 120 GB。当 N d N_d Nd较大时模型状态的内存需求从 2 Ψ 14 Ψ 16 Ψ 2Ψ 14Ψ 16Ψ 2Ψ14Ψ16Ψ字节降低到 2 Ψ 14 Ψ N d ≈ 2 Ψ 2Ψ \frac{14Ψ}{N_d} ≈ 2Ψ 2ΨNd14Ψ≈2Ψ字节得到8 倍的减少。
5.3 Pp参数分区
与优化器状态和梯度一样每个进程仅存储其对应分区的(模型)参数。当需要进行前向和反向传播时超出其分区的参数将通过广播从适当的数据并行进程接收。尽管乍一看这似乎会产生显著的通信开销但这种方法仅将基线数据并行系统的总通信量增加到 1.5 倍同时使内存减少与数据并行度 N d N_d Nd成正比。
内存节省通过参数分区将一个 Ψ Ψ Ψ参数模型的内存消耗从 16 Ψ 16Ψ 16Ψ降低到 16 Ψ N d \frac{16Ψ}{N_d} Nd16Ψ。如图1所示使用 P o s p g P_{ospg} Pospg 和 64路数据并行 N d 64 N_d 64 Nd64的情况下7.5 亿参数模型需要 1.9 GB 的模型状态内存而使用标准数据并行则需要 120 GB。ZeRO 可以使数据并行适应任意大小的模型只要有足够的设备共享模型状态。
5.4 对模型大小的影响 对 P o s P_{os} Pos、 P o s g P_{osg} Posg 和 P o s g p P_{osgp} Posgp三个阶段的分区分别将每个数据并行进程在模型状态的内存消耗减少至最多 4 倍、8 倍和 N d N_d Nd 倍。表1分析了在 ZeRO-DP 优化的 3 个阶段下几个示例模型在不同数据并行度下的模型状态内存消耗。如果没有 ZeRO内存消耗等于表中的第一行无论数据并行度如何。
当 N d 64 N_d 64 Nd64时分别使用 P o s P_{os} Pos、 P o s g P_{osg} Posg 和 P o s g p P_{osgp} PosgpZeRO 可以训练最多 7.5B、14B 和 128B 参数的模型当 N d 1024 N_d 1024 Nd1024时启用所有 ZeRO 优化 P o s g p P_{osgp} Posgp能够训练具有 1 万亿参数的模型或任何规模的模型而没有 ZeRO最大的数据并行模型只能运行少于 15 亿参数的模型。
6. 深入了解 ZeRO-R
6.1 Pa分区激活检查点
如 4.2 节所述模型并行设计上需要复制激活从而导致在模型并行 GPU 之间激活的冗余副本。ZeRO 通过对激活进行分区来消除这种冗余并且仅在计算过程中使用激活之前以复制形式逐层生成激活具体来说一旦计算完成模型某一层的前向传播输入激活便会在所有模型并行进程中进行分区直到在反向传播中再次需要此激活为止。在这一点上ZeRO 使用all-gather操作重新生成激活的复制副本。将这种优化称为 P a P_a Pa。它与激活检查点技术一起工作只存储分区的激活检查点而不是复制的副本。此外对于非常大的模型和极为有限的设备内存这些分区的激活检查点也可以卸载到 CPU以将激活的内存开销减少到几乎为零尽管会产生额外的通信成本将其称为 P a c p u P_{acpu} Pacpu。
内存节省通过分区激活检查点ZeRO 将激活的内存占用降低至与 MP 程度成正比的因子。考虑训练一个 100B 模型批量大小为 32序列长度为 1024MP 度为 16。如果对每个transformer层进行一个激活检查点那么每个 GPU 所需的内存将约为 33 GB仅用于存储激活检查点。但使用 ZeRO 的 P a P_a Pa 后这个内存需求可以降低至每个 GPU 大约 2 GB。此外这 2 GB 内存可以卸载到 CPU进一步将激活的内存占用减少到几乎为零。
6.2 CB恒定大小缓冲区
ZeRO 精心选择临时数据缓冲区的大小以平衡内存和计算效率。在训练过程中一些操作的计算效率可能高度依赖于输入大小较大的输入能够实现更高的效率。例如大规模的all-reduce操作比小规模操作实现更高的带宽。因此为了获得更好的效率高性能库如 NVIDIA Apex 或 Megatron 会在应用这些操作之前将所有参数融合到一个单一的缓冲区中。然而融合缓冲区的内存开销与模型大小成正比可能会造成限制。例如对于一个 3B 参数的模型32 位的融合缓冲区将需要 12 GB 的内存。为了解决这个问题在模型过大时简单使用高性能、高效的恒定大小的融合缓冲区。这样缓冲区大小不依赖于模型大小并且通过保持缓冲区大小足够大仍然可以实现良好的效率。
6.3 MD内存碎片整理
在模型训练中由于激活检查点和梯度计算而导致的内存碎片化现象。在进行激活检查点的前向传播过程中仅存储选定的激活用于反向传播而大部分激活被丢弃因为它们可以在反向传播期间重新计算。这导致了短期内存(被丢弃的激活)和长期内存(检查点激活)之间的交错从而造成内存碎片。同样在反向传播过程中参数梯度是长期存在的而激活梯度和计算参数梯度所需的其他缓冲区则是短期存在的。这种短期内存和长期内存的交错再次导致内存碎片。
有限的内存碎片通常不是问题尤其是在有充足内存的情况下但在有限内存下进行大模型训练时内存碎片会导致两个问题1)即使有足够的可用内存也会因缺乏连续内存而出现内存溢出2)由于内存分配器花费大量时间寻找满足内存请求的连续内存块导致效率低下。
ZeRO 通过为激活检查点和梯度预分配连续内存块并在生成时将它们复制到预分配的内存中实时进行内存碎片整理。 M D M_D MD不仅使 ZeRO 能够以更大的批量大小训练更大的模型还在有限内存下提高了训练效率。
7. ZeRO-DP 的通信分析
由于 ZeRO 通过消除内存冗余来提升模型大小ZeRO 驱动的 DP 方法的通信量与基线 DP 方法相比如何答案分为两部分1) 使用 P o s P_{os} Pos 和 P g P_g Pg 时ZeRO-DP 不会产生额外的通信同时实现最多 8 倍的内存减少2)在使用 P p P_p Pp 而不是 P o s P_{os} Pos 和 P g P_g Pg 时ZeRO-DP 的通信量最多为 1.5 倍同时进一步减少内存占用 N d N_d Nd 倍。
7.1 数据并行通信量
在数据并行训练过程中所有数据并行进程的梯度在反向传播结束时进行平均然后再计算下一步的更新。平均操作使用all-reduce集体通信。对于大模型all-reduce通信完全受限于通信带宽因此将分析限制在发送和接收每个数据并行进程的总通信量上。最先进的all-reduce实现采用两步法第一步是reduce-scatter操作它在不同的进程上减少数据的不同部分。下一步是all-reduce操作每个进程在所有进程上收集减少后的数据。这两个步骤的结果就是all-reduce。reduce-scatter和all-reduce都是通过管道化方法实现的这导致每个步骤的数据移动总量为 Ψ Ψ Ψ 元素(对于一个包含 Ψ Ψ Ψ 元素的数据)。因此标准 DP 在每个训练步骤中会产生 2 Ψ 2Ψ 2Ψ 的数据移动。
7.2 ZeRO-DP通信量
7.2.1 使用Posg的通信量
通过梯度分区每个进程仅存储更新其对应参数分区所需的梯度部分。因此ZeRO仅需对梯度执行scatter-reduce操作而不需要all-reduce这产生的通信量为 Ψ Ψ Ψ。在每个进程更新其负责的参数分区后会执行all-gather操作以从所有数据并行进程收集所有更新后的参数。这也产生了 Ψ Ψ Ψ的通信量。因此每个训练步骤的总通信量为 Ψ Ψ 2 Ψ Ψ Ψ 2Ψ ΨΨ2Ψ正好与基线DP相同。
7.2.2 使用Posgp的通信量
在参数分区后每个数据并行进程仅存储其更新的参数。因此在前向传播期间它需要接收所有其他分区的参数。然而这可以进行管道化以避免内存开销。在计算与特定分区对应的模型部分的前向传播之前负责该分区的数据并行进程可以将权重广播到所有数据并行进程。一旦该分区的前向传播完成参数可以被丢弃。因此总的通信量为 Ψ × N d N d Ψ \frac{Ψ×N_d}{N_d} Ψ NdΨ×NdΨ。
换句话说通过将其分布在整个前向传播的参数all-gather的过程下重新安排并在使用后丢弃参数。这个all-gather在反向传播中还需要按照反向顺序再次进行。因此总通信量是由这些all-gather操作产生的通信量加上梯度的scatter-reduce所产生的通信量之和。因此总量为 3 Ψ 3Ψ 3Ψ与基线相比为1.5倍。梯度和参数分区利用了这样一个见解——并非所有状态的梯度和参数都需要随时存在——通过明智地通信这些状态来优化内存。
8. ZeRO-R的通信分析
将分区激活检查点( P a P_a Pa)在ZeRO-R中的通信量与基线模型并行进行比较发现 P a P_a Pa引入的通信量增加通常小于基线MP的十分之一。分区激活检查点的通信量权衡取决于模型大小、检查点策略和MP策略。
在使用激活检查点的Megatron-LM中每个transformer块在前向传播中执行两次大小为batch × seq_length × hidden_dim维度的all-reduce操作在前向重新计算中又执行两次all-reduce在反向传播中再执行两次。每个块的总通信量为12 × seq_length × hidden_dim因为all-reduce的通信量为2 × message_size。
当ZeRO-R对激活检查点进行分区时需要在每个激活检查点的反向传播前进行一次额外的all-reduce操作。一般来说对每个transformer块检查输入激活因此每个transformer块需要一次all-reduce。因此 P a P_a Pa的通信开销为seq_length × hidden_dim因为all-reduce的通信量为message_size。因此 P a P_a Pa的总通信开销小于原始模型并行通信量的10%。
当MP与DP结合使用时 P a P_a Pa可以在增加10%的模型并行通信量的情况下将数据并行通信量减少一个数量级并在数据并行通信成为性能瓶颈时显著提升效率。 P a P_a Pa通过模型并行度减少激活内存消耗从而允许批量大小成比例增加。对于大型模型MP可以达到16这使得批量大小增加最多可达16倍。
数据并行训练的通信量与批量大小成反比。因此由于 P a P_a Pa引起的批量大小增加一个数量级可能导致数据并行通信量减少一个数量级。
最后如果应用 P a c p u P_{acpu} Pacpu分区激活检查点将被卸载到CPU激活内存需求几乎降低到零但相对于 P a P_a Pa往返于CPU内存的额外数据移动增加了2倍。在极端情况下如果由于小批量大小即使使用 P a P_a PaDP通信量仍是主要瓶颈那么 P a c p u P_{acpu} Pacpu可以通过增加批量大小来提高效率只要CPU数据传输开销少于DP通信量开销这通常对小批量大小成立。
9. 迈向万亿参数 目前发布的最大模型在100亿参数的范围内训练起来已经相当具有挑战性。达到万亿参数——大三个数量级的规模必然会发生。ZeRO并不声称能够知道或解决所有问题但ZeRO解决了从系统角度来看最根本的挑战之一在当前硬件上适配如此规模的模型同时确保良好的系统可扩展性。
从SOTA的跨越 目前SOTA框架Megatron能够以可接受的吞吐量训练的最大模型为160亿至200亿参数的模型运行在DGX-2系统中。而通过在多个DGX节点之间实现模型并行进一步扩展会由于节点间带宽的限制而导致显著的效率下降。
ZeRO大幅提升了可高效运行的模型规模。它使得当前一代硬件可以运行显著更大的模型而不需要细粒度的模型并行来跨越节点边界。如表1所示启用所有优化 P o s g p P_{osgp} Posgp的ZeRO能够仅使用数据并行在1024个GPU上容纳超过1万亿参数的模型。同时在结合模型并行时(如表2所示)ZeRO能够在1024个GPU上以16倍模型并行和64倍数据并行跨节点容纳超过1万亿参数的模型。
计算能力的差距 然而在可接受的时间范围内端到端地训练一个万亿参数的模型仍然可能需要大量计算能力而这在今天的AI集群中是缺乏的。
为了理解资源需求与Bert-Large进行了简要比较。Bert-Large可以在1024个GPU的DGX-2H集群上以67分钟的时间训练完成。一个一万亿参数的模型针对每个数据样本其计算量可能是Bert-Large模型的3000倍。即使我们假设序列长度和训练模型所需的样本总数相同训练一个一万亿参数的模型也需140天假设相同硬件和类似计算效率。实际上数据样本和序列长度很可能随模型规模的增加而增加训练可能需要超过一年。为了在合理的时间内训练一万亿参数的模型需要一个exa-flop级的系统。
总结
⭐ 主要介绍了DeepSpeed框架用到的零冗余优化器(ZeRO)技术使得训练万亿级别参数的大模型成为了可能。通过ZeRO-DP优化模型状态通过ZeRO-R优化剩余状态。
引用 Training Deep Nets with Sublinear Memory Cost ↩︎
