Efficient Pre-Training of LLMs via Topology-Aware Communication Alignment on More Than 9600 GPUs

作者/机构: Guoliang He, Youhe Jiang, Wencong Xiao, Kaihua Jiang, Shuguang Wang, Jun Wang, Zixian Du, Zhuo Jiang, Xinlei Zhang, Binhang Yuan, Eiko Yoneki (University of Cambridge, ByteDance Seed, HKUST)

A1 主要贡献

大规模语言模型（LLMs）的预训练是一个资源高度密集的过程，其性能取决于数据集大小、计算能力和模型参数。为了满足这些需求，公司不断构建大规模GPU集群并进行数千计算节点的大规模训练。然而，LLM预训练因其独特的通信模式——GPU在特定组内以稀疏但高容量的突发方式交换数据——而带来了独特的挑战。低效的资源调度会加剧带宽争用，导致训练性能不佳。

核心问题与研究目标：
现有的集群调度器（如【43, 44, 45, 46, 7, 5】）未能针对LLM工作负载集成网络拓扑感知的调度。它们的主要局限性在于缺乏对LLM训练中固有的高容量、稀疏活跃的分布式通信模式的认识。如图1a所示，生产环境中的LLM训练有30%-50%的时间用于通信，但如图1b所示，超过99%的GPU对之间没有直接流量，数据交换仅在特定的通信组内发生。现代GPU集群采用多层胖树网络拓扑（图2b），不合理的作业放置会导致显著的带宽损失和通信开销。为解决此问题，本文旨在开发一种高效的资源调度机制，以适应现代数据中心中资源密集且复杂的通信模式，有效对齐LLM通信模式与数据中心拓扑。


图 1: 生产环境中LLMs训练的通信特征。

关键挑战：
1. 通信模式与数据中心拓扑的错位：调度器通过装箱（bin-packing）优化GPU局部性，但缺乏对LLM预训练作业（LPJ）通信结构的感知。因此，即使作业被紧密打包，也可能因通信组跨越交换机而导致低效通信。
2. 不同通信维度间的权衡未被解决：在对齐数据并行（DP）和流水线并行（PP）通信模式之间存在根本性的权衡。由于GPU同时参与这两种通信组，无法同时完美对齐两者，调度器必须在放置时仔细平衡。

本文贡献：
本文提出了Arnold，一个协同设计训练框架和集群调度的系统，旨在有效对齐LPJ与现代数据中心网络拓扑。
1. 贡献 1：识别并描述了在大规模预训练中，将LLM通信模式与现代数据中心拓扑对齐的挑战。通过特征分析研究了物理网络拓扑对单个通信操作和端到端训练性能的影响。
2. 贡献 2：设计了一种调度算法，以有效对齐LPJ的通信模式与数据中心拓扑，并提出了一种资源管理策略来为放置预留节点。
3. 贡献 3：通过全面的仿真实验，评估了该调度算法在将通信组的最大扩展（maximum spread）降低多达1.67倍方面的有效性。在一次生产运行中，验证了所提出的调度器能将一个超过9600个GPU的LPJ性能提升10.6%。

A3 背景知识/关键Observation/设计原则

背景知识

分布式训练。LLMs是拥有数十亿参数的Transformer模型，必须使用多GPU系统进行训练【40, 36】。常见的训练框架【37, 35】采用混合并行策略来加速训练过程，包括：
* 数据并行（Data parallelism, DP）：ZeRO优化器【34】将模型权重和梯度分片到数据并行进程中，并在训练步骤结束时通过all-gather和reduce-scatter通信操作进行同步。
* 流水线并行（Pipeline parallelism, PP）：模型层被划分为几个阶段（PP size），每个阶段交替进行与相邻阶段的通信以及阶段内的计算。阶段间通信通过点对点（P2P）通信操作（如send-recv）执行。
* 张量并行（Tensor parallelism, TP）：一个PP阶段内的模型权重被进一步分片到多个GPU上，以减轻内存压力。在前向和后向传播过程中，需要all-gather和reduce-scatter来同步中间激活。

这些并行策略的组合形成了多样的GPU通信模式，训练框架使用通信组来管理这种复杂性。每个GPU在初始化时被分配到一个DP、TP和PP通信组中。不同并行策略和通信组的图示见图2a。先前的工作【24, 37】已指出，TP通信组应优先放置在同一节点内的GPU上，以利用高带宽的NVLink互连，因为其数据依赖性强。因此，节点间通信仅在DP和PP组内发生。由于只有节点间通信对物理网络拓扑敏感，本文的重点是DP和PP组的通信模式。


