【大模型】万亿级别的大语言模型训练，基础设施如何支持

前言

我们希望考虑哪些软件和系统配置可以使现有的大规模语言模型（LLM）在规模更大的模型上更有效地扩展，目前已经达到约1万亿参数。我们的分析受到LLM在各种应用中持续取得的成功和有效性的推动，并受到这样一个观察的启发，即模型浮点运算利用率（MFU）——评估专门用于模型训练期间人工智能（AI）加速器利用程度的常见效率指标——可能低至50%或更低。

要在类似于当前系统的体系结构上将模型大小增加10倍（10万亿参数）或更高，将需要显着提高MFU。以每个参数20字节的空间要求，仅存储模型的权重和优化器状态，我们将需要超过200 TB的内存。对于基于NVIDIA H100图形处理单元（GPU）的系统，具有80 GiB高带宽内存（HBM）内存，我们将需要2500个GPU和一个完全模型并行的实现来训练这样一个模型。在此规模上，没有已知的模型并行技术能够提供接近50%的MFU。

在这个例子的启发下，我们旨在确定阻止我们在使用8个相互连接的GPU集群构建的大型系统上培训多万亿参数模型的系统限制，类似于NVIDIA DGX和HGX。

我们首先提出了选择结构良好的多万亿参数LLM的方法。然后，使用我们自己的基于transformer的LLM训练的快速分析性能模型，我们搜索了数十亿个系统配置和执行策略。本文解释了我们的一些发现，可以概括如下。

1）培训百万亿参数的LLM是可行的，但需要每个GPU高达1 TiB的次级内存池，双向带宽为100 GB/s。

2）对于1T模型的强扩展在约12288个GPU左右停滞，因为矩阵乘法变得小而低效，并且无法与通信overlap。

3）超过10T模型需要更多的一级内存，其中HBM 大小与模型大小成比例。

4）将模型和系统大小增加到10T参数和10,000个GPU以上需要更大的快速网络域和更有针对性的软件优化。

总体而言，我们发现必须共同设计LLM、软件和硬件，以实现高性能和高效率。

一、实验方法论

为了性能估算，我们使用了Calculon，这是我们开发的LLM训练性能的快速开源分析模型。Calculon可以在约1毫秒内估算给定LLM、系统配置和软件执行策略的时间和资源使用情况，从而允许探索具有数十亿个这样的配置的大型设计空间。Calculon使用张量并行（TP）、流水线并行（PP）和数据并行（DP）模型对LLM进行建模，允许搜索确定最佳的分割并行配置。系统规格描述了一个基于加速器的分布式系统，具有连接到多个网络的两级内存层次结构，如下图所示。

为了验证其建模的准确性，Calculon与在NVIDIA的基于A100的Selene超级计算机上运行的Megatron LLM进行比较。Calculon实现了很高的准确性，在这些验证运行中的平均误差为3.4%，最大误差为7.25%，如表I所示。

我们进行的实验变化了系统大小、模型大小、内存容量、带宽和NVLink域大小，使用H100支持的FP8数据格式。对于每个系统，我们选择一个执行策略，考虑了多种最先进的软件优化并选择最佳执行策略。鉴于搜索空间很大，我们无法完全呈现实验结果，而是专注于我们发现的一些最重要的趋势。

我们的分析假设了一个计算节点的网络系统，其节点架构如图1所示。它的结构和连接方式类似于DGX或HGX。唯一的区别是在GPU上增加了卸载内存，除了HBM之外。这样的内存可以通过计算表达链接（CXL）连接，也可以由中央处理单元（CPU）托管，并直接从GPU访问，类似于NVIDIA的Grace-Hopper。

二、选择LLM配置

