WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training

发表时间: 2025-03 · arXiv:2503.17924 · OSDI 2025 (Meta)

文章标题：WLB-LLM: 面向大语言模型训练的工作负载均衡的4D并行
作者/机构：Zheng Wang (加州大学圣迭戈分校, Meta), Anna Cai (Meta), Xinfeng Xie (Meta), Zaifeng Pan (加州大学圣迭戈分校), Yue Guan (加州大学圣迭戈分校), Weiwei Chu (Meta), Jie Wang (Meta), Shikai Li (Meta), Jianyu Huang (Meta), Chris Cai (Meta), Yuchen Hao (Meta), Yufei Ding (加州大学圣迭戈分校, Meta)

A1 主要贡献

本文提出了 WLB-LLM，一种面向大语言模型（LLM）训练的工作负载均衡的4D并行方法。研究的核心问题是解决大规模LLM训练中存在的严重工作负载不平衡问题，该问题导致大量GPU资源未被充分利用。例如，在一个内部的8K-GPU训练任务中，最慢的GPU计算延迟比其他GPU长1.44倍。这种不平衡的根源在于注意力计算的输入依赖性以及训练数据中输入文档长度的可变性。具体来说，长文档尾部的token需要进行更多的注意力计算，导致包含长文档的输入序列即使总序列长度相同，其计算工作负载也显著更高。现有的4D并行（数据并行DP、流水线并行PP、上下文并行CP、张量并行TP）框架未能识别这种每个token计算强度的异质性，通过将所有token同质化处理并平均分配，不可避免地导致了工作负载不平衡。

为了解决上述问题，本文的研究目标是设计一种灵活且输入感知的文档打包和分片方法，使得每个GPU获得总计算和通信工作负载相等的输入token，同时最小化对数据随机性的影响。

本文的主要贡献如下：

  • 首次识别、分析并解决大规模LLM 4D并行训练中的工作负载不平衡问题。 本文深入分析了不平衡的两个主要来源：流水线并行（PP）层级的微批次（micro-batches）间不平衡，以及上下文并行（CP）层级的文档分片（document shards）间不平衡。
  • 在PP层级，设计了可变长度输入打包和自适应异常值延迟策略。 该方法允许将较短的文档组合成更长的序列，以匹配单个长文档的总计算工作负载。同时，通过自适应地延迟极端长文档的执行，在几乎不影响数据加载随机性和模型收敛的情况下，实现了接近最优的微批次间工作负载平衡。
  • 在CP层级，实现了一种细粒度的按文档分片方法和自适应分片选择机制。 该方法确保CP组内的每个工作单元（worker）负载相等。此外，通过启发式算法在运行时根据输入序列自适应选择最高效的分片策略，解决了内核效率和分片平衡之间的权衡。
  • 全面的评估表明，WLB-LLM在各种模型规模和上下文窗口大小下，平均实现了1.23倍的加速。


图1：在大规模LLM训练任务中观察到的工作负载不平衡问题及其原因。(a) 一个8K-GPU LLM训练任务中的归一化计算延迟。(b) 不平衡的原因：注意力计算的输入依赖性以及变化的输入文档长度。

A3 背景知识与关键观察

2. 背景知识

本节介绍了4D并行LLM训练的背景以及不同并行层级上工作负载不平衡的原因。

2.1 4D并行LLM训练

4D并行概览。训练拥有数十亿甚至数万亿参数的超大型LLM极具挑战性，需要大量的工程努力来调整多级并行【44，Lianmin Zheng, et al. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. OSDI 22, 2022】。最先进的LLM训练框架采用4维并行【8，Eric Harper, et al. NeMo: A Toolkit for Conversational AI and Large Language Models. 2019】【11，Sam Ade Jacobs, et al. DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. arXiv preprint arXiv:2309.14509, 2023】，包括数据并行（DP）【30，Samyam Rajbhandari, et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20, 2020】，流水线并行（PP）【21，Deepak Narayanan, et al. PipeDream: Generalized Pipeline Parallelism for DNN Training. SOSP, 2019】，上下文并行（CP）【25，NVIDIA. Megatron Core: Context Parallelism. 2023】，和张量并行（TP）【32，Mohammad Shoeybi, et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv:1909.08053, 2019】。图2展示了一个（TP=2, CP=2, PP=4, DP=4）的4D并行示例。


