Hugging Face高效训练技术二：大模型分布式训练策略——ZeRO、FSDP

• 《Hugging Face高效训练技术一：单 GPU 高效训练（Gradient Accumulation、Gradient Checkpointing、混合精度训练、优化其选择）》
• 《Hugging Face高效训练技术二：大模型分布式训练策略——ZeRO、FSDP》
• 《Hugging Face高效训练技术三：huggingface DeepSpeed文档》
• 《Hugging Face高效训练技术四：多GPU分布式训练（DP、PP、TP 、ZeRO）》
• 《Hugging Face高性能技术五：Transformer高效推断（bitsandbytes、FlashAttention、 BetterTransformer）》

一、ZeRO（零冗余优化器）

• DeepSpeed开源项目、DeepSpeed 官网 、huggingface DeepSpeed文档、知乎deepspeed入门教程
• 论文：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• 《论文解读系列第十三篇：ZeRO——面向万亿级参数的模型训练方法》

  现有普遍的数据并行模式下的深度学习训练，每一台机器都需要消耗固定大小的全量内存，这部分内存和并不会随着数据的并行而减小，因而，数据并行模式下机器的内存通常会成为训练的瓶颈。这篇论文开发了一种新颖的解决方案Zero Redundancy Optimizer (ZeRO)，主要用于解决数据并行状态下内存不足的问题。

  ZeRO通过跨数据并行进程划分模型状态（参数，梯度和优化器状态），而不是复制它们，从而消除了数据并行进程之间的内存冗余。它在训练期间使用动态通信方式，以在分布式设备之间共享必要的状态，以保持数据粒度的计算粒度和通信量。ZeRO 支持的数据并行性可以适应任意大小的模型，只要聚合的设备内存（the aggregated device memory）足够共享模型状态即可。

  按照作者的优化方案，在64个GPU的数据并行的情况下，我们可以将一个75亿个参数的模型的内存消耗从每GPU 120GB减少到每GPU 1.9GB，同时通信开销变为原来的1.5倍；或者每GPU 16.6GB而不增加通信开销。

  通过这样的方式，我们可以以同样的内存运行更大的模型，或者将原来必须使用模型并行才能运行的大模型使用数据并行方式进行训练，以减少模型并行的额外开销。

1.1 背景

  深度学习领域的模型越来越大，这显著地增加了模型的准确性。在NLP领域，像Bert-large（0.3B）、GPT-2（1.5B）、Megatron-LM（8.3B）、T5（11B）这样的transformer大型模型已经出现。然而，要继续扩展模型大小（从数十亿到数万亿参数），我们遇到了训练这些模型的挑战——它们无法容纳在单个设备（GPU或TPU）的内存中，而且仅仅添加更多设备也无法有效扩展训练。

  现有的解决方案存在一些限制。基本的数据并行（DP,Data Parallelism）不会减少每个设备的内存占用，对于具有超过1.4B参数的模型，它在当前32GB内存的GPU上会耗尽内存。其他解决方案，如流水线并行（PP,Pipeline Parallelism）、模型并行（MP,Model Parallelism）、CPU-Offloading等，虽然有一些作用，但它们都在功能性、可用性以及内存和计算/通信效率之间做出了权衡，而这些方面对于大规模的高速训练都是非常重要的。

  目前模型并行（MP）是其中一个有希望的解决方案，现在的文献中的大模型都采用了MP，但MP的扩展能力也有限。MP将模型在垂直方向上分割，将每层的计算和参数分配到多个设备上，这需要在每一层之间进行大量通信。因此，MP在单个节点内运行效果良好，但在单个节点之外的效率迅速下降。我们使用Megatron-LM在两个DGX-2节点上测试了一个40B参数模型，观察到每个V100 GPU的性能只有5Tflops（不到硬件峰值的5%）。

1.2 深度学习内存消耗分析

参考《大模型-LLM分布式训练框架总结》

  一个15亿参数的GPT-2模型在16位训练中的权重为3GB内存。但在使用Tensorflow或Pythorch这样的框架时，它无法在具有32GB内存的单个GPU上进行训练。那么在训练过程中，这些内存都消耗在了哪里呢？主要在是以下两个部分：

