作者/机构: Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li (Meta AI)

主要贡献

本文介绍了PyTorch的完全分片数据并行（Fully Sharded Data Parallel, FSDP），这是一个用于大规模模型训练的工业级解决方案。FSDP旨在解决现有分布式训练方法在通用性、与框架核心功能的耦合性以及对社区创新友好性方面的挑战。其核心研究目标是提供一个与PyTorch核心组件（如Tensor实现、分发器系统和CUDA内存缓存分配器）紧密协同设计的原生解决方案，以实现非侵入式的用户体验和高训练效率。

FSDP的核心算法源于DeepSpeed的Zero-RedundancyOptimizer技术，但经过了重新设计和实现。其工作原理是将模型实例分解为更小的单元，并将每个单元内的所有参数展平并分片。在计算前，分片的参数会按需被通信和恢复成未分片形式，计算后立即被丢弃。这使得FSDP在任何时候只需要在内存中物化一个单元的参数，从而显著降低峰值内存消耗。

为了应对大规模模型训练中的具体挑战，FSDP引入了多项关键创新和技术：
* 提升用户体验：针对大模型无法在单GPU上实例化的问题，FSDP引入了延迟初始化（deferred initialization）。该技术允许用户在虚拟设备上创建模型实例并记录初始化操作，随后在真实GPU上通过重放记录的操作来逐单元地初始化和分片模型，从而提供了与本地训练相似的用户体验。
* 应对硬件异构性：现代GPU集群通常具有非均匀的互连拓扑。FSDP提供了可配置的分片策略，用户可以根据集群的物理互连拓扑自定义分片方式，以优化通信效率。
* 优化资源利用率：为了减少通信开销导致的计算空闲（bubbles），FSDP通过激进的通信与计算重叠策略来最大化GPU利用率。这包括操作重排序和参数预取等一系列丰富工具。
* 精细化内存规划：为避免在接近GPU内存容量时频繁发生内存碎片整理从而拖慢训练，FSDP通过审慎地限制在途（inflight）的未分片参数所分配的内存块数量，并在必要时暂停CPU执行来进行优化。

实验结果表明，FSDP在小模型上能够达到与分布式数据并行（DDP）相当的性能，同时能够支持远大于DDP的模型，并在线性算力（TFLOPS）方面展现出近线性的扩展能力。

背景知识

PyTorch基础 PyTorch【24, PyTorch: An Imperative Style, High-Performance Deep Learning Library, 2019, Advances in Neural Information Processing Systems】已成为众多机器学习项目的基石。它将值存储在Tensor对象中，这是一种功能丰富的n维数组，支持多种数据操作。每个Tensor对象都有一个在特定设备上分配的关联存储。当Tensor仅表示简单的转换（如reshape和split）时，它们可以共享相同的底层存储。每个Module描述了从输入到输出值的转换，其前向传播行为由其forward成员函数指定。这样的模块可能包含作为参数的Tensor对象，例如Linear模块就包含权重和偏置参数。在前向传播过程中，Linear模块通过乘法和加法运算将这些参数应用于输入以产生输出。

分布式训练框架的需求 随着数据规模和模型复杂度的急剧增长，一个工业级的分布式训练框架对于构建在PyTorch之上的应用变得日益重要。本节阐述了PyTorch分布式训练能力的发展轨迹。

模型复制

DistributedDataParallel (DDP) 模型复制方法旨在通过横向扩展和跨多设备分布计算来处理海量数据集。DistributedDataParallel (DDP)【14, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv preprint】是PyTorch中属于此类别的第一个端到端分布式训练功能。DDP的采用已非常广泛，涵盖学术界和工业界。DDP在每个设备上维护一个模型副本，并在反向传播过程中通过AllReduce集合操作同步梯度，从而确保训练期间各副本之间的模型一致性。为了加速训练，DDP将梯度通信与反向计算重叠，促进了不同资源上的工作负载并发执行。然而，一个显著的限制是DDP要求所有模型参数、梯度和优化器状态都能装入单个GPU设备的内存中。因此，DDP不足以支持对于前沿机器学习突破至关重要的大型模型。例如，当使用40GB的GPU设备训练超过十亿参数的模型时，DDP很可能会在每个设备上遇到内存不足的错误。

模型分区

流水线并行与Tensor RPC 随着模型规模的增长，它们可能不再适合单个GPU设备。在这种情况下，一个可行的解决方案是将模型划分为更小的组件并将其分布在多个设备上。流水线并行【8, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019, Advances in neural information processing systems】和Tensor RPC【25, DISTRIBUTED RPC FRAMEWORK, 2023, pytorch.org】都属于这个方向。流水线并行涉及将一系列层分解为多个阶段，并以流水线方式将输入送入不同阶段以优化资源利用。另一方面，Tensor RPC提供了一个更底层的工具包，允许在远程设备上执行任意计算。尽管这两种技术都能够跨多设备扩展大型模型，但它们要么将模型限制为一系列阶段，要么需要修改模型编写代码以插入远程计算，这可能对用户的采用构成重大障碍。此外，许多工业训练基础设施仅支持单程序多数据（SPMD）范式，这需要一个更简单的入口点来处理大型模型。

模型分片