图2：LLM训练的4D并行概览。

  • 数据并行（Data Parallelism, DP）：在DP中，输入的全局批次被分区并分发到DP工作单元上，每个工作单元拥有一部分全局批次。默认情况下，模型参数在DP工作单元间是复制的。一些先进技术如ZeRO【30，Samyam Rajbhandari, et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20, 2020】和FSDP【27，Myle Ott, et al. Fully Sharded Data Parallel: Faster AI Training with Fewer GPUs, 2021】会将模型参数、梯度和优化器状态在DP工作单元间进行分区，以减少内存消耗。训练期间，每个DP工作单元独立计算其本地批次的参数更新，然后通过AllReduce（或使用FSDP时的ReduceScatter）与其他工作单元同步梯度。
  • 流水线并行（Pipeline Parallelism, PP）：在每个DP工作单元内部，设备通过流水线并行进一步划分为多个PP工作单元。在PP中，模型按层的方式被分割，每个PP工作单元拥有几个层的块。DP工作单元的输入批次也被划分为多个微批次。训练期间，一个微批次从第一个PP工作单元到最后一个，顺序遍历所有PP工作单元，然后在反向传播中反向进行。在前向和反向传播中，分别需要点对点（P2P）通信来发送激活值和梯度。
  • 上下文并行（Context Parallelism, CP）：CP旨在解决长上下文训练中激活值的大量内存消耗问题。如图2所示，CP复制模型参数，但将输入和激活值沿序列长度维度在CP工作单元间进行分片。CP可以以基于环的方式使用P2P通信【20，Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv preprint arXiv:2310.01889, 2023】运行，或者在前向阶段涉及AllGather通信来收集所有token的KV（键和值）张量，在反向阶段使用ReduceScatter通信处理KV张量的梯度【6，Abhimanyu Dubey, et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024】。所有其他操作符，如线性层和层归一化，在所有工作单元上独立运行，类似于DP。
  • 张量并行（Tensor Parallelism, TP）：TP将一个层内的模型参数（例如，注意力、FFN）分割并分发到多个TP工作单元。TP通常与序列并行（SP）【14，Vijay Anand Korthikanti, et al. Reducing Activation Recomputation in Large Transformer Models. MLSys, 2023】结合使用，以进一步分割输入张量和激活值。在本文中，当我们提到TP时，默认指的是TP和SP。使用TP时，每个GPU只有部分输入和参数，导致训练期间需要密集的AllGather和ReduceScatter通信。因此，TP通常应用于单个节点内，而其他级别的并行则应用于跨节点。

2.2 变化的输入文档长度

输入文档长度的统计特性。在大规模LLM训练中，输入文档的长度可能差异巨大，特别是在使用大上下文窗口时。为了说明这一点，我们分析了128K上下文长度训练任务中的训练数据特征，如图3所示。从每个文档的角度看，输入文档长度的分布高度倾斜。如图3左侧所示，大多数输入文档相对较短，但存在一些极长的文档，最长的达到完整的上下文窗口大小。输入批次中出现极长文档很容易导致微批次间的显著工作负载不平衡。从每个token的角度看，我们计算了上下文窗口内每个位置对应的平均文档长度。如图3右侧所示，不同token位置的平均文档长度差异约为1.4倍。中间位置的token更可能属于较长的文档。这是因为当数据加载器将输入文档打包到批次中时，偶然超出上下文窗口大小的长文档会在尾部被截断，以确保总序列长度与上下文窗口大小匹配。该图清楚地说明了长上下文LLM训练期间每个token计算强度的异质性。这些观察结果强调了需要一种输入感知的解决方案，该方案能够通过考虑文档长度和每个token计算强度的变化来动态平衡工作负载。


图3：输入数据统计：输入文档长度的分布（左）和每个token位置的平均文档长度（右）。

3. 设计动机

为了阐述WLB-LLM的设计动机，我们首先对4D并行LLM训练中不同并行层级上发生的工作负载不平衡进行了深入分析。然后，我们介绍了一种基线解决方案，该方案在批次间对输入文档进行洗牌和重打包。最后，我们讨论了为什么这种基线解决方案不能完全消除工作负载不平衡，并研究了输入打包平衡与模型收敛之间的权衡。

3.1 不平衡分析

不平衡的来源识别。在分析了我们内部8K-GPU、128K上下文窗口大小训练任务的性能追踪数据后，我们确定了两个主要的工作负载不平衡原因：（1）PP层级的微批次间工作负载不平衡，和（2）CP层级的序列分片间不平衡。为了证明这一点，我们累加了每个GPU上的注意力计算延迟。结果如图4所示。

  • PP层级不平衡：如图4(a)(1)所示，不同DP工作单元间的注意力计算延迟差异显著。在每个DP工作单元内部，我们观察到数据点形成的“垂直线”，每条线代表一个DP工作单元内的PP工作单元。每个PP工作单元显示出几乎相同的工作负载，因为一个DP工作单元内的所有PP工作单元处理相同的微批次集合。根据图4(a)(1)的结果，我们得出结论，DP和PP层级的工作负载不平衡源于微批次间的不平衡，这是由输入打包过程引起的。如图4(b)左侧所示，输入文档被打包成等长的序列（微批次）。虽然这种固定长度的打包策略确保了每个PP和DP工作单元处理相同数量的token，但每个token计算强度的变化导致了这些层级的工作负载不平衡。例如，一个只包含单个长文档的微批次（红色高亮）比由多个短文档组成的微批次（绿色高亮）具有更大的工作负载。