一个重要的参数是LLM的纵横比，定义为transformer块的隐藏维度与块数（也称为transformer层数）之间的比率。一些最近的研究声称理想的纵横比是一个恒定的128 ，而其他人声称随着块数的增加，纵横比应该呈指数增长 。这两种分析都是在比今天的生产LLM小2到5个数量级的LLM上进行的。在LLM专家之间没有共识的情况下，我们遵循表II建议的明显的当前实践，即线性外推纵横比transformer块数。尽管如此，我们的分析方法对于任何缩放函数都适用。

在扩展和塑造LLM时，一个挑战是将模型映射到可用的硬件上。一些模型，例如GPT-3 ，具有96个块，总共有1750亿个参数，被设计为2的幂次或2的倍数，使它们非常适合通常也是2的倍数的典型系统设计。其他模型则不那么容易映射。Turing-NLG 具有105个块，总共有5300亿个参数，导致可能的映射减少。PaLM 具有118个块，总共有5400亿个参数，一个2的质数乘以2，导致可能的映射进一步减少。

为了看到这些选择的影响，图2比较了两个大约11000亿个参数的相似大小的模型。一个有256个块，是2的幂次，另一个有254个块，是2的质数乘以2。当映射到4096个处理器时，256块模型产生15612832个可能的映射，而254块模型只产生842080个，少了18.5倍。因此，256块模型最终快了36%，MFU为75%，而254块模型的MFU为54%。

因此，我们建议使用2的幂次作为步长，扩展块数和注意力头数。这样做可以更容易配置张量和流水线并行，从而获得更好的整体性能。图3总结了各种纵横比的模型尺寸。这些模型都在各种常见的系统设计和各种系统大小上产生了数百万个映射解决方案。

图3中显示的隐藏步长是8192。然而，在找到最佳（最接近理想纵横比）时，我们使用了1024的步长。在本文的其余部分，我们使用表III中找到的模型配置。所有模型的序列大小为8192，前馈大小固定为隐藏大小的4倍，注意力头数等于块数。对于所有实验，我们将最大批处理大小限制为3072。

三、用于LLM扩展的张量卸载

使用具有8个NVLink连接的DGX/HGX H100标准扩展LLM是可能的，但要达到高性能并不容易。例如，图4a显示了在4096个GPU的固定系统大小上扩展模型大小时的培训效率。即使是最小的模型尺寸1T，也仅达到60%的效率，并迅速下降，直到18T，它无法运行。

主要的可扩展性问题是在训练期间存储权重和激活的内存不足。这反过来迫使使用激活重新计算和更高的模型并行性。具有大量流水线并行和缺乏备用内存的大型流水线并行会导致过多的时间开销，形成流水线泡沫。大规模的张量并行超出了8的NVLink大小，增加了通信时间，因为带宽不足。

这些问题可以通过次级内存池来解决，其中来自非活动transformer块的未使用张量可以根据需要传输和检索。这可以实现为CPU主机内存、一系列PCIe连接的SSD或CXL连接的内存。我们在图4b中考虑使用张量卸载时的培训效率，其中每个GPU的容量为1 TiB，带宽为无限。显然，具有足够的卸载容量和无限的卸载带宽，我们可以至少训练具有128T参数的模型。卸载容量的影响在图5中与256 GiB、512 GiB、1 TiB和2 TiB进行了比较。我们看到相对于无限带宽，使用每个方向的50 GB/s和100 GB/s的卸载带宽会导致相对减速。在4096个GPU系统上，以50 GB/s的速度，随着模型尺寸的增加，性能显著下降。以100 GB/s的速度，大多数系统几乎与无限带宽的性能相匹配，这表明这是一个足够的目标带宽。

重要的是，这些张量卸载内存需求在当前技术范围内。基于CXL 1.1和CXL 2.0的内存池，容量高达2 TiB，带宽高达89.6 GB/s已经可用 。基于NVIDIA的Grace-Hopper 的系统具有高达512 GiB的低功耗双数据速率（LPDDR）内存，带宽高达546 GB/s，远远超过我们的卸载需求估计。