图 2: LLMs并行化与数据中心拓扑。

数据中心拓扑。图2b概述了我们的HPC集群，其结构与其他现代数据中心相似。超过2000个节点通过三层交换机互连，形成一个类似CLOS的拓扑【8】。叶交换机表示为s0，连接同一机架内的节点。然后，多个s0交换机连接到一个汇聚交换机（s1），形成一个节点的minipod。最后，s1交换机连接到核心交换机，实现minipod之间的通信。每层交换机都有32个用于上行和下行链路的端口。当两个不同minipod的节点相互通信时，跳数最多。计算节点配备8个H800 GPU，每个GPU连接到一个InfiniBand【27】网卡。节点内的GPU通过高带宽链路（如NVLink【30】，带宽为400Gbps）连接，而节点间通信则通过InfiniBand网络实现。

关键观察与挑战

给定用户指定的GPU数量和LPJ的混合并行度，作业调度系统将作业入队并执行资源调度以在我们的GPU集群中找到最佳放置。然而，我们观察到现有调度系统在实践中未能将LLM通信模式与数据中心拓扑对齐。

观察1：作业放置的错位导致跨交换机通信增加。最先进的集群调度器【43, 2, 45】采用装箱（bin-packing）策略来增强LPJ的GPU局部性。然而，如图3a所示，即使调度器将一个4节点（32个GPU）的LPJ打包在一个minipod内，通信组仍可能未对齐，因为DP和PP组都涉及跨越核心交换机的通信，其通信距离更长。这种错位源于调度器在调度时缺乏对LPJ通信结构的感知，限制了其根据作业通信模式分配GPU资源的能力。


图 3: 通信模式的对齐。图中高亮显示了一个DP组和一个PP组。

观察2：DP和PP通信优先级之间的权衡未被解决。图3b和3c展示了LPJ的两种潜在对齐方式，一种优先考虑DP通信，另一种优先考虑PP。这揭示了两者之间的根本性权衡，因为DP和PP是LLM训练中广泛使用的正交并行策略。每个GPU都同时参与一个DP组和一个PP组，因此不可能同时完美对齐两种通信模式。一个精心设计的调度器必须考虑这种权衡，并在作业放置期间平衡两种组类型的对齐需求。

挑战：LPJ的通信与拓扑对齐调度。为了有效调度LPJ，调度器必须了解其多样的通信模式，并最小化它们在数据中心的扩展（spread）。此外，有效平衡DP和PP组的扩展至关重要，这需要对现代数据中心的通信模式进行深入的特征分析。

LPJ通信模式的特征分析

尽管先前的研究【45, 2】探讨了局部性和拓扑，但它们的范围受限于（1）关注数据并行（all-reduce）工作负载和（2）对节点间拓扑的考虑有限。为了弥补这些不足，我们使用NCCL测试（一个用于测量NCCL【28, 29】所用通信操作延迟和带宽的基准套件）来研究节点间拓扑的影响。我们专注于跨minipod的节点间拓扑，因为LPJ的规模通常需要跨多个minipod分配计算节点，其中最慢的通信路径往往决定了整体通信开销。

通信操作性能。图4研究了通信操作的性能。我们使用总线带宽（BusBw）作为性能度量，它通过考虑集合通信的rank数量来反映硬件峰值带宽。图4a展示了不同通信模式在不同消息大小下的性能。对于集合通信，消息大小必须大于 $2^8$ (≈ 256) MB才能充分利用带宽，而对于像send-recv这样的P2P通信，较小的消息大小 (≈ 2 MB) 就能使带宽饱和。使用一个广泛采用的分析模型（详见附录§C）并代入一个7B GPT模型的参数，我们可以得出DP和PP组的数据量分别为2GB和30MB，这表明带宽被充分利用。


图 4: 通信操作的性能。(AR: all-reduce, AG: all-gather, RS: reduce-scatter, SR: send-recv)

跨minipod通信的性能影响。图4b和4c说明了将通信组扩展到多个minipod时BusBw的下降情况。随着通信扩展到更多的minipod，集合操作的性能下降高达17%，而P2P操作的性能下降高达70%，这凸显了GPU局部性和对齐的关键重要性。此外，我们的发现表明，在多租户集群环境中，共同定位的作业可能会经历减少的带宽争用，如附录D中记录的作业间干扰模式所证实。

