的关系为: 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
DeepSpeed •微软开发的开源深度学习优化库 •PyTorch框架 零冗余优化器 ( Zero Redundancy Optimizer,ZeRO) •优化器状态分区(ZeRO-1) •梯度分区 (ZeRO
在去年 9 月份的博客中,微软这么介绍 ZeRO-Offload: ZeRO-Offload 继承了 ZeRO-2 的优化器状态和梯度分割。 但与 ZeRO-2 不同的是,ZeRO-Offload 不在每块 GPU 上保持优化器状态和梯度的分割,而是将二者卸载至主机 CPU 内存。
为了评估优化技术,该研究使用 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 进行模型评估,确保评估过程的高效性和准确性。
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的性能。
ZeRO优化器有三个主要版本:ZeRO-1、ZeRO-2和ZeRO-3,分别针对不同级别的内存优化。 ZeRO-1:优化器状态分片,每个GPU只保存部分模型参数的优化器状态 ZeRO-2:优化器状态和梯度分片,每个GPU只保存部分模型参数的优化器状态和梯度 ZeRO-3:优化器状态、梯度和模型参数分片, 参数高效微调则相对简单 生态兼容性:考虑与现有代码库和工具链的兼容性 根据不同的场景,可以推荐以下优化组合: 单卡微调小型模型:混合精度训练 + 梯度累积 + 梯度检查点 多卡微调中型模型:DDP + 混合精度训练 + ZeRO 原始配置(仅使用DDP): 模型:LLaMA-2-13B 每张GPU批处理大小:1 训练速度:约1.5样本/秒 显存使用:每张GPU约30GB 优化后配置: DDP + ZeRO-2:优化器状态分片,减少显存使用
添加梯度分区(Partition gradients,P_{os+g} ):又称为 ZeRO-2,在优化器状态分区的基础上,对梯度也进行分区。每个进程只需存储用于更新自身参数分区所需的梯度。 ZeRO-2 内存分析ZeRO-2 在 ZeRO-1 的基础上进一步优化,通过对梯度(Grad)也进行切分,减少显存占用并提高通信效率。 ZeRO-3 内存分析ZeRO-3 在 ZeRO-1 和 ZeRO-2 的基础上进一步优化,通过对优化器状态、梯度和权重进行全面切分,最大化显存节约,这种优化使得训练超大模型成为可能。
8experts_per_tok: 2expert_hidden_size: 4096# 训练超参learning_rate: 5e-4batch_size: 32grad_clip: 1.0梯度优化技巧:ZeRO