引言 — 许多超大规模企业都在竞相构建大型 GPU 集群，通常配备 64K 或更多 GPU，以适应所有类型的生成式 AI (genAI) 训练工作负载。虽然这些大型 Transformer 模型和数据集的大小可能需要数千个 GPU 进行训练，但在 GPU 之间提供任意到任意的非阻塞网络连接可能是一种过度设计。
了解 GenAI 训练工作负载的模型分区和流量模式可以帮助优化网络拓扑，并实现 GPU 和以太网交换机的有效利用。
本文的主要参考文献见文末，主要介绍了GPU 集群规模、模型分区以及训练工作负载的 GPU 之间的流量模式。
本文为第一部分，主要介绍了
GPU规模的数学计算
GPU集群的分区并行
I. 模型数学
GPU 集群大小和结构拓扑很大程度上取决于
工作负载的特征（参数大小、数据集大小和模型架构）
最佳训练时间和训练频率。换句话说，模型需要多久更新一次？
模型分区。当参数和数据集的状态空间很大时，必须有一个能够了解网络拓扑并能够迭代评估不同分区选择的编译器/优化器，以选择在计算和通信之间提供良好平衡的分区。
为了推导出用于训练 LLM 模型的 GPU 集群大小，我们可以安全地做出以下假设：
深度学习 Transformer 模型中大约 60-80% 的运算是矩阵乘法，可以使用张量核（TF32/TF16）运算。
由于内存和互连瓶颈，训练大型 Transformer 模型时的 GPU 利用率徘徊在 30-40% 左右。然而，利用率在很大程度上取决于有效的模型分区和底层 GPU 结构。Nvidia的论文声称，通过有效使用张量、管道和数据并行，具有一万亿个参数的大型模型的利用率达到 50%。
对于未发布实际训练 FLOP 的模型，可以使用如下估计：每个参数和每个token大约为~6-8 FLOPS。
表 1 显示了如果训练需要在一周内完成，使用旧版本 ( A100 ) GPU 训练各种模型所需的 GPU 规模。
OpenAI并未正式公布GPT-4参数和数据集大小。然而，许多论坛普遍猜测模型大小为 1.5 万亿个参数。使用与GPT-3中相同的数据集大小与参数的缩放比例，可以推断出GPT-4的数据集大小。上表假设训练采用 Tensor Float 32 (TF32) 格式，GPU 利用率通常为 30%。
从表中可以看出，275K A100 GPU 集群可以在大约一周的时间内训练GPT-4模型。据推测，OpenAI使用 24-32K A100 GPU 集群在三个月内（或 12 周）训练了他们的GPT-4模型。
Nvidia 的H100 GPU为 TF32 操作提供了比A100 GPU 好 2.5-3 倍的性能。此外，事实证明，混合精度训练是 TF32 操作的可行替代方案，可以加快训练性能。表 2 显示了具有H100和混合精度算法的 GPU 集群大小（TF16 和 TF32 之间的比率为 7:3）。
H100支持 FP8 算法。Nvidia的论文声称，当使用 FP8 算法进行卷积和矩阵乘法时，精度损失不到 1%，同时将输出张量保持在 FP32。微软最近发表了一篇论文，展示了 FP8 是训练 LLM 模型的可行选择。
使用彻底尝试不同方式对模型进行分区以减少通信开销的框架可以将 GPU 利用率提高到约 50%。
FP8操作和50%的利用率将进一步减少集群规模，如表3所示。
总而言之，使用最新的 GPU 硬件、混合精度和 FP8 浮点运算以及现代拓扑感知编译器可以将 GPT-4 的集群大小从275K缩小到大约34K GPU，并且该模型可以在一周内完成训练！而且，除了 GPT-4 之外，所有其他 LLM 模型都可以在不到一周的时间内使用不到 1K GPU 进行训练！
训练的总成本包括电力输送和散热解决方案。这在大型集群中可能非常昂贵，影响企业训练 LLM 模型的频率。如果企业可以容忍更长的训练时间（每月一次而不是每周一次），则有机会进一步优化硬件。例如，GPT-4模型需要约 8.5K（34K/4）个 GPU，训练周期为 4 周。
CSP通常使用 >64K GPU 构建更大的集群，以并行运行多个训练作业，并为下一代模型提供面向未来的结构（假设模型大小和数据集持续增长）。
Nvidia 发布了内存增加 75% 的H200 GPU。其路线图显示下一代B100将于 2024 年推出，X100将于 2025 年推出。鉴于 Nvidia 每年都会发布新 GPU，目前尚不清楚几年后面向未来是否具有明显优势。新版本的 GPU 肯定会比H100具有更好的内存容量/功率/性能和每 TFLOP（每秒万亿次浮点运算）成本。
构建多个合理大小（8K）的并行集群是否比构建一个具有 128K 节点的巨型集群更高效且更易于管理？
II. 并行
上一节中的集群数学假设模型参数和数据集分布在集群中的所有 GPU 上，并且训练并行进行。模型分区在优化网络拓扑和提高 GPU 效率方面起着关键作用。
让我们研究一下 GPU 服务器以及如何通过考虑服务器内和服务器间通信能力来对模型进行分区。
GPU 服务器
Nvidia 的基于H100的 GPU服务器由 8 个 GPU 和 8 个 ConnectX-7 NIC 组成。这些 GPU 可以使用连接到 NVSwitches 的高速 NVLinks 相互通信。每个 GPU 通过每个方向的 3600 Gbps NVLlinks（18 个链路，每个 200 Gbps）连接到一组 NVswitches。
服务器内部的 8 个 GPU 中的每一个都可以通过 8 个 ConnectX-7 NIC 与外部交换机通信，每个 GPU 能够向 NIC 发送 400Gbps。这 8 个 NIC 可以通过 4 个 800Gbps OSFP 光纤接口连接到外部 GPU 结构中的叶交换机。
与每个 GPU 通过结构连接到其他服务器所使用的 400Gbps 链路相比，服务器内的 GPU 可以通过 NVswitches 以 9 倍带宽相互通信。
此外，每台服务器都有两个单独的 OSFP 模块连接到存储网络。海量训练数据集（对于 26000 亿个Tokens 的GPT-4模型，约 10TB，其中有 40 亿个Tokens ）保存在存储文件系统中，并通过存储网络提取到系统内存中。此活动与训练重叠，理想情况下不应导致迭代之间出现停顿。
服务器内的每个H100 GPU 都有 80GB 内存。此内存存储模型参数（权重和偏差）、梯度、中间值（前向传递期间的输出激活和后向传递期间的梯度）、优化器状态、计算的暂存空间和其他杂项状态。对于大型语言模型，大多数内存都被模型参数、梯度和优化器状态占用。中间值所需的存储空间在很大程度上取决于模型架构和训练批次大小。
模型分区（Modle Partitioning）
典型的深度学习模型训练涉及将数据集分成相等数量 ( B ) 的批次。对每个批次进行一次训练迭代（运行前向传递以进行预测，并运行后向传递以计算每个参数的梯度），更新参数，然后使用下一个批次再次重复训练。此过程重复进行，直到完成所有批次。这称为一个时期。训练由数百个这样的时期组成，直到模型收敛。
数据并行（Data Parallelism）
加快大型数据集训练的最简单方法是将每个批次划分为Bm个小批次，并在每个小批次上训练模型的副本。这称为数据并行。
在数据并行中，每个模型副本都会对小批量进行一次完整的训练迭代。所有模型副本的每个参数的梯度都会在迭代结束时汇总。所有参数的更新梯度都会广播到所有模型副本，一旦模型副本使用更新的梯度调整其权重（参数），就会使用下一个批次再次重复训练。
对于 LLM 模型，批次中的数据通常包含序列长度（模型在训练期间预测下一个token时查看的连续token数）。不同模型的批次大小和序列长度有所不同。
当使用小批量训练模型时，从模型副本的一层传递到下一层的中间激活或结果的大小与最小批量大小成线性比例。
根据微软的论文，一个 175B 参数的模型需要大约 16 * 175B 内存来存储混合精度训练中的参数、梯度和优化器状态。显然，单个只有 80GB 内存的H100 GPU 只能拟合部分模型。仅拟合参数，GPT-3模型就需要约 35 个 GPU，而GPT-4模型大约为 300 个 GPU。
如果Nm是将模型划分成的 GPU 数量，Nd是数据并行组的数量（模型副本数量），则集群中的 GPU 总数为Nm x Nd ，与仅使用Nm 个GPU训练整个数据集相比，这将使训练速度提高Nd倍。
通过流水线实现模型并行
有两种方法可以在 Nm GPU 上对模型进行分区。在流水线模型并行中，模型分为几个阶段，每个阶段由一个或多个模型层组成。每个阶段都分配给不同的 GPU。当在具有重复transformer block的大型语言模型上使用时，可以为每个设备分配相同数量的transformer 层。流水线并行有时称为inter-operator并行。
流水线并行性通过将小批次（Bm）进一步划分为微批次（Bu）并以流水线方式在 GPU 中调度它们，使所有Nm 个GPU 保持活跃。
下图是流水线和数据并行的概念示例（每个模型副本有两个管道阶段）。
为了清晰起见，此图中的模型内部结构被简化了。
数据分布：在每个Nm GPU的模型副本中，
1 - 每个小批次进一步划分为“ Bu ”个微批次。第一个微批次在流水线的第一阶段（分区 #1）发送到 GPU。
训练
1 – 一旦第一阶段完成微批次的前向传递，它就会将其输出发送到下一阶段（管道中的下一个 GPU）并立即开始处理队列中的下一个微批次。
2 – 这个过程持续到所有微批次上的前向传递都完成
3 - 此时，开始对微批次进行反向传递。从最后一个阶段开始，每个流水线阶段都会计算所有微批次的梯度，并将其传递到上一个阶段。
4 - 每次反向传递结束时，每个流水线阶段中“每个参数”的梯度都会累积。在此模型副本中的所有反向传递结束时，所有参数都会有其累积梯度。
跨Nd 个数据并行组的梯度聚合：
1 – 然后将每个参数的梯度汇总到所有Nd模型副本中。
2 – 然后将每个参数的聚合梯度广播到保存相应参数的所有模型副本上的Nd GPU。
同步/重复：
1 - 一旦所有模型副本都更新了其参数，就会对训练中的下一个小批量重复该过程。当前批次在管道中执行时，会通过存储结构获取下一个小批量的数据集。
如下图所示，执行所有前向传递后再执行所有后向传递意味着 GPU 需要存储所有Bu微批次在前向传递期间生成的中间激活，以便它们可用于后向传递期间的梯度计算。
流水线并行性和气泡（灰色区域）的图示，其中前向和后向传递结合在一起。每个设备代表一个流水线阶段。四个流水线阶段和八个微批次。
图片来源：Nvidia 的 Megatron-LM 论文
相反，大多数流水线调度方案使用一个前向传递后再执行一个后向传递 ( 1F1B )，如下图 所示。这会将中间激活的存储减少到流水线的深度。
前向和后向传递的交错导致每个 GPU 中用于中间激活函数结果的存储空间减少。每个 GPU 在任何时候都存储来自 4 个微批次（流水线深度）的激活函数结果。
图片来源：Nvidia 的 Megatron-LM 论文
梯度聚合可以在一定程度上与前向/后向传递重叠。例如，由于模型副本流水线阶段中的 GPU 已积累了来自所有微批次的梯度，因此它可以与所有 Nd 模型副本中同一流水线阶段中的所有 GPU 一起参与梯度聚合。
发送所有Bu微批次后通常会有一个刷新期（同步期），以便所有模型副本中的 GPU 有足够的时间在下一次迭代开始之前更新其参数。刷新持续时间取决于管道深度和模型副本数量（ Nd）。模型副本越多，聚合梯度所需的时间就越长。
确定最佳流水线深度至关重要。目标是最大限度地利用每个 GPU 的计算和内存资源。一些方案试图以存储两对模型参数为代价来减轻流水线刷新。它们使流水线能够在下一对上工作，同时为当前对聚合和同步梯度。这可能会导致训练中的一些准确性损失，并且需要每个 GPU 中的内存增加一倍来存储参数。没有办法避免刷新/同步周期来获得准确的结果和更快的收敛。
除了参数/优化器状态之外，GPU 还必须在前向传递期间存储中间激活结果。如上图 所示，中间激活的存储与流水线级数成正比。为了节省用于参数的 GPU 内存空间，某些方案允许在后向传递中重新计算中间激活 - 从而在 GPU 计算和 GPU 内存之间进行权衡。
一些大型模型（例如GPT-3/GPT-4）需要更深的流水线才能适应每个流水线阶段中的参数/状态。使用更深的流水线，训练需要更长的时间，并且在迭代结束时需要刷新整个流水线的同步延迟也会占用大量的训练时间。
这就是张量或运算符内并行性有所帮助的地方。在张量并行性中，每个流水线阶段中的计算分布在多个 GPU 上，以便每个流水线阶段可以完成更多计算。这可以减少流水线分区并降低延迟。
后续文章将继续介绍张量并行，以及GPU网络的设计考虑。
参考文献：
Sharada Yeluri： GPU Fabrics for Gen AI Workloads
https://community.juniper.net/blogs/sharada-yeluri/2024/01/02/gpu-fabrics-for-genai-workloads
高阅读量文章