(a) 不平衡分析 (TP=8, CP=16, PP=16, DP=4): (1) 归一化计算延迟 (按DP和PP分组); (2) 一个CP组内的归一化计算延迟。


(b) PP层级的文档打包和CP层级的序列分片。
图4：工作负载不平衡来自PP层级的文档打包和CP层级的序列分片。

  • CP层级不平衡：为了更好地说明CP层级的不平衡，我们放大了一个特定的PP工作单元（也称为CP组）。如图4(a)(2)所示，CP工作单元之间观察到显著的工作负载不平衡（由圆圈表示），而每个CP工作单元内的TP工作单元表现出相似的计算延迟（每个圆圈内的数据点）。CP工作单元间的不平衡来自CP层级的序列分片。如图4(b)(2)所示，一个输入序列被划分为具有相等token数量的块，然后分发给CP工作单元。最先进的方法将输入序列划分为$2 \times CP\_size$个块。第$i$个CP工作单元被分配第$i$个和第$(2 \times CP\_size - 1 - i)$个块，以实现更好的负载均衡【6，Abhimanyu Dubey, et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024】。当序列只有一个文档时，这种分片策略效果很好。然而，如果序列由多个文档打包而成，它可能会导致CP工作单元之间显著的工作负载不平衡，如图所示。尽管序列块会进一步在TP工作单元间划分和分发，但在TP层级没有观察到不平衡。这是因为在计算之前，所有TP工作单元执行AllGather操作来收集整个序列块。因此，一个CP工作单元内的每个TP工作单元处理相同的序列块，消除了TP层级的不平衡。

不平衡的传播机制。在训练过程中，工作负载不平衡会从内层并行传播到外层并行。不平衡会累积和放大，最终对端到端训练延迟产生重大影响。在DP、CP和TP中，训练期间会执行诸如AllReduce、AllGather和ReduceScatter之类的集体通信【24，Nvidia. NVIDIA Collective Communication Library (NCCL)】。这些并行层级中的所有工作单元都以同步方式工作。因此，一个给定的DP、CP和TP组的训练延迟由该组内最慢的工作单元决定。相比之下，在PP层级，不同的PP工作单元互为生产者和消费者。如图5所示，PP的关键路径是最大微批次遍历所有PP工作单元的延迟，加上剩余微批次在第一个PP工作单元上的前向和反向传播延迟。PP工作单元之间独特的数据依赖关系放大了不平衡，导致对总延迟的影响更大。由于不平衡的传播，较高的微批次训练延迟可能由两个主要原因引起：（1）从内层并行层级（例如，CP分片不平衡）传播来的不平衡；（2）由于PP层级打包导致微批次本身具有更大的工作负载。这凸显了在所有并行层级消除不平衡的重要性。


图5：4D并行LLM训练中延迟在不同并行层级间的传播过程。工作负载不平衡的影响在PP层级被放大。

3.2 基线解决方案：固定长度打包

固定长度打包优化。解决工作负载不平衡问题的一个潜在方案是优化输入文档的打包。当前的4D并行框架要求所有微批次具有统一的序列长度【6，Abhimanyu Dubey, et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024】，等于上下文窗口大小，以便高效地批处理输入序列。基于这个固定长度的要求，我们实现了一个基线洗牌和打包优化，将输入文档打包成大小相等的微批次。形式上，给定来自一个或多个全局批次的一组输入文档，目标是将这些文档打包成几个总长度固定且注意力计算工作负载均衡的微批次。不失一般性，我们使用因果掩码作为计算注意力工作负载的例子。使用因果掩码时，一个微批次的注意力计算工作负载与$\sum_{i=1}^{N} d_i^2$成正比，其中$d_i$是微批次内每个文档的长度。寻找最优文档打包的问题是NP难的，因为它可以从经典的数划分问题【9，Juris Hartmanis. Computers and Intractability: A Guide to the Theory of NP-Completeness (Michael R. Garey and David S. Johnson). Siam Review, 1982】扩展而来，增加了数字总和约束和平方和目标。为了寻找最优打包，我们将其表述为一个整数线性规划（ILP）问题。假设我们有N个文档，每个文档长度为$d_i$，我们希望将这些文档打包成M个微批次，每个微批次的总长度为L。目标是最小化微批次间的最大工作负载：

其中，$x_{ij}$是一个二元变量，表示打包计划。具体来说，$x_{ij} = 1$表示文档$i$被打包到微批次$j$中。通过这种ILP表述，我们使用商业求解器【7，Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2024】来获得最优的打包计划。

3.3 权衡分析