1. model states（模型状态）：大部分内存用于存储这一部分对象，其中包括优化器参数(比如Adam中的动量和方差)、梯度、模型参数等。

  混合精度训练（mixed precision training）和Adam优化器基本上已经是训练语言模型的标配。前者在前向传播和反向传播过程中使用fp16（半精度）表示的参数和激活值进行计算；而在计算和应用参数更新时，通常需要将这些值转换回32位浮点数来进行计算，以避免数值精度损失，所以混合精度训练中需要同时保留fp16和fp32副本。

  Adam需要存储动量和方差两种优化器状态，以计算参数更新；此外，还需要存储梯度和模型权重本身的信息。所以对一个具有Ψ个参数的模型进行混合精度训练时，其内存需求为：

• fp16格式的参数和梯度，都是2Ψ字节
• fp32格式的优化器的状态信息，包括参数、动量和方差，都是4Ψ字节

  我们用KΨ来表示存储优化器状态信息所需要的内存数，则对于混合精度Adam优化器而言，总的内存需求为(2+2+k)Ψ=(2+2+4+4+4)Ψ=16Ψ字节，所以对于一个拥有15亿参数的模型（GPT-2），其内存需求至少为24GB，远远超过只保存fp16参数所需的3GB内存。

2. residual states（残余状态）：剩余内存用于存储这部分对象，其中包括激活函数、临时缓冲区、不可用的碎片化内存。

  例如对于一个拥有15亿参数的GPT-2模型，在序列长度为1K和批量大小为32的情况下，即使使用激活值检查点（Activation checkpointing）技术来减少其内存占用，但最终模型的激活值内存消耗还是有8GB左右。如果是1000亿参数的模型，这部分内存消耗将达到60G。另外用于存储中间结果的临时缓冲区会消耗6GB内存（GPT-2）。

  此外，如果没有足够的连续内存来满足内存请求，即使总可用内存大于请求的内存量，请求内存的操作也会失败，因为在训练非常大的模型时，会发生显著的内存碎片化，导致在一些极端情况下，即使仍然有超过30%的内存可用，也会出现内存不足的问题。

1.3 数据并行（Data Parallelism）

参考《Framework(二)：分布式训练》

  针对上述内存消耗巨大的情况，通常有两种解决方法：数据并行性（DP）和模型并行性（MP）。

  典型的数据并行实现（PyTorch的DDP和TF的tf.distribute.MirroredStrategy）来说，每个GPU通常存储的是完整的模型，包括权重和参数。然后将不同的数据批次分配给不同的GPU来进行处理，接着在所有GPU之间共享梯度信息，以更新模型的权重和参数。因为多个GPU可以同时处理不同的数据批次，而无需等待其他GPU完成它们的工作，所以这种方法可以提高训练速度，整个过程如下图所示：

整个数据并行训练过程可以用以下步骤来总结：

1. 初始化：在每个GPU上创建一个模型的副本（replica），每个副本都有相同的初始权重和参数。
2. 训练迭代：每个GPU接收其分配的数据批次，并在本地（每个GPU本身）计算损失函数的梯度（local gradients）。
3. 全体度求和（All-Reduce）：为了保持模型权重的一致性并允许模型更新，需要将来自各个GPU的local gradients汇总为一个global gradients，这个操作称为all-reduce，其包含两个操作，reduce-scatter和all-gather：
   

   • Reduce-Scatter：

     • 每个节点（replica）的local gradients都被分成不同的块或分片（blocks or shards，通常是均匀分配的），然后在N个 节点之进行 N-1 轮数据交换，这使得每个节点获得了其他节点的一部分梯度信息。
     • 每一轮的交换都会将部分梯度数据合并，以获得一部分全局梯度信息。最终每个 replica 都会持有来自所有其他 replica 的梯度分片的汇总（fully reduced data）

     reduced是指每一轮的合并操作，使得shards减少。fully reduced就是指最终所有的梯度shards都被合并成一个值。

   • All-gather：

     • 每个 replica 将在 Reduce-Scatter 阶段中获得的 fully reduced data再次分成多个小块，然后将其 广播其它节点。同样经过N-1轮数据交换，每个节点都获得了其他节点的全部梯度信息。
     • 每个节点根据收集到的全部梯度信息汇总为global gradients。
