ByteScale: Efficient Scaling of LLM Training with a 2048K Context Length on More Than 12,000 GPUs

发表时间: 2025-02 · arXiv:2502.21231 (ByteDance)

文章标题：ByteScale: 在超过12,000个GPU上以2048K上下文长度高效扩展LLM训练
作者/机构：Hao Ge (Peking University), Junda Feng (ByteDance Seed), Qi Huang (ByteDance Seed), Fangcheng Fu (Peking University), Xiaonan Nie (ByteDance Seed), Lei Zuo (ByteDance Seed), Haibin Lin (ByteDance Seed), Bin Cui (Peking University), Xin Liu (ByteDance Seed)

A1 主要贡献

为了满足现代LLM应用（如文档摘要、视频理解、智能体交互和代码补全）对长程依赖理解的需求，扩展模型的上下文长度至关重要。然而，扩展至长上下文面临着自注意力机制带来的内存和计算复杂度呈二次方增长的根本挑战。现有方法如Flash Attention将内存复杂度从$O(N^2)$降至$O(N)$，但要进一步扩展，则需跨设备分区序列。现有框架通常采用数据并行（DP，分发不同序列）和上下文并行（CP，切分单个序列）两种正交技术，并将设备组织成静态的2D网格。这种静态设计依赖于一个假设：所有序列长度相同，以保证负载均衡。

然而，在真实世界的训练场景中，无论是文本还是多模态数据，序列长度通常是可变的且分布倾斜。此外，强化学习中思维链推理过程的长度增加也加剧了长度异质性。现有框架为了处理最长序列，必须配置足够大的CP组，导致所有短序列也必须在整个CP组上进行不必要的分区和通信。这种数据异质性与静态系统设计之间的不匹配导致了两个核心问题：
1. 冗余通信：短序列被迫参与为长序列设计的复杂通信过程，即使它们本可以在单个或少数设备上处理。此外，对于短序列，用$O(N^2)$的计算来掩盖$O(N)$的通信非常困难。
2. 计算不均衡：尽管通过CP可以均匀分配令牌以平衡内存，但每个令牌的计算复杂度与原始序列长度相关（$O(N^2)$），导致不同设备上的执行时间不同，从而产生同步等待的空闲时间。

为解决上述挑战，本文提出了ByteScale，一个为大规模长短序列混合训练设计的高效、灵活且可扩展的训练框架。其主要贡献如下：

• C1: 提出混合数据并行（Hybrid Data Parallelism, HDP）：本文提出了一种新的并行策略HDP，它统一了数据间分区（DP）和数据内分区（CP），旨在将令牌均匀地分布在设备上。HDP能够灵活地使用[1, DP×CP]范围内的任意数量设备来处理可变长度的序列。
• C2: 通信优化：为了消除短序列的冗余通信，HDP提供了数据感知分片（data-aware sharding）能力，自动构建动态通信组，使每个序列能够以最少的设备数量进行处理。此外，HDP还提供选择性卸载（selective offloading），进一步压缩长序列的通信成本。
• C3: 平衡策略：为了缓解计算不均衡问题，本文设计了一种启发式算法，该算法根据数据和流水线并行的特性重新组织数据分配。此外，对于执行时间较短的设备，系统会分配更多的微批次（micro-batches），而不是像静态系统那样分配相同数量。
• C4: 评估：本文在一个拥有超过12,000个GPU的生产集群上进行了实验，模型规模从7B到141B，上下文长度从256K扩展到2048K。实验结果表明，与现有的最先进训练系统相比，ByteScale最高可实现7.89倍的加速。

A3 背景知识与关键观察

2. 背景知识

2.1 Transformer与大语言模型

Transformer架构【40，Attention is All you Need，2017，NeurIPS 2017】已成为当今大语言模型（LLMs）【5, 14, 32, 39】最流行和广泛使用的基础架构。它通常由一系列Transformer层组成，每层包含一个注意力模块和一个前馈网络（FFN）模块。如图1所示，自注意力机制需要序列中的所有令牌参与计算以捕获整个文本的上下文信息。相比之下，其他操作如归一化、线性投影和激活函数则执行令牌级计算，允许每个令牌独立处理。

2.2 分布式LLM训练

