[{"data":1,"prerenderedAt":58},["ShallowReactive",2],{"post-17582ba3":3},{"abbrlink":4,"title":5,"date":6,"categories":7,"tags":9,"excerpt":11,"html":12,"toc":13},"17582ba3","论文阅读: ZeRO-Offload: Democratizing Billion-Scale Model Training","2023-10-17",[8],"懵逼的深度学习",[10],"分布式训练框架","摘要 ZeRO-Offload 主要优化在于尽量减少数据在 GPU 与 CPU 之间的移动，并减少 CPU 计算时间，同时最大限度地节省 GPU 上的内存。","\u003Ch2 id=\"post-heading-摘要\">摘要\u003C\u002Fh2>\n\u003Cp>ZeRO-Offload 主要优化在于尽量减少数据在 GPU 与 CPU 之间的移动，并减少 CPU 计算时间，同时最大限度地节省 GPU 上的内存。\u003C\u002Fp>\n\u003Cp>原文链接：\u003Ca href=\"https:\u002F\u002Farxiv.org\u002Fpdf\u002F2101.06840.pdf\" target=\"_blank\" rel=\"noopener noreferrer\">https:\u002F\u002Farxiv.org\u002Fpdf\u002F2101.06840.pdf\u003C\u002Fa>\u003C\u002Fp>\n\u003Cp>开源代码：\u003Ca href=\"https:\u002F\u002Fgithub.com\u002Fmicrosoft\u002Fdeepspeed\" target=\"_blank\" rel=\"noopener noreferrer\">https:\u002F\u002Fgithub.com\u002Fmicrosoft\u002Fdeepspeed\u003C\u002Fa>\u003C\u002Fp>\n\u003Ch2 id=\"post-heading-背景\">背景\u003C\u002Fh2>\n\u003Cp>使用 Pytorch 训练内存不够，单个 V100 上性能为 30 TFlops\u003C\u002Fp>\n\u003Cp>目前的分布式训练框架动辄需要超过 10w 刀的设备，难以在单卡上发挥作用\u003C\u002Fp>\n\u003Cp>对于基于注意力机制的 LLM 训练，内存瓶颈主要在模型状态上\u003C\u002Fp>\n\u003Cp>现有的异构训练方法主要利用 CPU 内存，而忽略了 CPU 的计算潜力\u003C\u002Fp>\n\u003Ch2 id=\"post-heading-介绍\">介绍\u003C\u002Fh2>\n\u003Cp>ZeRO-Offload 优化：尽量减少数据在 GPU 与 CPU 之间的移动，并减少 CPU 计算时间，同时最大限度地节省 GPU 上的内存\u003C\u002Fp>\n\u003Cp>单个 V100 上性能提高到 40 TFlops，在 128 卡上实现接近线性加速，对比MP，模型规模增加了 4.5 倍\u003C\u002Fp>\n\u003Cp>ZeRO-Offload 利用了CPU内存和计算资源进行 Offload，并与 ZeRO-DP 相结合\u003C\u002Fp>\n\u003Ch3 id=\"post-heading-efficiency\">Efficiency\u003C\u002Fh3>\n\u003Cp>论文提出了一种名为Efficiency的offload策略，通过分析确定了CPU和GPU设备之间的最佳计算和数据划分策略，以在三个关键方面达到最优化：\u003C\u002Fp>\n\u003Cul>\n\u003Cli>在CPU上的计算量比GPU少多个数量级，防止CPU性能瓶颈；\u003C\u002Fli>\n\u003Cli>最小化CPU和GPU之间的通信量，防止通信瓶颈；\u003C\u002Fli>\n\u003Cli>在实现最小通信量的同时，可证明地最大化节约GPU内存。\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Cp>分析表明，要在上述方面达到最优，必须将梯度、优化器状态和优化器计算卸载到CPU上，同时在 GPU 上保留参数、前向和反向计算。这种策略使模型规模增加了10倍，并且通信量最小。能够在一块 V100 GPU上训练130亿参数模型，性能为 40 TFLOPS，而没有 CPU offload 的情况下只能训练12亿参数模型，且性能只有30 TFLOPS。\u003C\u002Fp>\n\u003Cp>offload 优化器计算要求CPU进行O(M)次计算，而GPU需进行O(MB)次计算，其中M和B分别为模型规模和 batch size 。在大多数情况下， batch size 较大，CPU计算量并不是瓶颈，但对于小 batch size，CPU计算量可能成为瓶颈。为了解决这个问题，采取了两种优化措施：\u003C\u002Fp>\n\u003Cul>\n\u003Cli>高效的CPU优化器，其速度比现有技术快6倍；\u003C\u002Fli>\n\u003Cli>延迟一步的参数更新，允许将CPU优化器步骤与GPU计算重叠，同时确保准确性。这两种措施共同保证了ZeRO-Offload在小 batch size 下也能保持效率。\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Ch3 id=\"post-heading-unique-optimal-offload-strategy\">Unique Optimal Offload Strategy\u003C\u002Fh3>\n\u003Cp>为了确定最佳的下载策略，ZeRO-Offload将深度学习训练建模为数据流图，并使用基于第一原理的分析方法将该图分割为CPU和GPU设备之间的部分。\u003C\u002Fp>\n\u003Cp>ZeRO-Offload 在以下三个关键方面优化了图的分割策略\u003C\u002Fp>\n\u003Cul>\n\u003Cli>要求更少数量级的CPU计算，以防止CPU成为性能瓶颈；\u003C\u002Fli>\n\u003Cli>确保最小化CPU和GPU内存之间的通信量；\u003C\u002Fli>\n\u003Cli>可证明地实现最小通信量的同时最大化内存节省。\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Cp>训练的计算复杂度通常为O(MB)，其中M为模型大小，B为有效batch size。为避免CPU计算成为瓶颈，只有那些计算复杂度低于O(MB)的计算才能转移到CPU上\u003C\u002Fp>\n\u003Cp>FWD 和 BWD 的计算复杂度都是O(MB)，必须在GPU上进行，而其余的计算，如范数计算、权重更新等，其复杂度为O(M)，可以转移到CPU上\u003C\u002Fp>\n\u003Cp>基于此，数据流图中的前向传播和反向传播节点合并为一个 FWD-BWD Super Node 并分配到 GPU 上\u003C\u002Fp>\n\u003Cp>还需要最小化 CPU 与 GPU 的通信带宽，如图中所示，最小通信量为 BWD后 GPU 发送到 CPU 的 2M 梯度与 CPU 发送到 GPU 的 2M 参数，只有将 fp32 模型状态（momentum 32、variance 32和p32），Param Update 和 float2half 计算放置在一起，为一个 CPU 上的 Update Super Node，才能达成最小通信量策略\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F1.webp\" width=\"800\"\u002F>\n\u003Cp>另外，上图中 2M 大小的 parameter 16 为什么分别到 FWD 与 BWD 有总和 4M 的数据传输，这个问题如有朋友知道，欢迎评论区讨论\u003C\u002Fp>\n\u003Cp>实验依据如下\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F2.webp\" width=\"500\"\u002F>\n\u003Cp>总之，ZeRO-Offload 策略就是将所有的 fp32 模型状态和 fp16 梯度存储在 CPU 内存中，并在 CPU 上计算参数更新。而 fp16 参数则保留在GPU上，前向和反向计算在GPU上进行\u003C\u002Fp>\n\u003Ch2 id=\"post-heading-zero-offload-schedule\">ZeRO-Offload Schedule\u003C\u002Fh2>\n\u003Ch3 id=\"post-heading-单卡策略\">单卡策略\u003C\u002Fh3>\n\u003Cp>ZeRO-Offload将数据进行分区，将fp16参数存储在GPU上，fp16梯度和所有优化器状态（如fp32动量、方差和参数）存储在CPU上。\u003C\u002Fp>\n\u003Cp>在训练过程中，首先通过 FWD 计算损失。由于fp16参数已经位于GPU上，因此这部分计算不需要与CPU进行通信。\u003C\u002Fp>\n\u003Cp>在 BWD 过程中，不同参数的梯度在后向调度的不同位置计算。ZeRO-Offload可以立即将这些梯度逐个或切分传输到 CPU 内存中。\u003C\u002Fp>\n\u003Cp>因此，在将梯度传输到CPU内存之前，只需要在GPU内存中临时保存少量的梯度。此外，每个梯度传输可以与反向计算重叠，消除大部分通信成本。\u003C\u002Fp>\n\u003Cp>在 BWD 之后，ZeRO-Offload在CPU上直接更新fp32参数和剩余的优化器状态（如动量和方差），并将更新后的 fp32 参数从 CPU 内存复制到 GPU 内存中的fp16参数中。\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F3.webp\" width=\"800\"\u002F>\n\u003Cp>上图展示了 ZeRO-Offload 每个步骤中的计算通信过程，下图伪代码展示了具体的调度过程\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F4.webp\" width=\"500\"\u002F>\n\u003Ch3 id=\"post-heading-多卡策略\">多卡策略\u003C\u002Fh3>\n\u003Cp>ZeRO-Offload 将梯度和优化器状态在不同的 GPU 之间进行 partition，并且每个 GPU 将自己的 part offload 到 CPU 内存中，存储持续整个训练过程\u003C\u002Fp>\n\u003Cp>BWD 过程中，在 GPU 上 reduce-scatter 计算梯度并平均，每个 GPU 仅将属于其 part 的平均梯度 offload 到 CPU 内存中\u003C\u002Fp>\n\u003Cp>一旦梯度在 CPU 上可用，优化器状态 part 对应的每个 DP 进程直接在 CPU 上并行更新对应的参数 part\u003C\u002Fp>\n\u003Cp>更新完成后，参数 part 发送到 GPU，在 GPU 上对参数进行类似 ZeRO-2 的 all-gather 操作\u003C\u002Fp>\n\u003Cp>下面是 ZeRO-Offload 的多卡数据布局示意图\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F5.webp\" width=\"800\"\u002F>\n\u003Cp>此外，ZeRO-Offload 还可以与 MP 和 Megatron-LM 同时使用，通过 offload MP 进程对应的梯度、优化器状态和优化器计算实现\u003C\u002Fp>\n\u003Ch2 id=\"post-heading-optimized-cpu-execution\">Optimized CPU Execution\u003C\u002Fh2>\n\u003Ch3 id=\"post-heading-implementing-the-cpu-optimizer\">Implementing the CPU Optimizer\u003C\u002Fh3>\n\u003Cp>优化 CPU Optimizer 性能\u003C\u002Fp>\n\u003Cul>\n\u003Cli>使用SIMD向量指令\u003C\u002Fli>\n\u003Cli>循环展开\u003C\u002Fli>\n\u003Cli>OMP 多线程并行计算\u003C\u002Fli>\n\u003Cli>混合精度训练\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Cp>在混合精度训练模式下，内存中的参数有两个版本\u003C\u002Fp>\n\u003Cul>\n\u003Cli>FP32 版本，用于优化器在 CPU 上进行更新\u003C\u002Fli>\n\u003Cli>浮点转换 FP32  后用于在 FWD 中计算激活的 FP16，用于GPU\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Cp>此外，梯度的动量和方差也以 FP32 保存在 CPU 上，防止在更新参数时出现精度丢失\u003C\u002Fp>\n\u003Cp>论文中的 adam 实现使用了 SIMD，通过将数据读入向量寄存器并使用多个融合乘加操作（FMA）构成主执行流水线，SIMD 矢量宽度由 simd_width 指定，并结合循环展开，自动调优的 unroll_width 为 8\u003C\u002Fp>\n\u003Cp>同时，对计算的数据按 tile_width 分区，将 CPU上计算好的 FP32 数据 part 转换为 FP16 拷贝到 GPU 中，同时在 CPU 上计算下一个 数据 part，实现通信计算 overlap\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F6.webp\" width=\"600\"\u002F>\n\u003Ch3 id=\"post-heading-one-step-delayed-parameter-update-dpu\">One-Step Delayed Parameter Update（DPU）\u003C\u002Fh3>\n\u003Cp>在 DPU 训练过程中，首先在前 N-1 个步骤中进行训练，避免在训练的早期阶段梯度变化较快时破坏训练的稳定性。\u003C\u002Fp>\n\u003Cp>在第N步，从 GPU 获取梯度，但跳过 CPU 优化器步骤，也不更新 GPU 上的 fp16 参数。\u003C\u002Fp>\n\u003Cp>在第 N+1 步，使用第N步的梯度在CPU上计算参数更新，在GPU上并行计算前向和后向传播，使用第 N-1 步更新的参数。\u003C\u002Fp>\n\u003Cp>从这一步开始，每个步骤的模型将使用从 i-1 步更新的参数进行训练，实现了 CPU 与 GPU 的计算重叠\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F7.webp\" width=\"1000\"\u002F>\n\u003Cp>作者通过对多个训练工作评估发现，如果在几十次迭代之后引入 DPU，不会对收敛产生影响\u003C\u002Fp>\n\u003Ch2 id=\"post-heading-实验\">实验\u003C\u002Fh2>\n\u003Ch3 id=\"post-heading-实验方法\">实验方法\u003C\u002Fh3>\n\u003Cp>硬件配置\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F8.webp\" width=\"500\"\u002F>\n\u003Cp>实验模型：GPT-2、BERT\u003C\u002Fp>\n\u003Cp>实验数据集：SQuAD\u003C\u002Fp>\n\u003Cp>对比框架：PyTorch DDP、Megatron、L2L、ZeRO\u003C\u002Fp>\n\u003Cp>模型配置\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F9.webp\" width=\"500\"\u002F>\n\u003Ch3 id=\"post-heading-实验结果\">实验结果\u003C\u002Fh3>\n\u003Cp>图7 在单个GPU，4个GPU和16个GPU（一个DGX-2节点）上可训练的最大模型参数量对比\u003C\u002Fp>\n\u003Cp>图8 在 batch size 为 512 的单个GPU上使用 PyTorch，L2L 和 ZeRO-Offload 训练的 TFLOPS 对比\u003C\u002Fp>\n\u003Cp>图9 使用 DPU 和不使用 DPU 的情况下，与 GPT-2 训练的 TFLOPS（batch_size = 8） 对比\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F10.webp\" width=\"800\"\u002F>\n\u003Cp>图10 针对不同模型大小，使用PyTorch、ZeRO-2、MegatronLM、ZeRO-Offload（无 MP）以及ZeRO-Offload（有 MP）训练的 TFLOPS 对比\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F11.webp\" width=\"600\"\u002F>\n\u003Cp>图11 针对不同 GPU 数量，ZeRO-2 和 ZeRO-Offload 的 TFLOPS 对比，从 64 卡开始 ZeRO-2 优于 ZeRO-Offload 是因为 ZeRO-2 没有 CPU-GPU 通信的额外开销\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F12.webp\" width=\"600\"\u002F>\n\u003Cp>本文优化过的 CPU-Adam 与 PT-GPU 性能差距并不大，不会成为性能瓶颈\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F13.webp\" width=\"600\"\u002F>\n\u003Cp>实验表明 DPU 策略在训练与微调上不会影响精度，最终的 F1 分数都是 0.92\u003C\u002Fp>\n\u003Cimg src=\"\u002Fimage\u002F论文阅读-ZeRO-Offload-Democratizing-Billion-Scale-Model-Training\u002F14.webp\" width=\"600\"\u002F>\n",[14,18,21,24,28,31,34,37,40,43,46,49,52,55],{"id":15,"text":16,"level":17},"post-heading-摘要","摘要",2,{"id":19,"text":20,"level":17},"post-heading-背景","背景",{"id":22,"text":23,"level":17},"post-heading-介绍","介绍",{"id":25,"text":26,"level":27},"post-heading-efficiency","Efficiency",3,{"id":29,"text":30,"level":27},"post-heading-unique-optimal-offload-strategy","Unique Optimal Offload Strategy",{"id":32,"text":33,"level":17},"post-heading-zero-offload-schedule","ZeRO-Offload Schedule",{"id":35,"text":36,"level":27},"post-heading-单卡策略","单卡策略",{"id":38,"text":39,"level":27},"post-heading-多卡策略","多卡策略",{"id":41,"text":42,"level":17},"post-heading-optimized-cpu-execution","Optimized CPU Execution",{"id":44,"text":45,"level":27},"post-heading-implementing-the-cpu-optimizer","Implementing the CPU Optimizer",{"id":47,"text":48,"level":27},"post-heading-one-step-delayed-parameter-update-dpu","One-Step Delayed Parameter Update（DPU）",{"id":50,"text":51,"level":17},"post-heading-实验","实验",{"id":53,"text":54,"level":27},"post-heading-实验方法","实验方法",{"id":56,"text":57,"level":27},"post-heading-实验结果","实验结果",1792849566659]