两种分片计算范式 除了分区，对模型参数进行分片也可以帮助减少其内存占用，并支持尺寸超出单个GPU设备内存容量的模型。在对模型进行分片后，每个rank只持有一部分模型参数，这使其无法执行与本地训练相同的计算。为保证正确性，训练过程需要采用以下一种或两种技术：
* 使用参数分片执行计算，并相应地通信激活值。通过这种方法，rank永远不需要完全物化任何参数。然而，每次通信都会出现在关键路径上，因为它被插入到两个连续且相互依赖的计算操作之间。因此，这种通信不容易与计算重叠，除非可以将不相关的计算或其他迭代的计算重新排序以与通信重叠。
* 通过在计算前按需通信参数来执行与本地训练相同的计算。由于参数通信与之前的计算没有任何数据依赖关系，它们可以与在相同前向或反向传播中执行的先前计算重叠。然而，这种方法要求按需通信的参数能够被完全物化并能装入单个GPU设备的内存中。

FSDP的选择 FSDP属于第二类，即通信参数。根据我们的观察和实验，这种方法足以支持当今和不久的将来绝大多数大型模型应用。值得注意的是，如果完全物化每个参数单元的需求成为瓶颈，我们可以进一步结合这两种技术来支持此类用例。

系统设计

FSDP核心机制 完全分片数据并行（FSDP）能够通过分片密集参数来扩展以适应可能无法放入单个GPU设备的大型模型。具体来说，FSDP将模型实例分解为更小的单元，并独立处理每个单元。在前向和反向计算期间，FSDP仅在需要时物化一个单元的未分片参数和梯度，而在其他时候，它保持参数和梯度为分片状态。在整个训练循环中，优化器状态始终保持分片。FSDP的内存需求与分片模型的大小加上最大完全物化FSDP单元的大小成正比。

FSDP工作流程概述 图1展示了使用一个简单的六层模型的整体工作流程。假设FSDP将模型分解为三个部分，即[layer0, layer3]、[layer1, layer2]和[layer4, layer5]。分解行为可以通过用户定义的函数来控制。然后，FSDP将这三个部分分别包装到一个FSDP单元中，并相应地对参数进行分片。为确保正确性，FSDP需要在相应的计算之前恢复未分片的参数。我们以包含[layer1, layer2]的FSDP unit1为例来解释这个过程。在前向计算进入layer1之前，FSDP通过从其他对等rank收集分片来恢复layer1和layer2的未分片参数。有了未分片的参数，FSDP运行这些层的本地计算，然后释放它刚刚收集的对等分片以减少内存占用。因此，在整个前向传播过程中，FSDP一次只需要完全物化一个单元，而所有其他单元都可以保持分片状态。同样，在反向计算期间，FSDP unit1在反向传播到达layer2之前恢复layer1和layer2的未分片参数。当自动求导引擎完成这两层的反向计算后，FSDP释放对等分片并发起ReduceScatter来规约和分片梯度。因此，在反向计算之后，每个rank只保留参数和梯度的一个分片。


图 1: FSDP算法概述

优化与定制 FSDP提供了广泛的优化和可调参数，以适应不同的模型结构和硬件能力。本节的其余部分将深入探讨模型初始化、分片策略、通信优化和内存管理的复杂性，这些都是FSDP底层设计的关键组成部分。

模型初始化

初始化挑战 在FSDP出现之前，PyTorch要求在单个设备上完全物化整个模型实例。尽管用户可以将不同的子模块分配到不同的设备上，但这需要修改模型源代码，这可能并不可行，特别是当模型作者和应用开发者属于不同方时。为了促进从本地训练到分布式训练的平稳过渡，FSDP必须有效地帮助物化和初始化一个巨大的模型，这带来了两个挑战：
* 如何在不物化任何张量存储的情况下创建一个模型实例，将初始化推迟到存储在具体设备上附加到张量时。
* 如何确保模型参数的初始化与用户的实现一致，即使模型大到无法装入单个GPU。

延迟初始化机制 为了克服第一个挑战，我们引入了一种称为延迟初始化（deferred initialization）的机制，它涉及在模拟的或“伪”（fake）设备上分配模型参数张量。在此过程中，对张量执行的所有初始化操作都会被记录下来。随后，当张量从“伪”设备移动到GPU设备时，所有记录的操作都会自动重放。通过采用这种技术，用户可以从任何第三方库生成一个模型实例，而无需分配任何GPU内存块，同时仍然准确地捕获其参数初始化实现。

逐单元初始化 如图1所示，一旦FSDP包装了模型，它就会被均匀地分布在所有GPU上，每个设备在其内存中只持有一个分片。因此，为了解决第二个挑战，每个rank理想情况下应该只物化和初始化它拥有的分片。然而，这并不总是实际的，因为我们无法预测用户将在模型的init方法中实现什么初始化逻辑。初始化逻辑可能依赖于设备上有一个未分片的参数，这使得对初始化过程进行分片变得不可能。因此，FSDP必须在执行Tensor初始化操作之前准备好未分片的参数，并同时减少内存占用。鉴于分片初始化是不安全的，FSDP采用了与处理模型前向和反向传播相同的方法，即一次初始化一个FSDP单元，并在移动到下一个单元之前对该单元进行分片。当与延迟初始化相结合时，FSDP遍历伪设备模型实例以将其分解为FSDP单元，一次将一个单元移动到GPU设备，并为该FSDP单元中的张量重放记录的初始化操作。


图 2: 通信效率与输入大小的关系

分片策略

分片策略的重要性 分片策略是FSDP中的一个重要元素，它在决定内存占用和通信开销方面扮演着重要角色。FSDP提供了多种分片策略，从完全复制到完全分片。为了概括这些分片策略，我们引入分片因子S作为参数分片所跨的rank数量。通过将分片因子设置为1，FSDP完全复制模型，简化为使用AllReduce进行梯度规约的普通数据并行。通过将分片因子设置为等于设备数量（即全局world size N），FSDP完全分片模型，每个设备只持有模型的1/N。当分片因子在1和N之间时，会发生混合分片。本节的其余部分将重点关注完全分片和混合分片，因为完全复制策略与现有的DDP【14, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv preprint】相似。