端到端训练性能。基于对单个通信操作的特征分析，我们按比例缩小了一个生产模型，并使用96个GPU（跨越2个minipod）运行该工作负载，以进一步了解网络拓扑对LLM训练的影响。图5a展示了LPJ在三种不同放置策略下的吞吐量。dp-aligned放置如图12所示。除了在第550步左右因垃圾回收有轻微波动外，吞吐量在200步后变得稳定。PP-aligned放置始终优于其他两种策略，表明优先考虑PP组通信能带来性能提升。对于密集模型和MoE模型，平均性能提升分别为2.3%和1.8%。对于密集模型，我们还观察到PP通信主导了通信开销，因为优先放置DP组并未带来加速。对于MoE模型，减少DP和PP组的扩展都有助于性能提升，但PP组的优化带来了更多改进。


图 5: 端到端训练性能。(a) 三种不同放置策略对两种LLM的比较。(b) 通过在优化对齐下扩展模型大小带来的性能提升。

模型规模的影响。图5b显示，如果我们通过增加更多层来扩展模型大小，性能提升会持续增加。我们认为这是因为通信是主要的性能瓶颈，而更大的模型会进一步放大通信量。因此，随着模型规模的增加，优化的好处变得更加显著。

minipod内部拓扑的影响。我们通过在单个minipod内改变节点放置来进一步检验训练性能对minipod内部网络拓扑的敏感性。对于一个24B参数的密集模型，观察到的最大性能变化为0.3%，对其他模型的影响可以忽略不计。由于一个组的通信开销通常由最慢的链路决定，并且LPJ通信组经常跨越多个minipod，我们得出结论，训练性能对minipod内部拓扑基本不敏感。

结论与推广。我们在另一个GPU集群（详见附录E）中重复了特征分析实验，发现最佳放置策略可能受到模型大小和GPU类型的影响。由于LPJ通常是提前调度并长期部署的，因此有必要在调度前进行特征分析，以识别组内的通信瓶颈。这使得可以相应地优化放置策略并平衡权衡，详见§5。

A2 方法细节

调度算法

Arnold的核心组件是其调度模块，该模块旨在根据用户指定的GPU数量和混合并行度，有效地为LPJ分配GPU。本节我们正式定义调度问题并提出我们的解决方案。

5.1 工作负载表示

通信矩阵表示。Arnold通过一个通信矩阵来表示一个LPJ，其中行代表一个PP组，列代表一个DP组。具体来说，给定一个作业规范，包括GPU总数、PP、TP的度，Arnold使用公式1计算通信矩阵的大小。

$$\begin{aligned} \begin{aligned} DP &= \#GPUs/TP/PP \\ \#row &= DP/(8/TP) \\ \#col &= PP \end{aligned} \end{aligned}$$

一个96个GPU，DP=6，PP=2的例子如图12所示。对于通信矩阵中的一个节点 $v_{ij}$，它附带一个向量 $[v_w, v_d, v_p]$，分别代表每个GPU的权重大小、DP和PP通信量。这些值使用附录§C中详述的分析模型计算，并用于平衡DP和PP组之间的权衡。

5.2 目标

调度目标函数。调度目标函数旨在最小化DP和PP组的最大扩展（spread）。对于通信矩阵中的一个节点 $v_{ij}$，有 $k$ 种可能的分配方案（即分配到 $k$ 个minipod中的一个）。因此，它有一个长度为 $k$ 的独热（one-hot）决策向量 $x_{ij}$，表示将其放置到 $k$ 个minipod之一的决策。然后，目标函数可以写成公式2。

$$MIN\ [\alpha \max_{i}(D(x_{i1}, x_{i2}, ...x_{in})) + \beta \max_{j}(D(x_{1j}, x_{2j}, ..., x_{mj}))]$$

其中，n个向量之间的距离D定义如下：

$$D(v_{1}, v_{2}, \ldots, v_{n}) = |\{i : \exists j, l \in \{1, 2, \ldots, n\}, j \neq l, v_{j}[i] \neq v_{l}[i]\}|$$

