Efficient Large-Scale Training with Pytorch FSDP and AWS

由莱斯·赖特、哈米德·肖贾纳泽里、吉塔·乔汗撰写

最先进的 AI 模型正变得极其庞大。训练这些模型的成本和开销正在迅速增加，需要大量的工程和猜测来找到正确的训练方案。FSDP 通过允许您使用相同数量的资源训练更大的模型，显著降低了这些成本。FSDP 降低了 GPU 上的内存占用，并且可以通过轻量级的配置使用，通常只需要几行代码。

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

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

在本系列博客文章中，我们将解释您可以使用 FSDP 运行的多项性能优化，以提升您的分布式训练速度和模型大小，同时考虑到您可用的服务器资源。在本系列中，我们将使用 HuggingFace T5 3B、11B 和 DeepVit 作为运行示例，在微调模式下进行操作。

作为本系列中讨论的一些优化的预览，我们以下列出了性能的“之前”和“之后”的 Flops 比例（请注意，这些结果可能因您的服务器资源和模型架构而异）。

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

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

在这篇第一篇帖子中，我们将快速概述 FSDP 及其如何使大规模 AI 模型的训练更加高效。我们将简要介绍可用的多种性能选项，并在后续帖子中深入探讨这些细节。然后，我们将总结如何利用 AWS 并行集群通过 FSDP 进行大规模训练。

优化	T5 模型	吞吐量提升
混合精度	3 B	5 倍
混合精度	11 B	10 倍
激活检查点（AC）	3B	10 倍
激活检查点（AC）	11B	100 倍
Transformer 包装策略	3B	2 倍
Transformer 包装策略	11 B	无法在没有 Transformer 包装策略的情况下运行实验。
完整分片策略	3 B	1.5 倍
完整分片策略	11B	无法与 Zero2 一起运行

在 T5 模型上相对于非优化模型的性能优化收益

在我们的 T5 3B 模型实验中，使用 transformer 封装策略相比默认封装策略，在 TFLOPS 指标上实现了超过 2 倍的吞吐量提升。通过将检查点释放的内存重新投资到更大的批次大小，激活检查点实现了 10 倍的性能提升。与 FP32 相比，使用 BFloat16 的混合精度实现了约 5 倍的性能提升。最后，与 zero2（DDP）相比，完整的分片策略实现了 1.5 倍的性能提升。

我们为更大的模型 T5 11B 进行了类似的实验，但更大的模型尺寸导致实验空间发生了一些变化。具体来说，我们发现需要两种优化，即 transformer 包装策略和激活检查点，才能使我们能够在 3 个节点上运行这些实验（每个节点有 8 个 A100 GPU 和 80 GB 的内存）。有了这些优化，我们可以将批大小设置为 50，并且与移除任何一个相比，获得更高的吞吐量。因此，与 3B 模型仅针对单个优化测试开/关运行不同，较大的模型实验是通过打开/关闭 3 个优化中的 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 通过消除优化器计算和状态存储的冗余，以及模型参数中存在的梯度内存存储，扩展了缩放模型。这种冗余减少，加上模型参数通信与模型计算同时进行的通信重叠增加，使得 FSDP 能够使用与 DDP 相同的资源训练更大的模型。

关键点在于，这种效率还允许训练比单个 GPU 更大的 AI 模型。现在可用于训练的模型大小已增加到所有 GPU 的聚合内存，而不是单个 GPU 的大小。（顺便提一下，FSDP 可以通过利用 CPU 内存来超越聚合 GPU 内存，但这里我们不会直接涉及这一方面）。

如前一篇博客文章所述，使用 DDP，我们能在 32 个 A100 GPU（4 节点，40GB 内存）上训练的最大模型参数量达到 30 亿，批大小为 128，借助激活检查点技术。相比之下，使用 FSDP，我们能够训练到 810 亿模型大小，结合了激活检查点、激活和参数卸载。在另一个实验中，我们使用 512 个 GPU 对 1T 参数模型进行了 FSDP 基准测试。

关于 FSDP 在参数层面的工作原理，以下动画展示了在两个 GPU 场景和简单的 8 参数模型中，模型参数是如何分片和通信的：

如上图所示，动画演示了模型在各个 rank 之间进行初始分片的过程，以及我们开始执行 all_gathers 和前向传递。

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

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

最后，每个 rank 将汇总的梯度通过优化器状态传递，以完成小批次的权重更新。

现在，模型已经分布到所有可用的 GPU 上，一个逻辑问题是，在这种参数划分的情况下，数据是如何在模型中移动的。

这是通过 FSDP 与所有 GPU 协调，有效地共享（通信）模型的相关部分来实现的。模型被分解为 FSDP 单元，每个单元内的参数被展平，然后跨所有 GPU 进行划分。在每个 FSDP 单元内，GPU 被分配交错的所有权，以管理单个模型参数。

通过交织，我们指的是以下情况——假设有两个 ID 为 1 和 2 的 GPU，FSDP 单元的所有权模式将是[12121212]，而不是连续的[111222]块。

在训练过程中，会启动一个 all_gather 操作，FSDP 单元内的本地拥有的模型参数由拥有 GPU 与其他非拥有 GPU 在需要时以“即时”的方式共享。FSDP 预取参数以重叠 all_gather 通信与计算。

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