完全分片

权衡与优化 完全分片策略导致最低的内存占用，但会产生最大的通信开销。例如，如果使用带宽最优的环形算法，完全分片的通信开销和通信量是DDP的1.5倍。因此，FSDP必须在此策略下仔细组织通信以最大化其效率。

通信效率的关键因素 我们进行了两组实验来理解输入大小对集合通信效率的影响。结果如图2所示，这帮助我们确定了效率的两个要素：
(1) 均匀的输入大小：Nvidia NCCL【22, The NVIDIA Collective Communication Library (NCCL), 2023, developer.nvidia.com】库为all-gather和reduce-scatter提供了高效的集合实现，这些实现要求跨rank的输入张量大小是均匀的。
(2) 更大的输入大小：对于固定的通信量，将数据批处理并发出更少的集合操作可以通过避免集合的启动开销和提高网络带宽利用率来提高性能。

针对均匀输入大小的优化 对于(1)，NCCL的AllGather API要求输入张量大小均匀，并将输出写入单个张量中。PyTorch的ProcessGroup包装了NCCL API，并通过支持跨rank不均匀的输入张量大小以及允许用户提供一个输出张量列表来增强它。这种灵活性带来了效率上的权衡，如图2 (a)所示。我们使用All-Gather Base来表示NCCL的AllGather行为，使用All-Gather来表示接受张量列表作为输出的行为。后者在通信前后会在各个输出张量和合并的单个大输出张量之间产生额外的拷贝。此外，对于不均匀的输入，ProcessGroup使用组内Broadcast来模拟AllGather的行为，这比All-Gather Base慢。在实验中，我们通过将1个元素和1/16个元素分别从rank 1移动到rank 0来人为制造不均匀性。结果表明，具有均匀输入大小的All-Gather Base实现了最高的效率。

针对更大输入大小的优化 对于(2)，图2 (b)将总通信量固定为 $2^{30} \approx 10亿$ 个FP32元素，并改变每次All-Gather的大小，即更小的All-Gather大小意味着更多的All-Gather调用。一旦All-Gather的大小减少到33M个元素以下，总通信时间开始迅速增加。

FlatParameter设计 因此，为了实现高效的通信，FSDP将一个FSDP单元内的所有参数组织成一个大的FlatParameter。FlatParameter合并了其内部各个参数的通信，并将它们均匀地分片到各个rank。具体来说，FlatParameter是一个一维张量，通过连接P个展平的原始参数并在右侧填充以达到一个可以被分片因子整除的大小来构建。为了分片FlatParameter，FSDP将其划分为大小相等的块，块的数量等于分片因子，并为每个rank分配一个块。FlatParameter的梯度继承了与FlatParameter相同的未分片和分片形状，FlatParameter及其梯度分别拥有原始参数及其梯度的底层存储。图3描绘了一个例子，我们使用一个FSDP单元在16个GPU上分片一个4×3的nn.Linear层。在这种情况下，每个GPU只持有FlatParameter中的一个元素，最后一个rank持有填充值。


图 3: 在16个GPU上进行完全分片

算法的通用性与内存分析 这种“展平-连接-分块”算法允许每个原始参数具有任意形状，同时最小化了所需的填充（最多为S−1），反映了其通用性。此外，在此算法下，分片的和未分片的FlatParameter及其梯度具有AllGather和ReduceScatter所期望的精确数据布局。这使得调用集合操作时无需对输入或输出张量进行任何额外的拷贝。

内存-吞吐量权衡 更正式地，假设一个模型有 $\Psi$ 个元素，FSDP构建了 K 个FlatParameter，其元素数量为 $p_1, \dots, p_K$，其中 $\sum_{i=1}^{K} p_i = \Psi$。对于分片因子 S，峰值参数内存贡献为 $O(\sum_{i=1}^{K} \frac{p_i}{S} + \max_{i=1}^{K} p_i)$，因为FSDP始终在GPU内存中保留每个大小为 $p_i/S$ 的本地分片FlatParameter，并且必须在前向和反向传播期间逐个物化每个大小为 $p_i$ 的未分片FlatParameter。由于第一项 $\sum_{i=1}^{K} p_i/S = \Psi/S$ 是固定的，峰值参数内存贡献由 $\max_{i=1}^{K} p_i$ 决定。同时，每次迭代的集合操作数量为 $O(K)$。这证明了FSDP的内存-吞吐量权衡：更细粒度的FlatParameter构建会降低峰值内存，但可能因需要更多集合操作而降低吞吐量。用户可以通过指定如何将子模块包装成FSDP单元来控制这种权衡。


图 4: 在16个GPU上的混合分片：GPU被配置为2个分片组和8个复制组

混合分片

定义与分组 我们将分片因子大于1但小于 N 的策略称为混合分片，因为它结合了分片和复制。对于全局world size N 和分片因子 S，参数在每个组 $G_1, \dots, G_{N/S}$ 内进行分片，并在每个互补组 $H_1, \dots, H_S$ 内进行复制，其中每个 $G_i, H_j \subseteq \{1, \dots, N\}$ 分别给出了分片组或复制组中的rank。

梯度规约 对于梯度规约，原来在所有rank上的单个reduce-scatter变成了在每个分片组内的reduce-scatter，然后是在每个复制组内的all-reduce，以规约分片的梯度。这种等价性来自于以下分解：