直观地说，这个距离衡量了一个通信组的扩展程度，即如果有任意两个向量在第 $i$ 个位置上不同，那么第 $i$ 个位置就对距离贡献1。目标函数旨在最小化DP和PP组最大扩展的加权和，取最大值是因为最慢的通信组是拖慢整个训练过程的“掉队者”（straggler）。权重参数 $\alpha$ 和 $\beta$ 代表亲和力，控制DP和PP组之间的权衡，且 $\alpha + \beta = 1$。由于离散距离的计算，公式2不能通过现成的求解器进行高效的在线调度，因此我们寻求简化。

领域特定简化。我们发现LPJ的通信组是同构和同步的，因为节点是“帮派调度”（gang-scheduled）的，并且必须在每个训练步骤结束时同步梯度。因此，每个PP组总是在大致相同的时间开始，并执行相同数量的计算和通信。同样，DP组也在同一时间执行梯度同步。这些特性使我们能够通过将调度单元粗化为一个通信组来简化调度目标函数。因此，我们将公式2转化为一个类似装箱问题的形式：

$$MIN \ [\alpha \sum_{j}(y_{j}) + \beta T]$$

约束条件如下：

$$\text { s.t. } \quad \forall i: \sum_{j} s_{i j} \leq T \quad(\text {Max Spread})$$

$$\forall j: \sum_{i} p_{i j} \leq c_{j} y_{j} \quad \text { (Capacity Const.) }$$
$$\forall i : \sum_{j} p_{i j}=1 \quad \text { (Allocation Const.) }$$
$$\forall i, j: p_{i j} \leq s_{i j} \quad \text { (Minipod Selection) }$$
$\forall j : y_j \in \{0, 1\}$
$$\forall i, j : s_{i j} \in\{0,1\}, p_{i j} \in[0,1]$$

其中，$y_j$ 表示第 $j$ 个minipod是否被使用，$c_j$ 是该minipod的归一化容量，根据可用节点数动态更新。$s_{ij}$ 表示第 $i$ 个通信组是否分配给minipod $j$，$p_{ij}$ 表示第 $i$ 个通信组分配给minipod $j$ 的百分比。$T$ 是一个引入的辅助变量，使我们能够最小化通信组的最大扩展。$\alpha$ 和 $\beta$ 是与之前相同的亲和力参数。总的来说，最小化 $T$ 能有效地将通信组整合到尽可能少的minipod中，而目标项 $\sum_j(y_j)$ 则控制了另一种通信组的扩展。例如，我们将PP组视为一个调度单元。通过设置 $\alpha = 0$，调度器可以将每个PP组放入一个minipod，但这会导致DP组产生更多的跨交换机通信；而通过设置 $\beta = 0$，调度器会最小化minipod的总体使用量，尽管这种放置可能导致PP组的跨交换机通信。在这个公式中，所有变量要么是整数，要么是分数。因此，目标函数可以使用现成的混合整数规划（MIP）求解器【4】进行高效的在线调度。求解MIP后，我们为属于同一minipod的节点分配连续的rank索引，以减少每个通信组内的跨交换机通信。

平衡权衡。公式4中的亲和力参数需要在DP和PP组之间进行权衡，这取决于模型配置和GPU类型(§4)。为了进行在线调度，我们将§4中的特征分析结果存储到一个数据库中，并通过搜索最佳匹配来确定亲和力参数的值。通信矩阵计算每个GPU的参数（$v_w$）和通信量（$v_d, v_p$）。然后，我们计算计算与通信的平均比率和DP与PP的通信量比率，即 $r_1 = v_d+v_p$ 和 $r_2 = \frac{v_d}{v_p}$，其中 $m_b$ 是微批量大小。这些比率用于通过比较欧氏距离在数据库中找到最匹配的作业，即 $\text{MIN}_{r_i,r_j} \sqrt{(r_1 - r_i)^2 + (r_2 - r_j)^2}$，因为如果GPU具有相似的计算负载和通信量，它们会表现出可比的性能特征。LPJ与元数据 $\langle \text{GPUtype}, j_{dp}, j_{pp} \rangle$ 相关联，其中 $j_{dp}$ 和 $j_{pp}$ 对应于DP对齐和PP对齐放置策略的改进。然后，亲和力参数 $\alpha$ 和 $\beta$ 根据 $j_{dp}$ 和 $j_{pp}$ 的相对性能改进得出，即 $\alpha = \frac{j_{dp}}{j_{dp}+j_{pp}}$ 和 $\beta = \frac{j_{pp}}{j_{dp}+j_{pp}}$。由于LLM训练的重要性和其统一的架构，LPJ通常是提前调度和预先进行特征分析的，因此数据库中的分析数据可以在在线调度时被查询。例如，对于H800 GPU集群中的一个24b密集模型，调度单元被设置为PP组，$\alpha$ 设置为零，因为PP组明显主导了通信开销。对于24b MoE模型，$\alpha = 0.3$ 且 $\beta = 0.7$。