打包平衡与模型质量的权衡。在更多的全局批次中优化输入文档打包可以帮助实现更高程度的工作负载平衡。然而，这也操纵了更多输入文档的执行顺序，影响了数据采样和加载的随机性。这可能对模型质量产生负面影响并影响模型收敛。为了评估打包平衡和模型质量之间的权衡，我们使用各种打包窗口大小对一个550M参数的模型进行了52K步的预训练。然后，我们评估了不同设置下打包后输入批次的工作负载不平衡程度。不平衡度计算为 $\frac{Max\_Attn}{Avg\_Attn}$，其中$Max\_Attn$表示全局批次中的最大注意力计算工作负载，而$Avg\_Attn$表示全局批次中所有微批次的平均注意力计算工作负载。如图6所示，当仅在一个全局批次内优化打包时，微批次间的工作负载不平衡仍然很高。如果全局批次的数量增加，固定长度打包优化可以实现更好的工作负载平衡。然而，由于重打包更多输入文档导致数据加载随机性降低，最终的训练损失会随着更多全局批次参与打包优化而增加。这些结果表明，朴素的固定长度打包优化无法在不损害模型质量的情况下实现良好的工作负载平衡，凸显了对更复杂解决方案的需求。

引出新方法。打包平衡和模型质量之间的权衡促使我们打破微批次的固定长度约束，设计一种更灵活的打包策略。在接下来的两节中，我们将介绍WLB-LLM的细节，包括PP层级的可变长度打包和启发式异常文档延迟优化（§4），以及CP层级的细粒度和自适应分片优化（§5）。


图6：更大的打包窗口改善了工作负载平衡，但导致训练损失增加。

A2 方法细节

4. 针对PP的可变长度打包和异常值延迟

在PP层级，我们专注于通过以工作负载感知的方式重打包输入文档来平衡微批次间的工作负载。首先，我们设计了一种可变长度打包策略，以在相同的打包窗口大小内实现比固定长度打包更高程度的工作负载平衡（§4.1）。其次，我们提出了一种异常文档延迟方法，以自适应地延迟极端长文档的训练。这种方法有助于最小化对数据随机性的影响，同时实现接近最优的微批次间工作负载平衡（§4.2）。最后，我们设计并实现了一种高效的启发式算法，以可忽略的开销在运行时优化打包（§4.3）。

4.1 工作负载感知的可变长度打包

固定长度打包的局限性与可变长度打包的提出。基线固定长度打包的主要局限性在于，如果一个全局批次中存在一个极端长的文档，它就无法实现平衡。例如，如果一个文档的长度等于上下文窗口大小，由于注意力计算的二次复杂度，就不可能创建另一个由较短文档组成且计算工作负载相等的微批次。为了解决这个局限，我们提出了一种可变长度打包策略，允许每个微批次具有不同的序列长度。我们设计的关键洞察是，一个微批次的工作负载不仅仅由注意力计算决定。其他操作，如GEMM计算、逐元素操作和集体通信（例如，AllGather和ReduceScatter），也对训练延迟有显著贡献，并且受到每个微批次内文档的影响。为了证明这一点，我们在图7中展示了操作延迟与文档长度之间的关系。可以观察到，注意力计算延迟随输入文档长度呈二次方增长，而其他操作，如GEMM、集体通信和逐元素操作，其操作延迟与文档长度呈线性关系。


图7：操作延迟与输入文档长度之间的关系。（Total Linear是GEMM、集体通信和逐元素操作的总和）

可变长度打包的优化目标。这种关系为超越固定长度打包、进一步改善工作负载平衡提供了机会。如果一个长文档的注意力计算延迟显著高于其他操作，我们可以将多个较短的文档打包在一起，以延长其他操作的延迟，从而匹配长文档的总延迟。具体来说，我们将固定长度打包扩展为可变长度方法。优化目标从仅平衡注意力计算工作负载转变为平衡包括所有操作在内的总工作负载：

其中，$W_a(\cdot)$和$W_l(\cdot)$是投影函数，分别根据文档长度计算注意力计算延迟和所有其他操作的延迟。$W_a(\cdot)$和$W_l(\cdot)$都可以通过离线剖析得出。而$L_{max}$表示GPU内存约束允许的最大序列长度。

4.2 异常文档延迟

异常文档延迟的动机。我们的可变长度打包策略通过允许将短文档打包成超过固定上下文窗口大小的更长序列，从而实现了更高水平的平衡。然而，在单个全局批次内，短序列的数量可能不足以均衡所有微批次的计算工作负载。为了应对这一挑战，我们提出自适应地延迟异常文档（即极端长的文档）的执行。这种方法基于我们的观察：尽管这些异常文档对工作负载不平衡有显著影响，但它们仅贡献了训练token的一小部分。通过选择性地延迟少数极端长文档的token训练，我们可以在最小化对数据随机性影响的同时，实现更均衡的微批次间工作负载。

设计细节。如图8所示，我们为异常文档实现了一个多级等待队列。假设我们有$n$个异常队列，每个队列关联一个超参数$L_i$，该参数指定了第$i$个队列中文档的最小长度，其中$L_1 < L_2 < \dots < L_n$。当一个新的全局批次到达时，长度大于$L_1$的文档被视为异常值，并被添加到相应的等待队列$i$中，其中$L_i \le \text{Doc\_len} < L_{i+1}$。这里，$\text{Doc\_len}$表示文档的长度。异常队列中文档的执行被延迟，直到该队列积累了足够的异常文档。一旦队列大小达到微批次的数量，异常文档就会从队列中弹出，并添加到每个微批次中。这确保了微批次间的工作负载平衡，因为每个微批次都包含一个异常文档。