4. 基于global gradients进行模型参数更新（weight update）
5. 重复步骤2至4，直到训练完成。

  在 Reduce-Scatter阶段，之所以要先将梯度数据被分成不同的shards，再分N-1次进行广播，而不是在一轮之内广播完成，是为了减少通信开销和内存占用，同时提高训练速度和可扩展性。

  具体来说，在大规模训练中，GPU的数量可能非常大，一次性将不同GPU的local gradients广播给所有GPU可能导致大量的数据流通信，也会占用大量的GPU内存。通过分桶操作，每次只传递local gradients的部分数据，可以减轻网络负载和内存压力，这样训练可以扩展到更多的GPU或更大的模型规模。

数据并行和模型并行的比较：

• 数据并行性（DP）：

  • 优点：具有良好的计算和通信效率，适用于多设备训练
  • 缺点：内存冗余严重。因为它在所有数据并行进程之间复制整个模型状态，导致冗余的内存消耗。

• 模型并行性（MP）：

  • 优点：内存效率较高，通过分割模型状态来减少内存消耗。
  • 缺点：计算和通信效率不佳，因为垂直切分模型会导致计算过于细粒度化，不利于规模化训练。

1.4 主要方法

1.4.1 ZeRO-DP优化模型状态内存

  数据并行和模型并行都保持了整个训练过程中所需的所有模型状态，但并不是所有时候这都是必需的。例如，仅在某个层的正向传播和反向传播期间才需要与每个层对应的参数。

  ZeRO-DP是一种改进的数据并行性方法，它通过对参数（包括优化器状态、梯度和参数）进行分区来消除内存冗余，使得每个GPU仅保存部分参数及相关状态，提高了内存效率；同时还通过在训练过程中使用动态通信来保持计算和通信效率。

图1：使用Adam优化器进行混合精度训练时，ZeRO-DP优化的三个阶段中，每个设备的model states内存消耗。

参数解释：

• Baseline：未优化的基线
• Ψ：模型大小，上图假设模型参数为Ψ=75亿
• K：存储优化器状态要消耗的内存倍数，上一节讲过，对于混合精度的Adam优化器而言，K=12
• $N_d$：数据并行度。基于Adam优化器的混合精度训练，数据并行度为Nd=64（即64个GPU）

上图展示了ZeRO-DP对数据并行优化的三个阶段：

1. 优化器状态分割（$P_{os}$）：
   在每个gpu中保存全部的参数和梯度，但是只保存1/Nd的优化器变量。通过将优化器状态进行分割，实现4倍的内存减少，同时保持与DP相同的通信量。

2. 梯度分割（$P_{os+g}$）：
   每个gpu中只保存1/Nd的梯度，实现8倍的内存减少，并保持与DP相同的通信量。

3. 参数分割（$P_{os+g+p}$）：
   每个gpu中只保存1/Nd的参数 ，实现64倍的内存减少，通信量会略微增加50%。作者通过用少量的计算的成本和通信成本换来了大幅的内存节省。

  上图显示了在不同数据并行度下（不同数量的GPU），每个设备内存消耗情况。可见，当启用了所有三个阶段时，ZeRO可以在仅使用1024个NVIDIA GPU的情况下训练具有万亿参数的模型（使用1024个GPU训练1T Model时，$P_{os}+g+p$ =15.6GB，当时的V100足有32GB）。

  上表是不同模型并行度和gpu下可实现的最大模型（以参数计量），最右边是作者的实现的测量值，左边是理论值。因此，这说明作者提出的内存计算是基本可靠的。

1.4.2 ZeRO-R优化残余状态内存

  在ZeRO-DP提高了模型状态内存效率之后，主要消耗在激活值、临时缓冲区以及无法使用的内存碎片这三个方面的剩余内存成为次要的内存瓶颈。为了解决这个问题，我们开发了ZeRO-R来进行优化：

1. 通过激活值分区来优化激活值内存
   我们使用Activation checkpointing技术来优化激活值内存，但这对于大模型仍旧是不足的。ZeRO-R通过激活值分割来识别和删除现有MP方法中的激活值复制，从而优化激活值内存。同时，它还会在适当的情况下将激活值卸载到CPU上，以释放GPU内存（CPU-offload）

     在训练过程中，激活值可能占用大量内存 。例如GPT-2有15亿参数，在序列长度为1K和批量大小为32的情况下，大约需要60GB的内存 。Activation checkpointing通过牺牲33%的重新计算开销来减少激活值内存占用（$N\to \sqrt{N}$），这将把该模型的激活内存消耗降低到约8GB左右。
   
     尽管Activation checkpointing显著减少了内存占用，但对于更大的模型，其内存开销依旧很大。例如，一个类似GPT的1000亿参数的模型，在批量大小为32的情况下，即使使用了Activation checkpointing，也需要约60GB的内存。

