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

Zheng Wang1,2, Anna Cai2, Xinfeng Xie2, Zaifeng Pan1, Yue Guan1, Weiwei Chu2, Jie Wang2, Shikai Li2, Jianyu Huang2, Chris Cai2, Yuchen Hao2, Yufei Ding1,2
University of California, San Diego1 Meta2

主要贡献

本文提出WLB-LLM，一种用于大型语言模型训练的工作负载平衡的4D并行策略。核心问题在于大型语言模型训练中存在的负载不均衡问题，导致GPU利用率低下。具体而言，在8K GPU的128K上下文窗口训练作业中，观察到计算延迟在GPU间显著差异，最慢GPU的计算延迟是其他GPU的1.44倍，这源于注意力计算的输入依赖性和输入文档长度的变异性。输入序列由多个文档组成，使用注意力掩码防止不同文档间的自注意力计算，导致长文档尾部标记需要更多注意力计算，从而引入每个标记的算术强度异质性。

研究目标是解决4D并行（数据并行、管道并行、上下文并行、张量并行）中的负载不均衡问题，提高训练效率。现有框架忽略了每个标记的算术强度异质性，导致负载不均衡，尤其在更大模型和更长上下文窗口下问题加剧。

创新点包括：首先，在管道并行级别引入工作负载感知的变长文档打包方法，以平衡微批次间的计算和通信负载；其次，在上下文并行级别引入细粒度逐文档分片策略，确保上下文并行组内每个工作者具有相同负载；此外，提出仅调整极长文档执行顺序以最小化对数据随机性的影响，并设计启发式算法自适应选择最有效分片策略。实验显示WLB-LLM在不同模型规模下平均加速1.23倍。

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

4D并行LLM训练概述。 训练具有数十亿甚至万亿参数的极大型LLM面临挑战，需要多级并行调优。先进LLM训练框架采用4维并行，包括数据并行（DP）【30，ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020, SC20, 无URL】、管道并行（PP）【21，PipeDream: Generalized Pipeline Parallelism for DNN Training, 2019, SOSP, 无URL】、上下文并行（CP）【25，Megatron Core: Context Parallelism, 2023, NVIDIA文档, https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/context_parallel.html】和张量并行（TP）【32，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019, arXiv:1909.08053】。图2展示了（TP=2, CP=2, PP=4, DP=4）的4D并行示例。

数据并行：在DP中，输入全局批次被分区并分布到DP工作者，每个工作者拥有全局批次的一部分。默认情况下，模型参数在DP工作者间复制。先进技术如ZeRO【30】和FSDP【27，Fully Sharded Data Parallel: Faster AI Training with Fewer GPUs, 2021, 无会议, 无URL】将模型参数、梯度和优化器状态分区以减少内存消耗。在训练中，每个DP工作者使用本地批次独立计算参数更新，然后通过AllReduce（或在使用FSDP时的ReduceScatter）与其他工作者同步梯度。

管道并行：在每个DP工作者内，设备进一步分区为多个PP工作者。在PP中，模型按层方式分割，每个PP工作者拥有几个层块。DP工作者的输入批次也被分为多个微批次。在训练中，一个微批次从第一个到最后一个PP工作者遍历，然后在反向传播中反向。点对点（P2P）通信用于发送激活和梯度。

上下文并行：CP旨在解决长上下文训练中激活的巨大内存消耗。如图2所示，CP复制模型参数，但沿序列长度维度分片输入和激活。CP以环状方式使用P2P通信【20，Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023, arXiv:2310.01889】，或在前向阶段使用AllGather收集所有标记的KV张量，并在反向中使用ReduceScatter处理KV张量的梯度【6，The Llama 3 Herd of Models, 2024, arXiv:2407.21783】。其他操作如Linear和LayerNorm在所有工作者上独立操作，类似于DP。

张量并行：TP在层内（例如注意力、FFN）分割模型参数并分布到多个TP工作者。TP通常与序列并行（SP）【14，Reducing Activation Recomputation in Large Transformer Models, 2023, MLSys, 无URL】结合，进一步分割输入张量和激活。本文中提及TP时默认包括TP和SP。由于每个GPU仅拥有部分输入和参数，导致训练中密集的AllGather和ReduceScatter通信。因此，TP通常应用于单个节点内，而其他并行级别应用于跨节点。