随着模型大小和训练数据的持续扩展，分布式训练技术在LLM训练中不可或缺。
- 数据并行（Data Parallelism, DP）：DP【9, 24, 37】将训练数据均匀分布在各个设备上，每个设备持有一个模型副本。在每个训练步骤中，设备独立处理其本地数据，然后全局同步梯度以更新模型。ZeRO系列方法【35，ZeRO: memory optimizations toward training trillion parameter models，2020，SC 2020】进一步增强了DP的可扩展性。
- 模型并行（Model Parallelism）：模型并行将模型分布在设备上，包括张量并行（TP）【38，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，2019，CoRR abs/1909.08053】和流水线并行（PP）【16, 28, 29】。TP执行操作内分区，将层内的操作和参数划分到不同设备（例如Megatron-LM中的行并行和列并行线性层），这需要通信中间结果（激活值），通常在单个节点内使用。PP采用操作间分区，将模型层分段为不同阶段，仅需在连续阶段间通过点对点（P2P）通信交换激活值，从而实现跨多节点的模型分区。
- 混合并行（Hybrid Parallelism）：混合并行结合多种并行策略以提高训练效率。特别是，Megatron-LM通过集成DP、TP和PP，采用了3D并行策略【21, 30, 38】，使其成为当今大规模模型训练的主流方法。
- 梯度累积（Gradient Accumulation）：为提高效率和收敛性，LLM通常需要大批量大小【6, 15, 39】（例如，在拥有10K GPU的集群中，每个批次处理近30-80M令牌是常见做法）。受硬件内存限制，一次性处理整个大批量是不可行的。梯度累积将每个全局批次（即每个训练步骤中采样的数据）分成多个微批次。这些微批次的梯度被累积起来，等同于一次性处理整个全局批次所产生的梯度。

2.3 填充与打包

为了在当前静态并行策略中支持可变长度的序列，需要使用填充（padding）和打包（packing）等技术。如图2所示，填充将同一批次中的序列填充到相同长度，但这会导致计算浪费。打包【22，Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance，2021，CoRR abs/2107.02027】将多个序列连接成一个单一序列，不含填充令牌。它采用一种特殊的分段注意力掩码，以确保每个序列在自注意力中被独立处理。

2.4 长上下文训练

由于自注意力的时间和内存复杂度均为$O(N^2)$，当上下文长度扩展时，这种二次复杂度成为瓶颈。Flash Attention【7, 8】通过优化内存I/O和采用分块（tiling）技术，将内存复杂度从$O(N^2)$降低到$O(N)$，但时间复杂度仍为$O(N^2)$。上下文并行（CP）【4, 23, 25, 31】进一步将序列划分到$P$个设备上，将每个设备的内存从$O(N)$减少到$O(N/P)$。根据图1，CP沿序列维度对QKV进行分片，跨令牌操作需要使用环形点对点（ring-style P2P）通信在设备间交换KV切片，该通信与计算重叠。该技术也适用于打包序列，其实现细节将在第7节详述。值得注意的是，每个子序列也必须在所有CP ranks上进行分片，如图2(c)和图3(a)所示。

3. 观察与动机

3.1 数据异质性

LLM在序列数据上进行训练。如第1节所述，训练数据通常包含可变长度的序列。存在两个观察和一个重大挑战：
- 观察1：真实世界数据集中序列长度呈偏斜分布。 如图4所示，我们分析了两个用于长上下文训练的数据集：开源的GitHub数据集和生产环境的Byted数据集。我们观察到两者在序列长度上都呈现偏斜分布。例如，在Byted数据集中，如果我们随机采样一个全局批次，近80%的样本是4K令牌或更短，而只有0.05%的样本能达到2M令牌。然而，从令牌分布的角度看，这0.05%的样本（>=2M）贡献了全局批次中12.1%的令牌，而1%的样本（>=128K）贡献了44.3%的令牌。尽管GitHub数据集中长序列的比例较低，但其16.2%的令牌来自超过128K的序列，显示出显著的数据异质性。
- 观察2：混合长短序列能提升模型性能。 现有工作【12，How to Train Long-Context Language Models (Effectively)，2024，CoRR】已证明，仅在长上下文数据上训练会导致短上下文性能下降。LLaMA3报告【11，The Llama 3 Herd of Models，2024，CoRR】指出，在训练一个128K上下文的模型时，将0.1%的长数据与原始短数据混合，可以优化短上下文和长上下文基准测试的性能。DeepSeek-R1【10，DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning，2025，CoRR abs/2501.12948】展示了RL过程中训练集上的平均响应长度，表明逐渐增加和多样化的响应长度有助于提高模型性能。

• 挑战：数据异质性导致效率下降。 尽管长短序列的混合训练对模型性能有益且常见，但它带来了新的挑战。现有系统中使用的静态并行策略不适合处理动态工作负载。这导致了冗余通信（§3.2）和计算不均衡（§3.3）的问题，我们将在下面更详细地讨论。

3.2 冗余通信

现有系统在整个训练过程中应用静态并行策略。通常，它们假设所有（打包的）序列长度相同，并设置一个固定的CP等级，以将它们分摊到足够多的设备上，从而避免内存溢出（OOM）错误。如§2.3所述，为了处理可变长度序列，通常会将序列打包至上下文长度。然而，如图3(a)-(b)所示，所有序列都必须在整个CP组中进行分区，即使对于较短的序列来说这是不必要的。

