[源码分析] Facebook如何训练超大模型---(1)

源码分析 Facebook如何训练超大模型---(1)

[源码分析] Facebook如何训练超大模型---(1)
- 0x00 摘要
- 0x01 简介
  - 1.1 FAIR & FSDP
  - 1.2 大规模训练计算能力需求
- 0x02 FSDP 如何工作
  - 2.1 全参数分片
  - 2.2 比对
  - 2.3 梳理
    - 2.3.1 思路
    - 2.3.2 流程步骤
- 0x03 How to use FSDP
  - 3.1 在语言模型中使用FSDP
  - 3.2 在计算机视觉模型之中使用FSDP
  - 3.3 在PyTorch Lightning使用FSDP
  - 3.4 直接从FairScale使用FSDP库
- 0x04 内存管理
- 4.1 Optimizer State Sharding (OSS)
  - - 4.1.1 训练流程
    - 4.1.2 最佳实践
    - 4.1.3 性能
- 4.2 Optimizer Gradient State Sharding
  - - 4.2.1 训练过程
    - 4.2.2 最佳实践
- 4.3 Optimizer Gradient Horizontal Model Sharding
  - - 4.3.1 训练过程
    - 4.3.2 最佳实践
    - 4.3.3 性能
- 0xFF 参考

0x00 摘要

我们在前文介绍过，微软 ZeRO 可以对一个万亿参数模型可以使用 8 路模型并行、64 路管道并行和 8 路数据并行在 4,096 个 NVIDIA A100 GPU 上进行扩展。

而FSDP（Fully Sharded Data Parallel）是Facebook 深度借鉴微软ZeRO之后提出的PyTorch DDP升级版本，可以认为是对标微软 ZeRO，其本质是 parameter sharding。Parameter sharding 就是把模型参数等切分到各个GPU之上。我们会以 Google，微软和 Facebook 的论文，博客以及代码来进行学习分析。

本系列其他文章如下：

[源码解析] PyTorch 分布式之 ZeroRedundancyOptimizer

[论文翻译] 分布式训练 Parameter sharding 之 ZeRO

[论文翻译] 分布式训练 Parameter Sharding 之 Google Weight Sharding

0x01 简介

1.1 FAIR & FSDP

大规模训练人工智能模型并不容易。除了需要大量的计算能力和资源外，训练非常大的模型背后还有相当大的工程复杂性。Facebook人工智能研究（FAIR）工程部一直致力于构建工具和基础设施，以使大型人工智能模型的培训变得更容易。

Fully Sharded Data Parallel（FSDP）是FAIR引入的最新工具。它将AI模型的参数在数据并行worker之间进行切分，并且可以选择将部分训练计算卸载到CPU。顾名思义，FSDP是一种数据并行训练算法。尽管参数被分片到不同的GPU，但每个微批次数据的计算对于每个GPU worker来说仍然是本地的。这种概念上的简单性使FSDP更易于理解，并且更适用于各种使用场景（与层内并行和流水线并行相比）。与optimizer state gradient sharding数据并行方法相比，FSDP在训练过程中通过通信和计算重叠对模型参数进行更均匀的切分，具有更好的性能。

FSDP可以使用更少的GPU更有效地训练数量级更大的模型。FSDP已在FairScale库中实现，允许工程师和开发人员使用简单的API扩展和优化其模型的培训。在Facebook，FSDP已经被整合并测试，用于训练一些NLP和Vision模型。

1.2 大规模训练计算能力需求

大规模模型训练需要大量的计算资源，比如OpenAI的GPT-3 拥有1750亿个参数。其训练估计需要355年的GPU时间，相当于1000个GPU连续工作4个月以上。

除了需要大量计算和工程资源外，大多数的训练扩展方法都会带来额外的通信成本，并且需要工程师仔细评估内存使用和计算效率之间的权衡。例如，典型的数据并行培训要求在每个GPU上维护模型的冗余副本，而模型并行培训为在worker（GPU）之间移动激活引入了额外的通信成本。