输入文档长度变异性。 在大规模LLM训练中，尤其在使用大上下文窗口时，输入文档长度可能显著变异。为说明此点，我们分析了128K上下文长度训练作业的训练数据特性，如图3所示。从每个文档视角，输入文档长度分布高度偏斜。如图3左侧所示，大多数输入文档相对较短，但存在一些极长文档，最长达到完整上下文窗口大小。极长文档出现在输入批次中容易导致微批次间显著负载不均衡。从每个标记视角，我们计算上下文窗口内每个位置对应的平均文档长度。如图3右侧所示，标记位置间的平均文档长度差异约为1.4倍。中间标记更可能属于较长文档。这是因为数据加载器打包输入文档时，超出上下文窗口大小的长文档在尾部被截断，以匹配总序列长度。该图清楚展示了长上下文LLM训练中每个标记计算强度的异质性。这些观察突显了对输入感知解决方案的需求，该解决方案能动态平衡负载，考虑文档长度和每个标记算术强度的变异。

不均衡分析。 在分析内部8K GPU 128K上下文窗口训练作业的性能跟踪后，我们识别出4D并行LLM训练中不同并行层次的负载不均衡的两个主要原因：（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工作者处理相同数量的标记，但每个标记计算强度的变异导致这些级别的负载不均衡。例如，包含单个长文档的微批次（红色突出）比由多个短文档组成的微批次（绿色突出）具有更大负载。

CP级别不均衡：为更好地说明CP级别的不均衡，我们放大特定PP工作者（也称为CP组）。如图4 (a)(2)所示，在CP工作者间观察到显著负载不均衡（圆圈表示），而每个CP工作者内的TP工作者显示相似的计算延迟（每个圆圈内的数据点）。CP工作者间的不均衡源于CP级别的序列分片。如图4 (b)(2)所示，输入序列被分为具有相同标记数量的块，然后分布到CP工作者。先前方法将输入序列分区为2×CP_size块。第i个CP工作者被分配第i个和(2×CP_size−1−i)个块，以实现更好负载平衡【6】。当序列仅有一个文档时，此分片策略效果良好。但如果序列打包多个文档，可能导致CP工作者间显著负载不均衡，如图所示。虽然序列块进一步分为TP工作者，但TP级别未观察到不均衡。这是因为在计算前，所有TP工作者执行AllGather操作收集整个序列块。因此，CP工作者内的每个TP工作者处理相同序列块，消除TP级别不均衡。

不均衡传播：在训练中，负载不均衡从内层并行传播到外层并行。不均衡被累积和放大，最终显著影响端到端训练延迟。在DP、CP和TP中，训练期间执行集体通信如AllReduce、AllGather和ReduceScatter【24，NVIDIA Collective Communication Library (NCCL), 无年份, developer.nvidia.com/nccl】。这些并行层次的所有工作者以同步方式工作。因此，给定DP、CP和TP组的训练延迟由组内最慢工作者决定。相反，在PP级别，不同PP工作者作为彼此的生产者和消费者。如图5所示，PP的关键路径是最大微批次遍历所有PP工作者的延迟，加上剩余微批次在第一个PP工作者上的前向和反向传播。PP工作者间独特的数据依赖关系放大不均衡，导致对总延迟更大影响。由于不均衡传播，更高微批次训练延迟可能源于两个主要原因：（1）从内层层次传播的不均衡（例如CP分片不均衡）；（2）由于PP级别打包而固有更大负载的微批次。这突显了消除所有并行层次不均衡的重要性。

基线解决方案：固定长度打包。 解决负载不均衡问题的一个潜在解决方案是优化输入文档打包。当前4D并行框架要求所有微批次具有统一序列长度【6】，等于上下文窗口大小，以启用输入序列的有效批处理。在此固定长度要求基础上，我们实现基线洗牌和打包优化，将输入文档打包成等大小的微批次。形式上，给定来自一个或多个全局批次的输入文档集，目标是将这些文档打包成几个具有固定总长度的微批次，并平衡所有微批次间的注意力计算负载。不失一般性，我们使用因果掩码作为示例计算注意力负载。因果掩码下，微批次的注意力计算负载正比于$\sum_{i=1}^N d_i^2$，其中$d_i$是微批次内每个文档的长度。搜索最优文档打包的问题是NP-hard，因为它可从经典数字分区问题【9，Computers and Intractability: A Guide to the Theory of NP-Completeness, 1982, SIAM Review, 无URL】扩展，通过添加数字和约束以及平方和目标。为搜索最优打包，我们将任务制定为整数线性规划（ILP）问题。假设有N个文档，每个长度$d_i$，我们希望将这些文档打包成M个微批次，每个总长度L。目标是最小化微批次间的最大负载：

其中$x_{ij}$是二元变量，表示打包计划。具体而言，$x_{ij}=1$意味着文档i打包到微批次j。在此ILP制定下，我们使用商业求解器【7，Gurobi Optimizer Reference Manual, 2024, 无会议, 无URL】获得最优打包计划。

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