例如，假设每个设备能处理8K令牌，要训练一个上下文长度为1M令牌的LLM，需要128的CP等级。此配置需要128个独立设备来处理一个1M令牌的序列。同时，大量较短的序列，如长度为4K、8K和16K的序列，被打包至1M令牌，并在一个有128个设备的CP组中处理。如图14所示，打包序列中的每个子序列都需要在CP ranks上被分区成128个块，并执行环形P2P通信。实际上，对于长度低于8K的序列，执行跨设备分区和通信是不必要的。对于16K令牌的序列，只需要两个CP ranks。对这些较短序列使用与最大序列长度相同的CP等级，会导致过度的冗余通信。当序列长度高度偏斜时，这个问题会更加严重。

3.3 计算不均衡

• 不均衡的FLOPs：尽管Flash Attention通过$O(N)$的内存复杂性实现了线性打包，但每个子序列的计算复杂度仍然是$O(S^2)$（其中S为子序列长度）。如图2(d)和图3(c)所示，即使两个打包序列包含相同数量的令牌，它们的实际计算工作负载也不同，这与注意力掩码的面积成正比。如图6(a)所示，当上下文长度短于8K令牌时，$O(S^2)$项相对不显著，允许打包有效地平衡内存和计算的工作负载。然而，对于长上下文训练任务，$O(S^2)$项成为计算的主要部分，导致不同打包序列之间出现显著的时间不平衡。为了直观解释，我们从GitHub数据集中采样了一个1.2M令牌的全局批次，并将其随机打包成最大32K令牌的微批次，以匹配模型的上下文长度。如图6(b)所示，我们记录了每个微批次的FLOPs（浮点运算次数），并观察到显著的可变性，这表明每个微批次的执行时间也不同。

• 不均衡的数据与流水线并行：微批次间不均衡的执行时间进一步降低了数据并行和流水线并行的效率。在数据并行中，所有DP ranks必须执行相同数量的微批次，然后在模型更新前同步梯度。如图3(c)所示，rank-2处理的令牌FLOPs比rank-0少，导致其在等待同步时出现空闲时间（即DP Bubble）。在流水线并行中，存在两种“气泡”：PP气泡出现在单个流水线内部，DP气泡出现在不同流水线（不同的DP组）之间。除了预热和冷却阶段的PP气泡，微批次间不均衡的FLOPs使得不同设备上的执行时间无法像理想流水线那样重叠。这导致了由阶段间等待引起的额外PP气泡，如图5所示。此外，由于每个微批次在流水线中的$P_{pp}$个不同阶段上顺序执行，任何DP气泡都将被放大$P_{pp}$倍。例如，考虑图5中所示的两个流水线，流水线(a)中的微批次0和7的前向和后向执行时间比流水线(b)中的要长。在$P_{pp}=4$的情况下，这个时间差被放大了四倍。因此，在执行8个微批次后，流水线(b)陷入了长时间的空闲期，等待梯度同步。这导致DP气泡占总执行时间的30%以上，远超正常的流水线气泡时间。

A2 方法细节

4. ByteScale 概览

我们提出ByteScale来解决这些挑战。如图7所示，它由三个主要组件构成。Profiler用于分析环境、模型配置、数据分布，并为其他组件构建成本模型。Communication Optimizer通过数据感知分片、动态通信和选择性卸载，为长短序列提升通信效率。Balance Scheduler通过并行感知的数据分配来解决计算不均衡问题。

5. 通信优化器

本节描述ByteScale如何优化通信开销。首先，它通过动态序列分片和通信减少短序列的冗余通信。其次，它通过选择性卸载进一步压缩长序列的通信成本。

5.1 数据感知分片与通信

• 混合数据并行（Hybrid Data Parallelism, HDP）：我们首先引入一种新的并行策略，即混合数据并行（HDP），以实现对不同序列长度的高效训练。DP和CP都将训练数据在设备间进行分区。DP通过将不同样本均匀分布到设备上执行数据间分区，而CP通过将单个样本分片到设备上执行数据内分区。HDP统一了数据间和数据内分区，其定义为将令牌均匀分布在设备上。它可以替代传统的DP和CP，HDP的并行度等于DP和CP并行度的乘积（即$P_{hdp} = P_{dp} \times P_{cp}$）。