5.3 资源管理

队列策略和资源预留。调度算法为GPU集群中的LPJ计算全局最优放置，这不可避免地会与其他作业发生冲突。为了解决这个问题，我们开发了一个排队策略来管理作业队列并为即将到来的LPJ预留资源。算法1阐述了我们的调度策略。一旦LPJ被规划，调度器就求解MIP方程并预留资源。此后，新来的作业将优先分配在预留区域之外。否则，为了提高资源利用率，如果一个新来作业的预测JCT（作业完成时间）早于LPJ的到达时间，它仍然可能被调度在预留区域内。如果这两个条件都不能满足，该作业将在调度间隔内被延迟。我们还采用了一个由机器学习驱动的JCT预测器来平衡排队延迟和资源利用率之间的权衡。其设置和评估分别在附录G和H中详述。

系统实现

系统实现概述。我们用超过3000行Python代码实现了一个Arnold原型。该原型包括调度模块和一个可以重放生产轨迹的轨迹驱动模拟器。我们还有一个与Kubernetes【21】集成的部署版本。对于训练框架，我们基于Megatron【24】进行构建，并对其进行修改以确保通信组遵循Arnold提供的放置策略。图6给出了Arnold的概览。


图 6: Arnold的架构概览。

A4 实验环境

硬件配置:
* 主集群: 超过2000个计算节点，每个节点配备8块H800 GPU。节点间通过三层CLOS-like拓扑的InfiniBand网络互联。节点内GPU通过高带宽NVLink连接。
* 备用集群: 节点配备L20 GPU（详见附录E）。
* 生产环境: 实验规模超过9600个GPU（即1200+个节点）。

软件配置:
* 代码实现: Arnold原型采用Python编写，代码量超过3000行。
* 依赖库与框架: 训练框架基于Megatron【24】进行修改。生产部署与Kubernetes【21】集成。
* 模拟器: 开发了一个轨迹驱动的模拟器用于低成本地评估调度算法。
* 对比基线:
* 模拟实验基线: Best-fit【32】, Random-fit【44】, GPU-packing【43, 45】, Topo-aware【2】。
* 生产实验基线: MegaScale【18】，一个用于LLM的业界先进生产系统。

数据集与模型:
* 特征分析: 使用按比例缩小的生产模型（包括密集型和MoE型）和24B参数的密集模型，以及NCCL测试套件【29】。
* 模拟实验: 使用从生产集群中提取的作业配置，分为小、中、大三种规模（详见表1）。
* 生产实验: 使用一个MoE变体的LLM，模型参数超过4000亿。实验分为208个GPU的中等规模和超过9600个GPU的全规模。

A4 实验结果

仿真实验

实验内容：我们使用轨迹驱动的模拟器，将Arnold的调度算法与四种基线算法（Best-fit, Random-fit, GPU-packing, Topo-aware）进行比较。评估指标为公式2中定义的通信组最大扩展的加权和。实验在三种不同规模的拓扑和作业配置下进行（见表1），并改变亲和力参数$\alpha$的值。

实验结果：
* 性能对比：如图7所示，Arnold的算法在所有设置下都持续优于其他基线，与最好的基线相比，最多可将加权扩展降低1.67倍，平均降低1.2倍。在简单的拓扑和作业配置下（设置(i)），Arnold的性能与其他打包策略相当，因为优化空间有限。但在中等和大型作业中，由于搜索空间巨大，Arnold的优势变得明显。
* 亲和力参数影响：我们观察到，随着$\alpha$（DP组的亲和力）增加，Arnold的性能逐渐接近其他基线。这是因为当$\alpha=1$时，目标函数退化为标准的装箱问题，与其他装箱算法没有区别。在实践中，根据我们的特征分析结果，$\alpha$通常不会大于0.5，因此Arnold通常表现更优。


表 1: 基准测试设置。网络拓扑{x}, {y}分别表示{x}个minipod和总共{y}个节点，作业配置中的数字分别是DP, TP, PP的度。调度单元是PP组。


图 7: 不同调度算法下通信组的加权扩展。