图8：异常文档延迟与可变长度打包相结合的过程。

4.3 启发式打包算法

启发式算法的设计。尽管ILP求解器可以为给定的输入文档集推导出最优打包方案，但求解时间过高，不切实际。为了解决这个问题，我们设计了一种启发式算法，该算法结合了可变长度打包和异常文档延迟优化，以高效地生成具有均衡工作负载的打包微批次。如算法1所示，该算法以数据加载器D、一组用于异常文档的FIFO等待队列Q、每次迭代的微批次数N以及序列长度上限$L_{max}$作为输入。它最终将为每个训练迭代输出一系列具有均衡工作负载的打包微批次。打包过程首先遍历数据加载器D中的输入批次，并将所有异常文档添加到Q中相应的等待队列以进行延迟处理（第3-10行）。如果Q中的任何队列q的大小达到N，q中的文档将被弹出并添加到当前批次的待处理文档集Doc_Set中（第11-15行）。打包以贪心方式进行：待处理文档按长度降序排序。算法首先尝试将每个文档打包到工作负载最小的微批次中（使用方程2中定义的$W_a(\cdot)$和$W_l(\cdot)$计算），前提是总长度保持在上限$L_{max}$内。如果这不可能，它会尝试长度最小的微批次（第16-25行）。无法在这些约束内容纳的文档将被保存到下一次迭代（第27行）。我们的启发式方法将异常文档延迟与可变长度文档打包相结合，以高效地平衡所有微批次的计算工作负载。

5. 针对CP的细粒度和自适应分片

在CP层级，我们的目标是通过实现一种细粒度的按文档分片策略来改善文档分片间的工作负载平衡，确保每个CP工作单元接收到相等的计算工作负载（§5.1）。此外，我们观察到注意力核效率和分片粒度之间存在权衡（§5.2）。为了最大化整体性能，我们进行了深入分析，并为给定的输入批次自适应地选择最优分片策略（§5.3）。

5.1 按文档分片设计

按文档分片的原理。在CP层级，微批次的序列在CP工作单元间进行分片。每个CP工作单元处理一个专属的序列分片。现有的CP实现采用按序列分片（Per-Sequence Sharding）策略，该策略将整个输入序列平均分片为$2 \times CP\_size$个块。当输入序列由多个文档打包而成时，这种方法很容易导致显著的注意力计算工作负载不平衡。为了消除CP层级的工作负载不平衡问题，我们提出以细粒度的方式对序列进行分片。具体来说，我们进行按文档分片（Per-Document Sharding），将每个文档划分为$2 \times CP\_size$个文档块。如图9所示，每个CP工作单元为每个输入文档获取一对对称的文档块。通过我们的细粒度按文档分片策略，每个CP工作单元不仅接收到相同数量的token（确保了GEMM计算和集体通信的工作负载平衡），而且还获得了相同的注意力计算工作负载。


图9：细粒度按文档分片概览。

避免填充。我们的细粒度按文档分片策略将每个输入文档划分为$2 \times CP\_size$个文档块。然而，文档长度并不总是能被$2 \times CP\_size$整除，这需要通过添加一些填充token来扩展文档长度。为了避免文档填充引入的冗余计算，我们设计了一种无填充的按文档分片方法。具体来说，我们将每个文档分成两部分：一部分可以被$2 \times CP\_size$整除，另一部分是剩余的token。假设第$i$个文档的长度为$L_i$，其中$L_i = D_i + R_i$，且$D_i = \lfloor \frac{L_i}{2 \times CP\_size} \rfloor \times (2 \times CP\_size)$。我们对$D_i$部分应用标准的按文档分片，而$R_i$中的token则以轮询方式分发给CP工作单元。由于$\sum_{i=1}^{n} L_i$和$\sum_{i=1}^{n} D_i$都可以被$2 \times CP\_size$整除，因此$\sum_{i=1}^{n} R_i = \sum_{i=1}^{n} (L_i - D_i)$也可以被$2 \times CP\_size$整除。这确保了每个CP工作单元接收到相等数量的token，从而消除了填充的需要。

5.2 内核效率与分片平衡