2. 恒定临时缓冲区大小，以在内存和计算效率之间取得平衡。
   缓冲区的大小通常会随着模型的规模变化而变化，但ZeRO-R采用了固定大小的缓冲区，这样可以防止随着模型规模的增加而导致缓冲区过大。同时ZeRO-R会保持缓冲区足够大，以保证计算效率。

3. 根据张量的不同生命周期来管理内存，以防止内存碎片化
   训练时产生的内存碎片化，是由于不同张量的生命周期差异引起的。内存碎片化可能导致内存分配失败，即使总内存足够也可能出现问题，因为无法获得足够的连续内存。

  有关训练更具体的内容，可以参考论文第7章Communication Analysis of ZeRO-DP和第8章Communication Analysis of ZeRO-R

1.5 总结

  ZeRO（Zero Redundancy Optimizer）是一种用于优化大规模深度学习模型训练的技术。它的主要目标是降低训练期间的内存占用、通信开销和计算负载，从而使用户能够训练更大的模型并更高效地利用硬件资源。

ZERO论文首先分析了模型训练中内存主要消耗在两个方面：

1. model states：模型状态，包括包括优化器参数（例如Adam的动量和方差）、梯度、模型参数
2. residual states：剩余状态，包括包括激活函数、临时缓冲区、内存碎片

  ZERO分别使用ZeRO-DP和ZeRO-R来优化model states和residual states。其中，ZeRO-DP包括三个阶段（以64GPU的混合精度训练举例，采用Adam优化器）：

1. 优化器状态分割（$P_{os}$）：在每个gpu中保存全部的参数和梯度，但是只保存1/Nd的优化器状态变量。此操作保持与DP相同的通信量，但是内存减少至原先的1/4；

2. 梯度分割（$P_{os+g}$）： 每个gpu中只保存1/Nd的梯度，通信量不变，内存减少至原先的1/8

3. 参数分割（$P_{os+g+p}$）： 每个gpu中只保存1/Nd的参数 ，通信量增加50%，但是内存大幅减少至原先的1/64（通过用少量的计算的成本和通信成本换来了大幅的内存节省）。

  ZeRO-R通过激活值分割来优化激活值内存；固定缓冲区大小以防止随着模型规模的增加而导致缓冲区过大；根据张量的不同生命周期来管理内存，以防止内存碎片化。

  通过上述一系列措施，ZeRO优化了模型训练过程中的内存消耗，并使用动态通信来保持计算和通信效率，使得我们可以用有限的计算资源来训练更大规模的模型。

1.6 官方视频：ZeRO & Fastest BERT，提高 DeepSpeed 深度学习训练的规模和速度（有空再补）

• YouTube视频：ZeRO & Fastest BERT: Increasing the scale and speed of deep learning training in DeepSpeed

  人工智能的最新趋势是更大的自然语言模型提供更好的准确性；然而，由于成本、时间和代码集成的难易程度，较大的模型很难训练。为了通过提高全球模型开发人员的规模、速度、成本和可用性来推进大型模型训练，微软于 2020 年 2 月开源了 DeepSpeed 库。过使用名为 ZeRO的内存优化系统，DeepSpeed 可以有效地训练具有 100-2000 亿个参数的模型，速度比最先进的技术快 10 倍。

在本次网络研讨会中，DeepSpeed 团队将讨论 ：

• DeepSpeed 功能、速度和规模优化以及未来路线图
• 如何使用 DeepSpeed 训练您自己的模型以及其他流行模型（例如 BERT 和 GPT-2）
• 深入探讨 ZeRO 优化器背后的技术和即将推出的功能，以及研究人员利用这些突破创建了Turing-NLG（最大的公开语言模型之一，拥有 170 亿个参数）
• 我们如何使用这项技术创造 BERT 训练世界纪录

资源列表：