打包平衡与模型质量间的权衡激励我们打破微批次的固定长度约束，并设计更灵活的打包策略。在接下来的两个部分，我们将呈现WLB-LLM的细节，包括PP级别的变长打包和启发式异常文档延迟优化（§4）以及CP级别的细粒度和自适应分片优化（§5）。

方法细节

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

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

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

其中$W_a(\cdot)$和$W_l(\cdot)$是投影函数，根据文档长度计算注意力计算延迟和其他所有操作的延迟。二者均可从离线剖析中派生。$L_{max}$表示GPU内存约束允许的最大序列长度。

异常文档延迟。 我们的变长打包策略通过允许短文档打包成超过固定上下文窗口大小的更长序列，实现更高平衡水平。然而，在单个全局批次内，短序列数量可能不足以均衡所有微批次的计算负载。为解决此挑战，我们提出自适应延迟异常文档，即极长文档。此方法基于我们的观察，虽然这些异常文档对负载不均衡有显著影响，但它们仅贡献训练标记的小比例。通过选择性延迟少数极长文档的标记训练，我们可实现微批次间更平衡负载，同时最小化对数据随机性的影响。

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

启发式打包算法。 虽然ILP求解器可为给定输入文档集派生最优打包，但求解时间不切实际高。为解决此点，我们设计启发式算法，将变长打包与异常文档延迟优化结合，以高效产生具有平衡负载的打包微批次。如算法1所示，算法输入数据加载器D、异常文档的等待队列集Q（以FIFO方式工作）、每个迭代的微批次数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）。我们的启发式方法结合异常文档延迟与变长文档打包，以高效平衡所有微批次的计算负载。

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

逐文档分片设计。 在CP级别，微批次的序列被分片到CP工作者。每个CP工作者处理独占序列分片。现有CP实现采用逐序列分片策略，将整个输入序列等分为2×CP_size块。此方法当输入序列打包多个文档时容易导致显著注意力计算负载不均衡。为消除CP级别的负载不均衡问题，我们提出以细粒度方式分片序列。具体而言，我们进行逐文档分片，将每个文档分为2×CP_size文档块。如图9所示，每个CP工作者为每个输入文档取一对对称文档块。通过我们的细粒度逐文档分片策略，每个CP工作者不仅接收相同数量的标记（确保GEMM计算和集体通信的负载平衡），还获得相同的注意力计算负载。

避免填充：我们的细粒度逐文档分片策略将每个输入文档分为2×CP_size文档块。然而，文档长度不总是可被2×CP_size整除，需要通过添加填充标记扩展文档长度。为避免文档填充引入的冗余计算，我们设计无填充逐文档分片方法。具体而言，我们将每个文档分割为两部分：可被2×CP_size整除的部分和剩余标记。假设第i文档的长度为$L_i$，其中$L_i = D_i + R_i$，$D_i = \lfloor L_i / (2 \times CP_{size}) \rfloor$。我们对$D_i$部分应用标准逐文档分片，而$R_i$中的标记以轮询方式分布到CP工作者。由于$\sum_{i=1}^n L_i$和$\sum_{i=1}^n D_i$均可被2×CP_size整除，因此$\sum_{i=1}^n R_i = \sum_{i=1}^n (L_i - D_i)$也可被2×CP_size整除。这确保每个CP工作者接收相等数量的标记，从而消除填充需求。

内核效率 vs. 分片平衡。 我们的逐文档分片策略完全消除CP级别的负载不均衡。然而，将每个文档分割为多个更短块可能损害内核效率。细粒度分片下内核效率可能下降有两个主要原因：（1）瓦片级计算浪费：注意力计算被分割为更小瓦片并分布到GPU上的不同线程块执行。例如，在FlashAttention【5，FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022, NeurIPS, 无URL】的注意力前向内核中，瓦片大小设置为128。如果标记数量少于瓦片大小，线程块仍对128个标记执行完整计算，这浪费大量计算。为说明此点，我们剖析查询标记长度从16到256的注意力前向延迟。如图10左侧所示，当查询标记数量（Q_len）从16增加到128时，内核延迟保持恒定。这是因为所有少于128标记的短文档在内核级别被填充到128标记。相反，当Q_len从128增加到256时，内核延迟显著上升。（2）低效张量内存加速器（TMA）使用：TMA是NVIDIA Hopper架构引入的特性，支持全局内存与共享内存间的异步内存复制【1，H100 GPU, 无年份, https://www.nvidia.com/en-us/data-center/h100/】。对于大文档长度（例如Q_len ≥ 256），多个线程块处理不同Q标记，同时共享文档块的相同KV标记。这允许KV张量加载通过TMA加载多播共享L2缓存，显著减少从全局内存到共享内存传输KV张量的延迟。如图10右侧所示，当Query张量长度从128增加到256时，实现TFLOPs显著增加，展示了利用TMA加载多播的影响。