细粒度分片的效率问题。我们的按文档分片策略完全消除了CP层级的工作负载不平衡。然而，将每个文档分割成多个较短的块可能会损害内核效率。细粒度分片可能导致内核效率下降的原因主要有两个：（1）Tile级别的计算浪费：注意力的计算被分割成更小的tile，并分发给GPU上的不同线程块执行。例如，在FlashAttention【5，Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and MemoryEfficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems, 2022】的前向注意力核中，tile大小设置为128。如果token数量小于tile大小，线程块仍将对128个token执行完整计算，这将浪费大量计算。为了说明这一点，我们对查询token长度从16到256的前向注意力延迟进行了剖析。如图10（左）所示，当查询token数（Q_len）从16增加到128时，内核延迟保持不变。这是因为所有少于128个token的短文档在内核级别都被填充到128个token进行计算。相反，当Q_len从128增加到256时，内核延迟显著上升。（2）张量内存加速器（TMA）使用效率低下：TMA是NVIDIA Hopper架构中引入的一项功能，它支持GPU上全局内存和共享内存之间的异步内存复制【1，H100 GPU. https://www.nvidia.com/en-us/ data-center/h100/】。对于长文档长度（例如，Q_len ≥ 256），多个线程块处理不同的Q token，同时共享文档块的相同KV token。这使得KV张量加载可以通过使用TMA加载多播经由L2缓存共享，显著减少了将KV张量从全局内存传输到共享内存的延迟。如图10（右）所示，当Query张量的长度从128增加到256时，实现的TFLOPs显著增加，证明了利用TMA加载多播的影响。

权衡的证实。这些剖析结果清楚地表明了注意力内核效率和CP分片平衡之间的权衡。如果将按文档分片应用于由短文档组成的输入序列，可能会在内核级别引入冗余计算并降低实现的TFLOPs，最终导致更长的注意力计算延迟。


图10：注意力内核性能剖析：（左）注意力前向延迟；（右）注意力前向内核实现的TFLOPs。

5.3 自适应分片选择

自适应选择策略。根据我们在5.2节的分析，尽管细粒度的按文档分片在CP层级实现了最优的工作负载平衡，但这并不一定保证更好的性能，因为对于按文档分片生成的更细粒度的文档块，注意力内核可能会变得效率较低。为了解决这个问题，我们建议在运行时为每个微批次自适应地选择最优的分片策略。如图11所示，在CP层级的前向传播过程中，会在CP工作单元间执行AllGather通信以收集完整的KV张量。然后，我们计算按序列分片和按文档分片两种情况下注意力内核的输入张量形状（查询token和键/值token的数量）。最后，我们预测注意力内核的延迟，并选择能产生较低注意力计算延迟的CP分片策略。

延迟预测方法。为了准确估计注意力内核的延迟，我们利用了§5.2中提供的见解。首先，我们计算注意力计算所需的总浮点运算次数。通过将文档块填充到tiling大小的倍数，也考虑了内核级别的填充。接下来，我们使用从离线剖析中收集的数据来估计给定张量形状下实现的TFLOPs，这些数据包括了TMA使用的影响。最后，通过将浮点运算量除以实现的TFLOPs来估计注意力内核的延迟。通过自适应地选择CP分片，WLB-LLM最小化了CP层级的训练延迟。


图11：自适应CP分片选择的过程。

A4 实验

实验环境

  • 硬件配置：实验在一个包含32个节点的集群上进行。每个节点配备8块通过高带宽NVLink互连的NVIDIA H100 SXM 80GB GPU。跨节点通信通过RoCE（RDMA over Converged Ethernet）进行。
  • 模型与并行配置：实验使用了内部一系列类LLaMA模型，涵盖550M、7B、30B和70B四种不同规模。7B模型与LLaMA2-7B模型【36，Hugo Touvron, et al. Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv preprint arXiv:2307.09288, 2023】架构相同，其他模型保持相同架构但按比例调整层数和模型维度。每个模型规模都评估了64K和128K两种上下文窗口大小。全局批次大小设为 PP_size × DP_size，所有评估均使用bfloat16精度。具体的训练设置和并行配置见表1。在硬件映射时，内层并行维度（如TP或CP）优先映射到节点内GPU以利用NVLink，外层并行维度（如DP）则映射到跨节点。

  • 软件配置：

    • 基线：

      1. Plain-4D：内部大规模LLM训练代码库，支持4D并行，可扩展至10万个GPU。直接使用数据加载器提供的输入批次，CP分片采用按序列分片。
      2. Fixed-4D：应用了基线固定长度打包优化（§3.2），使用贪心算法在单个全局批次内打包。CP分片在整个训练过程中固定使用按序列或按文档分片中的较优者。

    • 本文方法：WLB-LLM。

表1：模型和4D并行配置。

实验结果

6.2 训练性能

实验内容：在多种模型规模和上下文窗口大小上运行WLB-LLM及所有基线，评估端到端训练性能。

实验结果与分析：

  • Fixed-4D vs. Plain-4D：如图12所示，Fixed-4D相比Plain-4D仅取得了微小的改进，平均加速比约为1.03倍。原因是Fixed-4D仅在单个全局批次内调整打包且受限于上下文窗口大小，无法处理极端长文档（如长度等于上下文窗口）带来的不平衡。此外，其固定的CP分片策略也无法适应不同输入批次的特性，导致性能次优。
  • WLB-LLM vs. Baselines：如图12所示，WLB-LLM在所有配置下均显著优于基线，相比Plain-4D和Fixed-4D分别取得了1.23倍和1.19倍的平均加速。这得益于WLB-LLM的两项关键优化：PP层级的启发式可变长度打包与异常值延迟，以及CP层级的自适应分片选择。
  • 跨模型规模和上下文窗口大小：如图12所示，WLB-LLM在较大模型上加速比较低，因为大模型训练使用更多GPU，通信延迟占比增加，使得注意力层工作负载不平衡的影响相对减弱。另一方面，将上下文窗口从64K增加到128K，平均加速比从1.15倍提升至1.30倍，因为更长的上下文加剧了工作负载不平衡问题。