相比之下，FSDP相对而言没有做任何权衡。它通过在GPU上分割模型参数、梯度和优化器状态来提高内存效率，并通过分解通信并将其与前向和后向过程重叠来提高计算效率。FSDP产生与标准分布式数据并行（DDP）培训相同的结果，并提供易于使用的接口，该接口是PyTorch分布式数据并行模块的替代品。Facebook 的早期测试表明，FSDP可以扩展到数万亿个参数。

0x02 FSDP 如何工作

在标准DDP训练中，每个worker处理一个单独的批次，并使用all-reduce对worker之间的梯度进行汇总。虽然DDP已经变得非常流行，但它占用的GPU内存比它实际需要的要多，因为模型权重和优化器状态在所有DDP worker中都有一个副本。

2.1 全参数分片

减少副本的一种方法是应用全参数分片（ full parameter sharding）的过程，其中仅提供局部计算所需的模型参数、梯度和优化器的子集。这种方法的一个实现 ZeRO-3 已经被微软所普及。

解锁全参数切分的关键是：我们可以把DDP之中的all reduce操作分解为独立的 reduce-scatter 和 all-gather 操作。

图来自：https://engineering.fb.com/wp-content/uploads/2021/07/FSDP-graph-2a.png?w=1024

“All-reduce”是“reduce-scatter”和“all-gather”的组合。聚合梯度的标准 “All-reduce”操作可以分解为两个单独的阶段：“reduce-scatter”和“all-gather”。

“reduce-scatter”阶段，在每个GPU上，会基于rank 索引对 rank 之间相等的块进行求和。
“all-gather”阶段，每个GPU上的聚合梯度分片可供所有GPU使用。

通过重新安排reduce scatter和all gather，每个DDP worker只需要存储一个参数分片和优化器状态。

2.2 比对

下图显示了标准DDP训练（上半部分）和FSDP训练（下半部分）：

在标准的数据并行训练方法中，每个GPU上都有一个模型副本，向前和向后传递的序列只在自己的数据分片上进行运行。在这些局部计算之后，每个局部过程的参数和优化器与其他GPU共享，以便计算全局权重更新。
在FSDP中：

代码语言：txt复制

- **Model shard** ：每个GPU上仅存在**模型的分片**。
- **All-gather** ：每个GPU通过all-gather从其他GPU收集所有**权重**，以在本地计算前向传播。就是论文思路Pp下划线部分。
- **Forward（local）**：在本地进行前向操作。前向计算和后向计算都是利用完整模型。
- **All-gather** ：然后在后向传播之前再次执行此**权重**收集。就是论文思路Pp之中的下划线部分。
- **Backward（local）**：本地进行后向操作。前向计算和后向计算都是利用完整模型，此时每个GPU上也都是**全部梯度**。
- **Reduce-scatter** ：在向后传播之后，局部**梯度**被聚合并且通过 reduce-scatter 在各个GPU上分片，每个分片上的梯度是聚合之后本分区对应的那部分，就是论文思路Pg之中的下划线部分。
- **Update Weight（local）**：每个GPU更新其局部**权重**分片。

为了最大限度地提高内存效率，我们可以在每层向前传播后丢弃全部权重，为后续层节省内存。这可以通过将FSDP包装应用于网络中的每一层来实现（通过设置reshard_after_forward=True）。

下面是伪代码实现：

代码语言：javascript复制

FSDP forward pass:
    for layer_i in layers:
        all-gather full weights for layer_i # 权重
        forward pass for layer_i
        discard full weights for layer_i # 权重

FSDP backward pass:
    for layer_i in layers:
        all-gather full weights for layer_i # 权重
        backward pass for layer_i
        discard full weights for layer_i # 权重
        reduce-scatter gradients for layer_i # 梯度

2.3 梳理

我们结合论文思路再来梳理一下 FSDP。

2.3.1 思路

论文思路如下：