• HDP的异构行为：与DP和CP要求所有DP/CP ranks执行一致的计算或通信行为（例如，CP要求所有CP ranks参与同构的环形P2P通信）不同，HDP允许HDP ranks之间存在异构行为。它有两个关键特性：

  1. 更灵活的通信：HDP只要求不同的HDP ranks处理相等数量的令牌。这意味着一些HDP ranks可能被分配完整的序列（短序列），如S3和S5在图8(d)中所示，而其他一些ranks可能只处理序列的部分切片（长序列），如S4在图8(d)中所示。这需要建立更灵活的通信组。例如，在图8(d)中，仅在rank-[1~2]之间创建了一个大小为2的通信组来计算S4的分布式注意力，而rank-0和3可以执行本地计算而无需跨设备通信。在图8(b)中，序列S0被分片成四个切片，并在rank-[0~3]之间创建了一个大小为4的通信组。
  2. 更细粒度的通信：静态并行策略要求并行度的乘积等于集群中的设备数量，即$P_{dp} \times P_{cp} \times P_{tp} \times P_{pp} = N_{cluster}$，其中$P_{tp}$和$P_{pp}$实际上根据模型大小是固定的。为了利用所有设备并保持这种可除性，$P_{dp}$和$P_{cp}$只能按有限的因子进行缩放，导致粒度较粗（例如，假设每个rank能处理8K令牌，512K序列可以使用< $P_{dp}=2, P_{cp}=64$ >，而768K序列需要$P_{cp}=96$但必须使用< $P_{dp}=1, P_{cp}=128$ >）。同时，HDP可以使用[1, $P_{hdp}$]中的任意数量的ranks来处理一个序列，而无需考虑可除性约束（例如，当$P_{hdp} = P_{dp} \times P_{cp} = 128$时，HDP可以使用96个ranks处理一个768K的序列，同时使用剩下的32个ranks分别处理32个8K的序列）。

• NCCL缓冲优化：创建NCCL通信组会产生额外开销。首先，建立通信组的过程本身很慢，为每个序列动态创建新组会显著降低训练效率。其次，创建过多的通信组会为每个GPU消耗额外的5-10GB内存用于NCCL缓冲区，进一步减少可用内存。幸运的是，分布式注意力利用P2P通信。通过在所有HDP ranks上建立一个全局通信组，任意两个设备之间的P2P通信可以直接重用现有组，从而减轻了创建临时通信组带来的时间和内存压力。

• 优化器状态分片：HDP将令牌均匀地划分到设备上，并且既不分片模型参数也不分片梯度。这意味着HDP ranks像DP一样复制模型状态。因此，ZeRO系列技术也适用于HDP，如图8(a)所示，HDP在所有HDP ranks上利用ZeRO-1来最大程度地分片优化器状态，从而最小化内存使用。

• 损失与模型更新：尽管HDP ranks可能在不同的微批次中执行不同的异构通信，但参数的最终梯度与标准DP中获得的梯度是等效的。如图9所示，每个令牌都对参数$W$贡献一个梯度，最终梯度$grad_W$是全局批次（表示为B）中所有令牌梯度的总和。设$grad(t, W)$表示令牌$t$对参数$W$的梯度。那么$grad_W$可以表示为：
  
  由于参数是复制的，并且令牌在HDP ranks（表示为R）之间均匀分布，本地累积的梯度对应于分配给每个rank（表示为$B_r$，即rank $r$中的微批次）的令牌梯度的部分和。因此，与DP类似，将在所有HDP ranks上执行全局集体通信（如All-Reduce或Reduce-Scatter）来聚合部分梯度。这也得到了来自所有令牌的梯度$grad_W$：
  
  公式(2)等价于公式(1)，并确保HDP中梯度累积的结果与标准DP中的结果等效。此外，由于我们计算了全局批次中所有令牌的梯度$grad_W$，它也需要通过令牌总数进行缩放，我们通过令牌级损失来实现这一点，该损失按令牌数量而不是样本数量来缩放损失。

5.2 数据感知的选择性卸载

• 激活值卸载（Activation Offloading）：激活值的大小与序列长度成正比。受GPU内存限制，较长的序列需要更多的HDP ranks来分散激活值。例如，处理一个1M令牌的序列需要128个ranks（如果每个rank能处理8K令牌），这在当今昂贵的GPU资源下通常是无法承受的。实践中，现代GPU服务器通常配备的CPU内存远超GPU内存。因此，另一种方法是将激活值卸载到CPU，从而减少所需的ranks数量。支持这种方法可行性有两个特点：
  1. 激活值是后进先出（FILO）的：如图10所示，对于任何序列，在前向传播期间，它会按顺序由Transformer层处理，激活值会逐渐累积，直到最后一层后达到峰值。随后，在后向传播期间，这些激活值将从最后一层到第一层被消耗。由于早期层产生的激活值稍后才被使用（即FILO），因此在前向传播期间将这些激活值卸载到CPU，并在后向传播需要时重新加载回GPU是很有前景的。
  2. $O(N^2)$的计算可以掩盖$O(N)$的卸载：众所周知，由于PCIe带宽有限，GPU和CPU之间的数据传输通常效率不高。卸载时间通常远超计算时间，使其不切实际。幸运的是，如§2.4所述，注意力的计算复杂度是$O(N^2)$，而内存复杂度是$O(N)$。因此，对于足够长的序列，$O(N^2)$的计算时间将不可避免地超过$O(N)$的数据传输时间，从而允许卸载操作被完美地掩盖在计算之下。