可扩展性：我们通过与保证最优解的枚举方法比较调度延迟来评估算法的可扩展性。如图8所示，枚举方法不具备可扩展性，在简单拓扑下调度14个节点就需要30秒，在中等拓扑下调度10个节点需要超过100秒。相比之下，即使在拥有1000+个节点的集群中调度一个512节点的作业，Arnold的算法也保持了较低的延迟。


图 8: 不同配置下的调度延迟。

集群实验

实验内容：为了评估Arnold在真实环境中的有效性，我们在生产GPU集群中进行了实验。我们将Arnold与一个先进的LLM生产系统MegaScale【18】进行比较。实验首先在208个GPU的中等规模上验证加速效果，然后进行超过9600个GPU的全规模预训练。

端到端性能结果：
* 吞吐量提升：如图9a和9b所示，Arnold在中等规模和全规模实验中分别实现了5.7%和10.6%的平均吞吐量提升。
* 通信组扩展降低：在中等规模实验中，Arnold将DP组和PP组的最大扩展分别降低了3倍和2倍；在全规模实验中，则分别降低了1.5倍和1.3倍。全规模下降幅度较小是因为作业请求的GPU占集群总量一半以上，调度空间受限。
* 性能一致性：如图9c所示，在全规模实验的训练过程中，Arnold的性能持续稳定地优于MegaScale（除160步的torch profiling扰动外）。对于运行超过一个月的LPJ，这种提升对于下游任务和成本节约（GPU小时和人力资源）意义重大。


图 9: 集群实验。

分解分析：
* 意外发现：使用Torch profiler进行的内核级分析（图10）显示，虽然通信与拓扑对齐的放置策略如预期般提升了通信内核（如broadcast）的性能，但它同时导致其他内核（包括一个reduce-scatter通信核和一个GEMM计算核）性能下降。
* 根本原因：通过在附录I中进行的系统性消融研究，我们确定GPU流之间的资源争用是这一现象的根本原因。当通信和计算内核在多个流上并发执行时，就会出现资源争用。拓扑感知调度影响了通信效率，这反过来又影响了NCCL的动态自动调优，从而改变了计算和通信内核之间的资源（如SMs）分配。
* 结论：这些观察结果拓宽了拓扑感知调度的范围，表明其影响不仅限于通信效率，还延伸到影响计算内核的执行特性。


图 10: 分解分析。

A5 结论

本文介绍了Arnold，一个总结了我们在大规模有效调度LLM预训练作业（LPJ）方面经验的调度系统。通过深入的特征分析，我们开发了一种调度算法，以将LPJ与现代数据中心的拓扑对齐。通过基于模拟和在真实GPU集群中的实验，我们展示了该系统的有效性。

A6 附录

A 相关工作

LLMs训练。LLMs已成为机器学习领域的重要工作负载【31, 38, 50】。开发LLM的第一步需要用数万亿个token来训练一个大型Transformer模型，因此计算基础设施不断发展以适应这一挑战性工作负载。例如，通过模型并行器搜索高效的并行策略【52, 23, 39】，构建专门用于训练可扩展性的训练框架来协调大规模工作节点【18, 37, 24】，以及开发高性能算子以最大化硬件加速器的利用率【9, 6, 12】。然而，这些工作与本文提出的优化是正交的，因为物理网络拓扑仅在集群调度层可见。因此，我们的工作对底层基础设施代码是透明的，并且是在现有加速训练性能努力的基础上实现的加速。

深度学习作业调度器。用于深度学习任务的作业调度系统已被公司广泛部署【7, 45, 44, 2, 25】。然而，没有一个深度学习作业的规模和重要性接近LLM预训练作业。Arnold通过提供一个专为调度LLM预训练作业量身定制的解决方案来弥补这一差距，从而补充了现有调度器。最近的研究已经开始探索GPU集群中LLM工作负载的特性【13】。相比之下，我们的工作专门针对优化LLM预训练性能。

B 局限性与未来工作

故障恢复。在硬件发生故障时，LPJ的最优放置将不可避免地改变，但重新求解MIP然后迁移到新的放置方案成本太高。一种潜在的方法是在初始调度时增加通信组的GPU数量作为备份，这些备份GPU只运行可抢占的作业，在需要时可以替换故障节点。

