多GPU训练大语言模型，DDP， ZeRO 和 FSDP

在某些时候，我们可能需要将模型训练扩展到单个 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 代表零冗余优化器