• act_ctx组件：如图11(b)所示，我们设计了一个名为act_ctx的通用组件（列表1）来支持激活值卸载。该组件分别为D2H（设备到主机）和H2D（主机到设备）维护两个cuda流。它自动从计算图中捕获激活张量，并在前向传播的适当时间将它们卸载到CPU（使用asyncCudaMemcpy API），并在D2H流和计算流之间建立异步依赖关系。计算图中的原始张量被元数据{层id, 激活id}替换。同样，在后向传播期间，存储在计算图中的元数据用于索引和在H2D流中重新加载相应的激活值。图10展示了整个过程。act_ctx还支持一个名为offload_ratio的参数，提供令牌级的细粒度控制，以控制卸载到CPU的激活值比例。此功能在节省GPU内存与实现最佳计算重叠之间取得平衡。

# 列表1. act_ctx的用法
with act_ctx(offload_ratio=0.5):
    # 前向传播
    hidden_states = model_layer(hidden_states)
# 后向传播
loss.backward()

• 选择性卸载（Selective Offloading）：激活值卸载利用CPU内存来减轻GPU内存的负担。然而，只有对于长序列，计算才能完美地与卸载重叠。这意味着我们不能不加选择地卸载分配给每个rank的所有令牌。相反，我们必须根据FLOPs选择性地卸载每个令牌。假设每个rank的层数为$L$，每个rank的令牌容量为$C$。给定一个长度为$S_i \ge C$的序列，我们将每层的计算时间和激活大小分别定义为$T_{comp}(S_i)$和$Act(S_i)$。D2H和H2D的带宽被分析为$BW_{d2h}$和$BW_{h2d}$。我们的目标是找到一个卸载比率$\rho$，通过公式(3)最小化$S_i$所需的HDP ranks数量$n(S_i)$，其中$k_1, c_1, k_2, c_2$和$C_0$是我们为成本模型分析出的系数。
  
  由于不同的微批次具有相互独立的前向和后向传播，在列表1中，我们为每个微批次分配一个从公式(3)派生的独立offload_ratio。这种方法有效地将长序列所需的ranks数量从$S_i/C$压缩到$n(S_i)$，如图11(a)所示。它不仅显著减少了通信开销，还使更多可用的HDP ranks能够处理数据，从而提高效率。

• 重叠效率讨论：我们知道，NCCL通信需要占用一部分流式多处理器（SMs），以达到InfiniBand和NVLink的峰值带宽。因此，即使有通信-计算重叠，计算核心也无法完全利用所有张量核心，导致效率低下。幸运的是，D2H和H2D核心使用DMA引擎而不是SMs，使其能与计算和通信完美重叠。此外，我们使用缓存的固定主机内存来进一步减少CPU内存分配的开销，并加速设备和主机之间的数据交换。由于流水线并行交错了不同微批次的前向和后向传播，D2H和H2D核心可以同时执行，从而最大化PCIe的双向带宽。

5.3 整体流程

ByteScale的整体流程在算法1中概述。简而言之，该算法遍历全局批次中的每个序列$S_i$。对于长序列，它推导出卸载比率$\rho$并确定所需的ranks数量$n(S_i)$（第1-6行）。对于短序列，它将它们打包以填满每个rank的容量$C$（第7-9行）。处理后的序列随后被分配给$P_{hdp}$个ranks，算法返回微批次和offload_ctx以供执行（第10-12行）。

6. 平衡调度器

本节我们介绍平衡调度器，以解决DP和PP的不平衡问题。通过精心编排数据分配（替代算法1中的第10行），它在保持§5中实现的最小通信的同时，缓解了这些不平衡。我们将首先概述几个关键见解，然后提出我们的启发式解决方案。

6.1 重新定义微批次

梯度累积要求不同的DP ranks执行相同数量的微批次，这是基于所有微批次具有相同计算负载的假设。然而，如§3.3所述，不同微批次的执行时间可能显著不同。在ByteScale中，我们重新定义了一个更灵活的策略，允许不同的HDP ranks处理不同数量的微批次（大小相同但工作负载不同），以缓解不平衡问题。如图13所示，这使得所有ranks几乎同时完成计算。更重要的是，这个策略不影响模型收敛。无论序列如何分配给HDP ranks，我们最终都计算全局批次中所有令牌的梯度总和，如§5.1所讨论的，这确保了数学上的等价性。