其中 $g_r$ 表示rank r 上的梯度。

利用数据中心局部性 混合分片可以利用数据中心的局部性来加速训练，并可以减少跨主机流量以尽可能避免在超额订阅环境中的争用。同时，它在内存节省和吞吐量下降之间提供了渐进的权衡，这对于那些使用完全复制训练时所需内存足迹略高于设备容量且不希望完全分片的模型特别有帮助。图4展示了一个例子。

流量优化 具体来说，数据中心通常采用超额订阅的胖树网络拓扑【16, Incbricks: Toward in-network computation with an innetwork cache, 2017, Proceedings of the Twenty-Second International Conference on Architectural Support for Programming Languages and Operating Systems】，这导致了丰富的可利用局部性，并为减少跨主机流量提供了充分的理由【17, Plink: Discovering and exploiting locality for accelerated distributed training on the public cloud, 2020, Proceedings of Machine Learning and Systems】。混合分片可以提供一种自然机制，将设备网格映射到数据中心布局以利用这种局部性。例如，考虑一个集群由N个加速器组成，这些加速器分组到每个包含k个加速器的主机中（同一主机上的加速器之间的通信远快于跨主机的通信），我们可以设置S = k来限制AllGather（和ReduceScatter）操作在同一主机内，同时为跨主机的具有相同本地rank的加速器创建一个复制组。对于一个大小为M的模型，我们可以计算出混合设置中每个GPU的总跨主机流量为 $2 \frac{M}{S} \frac{k-1}{k}$，与完全复制的 $2M \frac{N-1}{N}$ 和完全分片的 $3M \frac{N-1}{N}$ 相比，这是一个巨大的减少。此外，由于混合分片中使用的AllReduce集合操作在更小的world size下运行，它们在经验上比在全局范围（完全复制和完全分片的情况）调用集合操作表现更好，这是由于掉队者效应和更大的网络干扰。

适应中等规模模型 混合分片的另一个重要设计动机是来自中等规模模型的需求。这些模型大到在使用完全复制训练时会导致内存不足问题，但又不够大到在使用完全分片时能充分利用加速器内存，导致运行时开销和内存浪费。混合分片策略通过简单地调整S创建了一个更丰富的内存-吞吐量权衡空间。

Autograd

与Autograd引擎的互操作 FSDP的FlatParameter必须与PyTorch的autograd引擎互操作，以确保(1)正确的梯度传播和(2)及时的梯度规约。对于(1)，回想一下FlatParameter及其梯度分别拥有原始参数及其梯度的底层存储。为了实现这一点，在前向计算之前，FSDP使用autograd可见的torch.split()和torch.view()调用，将原始参数设置为其未分片FlatParameter的视图。然后，autograd引擎自然地为未分片的FlatParameter分配梯度，并将每个原始参数的梯度写入由torch.split()的反向函数定义的适当偏移量。对于(2)，FSDP注册一个梯度钩子，该钩子仅在FlatParameter的梯度最终确定后运行。该钩子代表了后向传播逻辑，并包括梯度规约。值得注意的是，FSDP的方法建立在PyTorch的autograd引擎之上，而不是绕过它。因此，FSDP自动处理非传统情况，例如当并非所有参数都在前向传播中使用或在一次反向传播之前有多次前向传播时。


图 5: 重叠通信和计算

通信优化

内置优化技术 FSDP框架集成了一系列原生的通信优化技术。本节将揭示四个主要技术：重叠、反向预取、前向预取和累积。

重叠通信和计算

利用独立CUDA流 PyTorch的c10d库有一个ProcessGroup抽象，它代表一组可以一起运行集合操作的进程。对于NCCL后端，ProcessGroupNCCL实现为每个设备提供一个内部NCCL流，这个独立的内部流用于与当前流（通常是运行计算的默认流）进行异步执行。这些异步集合操作返回Work对象，调用Work.wait()会阻塞CPU线程直到集合操作完成。为了一般的正确性，ProcessGroupNCCL在运行集合操作之前会同步内部流与当前流。DistributedDataParallel利用async-collective-and-wait()方法来将梯度All-Reduce与反向计算重叠。然而，与DDP的反向传播中All-Reduce在要重叠的计算之后不同，FSDP的前向传播在要重叠的计算之后发出All-Gather，因为在即时执行（eager execution）模式下，FSDP无法知道下一个要All-Gather的FlatParameter以便将其重新排序到计算之前。这种内核发出顺序的差异使得遵循async-collective-and-wait()方法对FSDP来说是不可行的。也就是说，由于ProcessGroupNCCL与当前（默认）流同步，All-Gather要等到与之重叠的计算完成后才会运行。为了解决这个问题，FSDP使用一个独立的CUDA流来发出All-Gather，从而绕过了对默认流中先前计算的虚假依赖，并允许每个All-Gather进行重叠。因此，FSDP的集合同步是基于流操作的，而不仅仅是Work对象。图5展示了一个例子。请注意，反向传播不包括AG0的All-Gather，因为FSDP有意将最外层FSDP单元的参数保留在内存中，以避免在前向传播结束时冗余地释放，然后在开始反向传播时重新进行All-Gather。

反向预取

解决反向传播中的通信阻塞 FSDP强制每个rank使用单个CUDA设备，并为AllGather和ReduceScatter使用单个进程组，这意味着其集合操作在该进程组的内部NCCL流中顺序运行。在反向传播中，FSDP为当前的FlatParameter发出ReduceScatter，然后为下一个FlatParameter发出AllGather。因此，单个NCCL流迫使ReduceScatter阻塞下一个AllGather，这反过来又阻塞了下一个梯度计算，并可能暴露在关键路径上。