Pp: Parameter Partitioning，每个进程只存储与其分区对应的参数。当正向和反向传播需要其分区外的参数时，会通过broadcast操作从适当的数据并行进程接收这些参数。虽然乍一看，这可能会导致显著的通信开销，但我们发现，这种方法只会将基线DP系统的总通信量增加到1.5倍，同时实现与Nd成比例的内存减少。
Pos : Optimizer State Partitioning，对于一个NdNd 并行度的DP来说，我们将优化器状态分组到NdNd 个相等的分区中，这样第i个数据并行进程只更新与第i个分区对应的优化器状态。因此，每个数据并行过程只需要存储和更新总优化器状态的$ frac{1}{N_d} ，然后只更新，然后只更新 frac{1}{N_d}$ 个参数。在每个训练步骤结束时，我们会执行一个跨数据并行进程的all-gather操作，以获得跨所有数据并行进程的完全更新的参数。
Pg: Gradient Partitioning，由于每个数据并行进程只负责更新其相应的参数分区，因此，每个节点仅仅对自己负责的那部分参数的梯度进行规约。在归并之后，每个节点只需要自己参数分区对应的梯度，对于其他的梯度不再需要，所以它们的内存可以被释放。这将梯度的内存占用从2ψ字节缩减到 2ψNd2ψNd 。实际上，这是一种 Reduce-Scatter操作，不同参数的梯度被减少到不同的进程之中。

总结一下：因为模型参数被分区，所以参数梯度（在框架实现中，梯度往往是参数的成员变量）自然就被分区了。分区的参数被设置到优化器之中，所以优化器只会优化本分区的参数，所以优化器状态自然就是分区之后的。注意，在前向传播和后向传播时候，每个GPU都是用全部模型来计算，得到的梯度也是全部的梯度，只是存储时候只存储自己分区对应的部分。

2.3.2 流程步骤

我们再来展示一下具体流程。假设数据并行度为 n，则有 n 个GPU，那么每个GPU之上保存总模型参数的 1/n，同时梯度，优化器状态就自然被分区了，每个GPU之上还有数据并行。

起始状态：每个GPU之上是P_n, G_n, O_n 。注意，因为本GPU上模型是P_n ，所以O_n 自然就对应了P_n ，就自动分片了。
正向计算时候，每个 GPU_n 都把自己负责的参数P_n 广播给其他所有的 GPU，前向计算之后，每个 GPU_n 都得到自己输入训练数据 data_n 的损失 loss_n 。
反向计算时候，每个 GPU_n 也都把自己负责的参数 P_n 广播给其他所有的 GPU，最后计算得到对应于数据 data_n 的梯度 G_n 。
将梯度 G_n 聚合到对应的GPU_n 上，这时候GPU_n 上的梯度就是 reduce(G_0, ..., G_n) 之中自己rank对应的部分。注意，梯度聚合过程则使用了 reduce-scatter，因为每个gpu只需要更新自己负责的部分P_n,G_n,O_n ，所以不需要进行all-gather了。

0x03 How to use FSDP

目前，FAIR 提供四种解决方案来使用FSDP，以适应不同的需求。

3.1 在语言模型中使用FSDP

对于语言模型，可以在通过以下新参数，在 fairseq framework 之中支持 FSDP：

–ddp-backend=fully_sharded: 通过FSDP启用完全切分。
–cpu-offload: 将优化器状态和FP32模型副本卸载到cpu（与–optimizer=cpu_adam结合使用）。
–no-reshard-after-forward: 提高大模型训练速度 (1B params) ，类似于 ZeRO stage 2。
其他常见选项 (–fp16, –update-freq, –checkpoint-activations, –offload-activations, etc.) 还是继续正常工作。

具体请参阅fairseq教程。

3.2 在计算机视觉模型之中使用FSDP

对于计算机视觉模型， VISSL 中可以支持FSDP，并在regnet架构上进行了测试。像BatchNorm和ReLU这样的层已经被无缝地处理并已经测试过其收敛性。可以使用下面选项来启用 FSDP。

config.MODEL.FSDP_CONFIG.AUTO_SETUP_FSDP=True
config.MODEL.SYNC_BN_CONFIG.SYNC_BN_TYPE=pytorch
config.MODEL.AMP_PARAMS.AMP_TYPE=pytorch