图12：WLB-LLM和Fixed-4D相对于Plain-4D在各种配置下的训练性能加速比。

6.3 优化分析

实验内容：对WLB-LLM的各项优化技术进行分解分析和敏感性研究。

实验结果与分析：
* 加速比分解：在7B-128K配置下，单独应用WLB-LLM的各项优化技术（图13）。
* 仅应用CP层级的细粒度按文档分片，获得1.02倍加速，受限于内核效率与分片平衡的权衡。
* 加入自适应分片选择后，CP层级优化可将加速比提升至1.05倍。
* 仅应用PP层级的启发式可变长度打包与异常值延迟，获得了1.28倍的显著加速。
* 结合CP和PP层级的优化，最终实现了1.33倍的加速。

  • 上下文窗口大小敏感性分析：在7B模型上，评估不同上下文窗口大小（32K到160K）下WLB-LLM的加速比（图14）。随着上下文窗口增大，加速比也越显著，在160K上下文窗口下达到1.40倍。这是因为更大的上下文窗口增加了出现异常长文档的概率，并放大了注意力计算不平衡对训练延迟的影响。


图13：WLB-LLM在7B模型、128K上下文窗口下的性能分解。


图14：WLB-LLM在7B模型上跨不同上下文窗口大小的加速比。

6.4 消融研究

实验内容：分析WLB-LLM中打包和分片优化的有效性，并验证其系统优化不影响模型收敛。

实验结果与分析：
* 打包平衡与开销分析：在7B-128K训练任务中，使用指标 $\frac{\text{Max\_Latency} \times \text{PP\_size}}{\text{Total\_Latency}}$ 评估不同打包策略的不平衡度（越低越好），并记录打包开销（表2）。
* Fixed-Len Greedy（Fixed-4D使用）仅能轻微缓解不平衡。跨多全局批次打包可改善平衡度，但会增加训练损失（如图6）。
* Fixed-Len Solver能获得更好的平衡度，但打包开销巨大（例如，跨4个全局批次打包，每批次耗时超25秒）。
* WLB-LLM（使用2个异常队列时）实现了接近最优的1.05不平衡度，且每批次打包开销仅20毫秒，占单步训练时长的不到0.65%。

  • CP分片性能分析：在7B模型的单个Transformer层上（CP=4），比较不同分片策略的前向和反向延迟（图15）。

    • 按文档分片（Per-Doc）相比按序列分片（Per-Seq）在64K和128K上下文下分别取得1.01倍和1.07倍加速，证明了其在减少CP层级不平衡方面的有效性。
    • 然而，Per-Doc可能牺牲内核效率。WLB-LLM的自适应选择策略比静态的Per-Seq或Per-Doc策略分别提升了7.5%和3.4%的性能，且非常接近最优结果，证明了自适应选择的有效性。

  • 模型收敛分析：在550M模型上比较不同打包策略的训练损失曲线（图16）。

    • 跨8个全局批次打包导致训练损失明显增加（平均1.6%），因为它破坏了数据加载的随机性。
    • WLB-LLM的损失曲线与在单个全局批次上进行固定长度打包的曲线几乎重合。这是因为WLB-LLM只延迟了占少数token的异常文档，平均每个token仅被延迟0.5个迭代，最大程度地保留了原始数据分布，从而在不牺牲模型质量的情况下提升了训练效率。

表2：打包不平衡度与开销分析。


图15：CP分片性能比较。


图16：在550M模型上的训练损失比较。

A7 补充细节

7. 相关工作

分布式LLM训练框架。为了应对训练超大型LLM的挑战，先前的方法主要依赖于3D并行【14，Vijay Anand Korthikanti, et al. Reducing Activation Recomputation in Large Transformer Models. MLSys, 2023】【22，Deepak Narayanan, et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC, 2021】【32，Mohammad Shoeybi, et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv:1909.08053, 2019】【34，DeepSpeed Team and Rangan Majumder. DeepSpeed: Extreme-Scale Model Training for Everyone, 2020】【37，Colin Unger, et al. Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization. OSDI 22, 2022】【38，Guanhua Wang, et al. ZeRO++: Extremely Efficient Collective Communication for Giant Model Training. arXiv preprint arXiv:2306.10209, 2023】【44，Lianmin Zheng, et al. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. OSDI 22, 2022】，它集成了张量并行【14，Vijay Anand Korthikanti, et al. Reducing Activation Recomputation in Large Transformer Models. MLSys, 2023】【16，Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet Classification with Deep Convolutional Neural Networks. NeurIPS, 2012】【32，Mohammad Shoeybi, et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv:1909.08053, 2019】、流水线并行【10，Yanping Huang, et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS, 2019】【17，Joel Lamy-Poirier. Breadth-First Pipeline Parallelism. MLSys, 2023】【18，Zhuohan Li, et al. TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models. ICML, 2021】【22，Deepak Narayanan, et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC, 2021】和数据并行【30，Samyam Rajbhandari, et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20, 2020】【31，Jie Ren, et al. ZeRO-Offload: Democratizing Billion-Scale Model Training. USENIX ATC 21, 2021】【43，Yanli Zhao, et al. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. arXiv preprint arXiv:2304.11277, 2023】。