• DeepSpeed Website 、 DeepSpeed Library (GitHub)
• 论文： ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (publication)
• ZeRO 和 DeepSpeed：新的系统优化支持超过 1000 亿个参数的训练模型 (blog)
• ZeRO-2 和 DeepSpeed：打破深度学习速度和规模的障碍 (blog)
• DeepSpeed Fastest Bert deep dive (blog)
• Turing-NLG：微软的 170 亿参数语言模型 (blog)
• AI at Scale 大规模人工智能 (Project Page)
• ONNX runtime (GitHub)

二、ZeRO-Offload

论文：《ZeRO-Offload：Democratizing Billion-Scale Model Training.》

2.1 概述

  大规模模型训练通常需要复杂的模型重构和昂贵的GPU集群资源。ZeRO-Offload通过将数据和计算卸载到CPU来实现大规模模型训练（包括模型参数、梯度和优化器状态等模型状态），其优化目标有三个方面——最小化数据在GPU和CPU之间的移动、减少CPU计算时间以及最大化地节省GPU内存，从而提高训练效率。

  ZeRO-Offload能够在单个GPU上训练具有超过130亿参数的模型，比PyTorch的模型规模大了10倍，而且无需修改模型或牺牲计算效率。ZeRO-Offload还可以在多个GPU上扩展，最多支持128个GPU，提供接近线性的加速。此外，它还可以与模型并行一起工作，以在单个DGX-2计算机上训练具有超过700亿参数的模型，相比仅使用模型并行，可训练的模型规模增加了4.5倍。

  通过将计算和内存效率与易用性相结合，ZeRO-Offload使大规模模型训练变得更加平民化，即使你只有一个GPU，也可以进行大规模模型训练。

2.2 算法

  由于CPU的计算吞吐量比GPU的计算吞吐量慢多个数量级，因此，为了避免CPU计算成为瓶颈，我们必须避免将计算密集型组件转移到CPU。

  模型训练的每次迭代的计算复杂度主要为 O(MB)，其中M表示model size，B表示batch size。前向传播和反向传播，两者的计算复杂度都是O(MB)，必须在GPU上执行。而其他复杂度低于O(MB)的计算，例如 norm calculations和weight updates，其复杂度为O(M)，可以被转移到CPU上。基于此，我们将数据流图中的前向和反向节点合并成一个单一的超级节点（FWD-BWD）并将其分配到GPU上，示意图如下：

总结：将FWD,BWD放到GPU，norm calculations,weight updates放到CPU。

  ZeRO-Offload的策略是特别为使用Adam优化器进行混合精度训练的大型模型训练而设计的，这是目前大型模型训练的常用方法。

2.3 实验

图7显示了在单个DGX-2节点中，使用1个（32GB内存）、4个和16个GPU进行模型规模测试。

• 单个GPU：使用PyTorch DDP训练的最大模型大小为1.4B。ZeRO-Offload通过将昂贵的状态（如优化器状态和大部分梯度）卸载到CPU内存来最大化GPU内存的节省，最终可以在单个GPU上进行13B模型训练。
• 单个DGX-2节点中的多GPU：
  • PyTorch和L2L：最大可训练模型大小都保持不变，因为它们都不处理数据并行中的内存冗余，它们的可扩展性受限于单个GPU上的模型规模
  • MegaTron和ZeRO-2：都支持更多GPU的大模型训练，但即使使用16个GPU，它们也无法高效地扩展到超过15B参数的模型
  • ZeRO-Offload：通过将优化器状态和梯度分区和卸载到CPU内存，然后与模型并行结合使用，轻松实现了高达70B参数模型的训练。

  图8显示，ZeRO-Offload的吞吐量（TFLOPS）平均比L2L高出14%（最高可达22%）。这主要源于ZeRO-Offload在CPU与GPU之间的数据通信成本较低。

三、ZeRO-Infinity

论文：《 ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning》