在如下链接可以继续研究 this section 。

3.3 在PyTorch Lightning使用FSDP

为了更容易地与更通用的用例集成，PyTorch Lightning已经将FSDP作为beta功能。此教程包含一个关于如何将FSDP插件与PyTorch Lightning一起使用的详细示例。如下所示，添加plugins='fsdp'可以激活它。

代码语言：javascript复制

model = MyModel()
trainer = Trainer(gpus=4, plugins='fsdp', precision=16)
trainer.fit(model)

trainer.test()
trainer.predict()

3.4 直接从FairScale使用FSDP库

FSDP的主要开发库是FairScale.。您可以通过以下示例直接使用FairScale的FSDP，只需更换DDP。

代码语言：javascript复制

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
...
# sharded_module = DDP(my_module)
sharded_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()

FairScale中的FSDP库为大规模训练的许多重要方面提供了选项。当你希望使用FSDP的全部功能，你可以自行研究如下方面。

模型封装：为了最大限度地减少短期内的GPU内存需求，用户需要以嵌套方式封装模型。这增加了复杂性，但是在移植现有PyTorch模型代码时非常有用。
模型初始化：与DDP不同，FSDP不会在GPU工作进程之间自动同步模型权重。这意味着必须小心地进行模型初始化，以便所有GPU worker具有相同的初始权重。
优化器设置：由于分片和包装，FSDP只支持某些类型的优化器和优化器设置。特别是，如果模块被FSDP包装，并且其参数被展平为单个张量，则用户不能对此类模块中的不同参数组使用不同的超参数。
**混合精度 **：FSDP支持FP16主权重的高级混合精度训练，以及在梯度上FP16类型的reduce和scatter。但是，模型的某些部分可能只有在使用全精度时才收敛，在这些情况下，需要额外的wrapping，以便有选择地以全精度运行模型的某些部分。
状态检查点和推断：当模型规模较大时，保存和加载模型状态可能会变得很困难。FSDP支持多种方法使该任务成为可能，但这些方法是有代价的。
最后，FSDP通常与激活检查点函数一起使用，如checkpoint_wrapper 。用户可能需要仔细调整激活检查点策略，以便在有限GPU内存空间内容纳一个大型模型。

0x04 内存管理

我们接下来看看FSDP如何管理内存。

FairScale提供了受ZeRO <https://arxiv.org/pdf/1910.02054.pdf> 启发的算法：当使用数据并行训练时，您需要在计算/通信效率方面权衡内存的使用。另一方面，在使用模型并行训练时，需要为了内存而权衡计算/通信。

模型训练的内存使用通常分为两类：

模型状态：优化器状态、梯度、参数。
剩余状态：激活、临时缓冲区、碎片内存。

为了减少模型状态下的冗余，ZeRO提出了三种不同的算法。这些在FairScale中实现为优化器状态分片（Optimizer State Sharding，即OSS）、分片数据并行（Sharded Data Parallel，即SDP）和最终完全分片数据并行（Fully Sharded Data Parallel，即FSDP）。让我们深入了解每一个算法的实际机制，并理解它们为什么能够节省内存。

4.1 Optimizer State Sharding (OSS)

FairScale已经实现了与优化器内存相关的内存优化 OSS。

像Adam这样的优化器通常需要保持动量、方差。即使可以使用FP16精度的参数和梯度进行训练，参数和梯度也需要保存为FP32精度。当每个rank更新完整模型时，这意味着相当大一部分内存被优化器状态的冗余表示所占用。

为了克服这种冗余，优化器状态分片需要将模型优化步骤划分在不同的rank之间，以便每个rank只负责更新模型的对应分片。这反过来又确保优化器状态在每个rank上小得多，并且它不包含跨rank的冗余信息。

4.1.1 训练流程

训练流程可以从DDP的执行流程做如下修改：