6.2 解决PP不平衡

• 见解1：当不同长度级别的序列被分配到不同的流水线时，PP气泡较少。
  确保流水线处理的微批次具有相似的执行时间至关重要。如图13(b)所示，当$P_{pp}=4$时，时间轴上任意4个连续的微批次将由4个PP阶段同时执行。如果它们的执行时间差异很大，就会出现额外的PP气泡。由于全局批次中长序列数量有限，一些流水线不得不被分配多种长度级别的序列。幸运的是，只有在过渡阶段（例如，当4个连续的微批次属于不同长度级别时）才会导致额外的PP气泡。

• 策略：我们为平均执行时间较短的流水线分配更多的微批次。如图12(a)-(b)所示，pipeline-0处理平均执行时间较长的微批次，因此只分配了8个微批次。相比之下，pipeline-1被分配了18个微批次，以与pipeline-0同步。此外，由于微批次更多，气泡率进一步降低。

6.3 解决DP不平衡

• 见解2：当不应用流水线并行时，只需在每个时间步保持负载均衡。
  如果只应用DP而不应用PP，实现负载均衡只需要在任何给定时间，由不同HDP ranks执行的微批次具有相似的执行时间。无需考虑时间轴上不同时间步之间微批次的工作负载不平衡。

• 策略：一个直接的方法是在同一时间将相同长度级别的序列分配给不同的HDP ranks，如图13(a)所示。此外，我们仍然为处理较短序列的ranks分配比其他ranks更多的微批次。最终，这确保了所有HDP ranks几乎同时同步梯度。

6.4 平衡策略

算法2描述了平衡策略。首先，我们按长度降序对全局批次B中的序列进行排序。然后将这些有序序列划分为FLOPs总和近似相等的桶，因此平均长度较长的桶包含的序列较少（第3-5行）。其次，我们确定那些执行时间较短的ranks，以便后续分配（第7-9行）。第三，如果使用DP-Balance策略，我们从同一个桶中选择序列。否则，如果使用PP-Balance策略，我们从所有桶中顺序选择序列。实践中，执行时间较短的ranks会被分配更多的序列（第12-15行）。最后，我们重复第二和第三步，直到所有桶都为空。

7. 实现细节

ByteScale基于Python、C++和CUDA实现了约16000行代码，并已与MegaScale【18，MegaScale: scaling large language model training to more than 10,000 GPUs，2024，NSDI’24】集成，后者是一个用于LLM训练的高性能框架。为了支持大规模训练和通信，我们还应用了以下优化。

• GQA（Group Query Attention）：GQA已成为现代LLM（如LLaMA3和Mistral）中不可或缺的特性，它有助于减少KV头的数量，从而降低分布式注意力（dist-attn）的通信量。本文提到的所有系统都应用了GQA技术。

• 带打包的Dist-attn：由于工作负载与注意力掩码的面积成正比，按顺序将序列划分到设备上会导致工作负载不平衡。已有几种技术【4, 23, 31】被提出来解决这个问题。然而，它们不适用于打包序列的特殊分段因果注意力掩码。如图14所示，为了避免CP组内的异构计算和通信，我们优化了当前的dist-attn。打包序列的每个子序列被均匀地分成$2P$部分，并对称地分配给$P$个设备。这确保了每个设备持有所有子序列的$1/P$，并覆盖了注意力掩码面积的$1/P$。所有设备参与相同的环形P2P通信，数据交换量相同。

• 远程数据加载器（Remote Dataloader）：ByteScale在每个训练步骤需要全局批次信息来调度数据分配。然而，现有的数据加载器解决方案通常遵循SPMD（单程序，多数据）模式，其中每个rank只读取批次的部分数据。为了维护全局信息，所有HDP ranks必须同时读取整个全局批次，这对网络通信和CPU内存都造成了巨大压力。为了解决这个问题，我们使用Ray【26，Ray: a distributed framework for emerging AI applications，2018，OSDI’18】实现了一个远程数据加载器，它以全局视角提供实时调度和规划能力。如图15所示，考虑一个包含两个GPU节点（worker-0和worker-1）和一个作为Ray头节点的CPU节点的设置，存在三种由ray actors封装的角色。Server Roles是工作节点中的CPU进程，它们从HDFS获取并预处理原始数据，并生成元数据。Scheduler Role作为单一控制器，是worker-0中的一个CPU进程，它从所有服务器收集全局元数据，推导出加载计划，并将其广播给客户端。Client Roles是工作节点中的GPU进程，它们根据加载计划从服务器读取部分数据。