其他网络拓扑。虽然特征分析结果源自我们内部的数据中心环境，但它们对基于CLOS的网络拓扑具有广泛的适用性，这是现代数据中心部署中占主导地位的网络架构。通过改变亲和力参数，可以在自己的数据中心中有效地权衡DP和PP组的平衡。因此，特征分析方法和调度算法对于其他数据中心的大规模LPJ是可推广的。另一方面，专用于LLM训练的更新网络架构正在出现【33, 26, 41】，这需要仔细考虑调度算法。我们提出的调度方法通过将通信模式与数据中心拓扑对齐，为LPJ迈出了开创性的一步，并且其有效性已在真实的生产集群中得到评估。我们将这一方向的进一步探索留给未来的工作。

C 通信量的分析性估计

通信量分析模型。为了估计预训练作业的通信量，我们采用了一个针对GPT-based变体的分析模型。我们使用与先前工作【24】相同的符号，表示词汇表大小为V，全局批量大小为gb，微批量大小为mb，序列长度为s，隐藏维度为h，层数为l，DP大小为dp，PP大小为pp，VPP大小为vp，微批次数为m。我们有：

$$m = \frac{gb}{mb * dp}$$ $$DP-volume = h * (V + s) + l/pp * (4h^2 + 2h + \underbrace{8h^2 + 7h}_{\text{dense layer}})$$

对于MoE模型，我们可以相应地用MoE层的参数替换。

$$PP-volume = 2 * mb * s * h$$

D 对共享负载的敏感性


图 11: 带宽干扰。

共享负载下的带宽干扰。由于GPU集群通常是多租户的以提高资源利用率，我们也定量研究了节点间通信的干扰。在我们的集群上线之前，我们通过运行NCCL测试对集群进行了大规模压力测试。我们记录了总线带宽的时间序列，并显示在图11中。压力测试包含3个作业，每个作业请求数千个GPU，并跨越集群中的多个minipod。每个作业持续执行消息大小为2GB的all-to-all通信，这模拟了在繁忙集群上执行大量节点间通信的作业，因为all-to-all会在网络中产生大量流量。我们可以观察到所有三个作业的性能波动。例如，在01:22作业1启动后，作业3的性能有0.5GB/s（3%）的轻微下降。在压力测试期间，作业3的最大性能下降高达5%。这表明，跨越更多minipod的作业不仅会遭受更大的带宽损失，而且更容易受到集群中其他工作负载的干扰。

E Ada Lovelace GPUs

在L20 GPU上的验证。我们在另一个GPU集群中重复了特征分析实验，该集群的每个节点都配备了L20 GPU，以确保我们的发现不限于H800 GPU。我们在表2中简要总结了结果。我们观察到，对于某些模型配置，DP-aligned可以产生更大的加速。这可能是因为在训练期间，L20 GPU使用8位数据格式，这使得PP阶段之间的通信量减半。然而，DP组仍然使用32位进行参数和梯度同步，因此通信量保持不变。结果，DP组通信可能成为主导开销，使得优先考虑DP组的放置策略更有利。然而，随着模型规模的增长，瓶颈又转回到PP组的通信，因此PP-aligned变得更优。


表 2: L20 GPU集群上的结果。

F 通信矩阵

通信矩阵示例。图12中的放置示例具有DP组大小为6，PP组大小为2。不同颜色的节点表示它们被放置在不同的minipod中。数字是节点的rank。在这个例子中，DP组是对齐的，而PP组的通信必须跨越核心交换机。


图 12: LPJ的一个示例放置（96个GPU和12个节点）。

G 队列管理

作业调度策略。算法1展示了我们的排队策略。它管理作业队列并为即将到来的LPJ预留资源。它还采用了一个由机器学习驱动的作业完成时间（JCT）预测器，以平衡排队延迟和资源利用率之间的权衡。

# 算法 1 作业调度策略
1: J; # 作业的配置和元数据
2: V; # 集群的物理视图
3: Q; # 作业队列，按优先级和到达时间排序
4: O; # 延迟列表
5: function SCHEDULER(J, V, Q)
6:     while True do
7:         O ← ∅
8:         while Q is not empty do
9:             J ← Q.pop()
10:            if preemptable(J) then
11:                sched(J, V)
12:            else if J.request < V.free_resource then
13:                sched(J, V)
14:            else if pred_JCT(J) < arrival_time then
15:                sched(J, V)
16:            else
17:                O.add(J)
18:            end if
19:        end while
20:        Q ← O
21:        sleep(interval)
22:    end while
23: end function