wrapped optimizer根据参数大小（而不是使用顺序）以贪心算法方式来分割优化器状态。这是为了确保每个rank具有几乎相同的优化器内存占用。
训练过程类似于PyTorch的分布式数据并行（DDP）的过程。在每个rank上完成前向传播，然后是向后传播。在后向传播过程中，使用allreduce同步梯度。
每个rank只更新它负责的优化器分配状态参数，然后丢弃其余的。
更新后，将执行broadcast或allgather操作，以确保所有rank都收到最新更新的参数值。

当您使用具有附加状态的优化器（如Adam）时，OSS非常有用。如果您使用的是SGD或任何内存占用有限的优化器，那么在使用多个节点时，由于步骤4中的额外通信，您可能会看到速度减慢。在第2步的allreduce过程中，也有一些用于存储梯度的浪费内存，这些内存随后被丢弃。

4.1.2 最佳实践

OSS公开了一个broadcast_fp16 flag，您可能应该在多节点作业中使用它。在单节点实验中通常不需要这样做。
如果您的模型在大小方面极不平衡（例如，存在一个巨大的张量），那么这种方法将不会有很大帮助，而张量切分选项，如'fairscale.nn.FullyShardedDataParallel'将更可取。
3.OSS应该是DDP环境中的一个临时解决方案，其与大多数DDP功能保持兼容。

4.1.3 性能

在单个节点上，OSS应该总是比vanilla PyTorch快，内存节省会因使用的优化器而异
当使用多个节点时，OSS也可以比vanilla PyTorch快或慢，具体取决于所使用的优化器和可选标志（如上文提到的broadcast_fp16、梯度压缩、梯度累积）
如果您的实验可以使用更大的batch size，则采取更大的batch size并减少所涉及的rank数通常是有益的，或者使用梯度累积，因为这样可以降低通信成本。

4.2 Optimizer Gradient State Sharding

虽然OSS解决了优化器中的冗余问题，但依然存在梯度聚合计算的重复以及存在用于梯度的额外内存。为了克服冗余梯度内存，我们可以使用梯度分片或ZeRO-2。这已由FairScale中的分片数据并行（SDP）API实现。

为了启用梯度分片，每个 rank 都被分配一组参数，它们负责管理优化器状态以及梯度聚合。通过将一个模型分片分配给一个给定的rank，我们确保梯度被规约到特定的rank，而这些rank又负责相应的更新。因此这减少了通信和内存使用。

4.2.1 训练过程

训练过程如下：

与之前一样，包装的优化器在不同的列组中分割参数。
该模型现在使用分片数据并行（SDP）包装器进行包装，该包装器允许我们在训练过程中添加适当的hook并维护状态。
SDP关注于可训练的参数，并为每个参数添加了一个反向hook。
在反向传播过程中，梯度将规约到指定rank，rank是在 1 中作为切分过程的一部分指定的。使用reduce op代替allreduce op，从而减少通信开销。
每个rank更新其负责的参数。
更新后，将执行广播或allgather，以确保所有rank都收到最新更新的参数值。

OSS和SDPAPI都允许您减少用于梯度和优化器状态的内存，但是如果网络缓慢，则可能存在额外的通信成本。当遇到内存不足（OOM）问题时，可以把OSS和SDP作为第一步尝试。

4.2.2 最佳实践

如果使用多个节点，请通过指定reduce_buffer_size 参数确保SDP正在使用reduce buffers。改变它们的大小可能是一个优化目标，最佳配置可能取决于互连状况。
如果在单个节点上，通常最好不要使用'reduce_buffer_size'，因为它会带来延迟成本，但不会增加内存。将此值设置为0表示不使用此功能。
如果您的实验可以使用更大的batch size，则采取更大的batch size并减少所涉及的rank数通常是有益的，或者使用梯度累积，因为这样可以降低通信成本。

4.3 Optimizer Gradient Horizontal Model Sharding

为了进一步优化训练并实现更大的内存节省，我们需要启用参数切分。

参数切分类似于梯度和优化器状态，即，每个数据并行rank负责模型参数的一个分片。FairScale通过完全分片数据并行（FSDP）API实现参数分片，该API深受 ZeRO-3 <https://arxiv.org/pdf/1910.02054.pdf>的启发。