3.1 概述

  在过去的三年里，最大深度学习模型已经增长了1000多倍，达到了数千亿参数级别，而GPU内存仅增长了5倍（从16GB到80GB）。因此，模型规模的增长主要是通过系统创新来支持的。然而，我们已经接近了GPU内存的极限。为了训练一个拥有1万亿参数的模型，需要800个NVIDIA V100 GPU，而这样的集群对于大多数数据科学家来说根本不可及。此外，以这种规模训练模型需要复杂的并行技术组合，这对数据科学家来说是一项巨大的工作，需要对他们的模型进行重构。

  现代GPU集群在存储器方面非常异构。除了GPU内存外，它们还拥有CPU内存以及大规模的NVMe存储，后者的大小超过GPU内存的50倍，几乎比CPU内存大20倍。我们开发了ZeRO-Infinity，这是一种全新的异构系统技术，通过充分利用现代GPU集群中的这些异构内存系统（GPU、CPU和NVMe内存），在有限资源下能够训练前所未有规模的模型，而无需对模型代码进行重构。

  ZeRO-Infinity可以在单个NVIDIA DGX-2节点上微调具有1万亿参数的模型，比3D并行性增加了50倍。也可以在512个NVIDIA V100 GPU上维持超过25 petaflops的性能（达到峰值的40%）。

  NVMe内存是指使用NVMe接口的非易失性内存，通常是一种高速的、低延迟的存储介质，用于提供额外的内存资源，以帮助支持大规模深度学习模型的训练

3.2 算法

  下图是使用ZeRO-Infinity训练一个具有两层的模型，以及four data parallel (DP) ranks。图中下标表示层，上标表示DP ranks，例如$P_0^{(2)}$表示GPU2上的layer0参数。

在第一层的反向传播过程中，被分区的参数从较慢的内存移动到GPU上，并随后被收集以构建完整的层。在计算了梯度后，这些梯度会被聚合、重新分区，然后卸载到慢速内存中。

3.3 实验结果

图5：ZeRO-Infinity在训练几兆参数模型时的效率和可扩展性。

  图5a展示了ZeRO-Infinity在512个GPU上训练多达20万亿参数模型的性能。对于5000亿参数的模型（接近3D并行技术在这些资源上能够运行的最大规模），ZeRO-Infinity和3D并行技术实现了几乎相同的吞吐量，表明ZeRO-Infinity在训练效率上与最先进的技术不相上下。

  图5b显示，当训练1兆参数的模型时，ZeRO-Infinity从4个节点（64个GPU）到32个节点（512个GPU）实现了线性可扩展性（保持每个节点的批处理大小不变，并随着节点数量的增加增加总批处理大小）。

  图5c显示了在单个节点（16个GPU）上使用ZeRO-Infinity训练10亿到1兆参数模型的性能，而不需要模型并行。对于多达1000亿参数的模型，ZeRO-Infinity实现了每个GPU超过40 TFlops的出色性能，这使得只需一个NVIDIA DGX-2就可以微调GPT-3等大型模型成为可能。相比之下，3D并行技术无法扩展到超过200亿参数的模型。

  这些结果展示了ZeRO-Infinity的两个方面：

• 它可以在单个NVIDIA DGX-2节点上让用户轻松微调具有高达1兆参数的大型模型
• 易用性：在这种规模下，ZeRO-Infinity无需组合模型并行或管道并行，也无需对模型代码进行重构，这使得数据科学家可以轻松扩展他们的模型。

  a. 不同的训练策略可训练的模型规模。ZeRO-Infinity通过优化设备放置和分区策略，使模型规模从1.4B参数增加到1T参数，相对于数据并行技术提高了700倍。

  b. 最大隐藏层大小：内存中心瓷砖技术允许ZeRO-Infinity训练大型模型的大型隐藏层，最大可达64K，而无需模型并行。

  c. 梯度卸载：ZeRO-Infinity利用PCIe带宽卸载梯度，实现了比ZeRO-Offload更高的吞吐量。

  d. 预取和重叠：预取和重叠对于小批处理大小的性能至关重要，对大批处理大小的影响较小。

  e. 激活检查点卸载：CPU卸载激活检查点对小型隐藏层大小可能会降低吞吐量，但对大型隐藏层大小的影响较小。

  以上ZeRO的所有功能都在DeepSpeed中被实现。

四、FSDP: 使用更少的GPU进行更快的AI训练

• 论文《Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training》
• Fully Sharded Data Parallel: faster AI training with fewer GPUs、如何用数据并行训练万亿参数模型