然后通过 FSDP 单元处理数据，并丢弃从其他 GPU 接收到的参数以释放内存，为下一个单元做准备……这个过程会不断重复，贯穿整个模型以完成前向传递。然后（通常）重复这个过程进行反向传递。（注意——这是一个简化的版本，以便理解，实际上还有额外的复杂性，但这应该有助于构建 FSDP 过程的基本思维模型）。

这消除了 DDP 中存在的许多内存冗余，但同时也带来了更高的网络通信成本，需要在所有 GPU 之间来回传输这些请求的参数。将通信时间与正在进行的计算重叠是本系列中我们将讨论的许多性能改进的基础。关键收益通常基于通信可以经常与计算同时进行的事实。正如你可以推测的那样，拥有高通信速度对于 FSDP 性能至关重要。

我该如何优化我的 FSDP 训练？

我们将介绍四个主要性能改进 - 变压器包装器、激活检查点、混合精度以及选择合适的分片策略。下面的流程图将作为本篇帖子中我们将讨论的调整选项的清单。

包装策略 - 对于变压器，使用 Transformer 包装策略

首先的性能优化是利用 FSDP 转换器包装器来处理转换器模型。

预定义的包装策略之一是 size_based_autowrap_policy 。使用 size_based_autowrap_policy ，FSDP 将从底部到顶部遍历模块结构，一旦当前单元的大小策略（默认为 1e8，即 100M）中至少包含 min_num_params 指定的值，就会创建一个新的 FSDP 单元。如果无法将模块创建为 FSDP 单元，FSDP 将继续检查其父模块。这种基于大小的包装策略可能不适合某些模型结构，PyTorch 分布式团队正在积极开发下一个版本的新默认包装策略，该策略基于大小和模块执行顺序，用户可以简单地调整大小以实现优化的性能。

在当前版本中，您可以通过使用“转换器包装器”来显著提高运行转换器模型时的性能。您需要为您的模型提供适当的层类。这里的层类是包含多头注意力和前馈网络的类。

FSDP 将围绕层类形成 FSDP 单元，而不是基于参数大小的任意断点。通过在均匀重复于 transformer 中的层类周围分片模型，FSDP 可以创建更平衡计算和通信重叠的均匀 FSDP 单元。相比之下，基于大小的包装可能会为模型产生非常不均匀或倾斜的碎片，从而导致计算与通信重叠不匹配。如前所述，FSDP 高性能的主要驱动因素是通信和计算的叠加，这也是为什么 Transformer 包装器提供了改进性能的原因。请注意，Transformer 包装器也可以用于非 Transformer 模型，如果这些模型有一个均匀层的列表。

让我们比较在默认包装器和 Transformer 包装器下运行 T5，3B 参数模型时的性能差异。

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

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

在这种情况下，FSDP 将简单地使用单个 FSDP 单元包装整个模型。

在 NVIDIA A100-SXM4–40GB 和 8 个 GPU 上运行，我们能够达到 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 倍）！变压器包装策略导致 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%。

我们使用 BFloat16 在微调 3B HuggingFace T5 模型时获得了更大的加速，如图所示。我们观察到，由于精度较低，BFloat16 的验证损失在最初的几个 epoch 中略低于 FP32，但它能够赶上并最终达到与 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，它是一个即插即用的替代品。

任意精度优化器 - 超越混合精度，实现全 BF16 训练

混合精度训练，无论是在 FSDP 还是其他地方，都保持工作权重在降低的数据类型（BF16 或 FP16）中，同时保持主权重在完整的 FP32 中。主权重使用 FP32 的原因是，在纯 BF16 下运行会导致“权重停滞”，由于精度较低，非常小的权重更新会丢失，并且精度随着时间的推移而平坦，而 FP32 权重可以继续从这些小的更新中改进。

为了解决这个困境，我们可以使用 TorchDistX（Torch 分布式实验）中提供的 AnyPrecision 优化器，它允许您在纯 BF16 而不是 FP32 中成功训练并保持主权重。此外，与通常在 FP32 中存储优化器状态不同，AnyPrecision 还能够以纯 BF16 维护状态。

AnyPrecision 通过维护一个额外的缓冲区来跟踪权重更新过程中丢失的精度，并在下一次更新时重新应用，从而有效地解决了权重停滞问题，而无需 FP32。

为了比较使用 AnyPrecision 进行纯 BF16 训练的吞吐量增益，我们使用 FSDP 在 T5 11B 模型上进行了实验，包括常规 FP32 训练、混合精度训练（BF16）和纯 BF16 训练（使用 AnyPrecision 优化器），这些实验在之前提到的 3 个节点上使用 A100 GPU 进行。

如上图所示，使用 AnyPrecision 和纯 BF16 进行训练的吞吐量比混合精度提高了 2 倍，比 FP32 提高了 20 倍以上。

潜在的权衡是最终准确性的影响——在我们测试的案例中，由于略微降低的精度带来的正则化效应，准确度与 FP32 相当甚至更好，但你的结果可能会有所不同。

任何精度优化器可供您在此测试，并且是 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）。

集群工作流程

高级思想是拥有一个包含控制节点和计算节点的集群。实际的训练任务在计算节点上运行。在集群上运行训练任务的总体步骤如下：

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

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

下面是一个集群配置文件的示例。我们可以使用 AWS ParallelCluster CLI 创建集群。请注意，需要将之前获取的私有和公共子网 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 上提交一个 issue。

特别感谢：

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

资源：

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

我该如何优化我的 FSDP 训练？

文档

教程

资源