在某些时候，我们可能需要将模型训练扩展到单个 GPU 之外。当模型变得太大无法适应单个 GPU 时，需要使用多 GPU 计算策略。但即使模型适合单个 GPU，使用多个 GPU 来加速训练也是有好处的。即使您正在处理一个小模型，了解如何在多个 GPU 上分配计算也是有用的。
让我们讨论如何以一种高效的方式在多个 GPU 上进行扩展。您首先要考虑的情况是您的模型仍然适合单个 GPU。扩展模型训练的第一步是将大型数据集分布到多个 GPU 上，并并行处理这些数据批次。一种流行的实现这种模型复制技术的方法是 PyTorch 的分布式数据并行（Distributed Data Parallel，简称 DDP）。DDP 将您的模型复制到每个 GPU 上，并将数据批次并行发送到每个 GPU。每个数据集并行处理，然后在同步步骤中结合每个 GPU 的结果，从而更新每个 GPU 上的模型，使其在所有芯片上始终保持一致。这种实现允许在所有 GPU 上进行并行计算，从而加快训练速度。
请注意，DDP 要求您的模型权重以及所有其他参数、梯度和优化器状态都能适应单个 GPU 的内存。如果您的模型太大，无法做到这一点，您应该研究另一种技术，称为模型分片（model sharding）。PyTorch 中一种流行的模型分片实现是完全分片数据并行（Fully Sharded Data Parallel，简称 FSDP）。FSDP 的动机来源于微软研究人员在 2019 年发表的一篇论文，该论文提出了一种名为 ZeRO 的技术。ZeRO 代表零冗余优化器（Zero Redundancy Optimizer），其目标是通过在 GPU 上分布或分片模型状态来优化内存，做到零数据重叠。这使您能够在模型无法适应单个芯片内存时，将模型训练扩展到多个 GPU。
在回到 FSDP 之前，让我们快速了解一下 ZeRO 的工作原理。训练 LLMs 所需的所有内存组件，最大的内存需求是优化器状态，它占用了权重的两倍空间，其次是权重本身和梯度。之前展示的模型复制策略的一个限制是，您需要在每个 GPU 上保留完整的模型副本，这会导致冗余的内存消耗。您在每个 GPU 上存储了相同的数字。而 ZeRO 通过在 GPU 之间分布（也称为分片）模型参数、梯度和优化器状态，而不是复制它们，消除了这种冗余。同时，同步模型状态的通信开销接近前面讨论的 DDP。ZeRO 提供了三个优化阶段。ZeRO Stage 1 仅在 GPU 上分片优化器状态，这可以将内存占用减少最多四倍（分片的优化器状态包括Adam优化器（2个状态），激活和临时内存）。ZeRO Stage 2 还在芯片之间分片梯度。当与 Stage 1 一起应用时，这可以将内存占用减少最多八倍。最后，ZeRO Stage 3 在 GPU 之间分片所有组件，包括模型参数。当与 Stage 1 和 Stage 2 一起应用时，内存减少与 GPU 数量成线性关系。例如，在 64 个 GPU 之间分片可以将内存减少 64 倍。
让我们将这个概念应用于 DDP 的可视化，并用模型参数、梯度和优化器状态的内存表示替换 LLM。当您使用 FSDP 时，您会如同在 DDP 中看到的那样，在多个 GPU 之间分布数据。但在 FSDP 中，您还会使用 ZeRO 论文中指定的策略，在 GPU 节点之间分布或分片模型参数、梯度和优化器状态。使用这种策略，您现在可以处理那些过大而无法装入单个芯片的模型。
与 DDP 不同，每个 GPU 都具有处理每个数据批次所需的所有模型状态，FSDP 需要在前向和后向传递之前从所有 GPU 收集这些数据。每个 GPU 根据需求从其他 GPU 请求数据，以在操作期间将分片数据转化为未分片数据。操作后，您会将非本地的未分片数据释放回其他 GPU，恢复为原始分片数据。您也可以选择在后向传递期间保留这些数据以供将来操作使用。请注意，这再次需要更多的 GPU 内存，这是典型的性能与内存利用率之间的权衡决策。在后向传递的最后一步，FSDP 同步 GPU 之间的梯度，方式与 DDP 相同。通过 FSDP 描述的模型分片，您可以减少整体 GPU 内存利用率。可选地，您可以指定 FSDP 将部分训练计算卸载到 CPU，以进一步减少 GPU 内存利用率。为了在性能和内存利用率之间进行权衡，您可以使用 FSDP 的分片因子配置分片水平。分片因子为 1 基本上移除了分片，并完全复制模型，类似于 DDP。如果您将分片因子设置为最大可用 GPU 数，则启用完全分片。这具有最大的内存节约，但增加了 GPU 之间的通信量。任何介于两者之间的分片因子都启用部分分片。
让我们看看 FSDP 与 DDP 在每个 GPU 的每秒万亿次浮点运算（teraflops）中的性能比较。使用最多 512 个 NVIDIA V100 GPU（每个 80 GB 内存）进行这些测试。请注意，一个 teraflop 对应于每秒一万亿次浮点运算。第一个图显示了不同尺寸 T5 模型的 FSDP 性能。您可以看到 FSDP 完全分片（蓝色）、部分分片（橙色）和完全复制（绿色）的不同性能数据。供参考，DDP 性能显示为红色。对于第一个 611 百万参数和 2.28 十亿参数的 T5 模型，FSDP 和 DDP 的性能相似。现在，如果您选择超过 2.28 十亿参数的模型，例如 11.3 十亿参数的 T5，DDP 会遇到内存不足错误。另一方面，FSDP 可以轻松处理这种大小的模型，并在将模型精度降低到 16 位时实现更高的 teraflops。第二个图显示了当增加 GPU 数量从 8 到 512 时，每个 GPU 的 teraflops 减少了 7%，其中使用批次大小为 16（橙色）和批次大小为 8（蓝色）绘制了 11 十亿的 T5 模型。随着模型大小的增加并在越来越多的 GPU 之间分布，芯片之间通信量的增加开始影响性能，减慢计算速度。总而言之，这表明您可以将 FSDP 用于小型和大型模型，并无缝地在多个 GPU 之间扩展模型训练。
参考资料：
Coursera 课程 Generative AI with Large Language Models