参数分片有两个如下关键点：

Allreduce操作可以分为reduce和allgather，类似于以前的分片技术（优化器状态和梯度）。
可以使用FSDP API包装各个层，该API允许我们在给定实例中将单个层所需的所有参数引入给定GPU，计算前向传递，然后丢弃不属于该rank的参数。

使用FSDP很简单，只需要在代码中简单地替换原来的DDP即可。注意：FSDP目前要求模型是一个nn.Sequential模型。

代码语言：javascript复制

from torch.utils.data.dataloader import DataLoader
from torchvision.datasets import FakeData
from torchvision.transforms import ToTensor

from fairscale.experimental.nn.offload import OffloadModel

num_inputs = 8
num_outputs = 8
num_hidden =  4
num_layers =  2
batch_size =  8

transform = ToTensor()
dataloader = DataLoader(
    FakeData(
        image_size=(1, num_inputs, num_inputs),
        num_classes=num_outputs,
        transform=transform,
    ),
    batch_size=batch_size,
)

model = torch.nn.Sequential(
    torch.nn.Linear(num_inputs * num_inputs, num_hidden),
    *([torch.nn.Linear(num_hidden, num_hidden) for _ in range(num_layers)]),
    torch.nn.Linear(num_hidden, num_outputs),
)

4.3.1 训练过程

具体训练过程如下：

在开始计算特定层之前，allgather模型每个层的正向传播所需的参数。
计算向前计算。
在特定层开始反向传递之前，allgather模型每个层反向传播所需的参数。
计算向后传播。
规约梯度，以便在负责相应参数的rank上累积聚合梯度。
让每个rank使用聚合梯度更新已分配给它的参数。

有了FSDP，在使用API进行检查点设置和保存优化器状态时，需要做一些小的更改。鉴于优化器状态和参数的分片性质，任何旨在保存模型状态以供训练或推理的API都需要考虑保存所有worker的权重。FSDP实现所需的管道（required plumbing）以保存所有worker的权重、保存单个worker的权重以及保存所有worker的优化器状态。

FSDP还支持混合精度训练，其中计算和通信均以FP16精度进行。如果要减少在FP32中执行的操作（这是DDP的默认行为），则必须设置 fp32_reduce_scatter=True。

为了进一步节省内存，FSDP支持将当前未使用的参数和梯度卸载到CPU上。这可以通过将“move_params_to_cpu”和“move_grads_to_cpu”设置为True来启用。

4.3.2 最佳实践

对于FSDP，最好使用 model.zero_grad(set_to_none=True) ，因为它在单步执行后节省了大量内存。
torch.cuda.amp.autocast与FSDP完全兼容。您需要将'mixed_precision'arg设置为True。
如果与激活检查点相结合，则最好使用 FSDP(checkpoint_wrapper(module))而不是checkpoint_wrapper(FSDP(module)).。后者将导致更多的通信，速度也会变慢。
FSDP与使用pointwise优化器的DDP兼容，例如Adam、AdamW、ADADDelta、Adamax、SGD等。当使用non-pointwise优化器（例如Adagrad、Adafactor、LAMB等）时，sharding将导致略有不同的结果。

4.3.3 性能

为了获得最佳内存效率，请使用“auto_wrap”将网络中的每一层用FSDP进行封装，并将 reshard_after_forward 设置为True。这样速度会慢，但是显存开销最小。
为了获得最佳训练速度，请将 reshard_after_forward 设置为False（不需要包装每一层，但如果设置，则会进一步提高速度）。

支持，FSDP基本原理和如何使用我们已经介绍完毕，下一篇我们介绍其代码细节，看看究竟如何做到最大程度减少内存使用。

0xFF 参考

Fully Sharded Data Parallel: faster AI training with fewer GPUs

ZeRO & DeepSpeed：可以让训练模型拥有超过1000亿个参数的优化（微软）

Fully Sharded Data Parallel: faster AI training with fewer GPUs

https://github.com/microsoft/DeepSpeed

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training

NLP服务 mapreduce pytorch api

0 人点赞