使用卸载进行培训的效率显示在图6中，其中卸载带宽为100 GB/s，容量为256 GiB、512 GiB、1 TiB和2 TiB，分别跨越4096、8192、12288和16384个GPU。显示的主要趋势是：

• 在大系统上培训小模型和在小系统上培训大模型导致效率低。
• 在256 GiB时很少出现良好效率。
• 对于8k、12k和16k GPU，512 GiB通常足够。
• 1 TiB的容量几乎与2 TiB相同。

四、强力缩放

在本节中，我们分析了1T参数模型从4096个GPU强扩展到16384个GPU，检查了8和16的NVLink域大小。对于每个系统，我们使用Calculon对可能的配置进行详尽的搜索，通常是每个LLM-system对10-30百万个配置。结果见图7。我们分析了所有与配置相关的参数，如TP、PP、DP分割、微批大小、流水线交错等，并总结了以下值得注意的趋势。

扩展到12288个GPU表现良好，但在16384个GPU时受到影响。8的NVLink大小对12288个GPU已足够，但对16384个GPU需要16。添加额外的处理器需要将它们分配给张量、流水线或数据并行，但每种都会在时间或内存方面产生一些资源成本。我们确定了在16384个GPU上缺乏扩展性的几个原因。

1）在增加TP时，张量可能被分得太细，以在GPU上保持高计算效率。

2）在增加TP时，每个消息的大小可能变得足够小，以致成为延迟主导。

3）当试图重叠TP通信和计算时，增加TP会减小计算大小，但通信大小保持不变。在特定的FLOPs/带宽比下，通信隐藏减少，效率降低。

4）在重叠TP通信和计算时，为了维持NVLink的高带宽，GPU必须将许多核心专用于通信，降低了其计算速度。添加专门的直接存储访问（DMA）类似引擎用于通信将消除此开销，实现最佳重叠。

5）增加PP要么增加了流水线泡沫开销，要么需要更多的内存来实现更高级别的交错以减小流水线泡沫。

6）增加DP增加了内存，因为需要复制。

7）我们约束我们的模型具有最大批处理大小3072，以保持先前研究的收敛性质。此选择限制了最大可用的DP为3072，因此其余必须是TP或PP。

五、超过10T参数的比例模型

我们还分析了在GPU数量固定的系统中将模型大小增加到128T参数的效果。图8展示了在4,096个GPU上的结果，其中包括80 GiB和120 GiB的HBM，以及256 GiB、512 GiB、1 TiB和2 TiB的卸载容量。在4,096个GPU上，使用80 GiB的HBM对11T参数的模型进行尺度化效果良好，但要进一步扩展，HBM的大小必须增加到120 GiB，即使提供了额外的卸载内存。原因在于，必须有足够的HBM内存来容纳两个transformer块——一个用于计算，另一个用于卸载和预取，即使有卸载也是这样。在模型扩展期间，transformer块的大小主要由权重和激活引起。毫不奇怪，卸载内存容量也需要相应地进行扩展。

我们的实验表明，将HBM的大小扩展到120 GiB，卸载内存扩展到2 TiB足以实现100T参数的扩展。超过11T参数后，模型占用了大部分可用内存。这表明进一步的效率改进是可能的，可以通过提供更多内存或通过增加NVLink域的大小来减少每个GPU的权重空间并增加本地微批次大小。

这些实验显示，所提出的LLM在保持高于75%的MFU的同时，可以扩展到128T参数，这比目前系统上更小的LLM通常能够实现的要好。

六、结论

我们的共同设计分析显示，在选择合适的软件和硬件设置的情况下，包括用于张量卸载的次级内存池，可以有效地培训结构良好的多万亿参数LLM，其MFU为75%或更高。我们确定了强扩展（固定模型，增加GPU数量）下的最佳配置策略和基本限制。对于具有4,096个GPU的固定系统，我们展示了如何仅使用张量卸载来训练11T参数模型，以及如何使用120 GiB HBM和2 TiB卸载内存来扩展到128T参数。