📰 来源：Towards Data Science　|　📅 翻译日期：2026年7月10日
🔗 原文：查看原文
🤖 翻译：DeepSeek AI · 仅供参考

单GPU训练的局限与分布式需求

在单个GPU上训练模型很直接：加载权重、加载数据，然后等待完成。对大多数模型来说，这没问题，唯一成本是时间。当等待过长时，通常的解决办法是增加GPU。每个GPU并行处理不同数据切片，工作更快完成。模型状态不变，只是投入了更多人手。

这覆盖了常见情况——问题在于时间。大模型引入了第二个无法靠增加人手解决的问题：空间。过去一年，我的主要工作涉及在不同领域训练和微调大基础模型，包括单细胞模型和视觉-语言Transformer模型。一旦参数量超过几十亿，模型、梯度和优化器状态便无法放入单个GPU。你不能简单地跨GPU复制工作，因为任何单个GPU都放不下副本。你必须拆分模型本身。

两大策略：DDP对付时间，FSDP对付空间

这两个问题对应两种策略。DDP（Distributed Data Parallel） 针对时间。每个GPU保留完整模型副本，处理不同批次切片。FSDP（Fully Sharded Data Parallel） 针对空间。模型被分割成碎片，使任何GPU都无需持有整个模型。

但实际上不止这两种选择。它们只是同一个拨盘的两端。中间是DeepSpeed（微软的训练库）的ZeRO阶段，允许你逐步从一端滑向另一端，每一步以少量通信换取少量内存，而非直接从“人人持有全部”跳到“无人持有”。

策略如何工作：不同GPU持有什么？

每个分布式训练策略本质上回答一个问题：每个GPU各自持有什么副本？ 让我们看一个例子：

以Mistral-7B为例，使用Adam在BF16混合精度下微调。参数占 14 GB（70亿参数，每个2字节）。梯度又占 14 GB。优化器状态最昂贵：Adam为每个参数保存两个运行平均值（动量和方差），均以FP32存储，约 58 GB。总计达到 87 GB，模型还没见一个批次，也没计算一个激活值。

对于有 80 GB 显存的A100 GPU，模型放不下。因此我们需要分布。问题就是：每个GPU持有什么？

DDP：人人持有全部

简单答案是每个GPU持有全部。这是DDP，大多数人首先遇到的方式（我也是）。每个GPU保留模型的完整副本：所有参数、所有梯度、所有优化器状态。被分割的是数据。每个GPU取批次的不同切片，运行自己的前向和反向传播，并计算自己的梯度。

关键在于这些梯度各不相同，因为每个GPU看到的数据不同。如果每个GPU此时都执行优化器步骤，副本将漂移，你就不再拥有一个模型，而是四个。因此，在更新前，GPU会平均它们的梯度。每个GPU发送自己的梯度并接收所有其他梯度，最终得到相同的平均梯度。然后执行相同的优化器步骤，副本保持同步。这个平均化是DDP需要的唯一通信，每步发生一次。它被称为all-reduce，我们稍后会在讨论每个策略对网络的成本时回到这一点。

图1：DDP工作原理（作者绘制）

DDP快速且简单，正因为它通信极少。每步一次all-reduce，PyTorch甚至将其与反向传播重叠，所以在好硬件上几乎感觉不到。但其局限性内置于设计中：每个GPU持有全部，因此增加GPU对内存问题无济于事。如果 87 GB 放不进一个A100，那也放不进四个，因为每个仍然持有完整的 87 GB。为DDP增加GPU提高吞吐量，但不减少每个GPU需持有的内存。

FSDP：无人持有全部

如果问题是每个GPU都持有全部，显而易见的解决办法是确保没有GPU这样做。这就是FSDP。每个GPU不保留完整副本，而是将模型分割成碎片，每个GPU只拥有参数的一部分、梯度的一部分和优化器状态的一部分。在四个GPU上，每个持有约 87 GB 的四分之一，从而使模型适合。

但这立刻带来一个问题：你无法通过只拥有某层四分之一权重的GPU来运行该层。因此FSDP在需要之前重新组装每一层，并在之后立即拆解。当正向传播到达某一层时，GPU交换它们持有的碎片，使得每个GPU都能获得该层的完整权重。