使用 Pytorch FSDP 和 AWS 进行高效的大规模训练

作者：Less Wright、Hamid Shojanazeri、Geeta Chauhan

尖端 AI 模型正变得极其庞大。训练这些模型的成本和开销正迅速增加，并且涉及大量的工程和猜测以找到正确的训练方案。FSDP 通过使您能够使用相同数量的资源训练更大规模的模型，显著降低了这些成本。FSDP 降低了 GPU 上的内存占用，并且可以通过轻量级配置使用，该配置所需的精力大大减少，通常只需几行代码。

FSDP 的主要性能提升来自于最大化网络通信与模型计算的重叠，并消除传统数据并行训练 (DDP) 中固有的内存冗余。PyTorch FSDP 可以在与 DDP 相同的服务器资源上训练大约 4 倍大的模型，如果结合激活检查点和激活卸载，则可以训练 20 倍大的模型。

自 PyTorch 1.12 版本以来，FSDP 现在处于 beta 状态，并增加了一些可以调整的新功能，以进一步加速您的模型训练。

在本系列博文中，我们将解释您可以通过 FSDP 运行的多种性能优化，以在您可用的服务器资源范围内提高分布式训练速度和模型大小。我们在整个系列中都使用 HuggingFace T5 3B、11B 和 DeepVit 的微调模式作为运行示例。

作为本系列中讨论的一些优化的预览，我们在下方展示了按 Flops 衡量的优化前后的性能（请注意，这些结果可能因您的服务器资源和模型架构而异）。

*在 AWS A100 和 A10 服务器上测量的 T5 3B 性能。原始版本（未优化）和调整版本（应用了优化）

*在 A100 服务器上测量的 T5 11B 性能。原始版本（未优化）和调整版本（应用了优化）

在第一篇文章中，我们将快速概述 FSDP 以及它如何使大规模 AI 模型训练更高效。我们将简要介绍可用的多种性能选项，并在接下来的文章中深入探讨这些细节。最后，我们将概述如何利用 AWS parallel cluster 进行 FSDP 大规模训练。

优化	T5 模型	吞吐量提升
混合精度	3 B	5x
混合精度	11 B	10x
激活检查点 (AC)	3 B	10x
激活检查点 (AC)	11 B	100x
Transformer 包装策略	3 B	2x
Transformer 包装策略	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 模型那样仅针对单一优化进行开/关测试，而是在始终运行其他两种优化的同时开启/关闭 1 个优化，以便在每项测试的两种状态下都能使用可用的批量大小。

基于 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，我们在配备 40 GB 内存的 32 个 A100 GPU（4 个节点）上，在激活检查点的帮助下，可以训练的最大模型达到 3B 参数，批量大小为 128。相比之下，使用 FSDP，我们能够训练高达 81B 的模型，结合了激活检查点以及激活和参数卸载。在另一项实验中，我们使用 512 个 GPU 对 1T 参数模型进行了 FSDP 基准测试。

为了直观理解 FSDP 的参数级工作原理，下面我们展示一个动画，详细说明模型参数是如何分片和通信的，假设有两个 GPU 和一个简单的 8 参数模型

上图 - 动画展示了模型在各 rank 之间初始分片的步骤，并开始 all_gathers 和前向传播

我们继续模型的前向传播。在每个 FSDP 单元完成后，非本地拥有的参数会被丢弃以释放内存，并且可以选择性地对激活进行检查点。这个过程持续进行，直到完成前向传播并计算损失。

在反向传播期间，使用另一个 all_gather 来加载参数并计算梯度。然后将这些梯度 reduce_scattered，以便每个参数的本地所有者可以聚合并准备更新权重。

最后，每个 rank 将求和的梯度通过优化器状态传递，并更新权重以完成 mini-batch。

现在模型分布在整个可用的 GPU 集上，逻辑上的问题是数据如何通过模型传输，考虑到模型参数的分片。

这通过 FSDP 与所有 GPU 协调来有效地共享（通信）模型的各个部分来实现。模型被分解为 FSDP 单元，每个单元内的参数被展平，然后在所有 GPU 上进行分片。在每个 FSDP 单元内，GPU 以交错的方式分配对单个模型参数的所有权。

交错的意思是：假设有两个 GPU，ID 为 1 和 2，FSDP 单元的所有权模式将是 [12121212]，而不是连续的块 [111222]。

在训练期间，会启动一个 all_gather，当非本地所有者需要时，“及时地”由所有者 GPU 与其他非所有者共享 FSDP 单元内本地拥有的模型参数。FSDP 预取参数以重叠 all_gather 通信与计算。

当这些请求的参数到达时，GPU 会使用接收到的参数，结合它已经拥有的参数，创建一个完全填充的 FSDP 单元。因此，在持有完全填充的 FSDP 单元时，每个 GPU 会达到峰值内存使用。

然后它通过 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，即 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 参数模型的性能差异。

对于默认包装，我们无需采取任何行动——我们只需将模型传递给 FSDP，如下所示

model = FSDP(
      model,
      device_id=torch.cuda.current_device(),
  )

在这种情况下，FSDP 只会将整个模型包装在一个单一的 FSDP 单元中。

在具有 8 个 GPU 的 NVIDIA A100-SXM4–40GB 上运行，我们能够达到 2.3 TFlops，并且在批量大小为 14 时 GPU 内存利用率达到 95%。

然而，由于 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 模型，并说明在哪里可以找到层类。

混合精度 - 如果您的 GPU 是 Ampere 架构，请使用 BF16

FSDP 支持灵活的混合精度策略，使您能够对参数、梯度和缓冲区数据类型进行细粒度控制。这使您可以轻松利用 BFloat16 或 FP16 将训练速度提高多达 70%。

*请注意，BFloat 16 仅在 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 中还是其他地方，都会在保持主权重为完整的 FP32 的同时，将工作权重保持在降低的数据类型（BF16 或 FP16）。保持主权重为 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 训练，实验在之前提到的具有 A100 GPU 的 3 个节点上进行。

如上所示，使用 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 Accelerated Computing 实例、AWS ParallelCluster 和 Amazon Sagemaker。

在本系列博文中，我们使用了 Amazon EC2 p4d 实例的单实例多 GPU 配置以及使用 AWS ParallelCluster 和 SageMaker 的多实例配置来运行我们的训练作业。

在这里，我们将专门关注 AWS parallel cluster，并概述如何将其用于训练目的。

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-1c

CloudFormation 将代表我们部署新的 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 wrapper 和 activation checkpointing。

在接下来的文章中，我们将继续探讨 T5 模型，并深入研究上述每个主题，特别是分片策略和其他优化，以提供更多见解和详细信息。目前，关于分片策略的一个很好的参考可以在我们的视频教程此处找到。

如果您有任何问题或发现问题，请联系作者 Less、Hamid 和 Geeta，或者在 PyTorch Github 上提交问题。

特别感谢

PyTorch 分布式团队、沈力、Rohan Varma、Yanli Zhao、Andrew Gu、Anjali Sridhar、Ana Simoes、Pierre-Yves Aquilanti、Sundar Ranganathan，以及更广泛的 AWS 团队，感谢他们为我们提供运行大规模实验所需的基础设施和技术支持。

资源