• 融合的SoftmaxCrossEntropy：现代LLM通常使用具有大词汇表的tokenizer（例如，LLaMA3【11】中的128K，Mistral【2】中的130K，Qwen2.5【43】中的超过150K）。为了稳定精度，当前方法（如Megatron-LM中的VocabParallel）在计算SoftmaxCrossEntropyLoss之前将logits变量从BF16转换为FP32。然而，FP32的logits消耗大量内存。例如，在上下文长度为256K和词汇表大小为128K的情况下，在TP=8时需要16GB内存。此外，这些核函数是内存密集型且效率低下。如图16所示，我们开发了FusedSoftmaxCrossEntropy，它将众多操作融合成一个单一的核函数，接收BF16输入，但仍以FP32精度进行在线计算。与现有方法相比，它既节省了时间又节省了内存。

A4 实验环境与结果

8.1 实验设置

• 环境：我们的实验在一个拥有超过12,000个GPU的大规模生产GPU集群上进行。（由于商业和保密原因，有关生产集群的具体信息，如GPU的数量和类型，被隐藏。）

• 基线：我们的系统构建在MegaScale之上，这是一个用于大规模GPU集群的生产级LLM训练框架，已证明其性能优于DeepSpeed和Megatron-LM。因此，我们通过在三种情况下进行比较来展示ByteScale的优势：

  1. MegaScale (static)：使用静态并行策略（DP, TP, PP和CP），并为打包序列应用图14所示的dist-attn优化。
  2. HDP naive：MegaScale与朴素HDP，如算法1所述，仅应用通信优化。
  3. HDP balance：MegaScale与平衡HDP，如算法2所述，同时应用通信和平衡优化。
     为了公平比较，我们为所有三种情况设置了相同的$P_{tp}$和$P_{pp}$，并将情况②③中的$P_{hdp}$设为等于情况①中的$P_{dp} \times P_{cp}$，其中$P_{cp}$对应支持模型上下文长度所需的最小ranks数量。

• 模型和数据集：我们使用密集和稀疏LLM评估我们的工作，如表1所示。对于密集模型，我们选择了四种不同大小的LLaMA系列LLM：LLaMA-7B、LLaMA-13B、LLaMA-30B和LLaMA-70B。对于稀疏模型，我们选择了两种不同大小的Mistral系列LLM（MoE）：Mistral-8x7B（激活参数=13B/47B）和Mistral-8x22B（激活参数=39B/141B）。实验中使用了两个数据集，即GitHub和Byted，我们已在§3.1中介绍过。图4展示了这两个数据集的数据分布。

• 工作负载和指标：对于不同类型和大小的模型，我们将上下文长度从256K扩展到2M，集群规模从1024个GPU扩展到超过12,000个GPU，以更全面地评估ByteScale的性能。每个训练步骤的全局批次固定为32M令牌，这是大规模集群中的常见做法。我们使用吞吐量（每秒令牌数）作为评估性能的主要指标。所有结果都是在20次预热迭代后，对200次迭代取平均值。

Table 1. 用于评估的模型

8.2 端到端评估

我们首先通过测量每个训练步骤的平均吞吐量来评估三种方法的端到端性能，总体结果如图17所示。结果表明，朴素HDP和平衡HDP解决方案都优于基线，最大实现了7.89倍的加速。

• 可扩展性差异：随着上下文长度的增加，采用静态策略的基线必须增加CP等级以避免OOM错误。对于全局批次中的较短序列，我们必须打包它们并应用图14所示的dist-attn，这会遭受低效和冗余的通信。例如，对于GitHub数据集，全局批次中只有9.8%的令牌长于256K，当上下文长度从256K扩展到2M时，我们可以观察到基线的吞吐量几乎在上下文长度增加2倍时下降2倍。相比之下，在相同条件下，朴素HDP解决方案的吞吐量平均下降1.23倍，而平衡HDP解决方案的吞吐量平均仅下降1.08倍。朴素HDP解决方案减少了通信开销，但由于不平衡而使一些ranks空闲。同时，平衡HDP解决方案消除了这些气泡时间，并完全释放了由灵活和动态通信带来的性能。因此，ByteScale在GitHub数据集上比基线最多快7.89倍。

• 数据集差异：Byted数据集比GitHub数据集包含更多的长序列，全局批次中有37%的令牌长于256K。因此，平均吞吐量和加速比低于GitHub数据集。然而，由于ByteScale为长短序列都提供了通信优化，加速比仍可达到4.26倍。

• 并行策略差异：像LLaMA-7B、13B和30B这样的模型使用包括HDP和TP在内的并行策略，因此应用DP-Balance策略。相比之下，像LLaMA-70B、Mistral-8x7B和Mistral-8x22B这样的模型采用HDP、TP和PP，我们应用PP-Balance策略。可以观察到，与PP-Balance相比，带DP-Balance的HDP实现了更高的加速比。例如，在GitHub数据集和2M上下文长度下，带DP-Balance的HDP的加速比在6.21倍-7.89倍之间，而带PP-Balance的HDP的加速比仅在3.42倍-4.28倍之间。如图13所示，DP-Balance策略只需要在每个时间步平衡计算，这比PP-Balance策略要求的在所有时间步平衡计算更容易实现。