4.1 简介

  Facebook发布的FSDP（Fully Sharded Data Parallel）是一种用于分布式深度学习模型训练的并行计算策略，它对标的是微软在DeepSpeed中提出的ZeRO，可以看成PyTorch中的DDP优化版本。

  FSDP的核心思想是将模型参数分成多个小片段（shards），并在多个设备上并行处理这些参数片段。每个设备只处理模型参数的一个子集，而不是整个模型参数。FSDP的主要特点和原则包括：

1. 参数分片：模型的每个参数都被分成多个小片段（shards），每个片段存储在不同的设备上。这样，每个设备只需要处理自己负责的参数片段，而不需要存储整个模型。

2. 异步通信：FSDP采用异步通信模式，允许不同设备上的参数片段在不同的时间进行前向和反向传播，而无需等待其他设备。这提高了训练的效率。

3. 内存效率：由于模型参数被分片存储，FSDP可以更好地利用每个设备的内存，允许训练非常大的模型。

4. 通信策略：FSDP使用通信策略来控制参数片段之间的信息交换。通信可以定期发生，以确保参数片段的同步。

5. 容错性：FSDP具有一定的容错性，即使某个设备失败，训练也可以继续进行，因为参数片段可以在其他设备上备份。

  FSDP旨在解决大规模深度学习模型训练中的内存和计算瓶颈问题，允许研究人员和工程师训练超大规模的模型，同时有效地利用多个设备的计算资源。它通常用于大型深度学习训练集群中，以加速训练过程并提高训练效率。

4.2 sharding weight update

  本文1.3章节中讲解了数据并行的标准操作，其中每个分布式节点（replicas）都拥有完整的模型参数副本，并在每个训练步骤中都会做重复的update weight操作。对于大型模型，这种全局参数同步可能会成为性能瓶颈，因为需要大量的计算和通信资源。即使对于小模型而言，为防止global batch size过大，每个分布式节点都会采用较小的batch size，此时update weight也会成为训练中的重要耗时项。

为了解决这个问题，谷歌在2020年提出了sharding weight update。它的具体步骤如下：

1. 参数分片（Sharding）：模型参数被划分成多个分片（Shard），通常是均匀分布的。每个分布式节点只负责更新自己持有的参数分片。

2. 本地更新：在每个分布式节点上，只有该节点持有的分片被用于计算梯度和更新参数。这减少了每个节点需要处理的参数数量，从而提高了计算效率。

3. 全局同步：在一定的训练间隔之后，分布式节点之间进行全局同步。在全局同步中，各个节点交换它们的参数分片，以确保每个节点最终都具有完整的全局模型参数。

  这种方法的优点是，它减少了在每个训练步骤中进行全局参数同步的频率，因为只有在全局同步点才会进行参数交换，所以减少了计算和通信的开销，特别是在大型模型和小批次大小的情况下。同时，它允许分布式节点在本地更新参数，提高了计算效率。

  如下图所示，经过reduce-scatter后每个replica得到一个gradient shard，每个replica先更新自己的shard的weight，然后再进行all-gather，这样其实是和原始的all-reduce是等价的。但是经过这个调整，每个replica只是update weight shard，耗时就会降低了，相当于update weight也被各个replica给分担了。 ​

4.3 optimizer states all-gather

  optimizer往往包含额外的参数，比如SGD包含梯度的动量，而Adam包含动量和方差，，这些参数可以统称为optimizer states（优化器状态），它们也是需要同步更新的。如果也对optimizer states进行all-gather的话，通信成本就会比较大（原始的all-reduce并不需要）。

  optimizer states通常只在参数更新时需要，而在前向传播和反向传播中并不需要，所以只在需要时才对optimizer states进行all-gather，可以减少通信开销和内存使用，这样整个过程就变成：

图4：两种在循环中分片化优化器辅助变量（sharding auxiliary variables）的方式。左图：只在迭代之间保留分片化的auxiliary。右图：在迭代中保留auxiliary和权重分片，并在前向/后向传递之前对权重进行全局收集。

  左图weight的all-gather是在update后立即进行的，而右图是在需要的时候（forward和backward）才进行all-gather。右图方案有更大的优化空间，因为在前后向传播过程中，可以采用低精度fp16来all-gather来得到所需要的全部weight，这样就大大降低了内存使用和通信成本。另外weight和auxliary weight的生存周期也减少了。特别是optimizer的auxliary weight，这样就节省一部分内存空间，可以用来存储前后向传播中的activations和gradients。

  假定模型参数大小是W，auxliary weight大小是V，前后向传播中activations和gradients的峰值大小是P，共有N个shards，那么训练的峰值大小就从W+V+P降低为max(W+V/N+P,W+V)，这带来的一个好处是Adam将和SGD一样高效（Adam比SGD要多一份方差）。