预取策略 为了避免在反向传播中出现两个连续的暴露通信调用，FSDP的反向预取在当前的ReduceScatter之前发出下一个AllGather。然而，如前所述，即时执行模式的一个挑战是知道下一个要AllGather的FlatParameter是哪个。FSDP通过记录模块的前向执行逆序作为其反向执行顺序的代理来解决这个挑战。此外，前向顺序在每次迭代中都会重新记录，这意味着反向预取与跨迭代的动态性兼容。

前向预取

应对慢速CPU 对于一些CPU执行相对较慢的工作负载，CPU线程可能无法足够早地发出下一个前向AllGather以有效地填充NCCL流。如果模型在多次迭代中遵循静态计算图，那么FSDP可以假设来自前一次迭代的模块前向执行顺序，并在前向传播中显式地预取下一个AllGather。这种前向预取在当前FSDP单元的前向计算之前发出下一个AllGather。

梯度累积

两种累积方式 FSDP提供了两种梯度累积的变体：带通信和不带通信。带通信时，FSDP仍然在rank之间规约梯度，每个rank保存分片的梯度。只需运行多次迭代而不清除梯度即可实现。不带通信时，FSDP不在rank之间规约梯度，每个rank保存未分片的梯度。后一种变体以增加内存使用为代价换取减少通信，从而可以提高端到端的吞吐量。

内存管理

缓存分配器的角色 PyTorch使用一个CUDA缓存分配器作为中间层，为PyTorch程序提供GPU分配和释放请求。为了有效管理内存，FSDP使用一个速率限制器来考虑缓存分配器对使用多个CUDA流并运行快速CPU线程的程序的内存影响。

PyTorch缓存分配器如何影响内存

缓存机制与挑战 缓存分配器避免了频繁调用cudaMalloc和cudaFree，后者会引起昂贵的设备同步。具体来说，缓存分配器请求CUDA内存块，并在内部决定如何分割和重用这些块，而不将它们返回给CUDA，目标是达到一个稳定状态，不再需要调用cudaMalloc和cudaFree。缓存分配器从CPU线程运行，这意味着当CPU线程处理分配请求时，它必须决定使用哪个缓存分配器块。它不能等到需要该分配的GPU内核实际运行时才决定，而这可能要晚得多。

多流环境下的问题 对于单个流，缓存分配器可以根据流的顺序执行语义直接重用内存块。然而，对于独立的生产者和消费者流，没有流间的顺序保证，缓存分配器无法确定一个块是否可以安全重用，直到最后一个依赖该内存的GPU内核运行完毕。因此，如果CPU线程远超前于GPU执行，缓存分配器就无法为有待处理的消费者流GPU内核的生产者流重用块。此外，缓存分配器块是按流分配的，不能被其他流重用，这导致对生产者流的过度分配，而这些块本可以用于消费者流（例如用于激活值）。GPU本身可能有足够的内存来服务消费者流中的新分配，但对生产者流的过度分配可能导致缓存分配器无法服务它。这会强制执行一个阻塞的cudaFree序列来重置缓存分配器的内存状态，这被称为cudaMalloc retry，它会极大地降低训练吞吞吐量。

速率限制器

FSDP中的内存问题 FSDP在生产者流中分配代表未分片FlatParameter的AllGather目标张量，而使用AllGathered参数的前向和反向计算在消费者流（通常是默认流）中运行。对于一个快速的CPU线程，当缓存分配器必须服务下一个AllGather时，可能存在待处理的GPU计算内核，导致无法重用块。即使这些块在AllGather生产者流中不再活跃，这些保留的块也无法服务于默认计算流的分配请求，因此可能强制执行阻塞的cudaFree和cudaMalloc。

解决方案 FSDP提供了一个速率限制器，它有意地阻塞CPU线程以确保缓存分配器块的正确重用。它最多允许两个在途的AllGather，这是实现通信和计算重叠所需的最小数量。

实现

本节深入探讨了FSDP实现的复杂细节，这些细节虽然不改变FSDP的核心算法，但在采用FSDP之前理解它们至关重要。用户可以通过两个API访问FSDP：FullyShardedDataParallel模型包装器和fully_shard模块注解器。前者包装整个模型并将子模块替换为相应的FSDP单元。相比之下，后者将FSDP逻辑作为nn.Module的前向和反向钩子安装，保留了模型结构和参数的完全限定名称。

初始化

针对复杂情况的备选方案 第3.2.1节描述了FSDP高效初始化大型模型的解决方案，这在子模块初始化是自包含的情况下效果很好。在罕见的情况下，如果一个子模块的初始化依赖于另一个不同子模块的参数，按需物化和记录-重放的方法可能会失效，如果该参数属于不同的FSDP单元，因为该参数的未分片版本可能已经被丢弃以减少内存占用。因此，除了先进的延迟初始化外，FSDP还提供了另外两个选项：
* 在GPU上初始化未分片模型。模型初始化的内存需求可能小于训练时的需求，因为训练还涉及梯度、激活值和优化器状态。因此，如果训练步骤无法在单个GPU设备上执行，用户可能仍然能够在GPU上初始化整个模型，并将其传递给FSDP。然后，优化器应在FSDP对模型进行分片后实例化，以减少内存占用并与FSDP产生的已分片梯度对齐。
* 在CPU上初始化未分片模型。如果未分片模型的大小超过GPU内存容量，只能容纳在CPU内存中，那么在将其交给FSDP进行参数分片之前，将整个未分片模型移动到GPU是不切实际的。为了克服这个挑战，FSDP采用了一种流式方法，即模型被逐单元迁移到GPU。每个单元到达GPU后，其参数立即被分片，这在处理下一个单元之前减少了内存开销。即使在初始化过程中存在跨子模块的依赖关系，这种方法仍然可行，因为整个未分片模型的所有参数都存在于CPU内存中。

