尖端 AI 模型正变得极其庞大。训练这些模型的成本和开销正在迅速增加，并且需要大量的工程和猜测才能找到正确的训练方案。FSDP 通过让您使用相同数量的资源训练更大的模型来显著降低这些成本。FSDP 降低了 GPU 的内存占用，并且可以通过轻量级配置使用，通常只需几行代码即可，大大减少了所需的工作量。
FSDP 的主要性能提升来自最大化网络通信和模型计算之间的重叠，并消除传统数据并行训练 (DDP) 中固有的内存冗余。PyTorch FSDP 可以在与 DDP 相同的服务器资源上训练大约 4 倍大的模型，如果结合激活检查点和激活卸载，则可以训练 20 倍大的模型。
自 PyTorch 1.12 以来，FSDP 已进入测试阶段，并增加了许多新功能，可以进行调整以进一步加速您的模型训练。
在本系列博客文章中，我们将解释您可以使用 FSDP 运行的多种性能优化，以在可用服务器资源的背景下提高您的分布式训练速度和模型大小。我们将 HuggingFace T5 3B、11B 和 DeepVit 在微调模式下作为本系列中的运行示例。
作为本系列中讨论的一些优化的预览，我们展示了以下 Flops 缩放前后的性能（请注意，这些结果可能因您的服务器资源和模型架构而异）。
*在 AWS A100 和 A10 服务器上测量的 T5 3B 性能。原始未优化，经过应用优化后进行调整
*在 A100 服务器上测量的 T5 11B 性能。原始未优化，经过应用优化后进行调整
在这第一篇文章中，我们将快速概述 FSDP 以及它如何使大规模 AI 模型的训练更高效。我们将简要强调可用的多种性能选项，并在后续文章中深入探讨这些细节。然后，我们将总结如何利用 AWS 并行集群进行 FSDP 大规模训练。
| 优化 | T5 模型 | 吞吐量提升 |
| 混合精度 | 3 B | 5x |
| 11 B | 10x | |
| 激活检查点 (AC) | 3 B | 10x |
| 11 B | 100x | |
| Transformer 封装策略 | 3 B | 2x |
| 11 B | 无法在没有 Transformer 封装策略的情况下运行实验。 | |
| 全分片策略 | 3 B | 1.5x |
| 11 B | 无法与 Zero2 一起运行 |
T5 模型相对于未优化模型的性能优化增益。
在我们对 T5 3B 模型的实验中，使用 Transformer 封装策略 导致以 TFLOPS 衡量的吞吐量比默认封装策略高出 >2 倍。激活检查点 通过将从检查点释放的内存重新投入到更大的批大小中，实现了 10 倍的改进。使用 BFloat16 的混合精度 比 FP32 提高了约 5 倍，最后 全分片策略 比 Zero2 (DDP) 提高了 1.5 倍。
我们对更大的模型 T5 11B 进行了类似的实验，但更大的模型大小导致实验空间发生了一些变化。具体来说，我们发现需要两种优化：Transformer 封装策略和激活检查点，才能在 3 个节点（每个节点有 8 个 A100 GPU，80 GB 内存）上运行这些实验。通过这些优化，我们可以适应 50 的批大小，并获得比单独移除其中任何一个更高的吞吐量。因此，对于更大的模型实验，不再像 3B 模型那样只针对单一优化进行开关测试，而是始终运行另外两个优化，同时打开/关闭其中一个优化，以允许每个项目的两种测试状态都有可用的批大小。
基于 TFLOP 比较，对于 11B 模型，我们从优化中获得了更大的回报。混合精度（约 10 倍改进）和激活检查点（约 100 倍改进）对 11B 模型的影响比 3B 参数模型大得多。通过混合精度，我们可以适应约 2 倍大的批大小；通过激活检查点，可以适应 >15 倍批大小（从没有激活检查点的 3 增加到有激活检查点的 50），这转化为吞吐量的巨大提升。
我们还观察到，对于大于 3B 的模型，使用 Zero2 分片策略会导致内存中留给批处理数据的空间极小，并且必须使用非常小的批处理大小（例如 1-2），这实际上使得全分片策略成为启用更大批处理大小的必要条件。
注意——本教程假设您对 FSDP 有基本的了解。要了解 FSDP 的基础知识，请参阅入门和高级 FSDP 教程。
FSDP 是什么？它如何使大规模训练更高效
FSDP 在分布式数据并行的基础上进行扩展，不仅并行化数据，还并行化模型参数、优化器状态以及与模型相关的梯度。具体来说——每个 GPU 只存储整个模型的子集以及相关的优化器状态和梯度子集。
为了展示分布式训练的演变，我们可以从最初开始，当时 AI 模型只是在单个 GPU 上进行训练。
DDP (分布式数据并行) 是从仅使用单个 GPU 训练的初步提升，旨在解决数据和模型大小增长问题，其中多个 GPU 各自拥有自己的相同模型副本。这里的优势在于，每个批次的数据可以同时在每个 GPU 上独立分割和处理，从而并行化数据集的处理并通过增加 GPU 数量来提高训练速度。缺点是需要在反向传播后在每个 GPU 之间通信梯度以同步模型。
FSDP 通过消除优化器计算和状态存储的冗余以及 DDP 中存在的模型参数的梯度和内存存储（DDP = 分布式数据并行）来扩展模型缩放。这种冗余减少，加上模型参数通信与模型计算同时进行所增加的通信重叠，使得 FSDP 能够使用与 DDP 相同的资源训练更大的模型。
关键一点是，这种效率还允许训练比单个 GPU 更大的 AI 模型。现在可用于训练的模型大小已增加到所有 GPU 的聚合内存，而不是单个 GPU 的大小。（顺便提一下，FSDP 还可以通过利用 CPU 内存来超越聚合 GPU 内存，尽管我们在这里不会直接涵盖这方面）。
正如之前一篇博客文章所讨论的，使用 DDP，我们可以在 32 个 A100 GPU (40 GB 内存，4 个节点) 上训练的最大模型为 3B 参数，批处理大小为 128，借助激活检查点。相比之下，使用 FSDP，我们能够训练高达 81B 的模型，结合了激活检查点以及激活和参数卸载。在另一次实验中，我们使用 512 个 GPU 对 1T 参数模型进行了 FSDP 基准测试。
为了直观地了解 FSDP 在参数层面的工作原理，下面我们通过动画详细展示了在两个 GPU 场景和一个简单的 8 参数模型中，模型参数是如何分片和通信的。
上面——动画详细地展示了模型在各等级之间初始分片的步骤，我们开始进行 `all_gathers` 和前向传播。
我们继续通过前向传播遍历模型。在每个 FSDP 单元完成后，非本地拥有的参数将被丢弃以释放内存，并且可以选择性地对激活进行检查点。这个过程一直持续到我们完成前向传播并计算损失。
在反向传播期间，另一个 `all_gather` 用于加载参数并计算梯度。然后，这些梯度被 `reduce_scattered`，以便每个参数的本地所有者可以聚合并准备更新权重。
最后，每个等级通过优化器状态传递求和的梯度并更新权重以完成 mini-batch。
现在模型分布在所有可用的 GPU 上，一个合乎逻辑的问题是数据如何在模型参数被分片的情况下通过模型移动。
这通过 FSDP 与所有 GPU 协调，有效地共享（通信）模型的相应部分来实现。模型被分解成 FSDP 单元，每个单元内的参数被平坦化，然后分片到所有 GPU 上。在每个 FSDP 单元内，GPU 被分配单个模型参数的交错所有权。
交错是指以下内容——假设有 2 个 GPU，ID 为 1 和 2，FSDP 单元的所有权模式将是 [12121212]，而不是连续的 [111222] 块。
在训练过程中，当非本地 GPU 需要时，会启动 `all_gather`，并且 FSDP 单元中本地拥有的模型参数由拥有它的 GPU 与其他非拥有 GPU 以“即时”的方式共享。FSDP 会预取参数，以使 `all_gather` 通信与计算重叠。
当这些请求的参数到达时，GPU 会使用交付的参数，结合其已拥有的参数，创建一个完全填充的 FSDP 单元。因此，在持有完全填充的 FSDP 单元时，每个 GPU 都会达到内存使用峰值。
然后它通过 FSDP 单元处理数据，并丢弃从其他 GPU 接收的参数以释放内存供下一个单元使用……这个过程一遍又一遍地进行，通过整个模型完成前向传播。然后（通常）为反向传播重复此过程。（注意——这是为了理解而简化的版本……还有额外的复杂性，但这应该有助于构建 FSDP 过程的基本心理模型）。
这消除了 DDP 中存在的许多内存冗余，但也带来了更高网络通信量的开销，用于在所有 GPU 之间来回传输这些请求的参数。将通信时间与正在进行的计算重叠是我们将在此系列中讨论的许多性能改进的基础。 关键的收益通常基于通信可以与计算同时进行的事实。正如您可以推测的那样，拥有高速通信对 FSDP 性能至关重要。
我如何使用 FSDP 优化我的训练？
我们将介绍四种主要的性能改进——Transformer 封装器、激活检查点、混合精度以及选择合适的分片策略。下面的流程图将作为我们将在本文中讨论的调优选项的清单。
封装策略——对于 Transformer，使用 Transformer 封装策略
第一个性能优化是利用 FSDP Transformer 封装器用于 Transformer 模型。
预定义的封装策略之一是 `size_based_autowrap_policy`。使用 `size_based_autowrap_policy`，FSDP 将从下到上遍历模块结构，一旦当前单元至少达到大小策略中指定的 `min_num_params`（默认为 1e8，即 1 亿），就会创建一个新的 FSDP 单元。如果模块无法作为 FSDP 单元创建，FSDP 将继续检查其父模块。这种基于大小的封装策略可能不适用于某些模型结构，PyTorch 分布式团队正在积极开发下一个版本中基于大小和模块执行顺序的新默认封装策略，用户只需调整大小即可实现优化性能。
在当前版本中，通过使用“Transformer 封装器”，您可以极大地提高运行 Transformer 模型时的性能。您需要为您的模型提供适当的层类。在这里，层类是包含多头注意力 (Multi-Head Attention) 和前馈网络 (Feed Forward Network) 的类。
然后，FSDP 将围绕层类形成 FSDP 单元，而不是根据参数大小任意中断。通过围绕 Transformer 中统一重复的层类对模型进行分片，FSDP 可以创建统一的 FSDP 单元，从而更好地平衡计算和通信的重叠。相比之下，基于大小的封装可能会为模型生成非常不均匀或倾斜的分片，这会导致计算与通信重叠的不均匀匹配。如前所述，FSDP 高性能的主要驱动力是通信和计算的重叠，因此 Transformer 封装器提供了改进的性能。请注意，如果非 Transformer 模型具有统一层的列表，Transformer 封装器也可以用于非 Transformer 模型。
让我们比较一下 T5 3B 参数模型在默认封装器和 Transformer 封装器下运行时的性能差异。
对于默认封装，我们无需采取任何操作——我们只需将模型传递给 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.19 倍）！Transformer 封装策略产生了更细粒度、更均衡的 FSDP 单元，每个单元包含一个层类，从而实现了更有效的通信-计算重叠。
上图：基于封装器类型的 TFlops 图形比较
如果您正在训练 Transformer 模型，那么使用 Transformer 封装器配置 FSDP 训练是值得的。有关如何隔离层类的更多信息，请参阅我们关于 Transformer 封装的深入视频此处，其中我们演示了许多 Transformer 模型中层类所在的位置。
混合精度——如果您有 Ampere 架构 GPU，请使用 BF16
FSDP 支持灵活的混合精度策略，让您对参数、梯度和缓冲区数据类型进行细粒度控制。这让您可以轻松利用 BFloat16 或 FP16 将训练速度提高高达 70%。
*请注意，BFloat 16 仅在 Ampere 类型 GPU 上可用。在 AWS 上，p4dn 和 g5 实例提供此功能。
通过比较，我们可以看到在 8B DeepVit 模型上，完全调优的 BFloat16 与 FP32 相比，速度提高了 77%。
我们使用 BFloat16 在微调 3B HuggingFace T5 模型时获得了更大的加速，如下图所示。我们观察到，由于精度较低，BFloat16 的验证损失在最初几个时期略微落后，但它能够赶上并达到与 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 Distributed Experimental) 中提供的新 AnyPrecision 优化器，它允许您成功训练并将主权重保持在纯 BF16 而不是 FP32。此外，与通常在 FP32 中存储优化器状态不同，AnyPrecision 也能够将状态保持在纯 BF16 中。
AnyPrecision 通过维护一个额外的缓冲区来跟踪权重更新期间丢失的精度，并在下一次更新期间重新应用该精度，从而实现了纯 BF16 训练，有效解决了权重停滞问题，而无需 FP32。
为了比较使用 AnyPrecision 进行纯 BF16 训练可获得的吞吐量增益，我们使用 FSDP 对 T5 11B 模型进行了实验，分别使用常规 FP32 训练、BF16 混合精度训练以及使用 AnyPrecision 优化器进行纯 BF16 训练，实验在前面提到的 3 个带 A100 GPU 的节点上进行。
如上所示，使用 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 并行集群，并概述如何将其用于训练目的。
AWS ParallelCluster 设置
AWS ParallelCluster 是一款开源集群管理工具，可让您轻松地在 AWS 上部署和管理高性能计算 (HPC) 集群。AWS ParallelCluster 使用 yaml 配置文件来配置所有必要的资源。它还支持多种实例类型、作业提交队列、共享文件系统（如 Amazon EFS (NFS) 或 Amazon FSx for Lustre）以及作业调度器（如 AWS Batch 和 Slurm）。
集群上的工作流程
其高层思路是构建一个包含头节点来控制计算节点的集群。实际的训练任务在计算节点上运行。在集群上运行训练任务的总体步骤如下：
- 设置 AWS ParallelCuster（我们将在下面讨论）
- 连接到头节点，导入训练代码/设置环境。
- 拉取数据并将其放置在计算节点可以访问的共享文件夹中（FSx Lustre 驱动器）。
- 使用作业调度器（本例中为 Slurm）运行训练作业。
设置 AWS ParallelCuster
设置 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、子网、安全组和端点。完成后，您可以通过查询堆栈输出以及值PublicSubnet和PrivateSubnet来检索公共和私有子网的 ID。例如，使用 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 团队，感谢他们为我们提供基础设施和技术支持，以运行大规模实验。
资源