尖端的 AI 模型正变得极其庞大。训练这些模型的成本和开销正在迅速增加，并且需要大量的工程实践和反复摸索才能找到合适的训练方案。FSDP（全分片数据并行）通过让您在相同的资源下训练更大的模型，显著降低了这些成本。FSDP 降低了 GPU 上的内存占用，并且可通过轻量级的配置使用，通常只需几行代码即可完成，大大减少了工作量。
FSDP 的主要性能增益来自于最大化网络通信与模型计算之间的重叠，并消除了传统数据并行训练 (DDP) 中固有的内存冗余。在相同的服务器资源下，PyTorch FSDP 可以训练比 DDP 大约 4 倍的模型；如果结合激活检查点（Activation Checkpointing）和激活卸载（Activation Offloading），甚至可以训练大 20 倍的模型。
自 PyTorch 1.12 起，FSDP 已进入测试阶段，并增加了许多新功能，可以进行微调以进一步加速您的模型训练。
在本系列博文中，我们将解释您可以使用 FSDP 运行的多种性能优化方法，以在现有服务器资源范围内提高分布式训练速度和模型规模。我们使用处于微调模式下的 HuggingFace T5 3B、11B 和 DeepVit 作为贯穿本系列的运行示例。
作为本系列讨论的部分优化预览，我们在下方展示了按 Flops 缩放的优化前后性能对比（请注意，这些结果可能会根据您的服务器资源和模型架构而有所不同）。
*T5 3B 性能在 AWS A100 和 A10 服务器上测量。原始为无优化，调优后为应用了优化
*T5 11B 性能在 A100 服务器上测量。原始为无优化，调优后为应用了优化
在第一篇文章中，我们将快速概述 FSDP 以及它如何使大规模 AI 模型训练更高效。我们将简要强调可用的多种性能选项，并在后续文章中深入探讨这些细节。最后，我们将总结如何利用 AWS 并行集群（Parallel Cluster）通过 FSDP 进行大规模训练。
| 优化 | T5 模型 | 吞吐量提升 |
| 混合精度 | 3 B | 5倍 |
| 11 B | 10倍 | |
| 激活检查点 (AC) | 3 B | 10倍 |
| 11 B | 100倍 | |
| Transformer 封装策略 | 3 B | 2倍 |
| 11 B | 若无 Transformer 封装策略，无法进行实验。 | |
| 完全分片策略 | 3 B | 1.5倍 |
| 11 B | 无法在 Zero2 下运行 |
T5 模型相对于非优化版本的性能优化增益。
在我们针对 T5 3B 模型的实验中，使用 Transformer 封装策略 使得以 TFLOPS 衡量的吞吐量比默认封装策略高出 2 倍以上。激活检查点通过将检查点释放出的内存重新用于更大的批处理大小，带来了 10 倍的提升。BFloat16 混合精度与 FP32 相比提升了约 5 倍，最后，完全分片策略 (Full Sharding Strategy) 相对于 zero2 (DDP) 带来了 1.5 倍的提升。
我们对更大的模型 T5 11B 进行了类似的实验，但较大的模型尺寸导致实验空间发生了一些变化。具体来说，我们发现需要 Transformer 封装策略和激活检查点这两项优化，才能使我们在 3 个节点（每个节点有 8 个 80GB 内存的 A100 GPU）上运行这些实验。有了这些优化，我们可以容纳 50 的批处理大小，并获得比移除其中任何一个优化更高的吞吐量。因此，与 3B 模型仅测试单一优化开关不同，大模型实验是在开启/关闭这三个优化中的某一个时，始终运行另外两个，以便为每个测试项提供可用的批处理大小。
基于 TFLOP 比较，对于 11B 模型，我们从优化中看到了更大的回报。与 3B 参数模型相比，混合精度（约 10 倍提升）和激活检查点（约 100 倍提升）在 11B 模型上的影响要大得多。通过混合精度，我们可以容纳约 2 倍大的批处理大小，而通过激活检查点则可以容纳超过 15 倍大的批处理大小（从无激活检查点时的 3 增加到有激活检查点时的 50），这转化为巨大的吞吐量提升。
我们还观察到，对于这些大于 3B 的大模型，使用 Zero2 分片策略会导致内存中几乎没有剩余空间来存放批处理数据，并且必须使用非常小的批处理大小（例如 1-2），这使得完全分片策略成为实现更大批处理大小的必要条件。
注意：本教程假设您对 FSDP 有基本了解。要了解 FSDP 的基础知识，请参考 入门 和 FSDP 进阶 教程。
什么是 FSDP？它是如何使大规模训练更高效的？
FSDP 在分布式数据并行的基础上进行了扩展，不仅并行处理数据，还并行处理模型参数、优化器状态以及与模型相关的梯度。具体而言，每个 GPU 仅存储整个模型的一个子集，以及相关的优化器状态和梯度子集。
为了展示分布式训练的演变，我们可以从最初 AI 模型仅在单个 GPU 上训练开始说起。
DDP（分布式数据并行）是继单 GPU 训练之后的初步升级，旨在解决数据和模型尺寸的增长问题。在这种架构下，多个 GPU 每个都存放相同模型的完整副本。其优势在于每个批次的数据可以被拆分并在每个 GPU 上独立且同时处理，从而并行化数据集的处理并随着 GPU 数量的增加提高训练速度。代价是需要在反向传播后在每个 GPU 之间进行梯度通信以同步模型。
FSDP 通过消除 DDP 中存在的优化器计算和状态存储的冗余，以及模型参数的梯度和内存存储冗余，进一步扩展了模型缩放能力。这种冗余减少，加上增加了模型参数通信与模型计算同时进行的通信重叠，使得 FSDP 能够用与 DDP 相同的资源训练大得多的模型。
一个关键点是，这种效率也使得训练大于单 GPU 内存容量的 AI 模型成为可能。可用于训练的模型大小现在增加到了所有 GPU 的总内存，而不是单个 GPU 的大小。（顺便提一下，FSDP 还可以通过利用 CPU 内存超越总 GPU 内存，尽管我们在此不直接讨论这一方面）。
正如之前一篇博文所讨论的，使用 DDP 时，在 32 个 40GB 内存的 A100 GPU（4 个节点）上，在激活检查点的帮助下，我们能够训练的最大模型为 3B 参数，批处理大小为 128。相比之下，使用 FSDP，结合激活检查点以及激活和参数卸载，我们能够训练高达 81B 大小的模型。在另一个实验中，我们使用 512 个 GPU 对一个 1T 参数的模型进行了基准测试。
为了直观了解 FSDP 在参数层面的工作原理，下面我们展示了一个动画，详细说明了假设有两个 GPU 和一个 8 参数简单模型时，模型参数是如何分片和通信的。
上方动画展示了模型在各个 rank 之间初始分片的步骤，以及我们启动 all_gathers 和前向传播的过程。
我们继续进行模型的前向传播。每个 FSDP 单元完成后，非本地拥有的参数会被丢弃以释放内存，并可选择对激活进行检查点保存。此过程持续进行，直到完成前向传播并计算出损失。
在反向传播期间，使用另一个 all_gather 来加载参数并计算梯度。然后这些梯度被 reduce_scattered，以便每个参数的本地所有者可以聚合并准备更新权重。
最后，每个 rank 通过优化器状态传递求和后的梯度并更新权重，以完成小批次（mini-batch）训练。
既然模型现在分布在整个可用的 GPU 集合上，一个逻辑问题是，考虑到这种模型参数的分片，数据如何通过模型。
这是通过 FSDP 协调所有 GPU 有效共享（通信）模型各自部分来实现的。模型被分解为 FSDP 单元，每个单元内的参数被展平，然后在所有 GPU 之间进行分片。在每个 FSDP 单元内，GPU 被分配了单个模型参数的交错所有权。
通过交错，我们的意思是——假设有 ID 为 1 和 2 的两个 GPU，FSDP 单元的所有权模式将是 [12121212]，而不是连续的 [111222] 块。
在训练期间，启动一个 all_gather，FSDP 单元内的本地拥有的模型参数由所有者 GPU 在需要时以“即时”（just in time）的方式与其他非所有者 GPU 共享。FSDP 预取参数以将 all_gather 通信与计算重叠。
当请求的参数到达时，GPU 使用交付的参数以及它已经拥有的参数来创建一个完全填充的 FSDP 单元。因此，每个 GPU 在持有完全填充的 FSDP 单元时会出现内存使用高峰。
然后它通过 FSDP 单元处理数据，并丢弃从其他 GPU 接收到的参数，为下一个单元释放内存……该过程不断重复，贯穿整个模型以完成前向传播。该过程（通常）在反向传播中重复。（注：这是一个为了便于理解的简化版本，实际操作中存在额外的复杂性，但这应该有助于构建 FSDP 过程的基本心理模型）。
这消除了 DDP 中存在的大部分内存冗余，但付出了大量的网络通信成本，在所有 GPU 之间来回传输这些请求的参数。将通信时序与发生的计算重叠是本系列中我们将讨论的许多性能改进的基础。 关键增益通常基于通信往往可以在与计算相同的时间内发生这一事实。正如您可以推断出的，拥有高速通信对 FSDP 性能至关重要。
我该如何优化我的 FSDP 训练？
我们将涵盖四种主要的性能改进——Transformer 包装器、激活检查点、混合精度和选择正确的分片策略。下方的流程图将作为本系列讨论的调优选项的检查清单。
封装策略——对于 Transformer 模型，请使用 Transformer 封装策略
第一个性能优化是为 Transformer 模型利用 FSDP Transformer 包装器。
预定义的封装策略之一是 size_based_autowrap_policy。使用 size_based_autowrap_policy 时，FSDP 将从底向上遍历模块结构，一旦当前单元至少拥有该尺寸策略中指定的 min_num_params（默认为 1e8，即 100M），就会创建一个新的 FSDP 单元。如果模块无法创建为 FSDP 单元，FSDP 将继续检查其父模块。这种基于尺寸的封装策略可能不适合某些模型结构。PyTorch 分布式团队正在积极开发下一个版本中的新默认封装策略，该策略基于尺寸和模块执行顺序，用户只需调整尺寸即可获得优化性能。
在当前版本中，通过使用“Transformer 包装器”，您可以在运行 Transformer 模型时极大地提高性能。您需要为您的模型提供适当的层类（layer class）。这里，层类是指包含多头注意力（Multi-Head Attention）和前馈网络（Feed Forward Network）的类。
FSDP 然后会围绕层类而不是基于参数大小的任意断点来形成 FSDP 单元。通过围绕 Transformer 中均匀重复的层类对模型进行分片，FSDP 可以创建更均匀的 FSDP 单元，从而更好地平衡计算和通信的重叠。相比之下，基于尺寸的封装可能会为模型产生非常不均匀或倾斜的分片，进而导致计算与通信重叠的匹配不佳。正如前面讨论的，FSDP 高性能的主要驱动因素是通信和计算的重叠，这就是为什么 Transformer 包装器能提供更高性能的原因。请注意，如果非 Transformer 模型包含一系列均匀的层，也可以使用 Transformer 包装器。
让我们比较在默认包装器和 Transformer 包装器下运行 T5 3B 参数模型时的性能差异。
对于默认封装，我们不需要采取任何操作——我们只需将模型传递给 FSDP，如下所示
model = FSDP(
model,
device_id=torch.cuda.current_device(),
)
在这种情况下，FSDP 只是简单地将整个模型包装在一个 FSDP 单元中。
在 8 个 GPU 的 NVIDIA A100-SXM4–40GB 上运行，我们能够达到 2.3 TFlops 和 95% 的 GPU 内存利用率，批处理大小为 14。
然而，由于 T5 是一个 Transformer 模型，我们更好地利用该模型的 Transformer 包装器。
要使用它，我们需要隔离 Transformer 的层类，然后将其传递以创建我们的 Transformer 包装器。
from transformers.models.t5.modeling_t5 import T5Block
现在我们可以创建我们的 Transformer 包装器
transformer_auto_wrapper_policy = functools.partial(
transformer_auto_wrap_policy,
transformer_layer_cls={
T5Block, # < ---- Your Transformer layer class
},
)
模型感知包装器准备就绪后，我们可以初始化 FSDP
# invoke FSDP with your transformer wrapper policy:
model = FSDP(
model,
auto_wrap_policy=transformer_auto_wrapper_policy,
device_id=torch.cuda.current_device(), # streaming init
)
运行此封装后的模型，我们可以看到一些实质性的性能提升。我们可以容纳近两倍的批处理大小（达到 28），并且凭借更好的内存和通信效率，我们看到 TFlops 从 2.3 增加到 5.07。
因此，由于向 FSDP 提供了更多的模型信息，我们的训练吞吐量提高了 200% 以上 (2.19x)！Transformer 封装策略产生了更细粒度且平衡的 FSDP 单元，每个单元持有一个层类，从而导致更有效的通信-计算重叠。
上方：基于包装器类型的 TFlops 图形比较
如果您正在训练 Transformer 模型，使用 Transformer 包装器配置您的 FSDP 训练是非常值得的。有关如何隔离层类的更多信息，请查看我们关于 Transformer 封装的深度视频 这里，我们在其中演示了多个 Transformer 模型以展示在哪里可以找到层类。
混合精度——如果您有 Ampere 架构 GPU，请使用 BF16
FSDP 支持灵活的混合精度策略，让您可以精确控制参数、梯度和缓冲区数据类型。这使您可以轻松利用 BFloat16 或 FP16 将训练速度提高多达 70%。
*请注意，BFloat16 仅在 Ampere 架构 GPU 上可用。在 AWS 上，这可在 p4dn 和 g5 实例上使用。
作为比较，在 8B DeepVit 模型上比较完全调优的 BFloat16 与 FP32 时，我们可以看到 77% 的速度提升。
在微调 3B HuggingFace T5 模型时，使用 BFloat16 我们获得了更大的加速，如下图所示。我们观察到，由于较低的精度，BFloat16 的验证损失在前几个 epoch 中略微落后，但它能够追赶上来，并产生与 FP32 相同的最终准确度。
要使用混合精度，我们创建具有所需数据类型的策略，并在 FSDP 初始化期间传入它。
要创建我们的策略，我们需要导入 MixedPrecision 类，然后使用我们的自定义类定义自定义策略
from torch.distributed.fsdp import MixedPrecision
bfSixteen = MixedPrecision(
param_dtype=torch.bfloat16,
# Gradient communication precision.
reduce_dtype=torch.bfloat16,
# Buffer precision.
buffer_dtype=torch.bfloat16,
)
model = FSDP(
model,
auto_wrap_policy=transformer_auto_wrapper_policy,
mixed_precision=bfloatPolicy)
您可以根据需要混合和匹配参数、梯度和缓冲区的精度
comboPolicy = MixedPrecision(
# Param precision
param_dtype=torch.bfloat16,
# Gradient communication precision.
reduce_dtype=torch.float32,
# Buffer precision.
buffer_dtype=torch.float32,
)
对于 FP16 训练，您还需要使用 ShardedGradScaler，我们将在后续文章中介绍。对于 BFloat16，它是直接替换。
AnyPrecision 优化器——通过全 BF16 训练超越混合精度
混合精度训练（在 FSDP 和其他地方）将工作权重保持在减小的精度数据类型（BF16 或 FP16）中，同时将主权重保持在完整的 FP32 中。保持 FP32 主权重的原因是，以纯 BF16 运行会导致“权重停滞”，即由于较低的精度，极小的权重更新会丢失，准确度随时间趋于平稳，而 FP32 权重可以从这些小更新中继续改进。
为了解决这个困境，我们可以使用 TorchDistX（Torch 分布式实验性库）中提供的新 AnyPrecision 优化器，它允许您成功训练并将主权重保持在纯 BF16 中，而不是 FP32。此外，与通常将优化器状态存储在 FP32 中不同，AnyPrecision 也能够以纯 BF16 维护状态。
AnyPrecision 通过维护一个额外的缓冲区来跟踪权重更新期间丢失的精度并在下一次更新期间重新应用它，从而实现纯 BF16 训练……有效地解决了权重停滞问题，而无需 FP32。
作为使用 AnyPrecision 进行纯 BF16 训练所获得的吞吐量收益的比较，我们使用 FSDP 对 T5 11B 模型进行了实验，包括常规 FP32 训练、BF16 混合精度训练以及如前所述在 3 个节点 A100 GPU 上使用 AnyPrecision 优化器的纯 BF16 训练。
如上所示，使用 AnyPrecision 和纯 BF16 训练的吞吐量是混合精度的 2 倍，是 FP32 的 20 倍以上。
潜在的代价是对最终准确度的影响——在我们测试的情况下，由于略微降低的精度带来的正则化效应，准确度等于或优于 FP32，但您的结果可能会有所不同。
AnyPrecision 优化器可供您在 这里 进行测试，它是 AdamW 优化器的直接替换。
激活检查点——通过以计算换取内存来增加吞吐量
FSDP 在模型分片后支持激活检查点，并且易于实现。上图显示使用激活检查点可提高约 4 倍的吞吐量。
激活检查点是指在正向传播期间释放中间激活，并将一个检查点作为占位符留下。这通常使可用 GPU 内存增加 30% 以上。
权衡之处在于，在反向传播期间，这些先前删除的中间激活必须使用检查点中的信息重新计算（重复计算），但通过利用增加的 GPU 内存，可以增加批处理大小，使得净吞吐量大幅增加。
# verify we have FSDP activation support ready by importing:
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
checkpoint_wrapper,
CheckpointImpl,
apply_activation_checkpointing_wrapper,
)
实现激活检查点的步骤是首先导入 FSDP 检查点函数。我们需要声明我们的检查点包装器类型（不可重入），并创建一个检查函数来识别要包装的层，如下所示
non_reentrant_wrapper = partial(
checkpoint_wrapper,
offload_to_cpu=False,
checkpoint_impl=CheckpointImpl.NO_REENTRANT,
)
check_fn = lambda submodule: isinstance(submodule, T5Block)
apply_activation_checkpointing_wrapper(
model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
)
重要提示——这必须在模型通过 FSDP 初始化后运行。
总之，希望您已经看到一些 FSDP 选项的初步调优可以对您的训练性能产生重大影响。
至此，我们将注意力从如何利用 FSDP 进行内部扩展，转向如何使用 AWS 为 FSDP 扩展服务器硬件。
在 AWS 上使用 FSDP 进行大规模训练——对于多节点，优先考虑高速网络
AWS 提供了多种可用于通过 FSDP 进行分布式训练的服务：Amazon EC2 加速计算实例、AWS ParallelCluster 和 Amazon SageMaker。
在本系列博文中，我们使用了 Amazon EC2 p4d 实例，包括单实例多 GPU 配置，以及使用 AWS ParallelCluster 和 SageMaker 的多实例配置来运行我们的训练任务。
在这里，我们将特别关注 AWS ParallelCluster，并概述如何将其用于训练目的。
AWS ParallelCluster 设置
AWS ParallelCluster 是一款开源的集群管理工具，可让您轻松在 AWS 上部署和管理高性能计算 (HPC) 集群。AWS ParallelCluster 使用 YAML 配置文件来配置所有必要的资源。它还支持多种实例类型、任务提交队列、共享文件系统（如 Amazon EFS (NFS) 或 Amazon FSx for Lustre）以及任务调度器（如 AWS Batch 和 Slurm）。
集群上的工作流程
高层思路是拥有一个包含控制计算节点的头节点（head node）的集群。实际的训练任务在计算节点上运行。在集群上运行训练任务的整体步骤如下：
- 设置 AWS ParallelCluster（我们在下方讨论）。
- 连接到头节点，并导入训练代码/设置环境。
- 拉取数据并将其放在计算节点可以访问的共享文件夹中（FSx Lustre 驱动器）。
- 使用任务调度器（在此示例中为 Slurm）运行训练任务。
设置 AWS ParallelCluster
要设置 AWS ParallelCluster：
- 部署网络栈。 此步骤是可选的，因为您可以使用账户默认 VPC 并让 AWS ParallelCluster 创建您的子网和安全组。然而，我们更倾向于对期望的网络基础设施进行分区，并通过 CloudFormation 栈进行此部署。由于我们部署了一个公共子网和一个私有子网，我们希望在包含目标实例（在本例中为 p4d）的可用区内创建它们。我们通过以下 AWS CLI 命令查询它们在所使用的区域 (us-east-1) 的可用性：
aws ec2 describe-instance-type-offerings --location-type availability-zone \ --filters Name=instance-type,Values=p4d.24xlarge --region us-east-1 --output table我们看到三个包含 p4d 实例的可用区，我们在部署网络栈时选取其中一个 (us-east-1c，您的可能不同)。这可以通过 AWS 控制台或 AWS CLI 完成。在我们的案例中，我们使用后者，如下所示：aws cloudformation create-stack --stack-name VPC-Large-Scale --capabilities CAPABILITY_IAM --template-body file://VPC-Large-Scale.yaml --parameters ParameterKey=SubnetsAZ,ParameterValue=us-east-1cCloudFormation 将代表我们部署新的 VPC、子网、安全组和端点。完成后，您可以通过查询栈输出来检索公共子网和私有子网的 ID，即值PublicSubnet和PrivateSubnet。例如，使用 AWS CLI 获取私有子网：aws cloudformation describe-stacks --stack-name VPC-Large-Scale --query "Stacks[0].Outputs[?OutputKey=='PrivateSubnet'].OutputValue" --output text - 创建 ParallelCluster。 集群配置文件指定了我们集群的资源。这些资源包括头节点、计算节点的实例类型、S3 存储桶的访问权限、存放数据的共享存储。我们将使用 Amazon FSx for Lustre，它提供了一种完全托管的、带有 Lustre 的共享存储服务。这里 是集群配置文件的示例。我们可以使用 AWS ParallelCluster CLI 创建集群。请注意，私有和公共子网 ID 需要替换为您之前检索到的 ID。您将能够使用 AWS ParallelCluster CLI 来控制集群，以启动、停止、暂停等。
pcluster create-cluster --cluster-name my-hpc-cluster --cluster-configuration cluster.yaml - SSH 到头节点。 一旦集群准备就绪，我们就可以使用 SSH 协议连接到头节点，拉取我们的训练代码，并将数据放在集群配置文件中指定的共享存储中。
pcluster ssh --cluster-name cluster -i your-key_pair - 启动训练任务。 现在我们有了数据和训练代码，我们可以启动 Slurm 任务进行训练。这是使用 torchrun 启动任务的 Slurm 脚本示例。
设置集群的更多细节超出了本文章的范围，但我们将另设文章进行讨论。
接下来是什么？
在本文中，我们提供了 FSDP 的高层概述以及它如何高效地扩展分布式 AI 训练。随附的流程图将帮助您提供一份清单，以供回顾本文讨论的调优选项，例如 Transformer 包装器和激活检查点。
在接下来的文章中，我们将继续使用 T5 模型，深入探讨上述每个主题，特别是分片策略和其他优化，以提供更多的见解和细节。目前，分片策略的一个很好的参考资料在我们的视频教程 这里。
如果您有疑问或发现问题，请联系作者 Less、Hamid 和 Geeta，或在 PyTorch Github 上提交问题。
特别感谢
PyTorch 分布式团队、Shen Li、Rohan Varma、Yanli Zhao、Andrew Gu、Anjali Sridhar、Ana Simoes、Pierre-Yves Aquilanti、Sundar Ranganathan，以及更广泛的 AWS 团队，感谢他们为我们运行大规模实验提供基础设施和技术支持。
资源