FSDP(Fully Sharded Data Parallel)

news2025/3/15 7:36:50

在这里插入图片描述
完全分片数据并行 (FSDP) ，它将AI 模型的参数分片到数据并行工作器上，并且可以选择将部分训练计算卸载到 CPU。顾名思义，FSDP 是一种数据并行训练算法。尽管参数被分片到不同的GPU，但每个微批次数据的计算仍然是每个 GPU Worker 本地的。这种概念上的简单性使得 FSDP 更容易理解，并且更适用于广泛的使用场景（与层内并行和管道并行相比）。与优化器状态+梯度分片数据并行方法相比，FSDP 参数分片更均匀，并且能够通过训练期间的通信和计算重叠获得更好的性能。

借助 FSDP，现在可以使用更少的 GPU 更有效地训练数量级更大的模型。FSDP 已在FairScale 库中实施，允许工程师和开发人员使用简单的 API 扩展和优化模型的训练。在 Facebook，FSDP 已经被集成并测试用于训练我们的一些NLP和Vision模型。

大规模训练的计算成本高

NLP 研究是一个特殊领域，我们可以看到有效利用计算来训练人工智能的重要性。去年，OpenAI 宣布他们已经训练了GPT-3，这是有史以来最大的神经语言模型，拥有 1750 亿个参数。据估计，训练GPT-3 大约需要 355 个 GPU 年，相当于 1,000 个 GPU 连续工作四个多月。

除了需要大量的计算和工程资源之外，大多数像这样的扩展方法都会引入额外的通信成本，并要求工程师仔细评估内存使用和计算效率之间的权衡。例如，典型的数据并行训练需要在每个 GPU 上维护模型的冗余副本，并且模型并行训练会引入额外的通信成本以在工作线程 (GPU) 之间移动激活。

相比之下，FSDP 相对来说不需要权衡。它通过跨 GPU 分片模型参数、梯度和优化器状态来提高内存效率，并通过分解通信并将其与前向和后向传递重叠来提高计算效率。FSDP 产生与标准分布式数据并行 (DDP) 训练相同的结果，并且可在易于使用的界面中使用，该界面是 PyTorch 的 DistributedDataParallel 模块的直接替代品。我们的早期测试表明，FSDP 可以扩展到数万亿个参数。

FSDP的运作方式

在标准 DDP 训练中，每个工作人员处理一个单独的批次，并使用全归约操作对各个工作人员的梯度进行求和。虽然 DDP 已经变得非常流行，但它占用的 GPU 内存超出了其需要，因为模型权重和优化器状态会在所有 DDP 工作线程之间复制。

减少重复的一种方法是应用称为全参数分片的过程，其中仅提供本地计算所需的模型参数、梯度和优化器的子集。这种方法的实现 ZeRO-3 已被 Microsoft 推广。

解锁全参数分片的关键在于我们可以将DDP 中的全归约操作分解为单独的归约分散和全收集操作：

在这里插入图片描述
然后，我们可以重新安排reduce-scatter和all-gather，以便每个DDP工作线程只需要存储参数和优化器状态的单个分片。下图说明了标准 DDP 培训（上）和 FSDP 培训（下）：

为了最大限度地提高内存效率，我们可以在每层前向传递后丢弃全部权重，为后续层节省内存。这可以通过将 FSDP 包装器应用到网络中的每一层来实现（使用reshard_after_forward=True ）。