LLMs高效的多 GPU 计算策略Efficient multi-GPU compute strategies

很有可能在某个时候，您需要将模型训练工作扩展到超过一个GPU。在上一个视频中，我强调了当您的模型变得太大而无法适应单个GPU时，您需要使用多GPU计算策略。但即使您的模型确实适合单个GPU，使用多个GPU加速训练也有好处。即使您正在使用小型模型，了解如何跨GPU分配计算也会很有用。让我们讨论如何有效地跨多个GPU进行这种扩展。

您将从考虑模型仍然适合单个GPU的情况开始。扩展模型训练的第一步是将大型数据集分布到多个GPU，并并行处理这些数据批次。这种模型复制技术的一个流行实现是PyTorch的Distributed Data Paraller分布式数据并行，或简称DDP。DDP将您的模型复制到每个GPU，并将数据批次并行发送到每个GPU。每个数据集并行处理，然后一个同步步骤组合每个GPU的结果，从而更新每个GPU上的模型，这在芯片之间始终是相同的。这种实现允许跨所有GPU进行并行计算，从而加快训练速度。请注意，DDP要求您的模型权重和所有其他训练所需的参数、梯度和优化器状态都适合单个GPU。

如果您的模型对此过大，您应该考虑另一种称为模型分片Model Sharded的技术。模型分片的一个流行实现是PyTorch的Fully Sharded Data Parallel完全分片数据并行，或简称FSDP。FSDP的动机是2019年由Microsoft研究人员发表的一篇论文，该论文提出了一种称为ZeRO的技术。ZeRO代表零冗余优化器，ZeRO的目标是通过分布或分片模型状态跨GPU与ZeRO数据重叠来优化内存。这允许您在模型不适合单个芯片的内存时，跨GPU缩放模型训练。

在回到FSDP之前，让我们快速看看ZeRO是如何工作的。

本周早些时候，您查看了训练LLM所需的所有内存组件，最大的内存需求是优化器状态，它占用的空间是权重的两倍，然后是权重本身和梯度。

让我们将参数表示为这个蓝色框，梯度为黄色，优化器状态为绿色。

我之前展示的模型复制策略的一个限制是您需要在每个GPU上保留一个完整的模型副本，这导致了冗余的内存消耗。您在每个GPU上存储相同的数字。

另一方面，ZeRO通过分布也称为分片模型参数、梯度和优化器状态跨GPU，而不是复制它们，从而消除了这种冗余。与此同时，沉没模型状态的通信开销接近之前讨论的Distributed Data Paraller (DDP)。

ZeRO提供了三个优化阶段。ZeRO阶段1仅跨GPU分片Optimizer States优化器状态，这可以将您的内存占用减少到四分之一。

ZeRO阶段2还将Gradient梯度分片到芯片上。与阶段1一起应用时，这可以将您的内存占用减少到八倍。

最后，ZeRO阶段3将所有组件（包括模型参数Parameters）分片到GPU上。与阶段1和2一起应用时，内存减少与GPU数量成线性关系。例如，跨64个GPU的分片可以将您的内存减少64倍。

让我们将这个概念应用到GDP的可视化
，

并用模型参数、梯度和优化器状态的内存表示替换LLM。当您使用FSDP时，您将数据分布到多个GPU，如您在DDP中看到的那样。

但是，使用FSDP，您还可以使用ZeRO论文中指定的策略之一，将模型参数、梯度和优化状态分布或分片到GPU节点上。使用这种策略，您现在可以使用太大而无法适应单个芯片的模型。

与DDP相反，其中每个GPU都有本地化处理每批数据所需的所有模型状态，FSDP要求您在前向和后向传递之前从所有GPU收集此数据。

每个CPU按需从其他GPU请求数据，将分片数据转化为非分片数据以供操作使用。操作后，您将非分片的非本地数据释放回其他GPU作为原始分片数据。您还可以选择在后向传递期间为未来的操作保留它。注意，这需要更多的GPU RAM，这是一个典型的性能与内存权衡决策。

在后向传递后的最后一步，FSDP与DDP相同地跨GPU同步梯度。

如FSDP所描述的模型分片

1. 允许您减少整体GPU内存使用。
2. 您还可以选择让FSDP将部分训练计算卸载到GPU，以进一步减少GPU内存使用。
3. 为了管理性能与内存使用之间的权衡，您可以使用FSDP的sharding factor分片因子配置分片级别。

分片因子为1基本上删除了分片并复制了与DDP类似的完整模型。

如果您将分片因子设置为可用GPU的最大数量，您将打开完整的分片。这节省了最多的内存，但增加了GPU之间的通信量。

中间的任何分片因子都启用了超分片。

让我们看看FSDP与DDP在每个GPU的teraflops上的性能如何。这些测试使用最多512个NVIDIA V100 GPU执行，每个GPU有80GB的内存。注意，一个teraflop对应于每秒一万亿次$10^{12}$浮点运算。第一个数字显示了不同大小T5模型的FSDP性能。您可以看到FSDP的不同性能数字，完整分片为蓝色，超分片为橙色，完整复制为绿色。作为参考，DDP性能以红色显示。

对于首先有611百万参数和22.8亿参数的25个模型，FSDP和DDP的性能相似。现在，如果您选择超过22.8亿的模型大小，例如25个模型有113亿参数，DDP会遇到内存不足的错误。另一方面，FSDP可以轻松处理这种大小的模型，并在将模型的精度降低到16位时获得更高的teraflops。

第二个数字显示了当增加GPU数量从8-512为11亿T5模型时，每个GPU teraflops减少了7%，

这里使用了批量大小为16的橙色和批量大小为8的蓝色。随着模型在大小上增长并分布到越来越多的GPU上，芯片之间的通信量增加开始影响性能，减慢计算。

总之，这表明您可以使用FSDP进行小型和大型模型，并无缝地跨多个GPU扩展模型训练。

我知道这次讨论非常技术性，我想强调的是，您不需要记住所有的细节。最重要的是，当训练LLM时，了解数据、模型参数和训练计算如何跨进程共享。鉴于跨GPU训练模型的费用和技术复杂性，一些研究人员一直在探索如何使用较小的模型实现更好的性能。在下一个视频中，您将了解有关计算最佳模型的研究。让我们继续看下去。

参考

https://www.coursera.org/learn/generative-ai-with-llms/lecture/e8hbI/optional-video-efficient-multi-gpu-compute-strategies