的关系为: zero 三种优化策略 对显存的进一步优化也就从这三方面下手,即零冗余优化,分为三个层次: • zero-1 仅对优化器分片 优化后的内存占比为: 当N比较大时, 显存占用相当于原来的 • zero zero-1 和zero-2 相对于baseline 数据并行,通信量是没有变化的,zero-3变为1.5倍。 。
• zero-2 对优化器和梯度分片 优化后的内存占比为:,当N比较大时, 显存占用相当于原来的 。 2,ZeRo-2参数更新过程 策略:将模型的梯度和优化器状态进行分片。每个GPU上,保存一份完整的参数副本,以及分片后的梯度和优化器状态。 内存:相对于数据并行,内存最大能减少到原来的1/8倍。 梯度同步: 与ZeRo-2一样,3个GPU设备前向计算获得损失后,同时反向传播,对梯度进行 Reduce-Scatter通信,每个GPU设备获得负责层的平均梯度,通信量为M。 以上流程中总的通信量: 前向计算 M + 反向传播 2M + 参数更新同步 0 = 3M ZeRo策略总结 ZeRo通过分片策略实现内存与通信的权衡,其核心思想是 “以通信换内存” 维度 ZeRo-2
我们将此版本称为 ZeRO-2。 ZeRO-2显着降低了训练大型模型的内存占用,这意味着可以使用更少的模型并行度和更大的批量大小来训练大型模型。 对于相同大小的模型,与仅使用Megatron-LM相比,ZeRO-2的训练速度快10倍,与ZeRO-1相比快5倍。 ZeRO-2降低了模型状态的内存占用,使我们可以在每个GPU上容纳更大的批量大小,从而提高了性能。随着数据并行度的增加,ZeRO-2减少了模型状态的内存占用,这也导致了超线性加速的效果。 平民化大模型训练: ZeRO-2使模型科学家能够高效地训练高达130亿个参数的模型,而无需进行通常需要模型重构的模型并行(图2,右下角)。 你也可以使用此方法保存ZeRO-2权重。 如果你想获取fp32权重,我们提供了一种特殊的脚本,可以进行离线合并。它不需要配置文件或GPU。
1,ZeRo-2 参数更新和通信量计算 策略:将模型的梯度和优化器状态进行分片。每个GPU上,保存一份完整的参数副本,以及分片后的梯度和优化器状态。 与ZeRo-2一样,3个GPU设备前向计算完成后,同时反向传播,对梯度进行 Reduce-Scatter通信,每个GPU设备获得负责层的梯度平均,通信量为M。注意! 以上流程中总的通信量: 前向计算 M + 反向传播 2M + 参数更新同步 0 = 3M 3,ZeRo-2&3策略总结 ZeRo通过分片策略实现内存与通信的权衡,其核心思想是 “以通信换内存” 维度 ZeRo
ZeRO-Offload 核心技术是在 ZeRO-2基础之上将优化器状态和梯度卸至 CPU 内存。 其背后的关键技术是:在 ZeRO-2 基础上将优化器计算,优化器状态和梯度卸载到 CPU 内存。 ZeRO-2除了对优化器状态进行分区外,还对梯度进行分区, ZeRO-3对所有模型状态进行分区。 ZeRO-Offload 与ZeRO-2协同工作,因此我们将对其进行进一步讨论。 ZeRO-2的计算和通信具体描述如下。 在前向传播过程中,每个GPU计算不同mini-batch的损失。 2.6.1 单机计划 ZeRO-2 在每个 GPU 上保存一部分优化器状态量和梯度,ZeRO-Offload 继承了 ZeRO-2 的划分优化器状态量和梯度的方法。
DeepSpeed ZeRO Stage-2(后文简称 ZeRO-2)是该框架中一项节约显存占用且不增加额外计算量和通信量的技术。如上图所示,由于计算需要,每个 Rank 都需要拥有全部的参数。 于是 ZeRO-2 提出将优化器状态均匀地切分在每个 Rank 上(注意,这里并不需要保证每个变量被均分或完整保留在某个 Rank 上),每个 Rank 在训练进程中只负责对应部分的优化器状态和模型参数的更新 默认情况下,ZeRO-2 在反向时在所有 Rank 间使用 Reduce 方式聚合梯度,而后每个 Rank 只需要保留自身所负责的参数的部分,既消除了冗余的重复计算,又降低了显存占用。 Megatron-LM DistributedOptimizer 原生 Megatron-LM 通过 DistributedOptimizer 实现了类似 ZeRO-2 的梯度和优化器状态切分,以减少训练中的显存占用
其次,使用比较多的是 Zero-2 和Zero-3,但目前在 Megatron 和 Nemo 上还没有做支持。 所以当同时使用 Zero-2 或 Zero-3 加上流水线并行时,通信量会大幅上升,因此这里不推荐大家同时使用流水线并行和 Zero-2 或 Zero-3,但流水线并行可以和 Zero-1 同时使用。 若在 Zero-1 基础上再做 Zero-2,可以将梯度也进行 DP 维度的拆分,梯度的显存占用由 2.74GB 降低到 0.34GB。 相对而言,使用 Zero-2、Zero-3 对梯度和参数做拆分时的显存开销优化不明显,同时会增加非常多的通信 overhead。所以建议如果不是显存特别紧张的条件下,尽量不要用 Zero-3。 另外,值得注意的是,我们并不推荐使用 PP 后再叠加 ZeRO-2/3,主要原因是一方面它对显存的进一步优化不是很明显;另一方面是它会带来一定的通信冲突,所以建议大家使用 TP 和 PP 后用 Zero
ZeRO-2:每个设备保存完整模型参数。梯度计算完一层后立即通过Reduce-Scatter分发,每个设备只保留自己负责的梯度分片。ZeRO-3:模型参数、梯度、优化器状态全部切分。 ZeRO-2=字典还是每人一本,但"笔记本"和"错题集"(梯度)都拆成4份。ZeRO-3=连字典也拆成4份!每人只带1/4字典+1/4笔记+1/4错题集。需要查某个字时,临时向拥有该部分的同学借阅。 【进阶】答案:ZeRO通过逐级切分实现显存优化:ZeRO-1切分优化器状态(单卡=4Φ+12Φ/N),ZeRO-2额外切分梯度(2Φ+14Φ/N),ZeRO-3全部切分(16Φ/N)。 详细说明:要点说明ZeRO-1只切分优化器,通信开销最小ZeRO-2梯度Reduce-Scatter后只保留自己的分片ZeRO-3前向时All-Gather临时收集参数,计算后丢弃权衡级别越高,显存越省 ZeRO-3虽然最省显存,但通信开销显著增大(每次前向/反向都需All-Gather),在机内NVLink带宽下ZeRO-2通常是更好的平衡点。
5.2原文核心要点ZeRO(ZeroRedundancyOptimizer)分三个级别:ZeRO-1切分优化器状态,ZeRO-2额外切分梯度,ZeRO-3进一步切分模型参数。 ZeRO-2:词典还是每人一本,但笔记本和参考书都只各买1套分着保管。书包更轻了,沟通量还是一样。ZeRO-3:连词典都只买1套,700页拆成8份每人87页。 ZeRO-3相对ZeRO-2的额外通信量比例:CZeRO-3−CZeRO-2CZeRO-2=3Φ−2Φ2Φ=50%\frac{C_{\text{ZeRO-3}}-C_{\text{ZeRO-2}}}{C (Φ=7B)在8张GPU上训练为例:CZeRO-1=CZeRO-2=2×7B×2B/param=28GBC_{\text{ZeRO-1}}=C_{\text{ZeRO-2}}=2\times7\text (进阶)答案:ZeRO-1切分优化器状态,ZeRO-2额外切分梯度,ZeRO-3进一步切分参数。
在去年 9 月份的博客中,微软这么介绍 ZeRO-Offload: ZeRO-Offload 继承了 ZeRO-2 的优化器状态和梯度分割。 但与 ZeRO-2 不同的是,ZeRO-Offload 不在每块 GPU 上保持优化器状态和梯度的分割,而是将二者卸载至主机 CPU 内存。
DeepSpeed •微软开发的开源深度学习优化库 •PyTorch框架 零冗余优化器 ( Zero Redundancy Optimizer,ZeRO) •优化器状态分区(ZeRO-1) •梯度分区 (ZeRO
为了评估优化技术,该研究使用 DeepSpeed 逐一启用如下优化:ZeRO-2、ZeRO-3、offloading、激活重计算、量化和 FlashAttention,以衡量性能改进以及时间和内存消耗方面的下降 对于具有卸载功能的 ZeRO-2 和 ZeRO-3,他们分别将优化器状态和优化器状态 + 模型卸载到 CPU RAM。对于量化,他们使用了具有双重量化的 4bits 配置。
(二)ZeRO-2:梯度与参数分割ZeRO-2 在 ZeRO-1 基础上,进一步对梯度和参数进行分割处理。不仅优化器状态,模型的梯度和参数也被分散存储在不同 GPU 上。
3.1.2 ZeRO三个级别 ZeRO-1:优化器状态分片 ZeRO-2:梯度分片 ZeRO-3:参数分片 3.2 ZeRO-1:优化器状态分片 3.2.1 工作原理 ZeRO-1将优化器状态(如Adam -2:梯度分片 3.3.1 工作原理 ZeRO-2在ZeRO-1的基础上,进一步对梯度进行分片存储,减少每个进程存储的梯度数量。 3.3.2 内存节省计算 ZeRO-2:每个进程仅存储1/N的优化器状态和1/N的梯度 内存节省:与ZeRO-1相比,额外节省约(N-1)/N的梯度内存 总体内存节省:对于Adam优化器,总内存使用可减少约 4/3*(N-1)/N倍 3.3.3 实现示例 # DeepSpeed ZeRO-2配置 zero_config = { "zero_optimization": { "stage , training_data=dataset, config=zero_config ) 3.4 ZeRO-3:参数分片 3.4.1 工作原理 ZeRO-3是ZeRO的最高级别,它在ZeRO
zero-offload 主要结合了zero-2的训练策略,每个GPU上保留一份参数副本,将梯度和优化器分片后,由不同的GPU进行通信和更新。 GPU 仅将其分区中经过平均的梯度卸载到 CPU 内存 • 梯度在 CPU 上可用后,各数据并行进程直接在 CPU 上并行更新优化器状态分区 • 更新后,参数分区移回 GPU,随后在 GPU 上执行类似 ZeRO
技术特点 并行训练:支持 DDP 或 DeepSpeed ZeRO-2/3 进行训练,并支持数据并行和张量并行。 模型评估:使用 vLLM 进行模型评估,确保评估过程的高效性和准确性。
修复 ZeRO-1 / ZeRO-2 CPU-offloaded gradient norm 问题 更新中明确提到: • 修复 ZeRO-1 和 ZeRO-2 在 CPU offload 场景下的 gradient
40% 训练优化 FlashAttention-2 + 梯度累积 显存占用降低60% 数据处理 SentencePiece + 中文清洗策略 中文分词准确率提升35% 分布式训练 DeepSpeed ZeRO
ZeRO有三个优化阶段: ZeRO-1:优化器状态分区 ZeRO-2:优化器状态+梯度分区 ZeRO-3(也称为FSDP):优化器状态+梯度+参数分区 ZeRO的理念是在数据并行进程之间,对这些对象进行分区 ZeRO-2:增加梯度分区 在反向传播中,不再对梯度执行all-reduce操作,而是仅执行reduce-scatter操作。只传播1/N_d的梯度,与ZeRO-1相比,节省了更多内存。 反向传播与之类似,只是流程相反: 由于需要在前向传播和反向传播中持续进行all-gather操作,与ZeRO-2相比,每个训练步骤会增加(2×层数-1)次额外的操作。
zero优化——Model States k=12 Nd=64(64卡的数据并行)model=8.3B zero-0 8.3*16=133GB zero-1 4*8.3+12*8.3/64=35GB zero 3、百亿参数规模 腾讯云TI平台团队使用黑石2.0 64卡训练83亿和100亿参数规模 GPT-2, 通过使用zero-2纯数据并行,每卡可以达到近40TFLOPs的性能。