方案的局限性 请注意，上述两种方法都有其自身的局限性。第一种方法要求整个模型能装入单个GPU设备，因此对于更大的模型变得不可行。第二种方法可以处理更大的模型，因为CPU的内存要大得多。然而，由于CPU有限的内存带宽和并行化能力，这种方法与延迟初始化相比可能会经历显著的减速。鉴于这些观察，即使用户处理的模型大小在前两种方法所能覆盖的范围内，他们可能仍倾向于使用延迟初始化。

FSDP单元的界定 为了界定每个FSDP单元的范围，用户可以选择通过在模型源代码中侵入式地将其应用于子模块来使用FullyShardedDataParallel包装器，或者在实例化时为auto_wrap_policy参数提供一个自定义函数。选择最佳的包装方法通常需要一些实验和测量。

Flat Parameters

FlatParameter的管理 FlatParameter类继承自nn.Parameter，其行为类似于一个nn.Parameter。FSDP实现了一个配套的FlatParamHandle类，负责管理单个FlatParameter实例。前端，无论是FullyShardedDataParallel还是fully_shard，仅通过FlatParamHandle与FlatParameter交互。

FSDP单元边界的重要性 一个FlatParameter容纳一个FSDP单元内所有参数张量的存储。FSDP单元的边界控制着AllGather和ReduceScatter的时机，这对整体FSDP性能有直接影响。在理想情况下，FSDP单元边界应与模型执行顺序对齐。

FlatParameter的构建 FSDP在构建时可以访问模型的静态nn.Module结构。幸运的是，尽管这种结构不保证忠实地代表模型执行顺序，但模型作者通常会将层和更广泛的块转换为嵌套的nn.Module定义，这可能自然地具有所需的参数局部性。FSDP可以利用该结构来选择FlatParameter的构建方式。实际上，FSDP支持对nn.Module进行注解，并遵循一个简单的规则：被注解的nn.Module中的所有参数都被分配给一个FlatParameter，不包括那些已经被分配的参数。这个规则很自然地适用于嵌套注解，其中块被注解，形成大小合适的FlatParameter，任何剩余的参数则被分配给它们的父级。

动态重构 我们探索的另一种方法是使用执行顺序并动态地重构FlatParameter。这种方法从一个初始的小FlatParameter构建开始，运行一个可能效率不高的第一次迭代，同时观察执行顺序，然后根据观察到的顺序通过合并现有的小FlatParameter来重构它们。

运行时

通信操作的插入 FSDP通过加入通信操作来增强本地模型实例，以规约梯度和收集参数。及时启动这些操作对于确保正确性和效率至关重要。过早启动通信会导致消耗待更新的参数或梯度，而过晚启动通信则会浪费网络带宽并延迟后续计算。

利用PyTorch钩子 为了将与通信相关的代码插入到模型的前向传播中，FullyShardedDataParallel nn.Module包装器重写了nn.Module的forward()方法来安装前向传播前和后的逻辑，而函数式的fully_shard则通过register_forward_pre_hook()和register_forward_hook()等方法注册nn.Module钩子来实现它们。从反向传播中捕获适当的信号更具挑战性，因为PyTorch会自动且透明地处理反向传播。幸运的是，autograd引擎暴露了多种钩子，使得能够以精确的粒度安装自定义逻辑。
* 通过register_hook()在Tensor上注册的钩子允许在Tensor的梯度生成时运行自定义函数。这可以帮助将FSDP逻辑锚定到反向传播中某个激活值的梯度计算上。FSDP在每个FSDP单元的前向输出张量上注册这种类型的钩子，以便在反向传播进入该FSDP单元之前插入通信。
* 通过queue_callback()在backward()上注册的钩子在退出当前autograd GraphTask之前立即运行，这通常是整个反向传播结束的时候。FSDP依赖此钩子等待待处理的通信，以便后续的优化器步骤不会过早地消耗梯度。
* 在AccumulateGrad autograd函数上注册的钩子在参数的梯度在当前反向传播中完成累积时触发。FSDP将这种类型的钩子附加到每个FlatParameter的AccumulateGrad函数上，以便在梯度准备好时立即启动ReduceScatter。请注意，上面提到的Tensor钩子可能可以实现相同的行为，但可能会引入不必要的延迟，因为它还需要等待输入激活值的梯度计算。

集成方式 上述方法共同将FSDP算法以一种非侵入性且高效的方式与PyTorch nn.Module和autograd引擎集成在一起。

原生混合精度

灵活的混合精度机制 FSDP提供了一种多功能的原生混合精度机制。在参数管理方面，它遵循标准的混合精度技术，即同时维护参数的低精度和全精度副本【18, Mixed Precision Training, 2017, arXiv】。前向和反向计算使用低精度，而优化器步骤使用全精度。FSDP允许用户为参数、梯度规约和非可训练缓冲区指定精度，如果需要，每个都可以独立设置。