这些剖析结果清楚展示了注意力内核效率与CP分片平衡间的权衡。如果对由短文档组成的输入序列应用逐文档分片，它可在内核级别引入冗余计算并减少实现TFLOPs，最终导致更长的注意力计算延迟。

自适应分片选择。 根据§5.2的分析，虽然细粒度逐文档分片实现CP级别的最优负载平衡，但它不一定保证更好性能，因为注意力内核可能随着逐文档分片生成的更细粒度文档块而变得低效。为解决此问题，我们提出在运行时为每个微批次自适应选择最优分片策略。如图11所示，在CP级别的前向传播期间，跨CP工作者执行AllGather通信以收集完整KV张量。然后，我们计算逐序列和逐文档分片情况下的注意力内核输入张量形状（查询标记数量和键/值标记数量）。最后，我们预测注意力内核延迟，并选择产生更低注意力计算延迟的CP分片策略。

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

实验环境

数据集：实验使用内部LLaMA-like模型的训练数据，上下文窗口大小为64K和128K。数据包括高度偏斜的文档长度分布，用于长上下文预训练。数据集规模未明确指定，但用于从550M到70B参数模型的预训练，全球批次大小设置为PP_size × DP_size。

模型架构关键参数：评估了550M、7B、30B和70B参数的LLaMA-like模型。7B模型与LLaMA2-7B架构相同【36，Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023, arXiv:2307.09288】。其他模型保持相同架构，但按比例调整层数和模型维度大小。使用bfloat16精度。上下文窗口大小为64K和128K。

硬件配置：部署在32节点集群上，每个节点配备8× NVIDIA H100 SXM 80GB GPU，通过高带宽NVLink互连，跨节点通信通过RDMA over Converged Ethernet (RoCE)实现。4D并行映射优先将内层并行（如TP或CP）映射到节点内GPU，利用NVLink高效通信；外层并行如DP映射到多个节点。

软件配置：基于内部4D并行LLM训练代码库，支持扩展到100K GPU。使用Python语言，实现依赖DeepSpeed、Megatron等库（隐含于引用中），操作系统未指定但兼容NVIDIA GPU环境。基线包括Plain-4D（无优化打包，使用逐序列分片）和Fixed-4D（固定长度打包，使用贪婪算法）。

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

实验结果

端到端训练延迟改进。 在不同模型大小（550M到70B）和上下文窗口大小（64K和128K）下运行WLB-LLM和基线。Fixed-4D相对于Plain-4D平均加速1.03倍，受限于单个全局批次打包和固定上下文窗口，无法处理极长文档。WLB-LLM相对于Plain-4D和Fixed-4D分别加速1.23倍和1.19倍，通过变长打包、异常延迟和自适应分片实现更高平衡。更大模型加速略低，因为通信比例增加；更大上下文窗口从1.15倍提高到1.30倍（图12）。

性能分解（7B-128K）。 单独应用优化：逐文档分片加速1.02倍；自适应分片加速1.05倍；变长打包结合异常延迟加速1.28倍；组合CP和PP优化加速1.33倍（图13）。

上下文窗口大小敏感性（7B模型）。 上下文窗口从32K到160K，加速从约1.1倍增加到1.40倍，更大窗口加剧不均衡（图14）。

打包平衡和开销分析（7B-128K）。 使用不均衡度指标$ \frac{Max_{Latency} \times PP_{size}}{Total_{Latency}} $。原始打包不均衡度1.39；固定长度贪婪（1全局批次）1.35；固定长度求解器（1全局批次）1.23但开销高（>25s/批次）；WLB-LLM（2队列）1.05，开销20ms/批次（<0.65%步长延迟）（表2）。

表2：打包不均衡度和开销分析。

CP分片性能分析（7B单层，CP=4）。 逐文档分片相对于逐序列加速1.01倍（64K）和1.07倍（128K）；WLB-LLM自适应选择相对于静态分片改善7.5%（64K）和3.4%（128K），接近最优（图15）。

模型收敛分析（550M模型）。 WLB-LLM损失曲线与单个全局批次固定长度打包几乎相同，平均延迟0.5迭代/标记，不影响收敛；跨8全局批次打包导致损失增加1.6%（图16）。

结论

WLB-LLM是一种用于LLM训练的工作负载平衡的4D并行框架，识别并解决不同并行层次的负载不均衡。在管道并行级别，引入启发式变长文档打包算法，有效缓解微批次间负载不均衡。在上下文并行级别，提出细粒度和自适应逐文档分片策略，实现最优训练性能。实验显示WLB-LLM在各种模型大小和并行配置下平均加速1.23倍。未来工作可扩展到更大规模训练和更多并行维度。

附录

论文无明确附录内容。