4D并行与上下文并行。尽管3D并行在扩展模型规模方面证明了其有效性，能够训练万亿参数模型【34，DeepSpeed Team and Rangan Majumder. DeepSpeed: Extreme-Scale Model Training for Everyone, 2020】，但它在有效扩展上下文窗口大小方面存在困难。为了解决这个限制，引入了一个新的并行维度——上下文并行【25，NVIDIA. Megatron Core: Context Parallelism. https://docs.nvidia.com/megatron-core/ developer-guide/latest/api-guide/context_ parallel.html, 2023】，形成了4D并行训练范式。上下文并行将输入上下文分割成块，有效减少了与极长序列输入相关的内存瓶颈。最初，上下文并行采用基于环的方法，通过P2P通信重叠通信和计算【20，Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv preprint arXiv:2310.01889, 2023】。更新的方法利用集体通信方法（如AllGather或AlltoAll）来聚合键和值张量，提供了更大的灵活性并更好地支持各种类型的注意力掩码【6，Abhimanyu Dubey, et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024】【11，Sam Ade Jacobs, et al. DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. arXiv preprint arXiv:2309.14509, 2023】【25，NVIDIA. Megatron Core: Context Parallelism. https://docs.nvidia.com/megatron-core/ developer-guide/latest/api-guide/context_ parallel.html, 2023】。上下文并行的引入使得具有长上下文窗口的LLM的高效训练成为可能。然而，现有的4D并行框架忽略了每个token计算强度的异质性，导致GPU之间存在显著的工作负载不平衡。

LLM训练的输入填充和打包。LLM的输入由不同长度的样本组成。为了将输入文档批处理在一起并优化LLM训练期间的GPU利用率，输入文档必须被组织成相同长度的张量。这可以通过两种主要方法实现：填充（Padding）【32，Mohammad Shoeybi, et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv:1909.08053, 2019】和打包（Packing）【15，Mario Michael Krell, et al. Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance. arXiv preprint arXiv:2107.02027, 2021】。填充涉及在微批次内向较短的文档添加零填充。然而，这种方法不可避免地引入了冗余的计算、通信和内存开销。为了解决这个问题，先前的工作集中于设计更高效的内核以减少冗余计算【40，Jinle Zeng, et al. Boosting Distributed Training Performance of the Unpadded BERT Model. arXiv preprint arXiv:2208.08124, 2022】【41，Yujia Zhai, et al. ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs. IPDPS, 2023】或优化批处理策略以最小化填充【12，Chenyu Jiang, et al. DynaPipe: Optimizing Multi-task Training through Dynamic Pipelines. EuroSys, 2024】。为了完全消除冗余计算，最近的工作提出将短文档打包在一起形成一个长的输入序列【15，Mario Michael Krell, et al. Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance. arXiv preprint arXiv:2107.02027, 2021】。打包后，必须应用额外的注意力掩码以确保token只关注同一文档内的其他token【6，Abhimanyu Dubey, et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024】【29，Pytorch. FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention. https://pytorch. org/blog/flexattention/, 2024】。由于其在避免冗余计算方面的效率，打包已成为LLM训练的主流选择。例如，LLaMA-3在其训练过程中采用了输入打包【6，Abhimanyu Dubey, et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024】，而最先进的高性能注意力实现（例如，FlashAttention【4，Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv preprint arXiv:2307.08691, 2023】【5，Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and MemoryEfficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems, 2022】）也支持带文档打包的高效注意力计算。WLB-LLM专注于输入打包，并解决了当输入文档被打包在一起时出现的工作负载不平衡问题。

A5 结论

本文介绍了WLB-LLM，一个用于LLM训练的工作负载均衡的4D并行框架。WLB-LLM识别并解决了不同并行层级的工作负载不平衡问题。在流水线并行层级，WLB-LLM引入了一种新颖的启发式可变长度文档打包算法，有效缓解了微批次间的工作负载不平衡。在上下文并行层级，WLB-LLM提出了一种细粒度且自适应的按文档分片策略，以实现序列分片的最佳训练性能。全面的实验表明，WLB-LLM在各种模型规模和并行配置下均优于现有的4D并行框架，平均实现了1.23倍的加速。