内存优化 对于一个拥有 $\Psi$ 个参数元素（torch.numel）的模型，每个低精度元素 b_low 字节，每个全精度元素 b_full 字节，这种混合精度方法通常会将内存开销从 $b_{full}\Psi$ 增加到 $(b_{low} + b_{full})\Psi$，因为需要维护两种精度的副本。然而，FSDP可以避开这个问题，因为我们的设计是始终将每个本地分片的FlatParameter保留在GPU内存中，并只动态分配未分片的FlatParameter。对于具有元素数量为 $p_1, \dots, p_K$ 的K个FlatParameter，FSDP的参数峰值内存贡献实际上从 $\frac{b_{full}}{S} \sum_{i=1}^{K} p_i + b_{full} \max_{i=1}^{K} p_i$ 字节减少到 $\frac{b_{full}}{S} \sum_{i=1}^{K} p_i + b_{low} \max_{i=1}^{K} p_i$ 字节。换句话说，FSDP直接将第二项 $b_{full} \max_{i=1}^{K} p_i$ 减少到 $b_{low} \max_{i=1}^{K} p_i$。

实现优势 与在操作符级别执行即时类型转换的torch.amp.autocast相比，FSDP的原生混合精度仅在其前向传播前为每个FlatParameter引入一次从全精度到低精度的转换，并且如果在前向传播后重新分片，则在其反向传播前也进行一次。此外，FSDP的混合精度允许所有集合操作在低精度下运行，这节省了通信量。

梯度缩放 用户最常选择FP16或BF16作为低精度，FP32作为全精度。与FP32相比，FP16的动态范围更小，这使得FP16面临更大的数值下溢和溢出风险。标准的解决方案包括一个梯度缩放器【1, torch.amp Gradient Scaling, 2023, pytorch.org】，它将梯度缩放到一个安全的幅度。然而，由于FSDP跨rank对梯度进行分片，一个普通的本地梯度缩放器实现会破坏数学等价性，因此，FSDP提供了自己的分片梯度缩放器。

实验环境

  • 模型:
    • 语言模型: HuggingFace T5-11B transformer【26, Exploring the limits of transfer learning with a unified text-to-text transformer, 2020, The Journal of Machine Learning Research】, minGPT-175B transformer【3, Language models are few-shot learners, 2020, Advances in neural information processing systems】, T5-611M, T5-2B。
    • 推荐模型: DHEN【33, DHEN: A Deep and Hierarchical Ensemble Network for Large-Scale Click-Through Rate Prediction, 2022, arXiv】，包含7680亿稀疏参数和5.5亿密集参数。
    • 其他模型: RegNet【29, RegNet: Multimodal sensor registration using deep neural networks, 2017, 2017 IEEE intelligent vehicles symposium (IV)】 (9B参数), DeepViT【36, Deepvit: Towards deeper vision transformer, 2021, arXiv preprint】 (8B参数)。
  • 硬件配置:
    • GPU: 最多512块 A100 80GB GPU。
    • 网络: 2Tb/s RoCE网络互连。
  • 软件配置:
    • 框架: PyTorch 2.0。
    • 实现: 使用PyTorch FSDP进行训练。优化器为Adam。
    • 训练设置: 针对大型模型实验，应用了激活检查点（Activation Checkpointing）和BF16混合精度。

实验结果


图 6: 模型规模和训练效率

模型规模对性能的影响

实验内容：我们研究了FSDP在处理从611M到175B不同规模模型时的性能，并与DDP【14, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv preprint】进行了比较。同时，我们使用GPT-175B模型测量了反向预取带来的加速效果。
实验结果与分析：
* 与DDP对比（图6 (a)）：在611M和2.28B的T5模型上，FSDP和DDP的性能相似。然而，当模型规模超过2.28B时，DDP会遇到内存不足（OOM）错误。相比之下，FSDP可以轻松支持11B模型，并通过启用BF16混合精度获得显著更高的TFLOPS。这表明FSDP可以无缝地用于从小到大的各种模型。
* 反向预取效果（图6 (b)）：在使用GPT-175B模型（通信开销更显著）的实验中，开启反向预取带来了大约18%的速度提升，并且这种TFLOPS增益在不同GPU集群规模下都得以保持。因此，在后续实验中，我们默认开启反向预取。

通信节流的影响

实验内容：我们研究了对FSDP通信进行节流（throttling）的影响。如第3.4节所述，过于激进地发起AllGather可能导致不必要的内存占用和性能问题。我们对三种不同类型的模型（RegNet-9B, T5-11B, DeepViT-8B）应用了速率限制器，并使用了各自可行的最大批次大小。
实验结果与分析（图6 (c)）：
* 速率限制器的效果并非在所有情况下都是正向的。在RegNet实验中没有获得任何加速，在DeepViT实验中甚至导致了性能下降（约5%的开销）。
* 这种行为是预期的，因为节流通信只有在CPU线程过于激进地分配GPU内存块并导致碎片整理时才能提升训练效率。torch.cuda.memory_stats()中的num_alloc_retries可以作为判断是否发生碎片整理的有效指标。
* 在T5模型实验中，速率限制器带来了高达5倍的速度提升，证明了其在特定场景下的巨大价值。然而，在通信主导的情况下，延迟AllGather可能会阻塞后续计算，从而产生开销（如DeepViT）。因此，在启用速率限制器之前，应确认训练中是否发生了内存碎片整理。


图 7: 训练吞吐量：为符合DHEN惯例，我们对DHEN使用样本/GPU/秒（QPS）作为单位。


图 8: 内存占用

大规模模型的高效训练