8.3 案例研究

为了更深入地剖析ByteScale的卓越性能，我们选择Byted数据集，并在一个有1024个GPU的集群上训练LLaMA-7B，上下文长度为2M。图18展示了单个训练步骤中不同ranks的详细运行时状态。

• 通信密集型案例：首先，我们从集群中随机选择4个ranks，并记录它们在每个方法的训练步骤中的前向和后向时间。如图18(b)所示，对于基线，微批次数量设为8，我们必须设置$P_{cp}=256$以支持2M的序列长度。可以观察到这4个ranks表现出相似的执行时间。这是因为大多数微批次（除了第三个）不具有$O((2M)^2)$的计算复杂度，但必须处理2M的通信量。如图18(a)所示，P2P通信时间远超计算时间，导致一个微批次的执行时间几乎由通信决定（占总时间的97.6%）。

• 计算不平衡案例：在朴素HDP解决方案下，全局批次内的序列按所需的最小ranks数量进行分片。如图18(a)所示，一个312K的序列仅由39个HDP ranks分片作为微批次，因此计算时间可以与通信开销重叠。然而，由于ranks之间的不平衡，训练效率仍然存在问题。如图18(b)所示，尽管第三个rank在1分41秒内完成了它的8个微批次，但它必须等待第一个rank在4分32秒时完成，导致171秒的空闲时间。即便如此，朴素HDP解决方案相比基线节省了4分8秒。

• 平衡案例：在平衡HDP解决方案下，所有ranks几乎同时完成执行。如图18(b)所示，在任何时间步，每个rank都被分配了具有相似FLOPs的微批次，而执行时间较短的ranks（如第三和第四个rank）将被分配更多的批次。因此，该步骤的总时间进一步减少到2分37秒，相比基线节省了6分3秒。

• 总体比较：如图18(c)所示，我们记录了所有1024个GPU在单个步骤中的有效计算时间。可以发现，朴素HDP解决方案相比基线将峰值执行时间减少了1.7倍，但ranks之间存在显著的时间差异，最大值和最小值之间相差4.7倍（min=60s, max=279s, std=68s）。平衡HDP解决方案消除了时间差异，从而相比朴素解决方案进一步将执行时间减少了2.3倍。

8.4 消融研究

为了探究ByteScale中每个组件的有效性，我们使用与§8.3相同的配置进行了消融实验，如图20所示。

• 动态通信的有效性：虽然图18提供了单个训练步骤中的运行时快照，我们进一步分析了两个小时内的网络流量和张量核心利用率，如图19所示。可以观察到，基线表现出非常繁重的RDMA流量，但相应的张量核心利用率很低。这是因为大多数ranks受通信限制，计算单元由于等待冗余通信而大部分时间处于空闲状态。这一观察与图18(a)中描述的情况一致。当我们应用朴素HDP解决方案时，峰值RDMA流量几乎减半，表明大量不必要的通信已被消除。此外，张量核心利用率也从10%增加到40%。因此，它相比基线实现了1.59倍的加速。然而，由于不平衡问题，这些改进不稳定，通信和计算硬件单元偶尔会经历停顿或空闲期。

• 选择性卸载的有效性：选择性卸载作为朴素HDP解决方案的补充。如图21所示，当卸载比率$\rho=1.0$时，激活值卸载节省的内存与重计算相同。当上下文长度设为64K时，计算无法完全与卸载重叠。然而，当我们将比率降低到$\rho=0.5$时，它可以在不影响吞吐量的情况下节省32.3%的内存。此外，卸载比率是根据公式3自动推导的，随着上下文长度的增加，可以设置更高的比率以节省更多内存（例如，为256K设置$\rho=1.0$不会降低吞吐量）。该方法减少了长序列所需的ranks数量，使得相同数量的ranks可以同时处理更多序列，从而将加速比从1.59倍提高到2.01倍。

• 平衡策略的有效性：如图19所示，平衡策略稳定了RDMA流量，并使张量核心利用率持续保持在40%左右。这表明计算和通信的硬件单元在超过两个小时内持续满负荷工作而没有空闲。因此，平衡HDP解决方案将加速比从2.01倍提高到3.69倍，超过了任何其他策略带来的改进。

• 远程数据加载器的有效性：我们采用图15所示的远程加载器，并使用CPU预取来将数据读取与计算重叠。这种方法进一步将加速比从3.69倍提高到3.89倍。