JCT预测器。JCT预测使得可以机会性地将短时作业调度到预留资源上，只要它们能在LPJ到达之前完成。这有助于提高资源利用率并减少排队延迟。预测基于与作业相关的元数据，如请求的CPU和GPU数量、请求的驱动器数量、任务所有者的部门等。虽然估计确切的JCT本质上是困难的，但我们采用了一种粗粒度的预测策略，将JCT分为不同的时间间隔。为了训练JCT预测器，我们从数据库中检索历史轨迹数据。然后，我们对数据进行预处理，例如移除被用户提前杀死的作业，并将JCT划分为10分钟的间隔。接着，我们训练模型，根据与作业相关的元数据来预测新来的作业可能落入的间隔。我们尝试了深度神经网络（DNN）和梯度提升预测器（GBM）【20】，发现GBM性能更高，这可能是因为它能更好地处理分类变量。


图 13: JCT预测。

为了展示有效性，我们提取了4个月的轨迹数据，并将其分为训练集（90%）和测试集（10%）。图13展示了在测试集上的一个预测示例。我们应用随机网格搜索来优化超参数，并使用bagging来确定不确定性估计。我们观察到测试集上的RMSE为1.61，最近的研究表明这种预测可以帮助调度决策【43, 3】。

H 队列管理评估

队列管理策略评估。我们收集作业轨迹并在轨迹驱动的模拟器中重放。图14显示了自从一个LLM预训练作业被规划以来，分配率和保留率随时间的变化。分配率由已分配作业的节点数除以总节点数确定。保留率衡量为LLM作业规划的节点中有多少被其他作业占用，这在LLM作业到达时不可避免地需要手动抢占。在18:00，Arnold被告知LLM作业将于22:00到达，因此它触发代码路径来相应地规划和预留资源。我们使用默认的装箱算法来调度其他作业。


图 14: 自一个LLM预训练作业被规划以来，分配率和保留率随时间的变化。

初始分配率为0.9，然后由于需要为集群中的1200多个节点预留资源而逐渐下降到0.5以下。在临近周期的开始，保留率相对较高并与分配率匹配，因为之前分配的作业不知道即将到来的LLM预训练作业。然后，保留率比分配率下降得更快，因为节点已被预留，并在周期结束时接近0，显示了调度策略的有效性。对于“预留并打包”（reserving-and-packing）策略【43】，它不提供强有力的预留语义（即尽力而为）。因此，调度器在LLM作业到达时将无法生成可行的解决方案，更不用说最优放置了（橙色线）。JCT预测通过机会性地将短时作业调度到预留区域，在资源利用率和保障之间导航权衡空间。在没有JCT预测的情况下，排队延迟和资源空闲时间都会增加，如绿线所示。

I 分解分析


图 15: 消融实验的分解分析。

内核级性能分解。我们通过对每种类型的内核持续时间求和来比较聚合的内核级指标。图10说明了Arnold在全规模实验中实现的加速（绿色）和减速（红色）。我们只报告差异显著的内核。最显著的加速是广播内核（10%），这是我们通信库优化的P2P实现。然而，这种加速被一个reduce-scatter内核甚至一个计算内核的减速所部分抵消。这种减速与我们的预期相悖，因为Arnold的调度也减少了DP组的扩展。此外，我们没有改变其他配置，所以GEMM内核的减速是出乎意料的。

原因分析：GPU流干扰。经过深入调查，我们怀疑减速是由于GPU流之间的干扰。由于混合并行，GPU在训练期间维护多个并发发出操作的流。虽然计算与通信的重叠表明了良好的性能优化，但它也导致了资源争用和干扰。

网络拓扑对计算内核的影响。为了调查这一反直觉的结果，我们通过修改NCCL来隔离流的影响。例如，我们添加了额外的环境变量，如NCCL_DP_MIN_NCHANNELS，以对DP流进行细粒度控制。我们禁用了通道自动调优，并在设置和不设置NCCL变量的情况下重新运行作业。图15显示了分解分析。通过设置NCCL变量，通信内核获得了加速，而计算内核则出现了减速。由于唯一的改变是NCCL变量，这表明如果我们为通信分配更多的GPU SM，计算内核会因可用SM减少而遭受性能损失。在生产训练中，NCCL变量是动态自动调优的，因此，考虑到网络拓扑优化的调度会影响DP和PP组的通信，最终它会导致计算内核的性能变化。