实验内容：我们使用完全分片（Full Sharding）策略，并开启预取和速率限制器，对三种大型模型（DHEN推荐模型、minGPT-175B、T5-11B）进行了评估。实验中还应用了激活检查点和BF16混合精度。
实验结果与分析：
* DHEN推荐模型（图7 (a) 和图8 (a)）：实验表明FSDP能够在大规模GPU集群上支持DHEN模型。实验对比了两种配置：RAF（前向传播后重分片）和NRAF（前向传播后不重分片）。
* RAF配置（完全分片）产生了最小的内存占用，但代价是QPS（吞吐量）较低。
* NRAF配置（混合分片）则相反，内存占用更高但QPS也更高，因为它分片组更小且跳过了一次重分片。
* 随着集群中GPU数量的增加，由于每个rank的模型分片大小减小，峰值内存使用量持续下降。
* minGPT-175B模型（图7 (b) 和图8 (b)）：
* 在批次大小为1和2的情况下，每个GPU分别实现了超过173和186 TFLOPS的性能，相当于A100 GPU BF16峰值算力（312 TFLOPS）的55%和60%。
* 模型在从128个GPU扩展到512个GPU时，在TFLOPS方面表现出线性扩展性，证实了FSDP在处理计算密集型或高速网络互连的大模型时的有效性。
* 一个特例是，在128个GPU上使用批次大小为2时，单GPU TFLOPS显著降低。这是由于反向传播过程中的CUDA内存碎片整理所致，该情况下的反向传播时间占迭代延迟的85.56%（正常约为67%）。图8的左上角图证实了这一点，PyTorch CUDA缓存分配器耗尽了全部80GB的CUDA内存。
* T5-11B模型（图7 (c) 和图8 (c)）：
* 所有实验都在远低于GPU内存容量的情况下舒适运行，不太可能发生碎片整理。
* 尽管如此，当GPU数量从8个增加到512个时，仍然观察到单GPU TFLOPS有7%的下降。这表明在大型集群上，通信开销开始超过计算，通信和计算之间的近乎完美的重叠不再能实现。

讨论

本节讨论FSDP如何与其他并行范式结合，以及采用FSDP时已知的局限性。

FSDP的互操作性

与其他并行范式的结合 为了进一步提高分布式训练的可扩展性和效率，需要将FSDP与其他范式相结合。本节简要介绍了FSDP的设计如何使其能够与其他类型的并行性进行混合和匹配。

流水线并行

功能集成与挑战 流水线并行可以通过使用FSDP包装每个独立的流水线阶段来与FSDP进行功能性集成。然而，由于流水线并行将输入的mini-batch划分为更小的micro-batch，FSDP中的默认完全分片策略将不得不为每个micro-batch取消分片模型参数。因此，将这些方法与默认FSDP配置相结合可能会导致显著的通信开销。幸运的是，FSDP提供了替代的分片策略，可以在前向传播后保持参数不被分片，从而避免了每个micro-batch不必要的AllGather通信。诚然，这需要在GPU设备上存储整个流水线阶段的参数，但FSDP仍然可以减少内存使用，因为它仍然对梯度和优化器状态进行分片。

张量并行

2D并行实现 与FSDP不同，张量并行在计算过程中保持参数分片，这在任何子模块大到无法装入GPU内存时是必要的。目前，PyTorch提供了一个名为parallelize_module的原型功能，可以与FSDP结合构建2D并行。它的工作原理是将设备组织成一个2D网格，其中PyTorch的分布式张量DTensor在一个维度上管理张量并行，而FSDP在另一个维度上应用分片数据并行。这两个维度分别通信激活值和参数。我们通常将阻塞后续计算的张量并行通信保持在节点内以利用更高的网络带宽，并允许FSDP通信在另一个网格维度上跨节点操作。

局限性

实践中遇到的挑战 在我们与生产和研究应用的合作中，我们遇到了与FSDP相关的一些局限性。本节旨在讨论两个不易察觉且在故障排除时构成重大挑战的棘手问题。

数学等价性

优化器计算的差异 FSDP不能保证它总是能实现与本地训练相同的数学等价性，特别是在优化器计算方面。这是因为优化器步骤是在分片参数上操作的，其数据布局是FSDP的FlatParameter分片算法的函数，该算法不尊重单个参数的边界。因此，任何依赖于原始参数未分片值（例如，向量范数）、其张量结构（例如，近似二阶优化器）或需要所有参数的全局状态的优化器计算都将变得无效。解决这个问题需要不均匀分片、填充或额外的通信，所有这些都会损害性能。如何将这类优化器计算与分片协同设计是一个开放的研究问题。

共享参数

处理共享参数的复杂性 对于共享参数，FSDP必须确保不将它们展平到多个FlatParameter中，并确保在所有使用场景需要时它们能被正确地取消分片。如果处理不当，PyTorch可能会引发关于张量存储丢失或大小不匹配的错误，这可能发生在某个FSDP单元试图使用一个已经被前一个FSDP单元重新分片的共享参数时。当前的建议是构建FSDP单元，使得共享参数属于最低共同祖先单元，以确保共享参数在所有使用期间都保持未分片状态。这可能需要对模型结构进行一些检查才能正确完成，并且可能会不希望地使FlatParameter在很长一段时间内保持未分片状态，因此我们正在研究改进共享参数处理的方法。

结论

本手稿阐述了截至PyTorch 2.0版本中FullyShardedDataParallel的底层原理、设计理念和实现。FSDP通过一系列先进技术实现了可用性和效率，包括延迟初始化、灵活的分片策略、通信重叠与预取，以及对通信集合操作的速率限制。所有这些技术都与其他关键的PyTorch组件紧密协同设计，以确保解决方案的健全性和鲁棒性。评估表明，FSDP能够支持大型语言模型和推荐模型，并具有近线性的可扩展性。