下图对比了standard DDP training 和 FSDP training ：

  在标准数据并行训练中，每个 GPU 上都有完整模型的一份拷贝。然后，在每个 GPU 上，仅对数据的一个分片进行一系列前向传播和反向传播的计算。在本地计算完成后，每个 GPU 的参数和优化器状态需要与其他 GPU 共享，以计算全局模型的权重更新。

  而在完全分片数据并行训练中，每个 GPU 上只有模型的一个分片。然后，在本地计算前向传播时，需要从其他 GPU 收集所有分片的权重信息，以完成前向传播计算。同样，在反向传播之前，再次执行权重的收集。在完成反向传播后，本地梯度需要通过减少散布步骤进行平均并分片到各个 GPU 上，以允许每个 GPU 更新其本地权重分片。

  这两种方法的区别在于参数和权重信息的共享方式以及计算的分布方式。完全分片数据并行训练试图减少每个 GPU 上的模型复制，但需要更复杂的权重信息的收集和分布，以确保每个 GPU 具有必要的信息来执行前向和反向传播。这个方法的选择通常取决于训练环境和硬件资源的特性。

4.4 FSDP的使用

FSDP特点如下：

• FSDP对parameters 和optimizer state进行sharding；
• 当reshard_after_forward=False，和PyTorch DDP通信成本一样，类似ZeRO-DP-2（optimizer state+gradient sharding DP，即两阶段的$P_{os+g}$）；
• 当reshard_after_forward=True，通信成本增加50%，类似ZeRO-DP-3（$P_{os+g+p}$），速度会慢，但是显存开销最小（ZeRO-DP详见本文1.3.1章节）。
• FSDP用8 GPUs可以训练13B parameter models，用128 GPUs可以训练175B parameter models。当设置cpu_offload=True，可以用256 GPUs训练 1T parameter models。
• FSDP只兼容pointwise Optimizers（Adam, AdamW, Adadelta, Adamax, SGD等），如果是non-pointwise Optimizers（Adagrad, Adafactor, LAMB等），sharding将得到稍微不一样的结果。

可以通过以下示例直接使用 FairScale 的 FSDP，只需替换 DDP(my_module)：

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
...
sharded_module = DDP(my_module)FSDP(my_module)
optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
for sample, label in dataload.next_batch:
  out = sharded_module(x=sample, y=3, z=torch.Tensor([1]))
  loss = criterion(out, label)
  loss.backward()
  optim.step()

4.5 FSDP和ZeRO的异同

  FSDP（Fully Sharded Data Parallel）和ZeRO（Zero Redundancy Optimizer）都是用于大规模深度学习模型训练的优化策略，都采用了内存优化，用于解决大型深度学习模型训练中的内存和计算瓶颈，以便在有限的硬件资源上训练更大规模的模型；但它们的焦点和目标略有不同：

1. 关注点：

   • FSDP：FSDP的主要关注点是模型参数的并行存储和计算，它通过将参数分成小片段并在多个设备上并行处理来降低内存占用。
   • ZeRO：ZeRO的主要关注点是优化训练中的冗余内存，特别是优化器状态、梯度和参数的内存占用。它的目标是在有限的设备内存上训练大规模模型。

2. 通信策略：

   • FSDP：FSDP使用异步通信来实现参数片段之间的信息交换，以提高训练效率。
   • ZeRO：ZeRO使用动态通信调度来减少通信量，同时保持计算粒度，以提高训练速度。它特别关注了通信中的内存重叠。

3. 组合使用：

   • FSDP：通常与数据并行一起使用，以实现参数的全分片和计算的全并行。
   • ZeRO：通常与数据并行一起使用，以减少模型参数的内存冗余，但也可以与模型并行结合使用，以降低激活内存。

4. 适用性：

   • FSDP：主要用于分布式训练和大规模模型的内存优化，适用于需要高度并行化和大规模训练的情况。
   • ZeRO：主要用于在有限设备内存上训练大规模模型，适用于需要降低内存占用的情况。