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

作者/机构: 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)长短序列混合训练场景,提出了一个名为 ByteScale 的高效、灵活且可扩展的训练框架。现有框架通常采用静态的通信网格(如二维网格)来组织设备,将数据并行(DP)和上下文并行(CP)视为正交的技术,但这导致在处理长度可变的序列时,出现冗余通信和计算不平衡的问题,从而降低了训练效率。

核心问题:
1. 冗余通信: 静态并行策略强制所有序列(无论长短)都使用为最长序列配置的上下文并行(CP)组进行分区和通信,即使短序列并不需要跨设备分区,也产生了不必要的通信开销。
2. 计算不平衡: 尽管通过打包(packing)技术可以在设备间均匀分配令牌数量,但由于自注意力机制的计算复杂度为 $O(L^2)$($L$为序列长度),不同打包序列的实际计算负载(FLOPs)差异巨大,导致设备间出现计算不平衡,产生空闲等待时间(气泡)。

研究目标与创新点:
为了解决上述挑战,ByteScale 提出了以下核心贡献:

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

A3 背景知识与关键洞察

2.1 Transformer 与大语言模型

Transformer 架构【40, Attention is All you Need, 2017, NeurIPS 2017】是当前大语言模型(LLM)的主流基础架构。它由一系列 Transformer 层堆叠而成,每个层包含一个注意力模块和一个前馈网络(FFN)模块。如图 1 所示,自注意力机制需要在整个序列的所有令牌间进行计算以捕获上下文信息,而其他操作(如归一化、线性投影和激活函数)则是逐令牌计算的,每个令牌可以独立处理。

图 1. Transformer 层的架构
图 1. Transformer 层的架构

2.2 分布式 LLM 训练

  • 数据并行 (Data Parallelism, DP):【9, Large Scale Distributed Deep Networks, 2012, NeurIPS 2022】【24, PyTorch Distributed: Experiences on Accelerating Data Parallel Training, 2020, Proc. VLDB Endow.】【37, Horovod: fast and easy distributed deep learning in TensorFlow, 2018】将训练数据均匀分布到不同设备上,每个设备持有一份模型副本。设备独立处理本地数据,然后全局同步梯度来更新模型。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】和流水线并行(PP)【16, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019, NeurIPS 2019】【28, PipeDream: generalized pipeline parallelism for DNN training, 2019, SOSP 2019】【29, Memory-Efficient Pipeline-Parallel DNN Training, 2021, ICML 2021】。TP 在层内划分操作和参数,需要通信中间结果(激活值),通常用于单节点内。PP 将模型层切分为不同阶段,只需在相邻阶段间通过点对点(P2P)通信交换激活值,可跨多节点。
  • 混合并行 (Hybrid Parallelism):结合多种并行策略提升效率。Megatron-LM【21, Reducing Activation Recomputation in Large Transformer Models, 2022】【30, Efficient large-scale language model training on GPU clusters using megatron-LM, 2021, SC 2021】【38, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019】提出的 DP、TP 和 PP 结合的 3D 并行策略是当前大规模模型训练的主流方法。
  • 梯度累积 (Gradient Accumulation):为了提升效率和收敛性,LLM 通常需要大批量(large batch size)。梯度累积将每个全局批次(global batch)划分为多个微批次(micro-batches),累积这些微批次的梯度,其效果等同于一次性处理整个全局批次。

2.3 填充与打包

为了在静态并行策略中支持可变长度序列,需要使用填充(padding)和打包(packing)技术。如图 2 所示,填充将同一批次中的序列补齐到相同长度,但这会造成计算浪费。打包【22, Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance, 2021】则将多个序列拼接成一个长序列,不使用填充令牌,并采用一种特殊的分段注意力掩码(segmented attention mask)来确保每个序列被独立处理。

图 2. 序列填充与打包
图 2. 序列填充与打包

2.4 长上下文训练

自注意力的时间和内存复杂度均为 $O(L^2)$,这成为扩展上下文长度的瓶颈。Flash Attention【7, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023】【8, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022, NeurIPS 2022】通过优化内存 I/O 和使用分块技术,将内存复杂度从 $O(L^2)$ 降至 $O(L)$,但时间复杂度仍为 $O(L^2)$。上下文并行(Context Parallelism, CP)【4, Striped Attention: Faster Ring Attention for Causal Transformers, 2023】【23, LightSeq: Sequence Level Parallelism for Distributed Training of Long Context Transformers, 2023】【25, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023】【31, NVIDIA: Context Parallelism, 2024】进一步将序列划分到 $C$ 个设备上,将内存从 $O(L)$ 降至 $O(L/C)$。CP 沿序列维度对 QKV 进行分片,跨令牌操作需要设备间通过环形点对点通信交换 KV 切片,并与计算重叠。该技术也适用于打包序列,如图 2(c) 和 3(a) 所示,每个子序列也必须被划分到所有 CP rank 上。

图 3. 带打包的上下文并行
图 3. 带打包的上下文并行

3.1 数据异构性

  • 观察 1:真实世界数据集中的序列长度呈偏态分布。如图 4 所示,我们分析了开源数据集 GitHub 和一个用于长上下文训练的生产数据集 Byted。两者都显示出序列长度的偏态分布。例如,在 Byted 数据集中,随机抽样的全局批次中近 80% 的样本长度小于等于 4K,而只有 0.05% 的样本能达到 2M。然而,从令牌分布来看,这 0.05% 的长样本(>=2M)贡献了全局批次中 12.1% 的令牌。GitHub 数据集中,超过 128K 的序列贡献了 16.2% 的令牌,显示出显著的数据异构性。

    图 4. 两个数据集中的样本和令牌分布
    图 4. 两个数据集中的样本和令牌分布

    * 观察 2:混合长短序列能提升模型性能。现有工作【12, How to Train Long-Context Language Models (Effectively), 2024】表明,仅在长上下文数据上训练会导致短上下文性能下降。LLaMA3 报告【11, The Llama 3 Herd of Models, 2024】指出,在训练 128K 上下文模型时,混合 0.1% 的长数据能优化短长上下文的性能。DeepSeek-R1【10, DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning, 2025】也展示了在强化学习过程中,逐步增加和多样化的响应长度有助于提升模型性能。
    * 挑战:数据异构性导致效率下降。虽然混合训练有益,但现有系统的静态并行策略不适应动态工作负载,导致了冗余通信和计算不平衡问题。

3.2 冗余通信

  • 静态并行策略的局限性。现有系统在整个训练过程中应用静态并行策略。它们通常假设所有(打包的)序列长度相同,并设置固定的 CP 并行度以避免内存溢出(OOM)。如图 3(a)-(b) 所示,所有序列,无论长短,都必须在整个 CP 组上进行分区,即使这对于短序列来说是不必要的。
  • 具体示例。假设每个设备能处理 8K 令牌,训练一个 1M 上下文的模型需要 CP 并行度为 128。这意味着处理一个 1M 令牌的序列需要 128 个设备。同时,大量较短的序列(如 4K, 8K, 16K)被打包成 1M 长度,也在这 128 个设备上处理。如图 14 所示,打包序列中的每个子序列都需要被切成 128 块,在 CP ranks 间进行环形 P2P 通信。实际上,对于长度小于 8K 的序列,跨设备分区和通信是完全不必要的;对于 16K 的序列,也仅需 2 个 CP rank。使用与最大序列长度相同的 CP 并行度,导致了短序列的过度冗余通信。

3.3 计算不平衡

  • 不平衡的 FLOPs。尽管 Flash Attention 使得打包后的序列在内存上是线性复杂度 $O(L)$,但每个子序列的计算复杂度仍然是 $O(L^2)$。如图 2(d) 和 3(c) 所示,即使两个打包序列包含相同数量的令牌,它们的实际计算工作负载也不同,这与注意力掩码的面积成正比。如图 6(a) 所示,当上下文长度较短(<8K)时,$O(L^2)$ 项不显著,打包能有效平衡内存和计算负载。但在长上下文训练中,$O(L^2)$ 项成为主导,导致不同打包序列间存在显著的时间不平衡。如图 6(b) 所示,从 GitHub 数据集采样打包的微批次,其 FLOPs 存在巨大差异。

    图 6. 不平衡的计算
    图 6. 不平衡的计算

    * 不平衡数据与流水线并行。微批次间执行时间的不平衡进一步降低了数据并行和流水线并行的效率。在数据并行中,所有 DP rank 必须执行相同数量的微批次,然后同步梯度。如图 3(c) 所示,计算量较小的 rank-2 会提前完成,进入空闲等待(DP 气泡)。在流水线并行中,除了固有的 PP 气泡外,微批次间的 FLOPs 不平衡会导致各阶段执行时间无法理想重叠,产生额外的 PP 气泡(级间等待),如图 5 所示。更严重的是,每个微批次需在 $P_{pp}$ 个阶段上顺序执行,任何 DP 气泡都会被放大 $P_{pp}$ 倍。如图 5 所示,流水线(b)在执行完 8 个微批次后,会长时间空闲等待流水线(a),导致 DP 气泡占总执行时间的 30% 以上。

    图 5. 不平衡的数据和流水线并行
    图 5. 不平衡的数据和流水线并行

A2 方法细节

4. ByteScale 概览

为了应对上述挑战,我们提出了 ByteScale。如图 7 所示,它包含三个主要组件:
1. Profiler (分析器): 用于分析环境、模型配置、数据分布,并为其他组件构建成本模型。
2. Communication Optimizer (通信优化器): 通过数据感知分片、动态通信和选择性卸载,来提高长短序列的通信效率。
3. Balance Scheduler (平衡调度器): 通过并行感知的数据分配,来解决计算不平衡问题。

图 7. ByteScale 概览
图 7. ByteScale 概览

5. 通信优化器

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

5.1 数据感知分片与通信

  • 混合数据并行 (Hybrid Data Parallelism, HDP)。我们引入了一种新的并行策略 HDP,以实现对不同长度序列的高效训练。DP 和 CP 都将训练数据划分到设备上,但 DP 是数据间划分,CP 是数据内划分。HDP 统一了这两种划分方式,其目标是均匀地将令牌分布到设备上。HDP 可以替代传统的 DP 和 CP,其并行度 $D_{hdp} = D_{dp} \times D_{cp}$。与 DP 和 CP 不同,HDP 允许其 ranks 之间存在异构行为,具有两个关键特性:

    • 更灵活的通信: HDP 只要求不同 rank 处理相同数量的令牌。这意味着某些 rank 可能被分配完整的短序列(如图 8(d) 中的 $s_3$ 和 $s_5$),而其他 rank 可能只处理长序列的一个切片(如图 8(d) 中的 $s_4$)。这需要建立更灵活的通信组。例如,图 8(d) 中仅在 rank-[1~2] 之间创建了一个大小为 2 的通信组来计算 $s_4$ 的分布式注意力,而 rank-0 和 3 可以进行本地计算。
    • 更细粒度的通信: 静态并行策略要求 $D_{dp} \times D_{cp} \times D_{tp} \times D_{pp} = N_{cluster}$,这限制了 $D_{dp}$ 和 $D_{cp}$ 的选择,粒度较粗。而 HDP 可以在 [1, $D_{hdp}$] 范围内使用任意数量的 rank 处理一个序列,不受整除性约束。例如,当 $D_{hdp}=128$ 时,HDP 可以用 96 个 rank 处理一个 768K 的序列,用剩下的 32 个 rank 单独处理 32 个 8K 的序列。

    图 8. HDP 示意图
    图 8. HDP 示意图

    * NCCL 缓冲优化。动态创建 NCCL 通信组会带来巨大开销:创建过程本身很慢,且大量通信组会额外消耗每个 GPU 5~10GB 的内存用于 NCCL 缓冲。幸运的是,分布式注意力使用 P2P 通信。通过在所有 HDP rank 间建立一个全局通信组,任意两个设备间的 P2P 通信可以直接复用该组,从而避免了创建临时通信组的时间和内存压力。
    * 优化器状态分片。HDP 均匀地将令牌划分到设备上,但模型参数和梯度是复制的,这与 DP 类似。因此,ZeRO 系列技术也适用于 HDP。如图 8(a) 所示,HDP 在所有 ranks 间使用 ZeRO-1 来最大程度地分片优化器状态,以最小化内存使用。
    * 损失与模型更新。尽管 HDP ranks 在不同微批次间可能执行异构通信,但参数的最终梯度与标准 DP 相同。如图 9 所示,每个令牌都对参数 $\theta$ 贡献一个梯度,最终梯度 $G_{\theta}$ 是全局批次 B 中所有令牌梯度的总和。令 grad($t$, $\theta$) 为令牌 $t$ 对参数 $\theta$ 的梯度,则 $G_{\theta}$ 可表示为:
    公式1
    公式1

    由于参数是复制的,令牌被均匀分布到 HDP ranks (R) 上,每个 rank $r$ 上的局部累积梯度对应于分配给它的令牌子集 $B_r$ 的梯度之和。因此,与 DP 类似,在所有 HDP ranks 间执行一次全局集合通信(如 All-Reduce 或 Reduce-Scatter)来聚合部分梯度,同样可以得到来自所有令牌的梯度 $G_{\theta}$:
    公式2
    公式2

    公式(2)等价于公式(1),确保了 HDP 中的梯度累积结果与标准 DP 在数学上等价。我们通过令牌级损失(token-level loss)来实现这一点,即用总令牌数而不是样本数来缩放损失。
    图 9. 令牌级梯度
    图 9. 令牌级梯度

5.2 数据感知的选择性卸载

  • 激活值卸载。激活值的大小与序列长度成正比,长序列需要更多的 HDP rank 来分散激活值,这在资源上是昂贵的。现代 GPU 服务器通常配备远超 GPU 显存的 CPU 内存,因此可以将激活值卸载到 CPU。这种方法的可行性有两点支撑:
    • 激活值是“先进后出”(FILO): 如图 10 所示,前向传播时,激活值逐层累积;反向传播时,从最后一层开始消耗。较早层产生的激活值被使用得更晚,因此适合在前向传播时卸载到 CPU,在反向传播需要时再加载回 GPU。
    • $O(L^2)$ 计算可以重叠 $O(L)$ 卸载: GPU 和 CPU 之间的数据传输通常因 PCIe 带宽有限而效率不高。但幸运的是,注意力的计算复杂度是 $O(L^2)$,而内存(即传输数据量)复杂度是 $O(L)$。对于足够长的序列,$O(L^2)$ 的计算时间必然会超过 $O(L)$ 的数据传输时间,使得卸载可以被计算完美掩盖。
      我们设计了一个名为 act_ctx 的通用组件(如代码清单 1)来支持激活值卸载。该组件为 D2H(设备到主机)和 H2D(主机到设备)维护两个独立的 CUDA 流,自动捕获计算图中的激活张量并异步卸载到 CPU,同时在计算流和 D2H 流之间建立依赖。act_ctx 还支持一个 offload_ratio 参数,提供令牌级的细粒度控制,以平衡 GPU 内存节省和计算重叠效率。
# 代码清单 1. act_ctx 的用法
from activation_offload import act_ctx

# 启用激活值卸载
with act_ctx(offload_ratio=ratio):
    # Transformer 前向传播
    output = model(input)
# 反向传播
loss.backward()
图 10. 逐层激活值卸载
图 10. 逐层激活值卸载
  • 选择性卸载。激活值卸载利用 CPU 内存减轻 GPU 内存负担,但只有长序列的计算才能完美重叠卸载开销。因此,我们必须根据 FLOPs 选择性地卸载每个令牌。我们建立了一个成本模型(公式 3),旨在找到一个卸载比例 $\alpha$,以最小化处理序列 $S_i$ 所需的 HDP rank 数量 $C(S_i)$。
    公式3
    公式3

    由于不同微批次的前向和反向传播是独立的,我们为每个微批次分配一个由公式 3 推导出的独立 offload_ratio。这种方法有效地将长序列所需的 rank 数量从 $L_i/T$ 压缩到 $C(S_i)$,如图 11(a) 所示,显著减少了通信开销,并使更多 HDP rank 可用于处理数据,从而提高效率。
    图 11. 数据感知的选择性卸载
    图 11. 数据感知的选择性卸载
  • 重叠效率讨论。D2H 和 H2D 内核使用 DMA 引擎而非 SMs,使其能与计算和通信完美重叠。我们还使用缓存的固定主机内存(cached pinned host memory)来减少 CPU 内存分配开销并加速数据交换。在流水线并行中,D2H 和 H2D 内核可以同时执行,从而最大化 PCIe 的双向带宽。

5.3 整体流程

ByteScale 的整体流程如算法 1 所示。简而言之,算法遍历全局批次中的每个序列 $S_i$。对于长序列,它推导出卸载比例 $\alpha$ 和所需的 rank 数量 $C(S_i)$(1-6行)。对于短序列,它将它们打包以填满每个 rank 的容量 $T$(7-9行)。处理后的序列被分配给 $D_{hdp}$ 个 rank,算法返回每个 rank 的微批次和 offload_ratio 用于执行(10-12行)。

算法 1: 朴素 HDP 解决方案
算法 1: 朴素 HDP 解决方案

6. 平衡调度器

本节介绍平衡调度器如何解决 DP 和 PP 的不平衡问题。通过精心设计数据分配(替代算法 1 的第 10 行),它在保持最小通信开销的同时缓解了这些不平衡。

6.1 重新定义微批次

梯度累积要求不同 DP rank 执行相同数量的微批次,这是基于所有微批次计算负载相同的假设。然而,实际执行时间差异很大。在 ByteScale 中,我们重新定义了一个更灵活的策略,允许不同的 HDP rank 处理不同数量的微批次(大小相同但工作负载不同),以缓解不平衡问题。如图 13 所示,这使得所有 rank 能在同一时间完成计算。更重要的是,该策略不影响模型收敛,因为我们最终计算的是全局批次中所有令牌的梯度总和,保证了数学等价性。

图 13. 平衡策略
图 13. 平衡策略

6.2 解决 PP 不平衡

  • 洞察 1:将不同长度级别的序列分配到不同的流水线可以减少 PP 气泡。关键是确保一个流水线处理的微批次具有相似的执行时间。如图 13(b) 所示,如果时间线上连续的 4 个微批次执行时间差异很大,就会产生额外的 PP 气泡。我们应尽量将相似长度的序列分配给同一个流水线。
  • 为执行时间较短的流水线分配更多微批次。如图 12(a)-(b) 所示,pipeline-0 处理平均执行时间较长的微批次,因此只分配了 8 个微批次。相比之下,pipeline-1 被分配了 18 个微批次,以与 pipeline-0 同步完成。此外,由于微批次更多,其气泡率也进一步降低。

    图 12. 平衡的数据和流水线并行
    图 12. 平衡的数据和流水线并行

6.3 解决 DP 不平衡

  • 洞察 2:当不应用流水线并行时,只需在每个时间步保持负载均衡即可。如果只应用 DP,实现负载均衡只需要在任何给定时间点,不同 HDP rank 执行的微批次具有相似的执行时间。
  • 实现方法。一个直接的方法是在同一时间将相同长度级别的序列分配给不同的 HDP rank,如图 13(a) 所示。同时,我们仍然为处理较短序列的 rank 分配更多的微批次。最终,这确保了所有 HDP rank 几乎同时同步梯度。

6.4 平衡策略

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

算法 2: HDP 的平衡策略
算法 2: HDP 的平衡策略

7. 实现细节

ByteScale 基于 Python, C++ 和 CUDA 实现,代码约 16K 行,并已集成到高性能 LLM 训练框架 MegaScale【18, MegaScale: scaling large language model training to more than 10,000 GPUs, 2024, NSDI’24】中。

  • GQA (Group Query Attention): GQA 减少了 KV头的数量,从而降低了分布式注意力(dist-attn)的通信量。
  • 带打包的 Dist-attn: 为避免打包序列在 CP 组内造成异构计算和通信,我们优化了 dist-attn。如图 14 所示,我们将打包序列的每个子序列均匀划分为 $2C$ 份,并对称地分配给 $C$ 个设备。这确保了每个设备持有所有子序列的 $1/C$ 部分,并覆盖 $1/C$ 的注意力掩码区域,所有设备参与相同的环形 P2P 通信,数据交换量也相同。

    图 14. 针对打包序列优化的 Dist-attn
    图 14. 针对打包序列优化的 Dist-attn

    * 远程数据加载器 (Remote Dataloader): ByteScale 需要在每个训练步骤获取全局批次信息以进行调度。为避免所有 rank 同时读取整个全局批次带来的网络和 CPU 内存压力,我们使用 Ray【26, Ray: a distributed framework for emerging AI applications, 2018, OSDI’18】实现了一个远程数据加载器。如图 15 所示,它包含三种角色:Server Roles(在工作节点上预处理数据)、Scheduler Role(作为单一控制器收集元数据并制定加载计划)和 Client Roles(GPU 进程根据计划从服务器读取部分数据)。

    图 15. 远程数据加载器
    图 15. 远程数据加载器

    * 融合的 SoftmaxCrossEntropy: 现代 LLM 的词汇表很大,计算 SoftmaxCrossEntropyLoss 前将 logits 从 BF16 转换为 FP32 会消耗大量内存。如图 16 所示,我们开发了 FusedSoftmaxCrossEntropy,它将多个操作融合成一个内核,接收 BF16 输入,但在内部以 FP32 精度进行在线计算,从而节省了时间和内存。

    图 16. 融合的 SoftmaxCrossEntropy
    图 16. 融合的 SoftmaxCrossEntropy

A4 实验环境

  • 硬件配置: 实验在一个拥有超过 12,000 个 GPU 的大规模生产集群上进行。具体的 GPU 型号和数量因商业机密而未透露。
  • 软件配置: 系统基于 MegaScale【18, MegaScale: scaling large language model training to more than 10,000 GPUs, 2024, NSDI’24】框架实现,这是一个性能优于 DeepSpeed 和 Megatron-LM 的生产级 LLM 训练框架。
  • 模型架构:

    • 密集模型: LLaMA 系列,包括 LLaMA-7B, 13B, 30B, 70B。
    • 稀疏模型 (MoE): Mistral 系列,包括 Mistral-8x7B 和 Mistral-8x22B。
    • 具体模型参数见表 1。

    表 1. 用于评估的模型
    表 1. 用于评估的模型

    * 数据集:
    * GitHub: 开源数据集,长序列比例较低。
    * Byted: 生产数据集,长序列比例较高。
    * 数据分布如图 4 所示。
    * 基线:
    1. MegaScale + 静态并行: 使用 DP, TP, PP, CP 策略,并对打包序列应用了优化的 dist-attn(图 14)。
    2. MegaScale + 朴素 HDP: 只应用通信优化(算法 1)。
    3. MegaScale + 平衡 HDP (ByteScale): 同时应用通信和平衡优化(算法 2)。

A4 实验结果

8.2 端到端评估

  • 总体性能: 如图 17 所示,朴素 HDP 和平衡 HDP (ByteScale) 均优于基线,ByteScale 最高取得了 7.89 倍的加速比。
  • 可扩展性差异: 随着上下文长度增加,基线的吞吐量几乎在上下文长度翻倍时减半。相比之下,朴素 HDP 的吞吐量平均下降 1.23 倍,而平衡 HDP (ByteScale) 仅下降 1.08 倍。这是因为 HDP 极大地减少了冗余通信,而平衡策略消除了因计算不平衡造成的空闲时间。
  • 数据集差异: 在长序列更多的 Byted 数据集上,吞吐量和加速比低于 GitHub 数据集,但 ByteScale 依然能达到最高 4.26 倍的加速比。
  • 并行策略差异: 使用 DP-Balance 策略的模型(如 LLaMA-7B/13B/30B)比使用 PP-Balance 策略的模型(如 LLaMA-70B)获得了更高的加速比(6.21x-7.89x vs 3.42x-4.28x)。这是因为 DP-Balance 只需平衡每个时间步的计算,比 PP-Balance 平衡所有时间步的计算更容易。

图 17. 端到端评估(单位:令牌/秒)
图 17. 端到端评估(单位:令牌/秒)

8.3 案例研究

  • 实验设置: 在 1024 个 GPU 上训练 LLaMA-7B,上下文长度 2M,使用 Byted 数据集。图 18 展示了一个训练步骤中不同 rank 的运行时状态。
  • 通信受限情况 (基线): 基线方法中,P2P 通信时间占总时间的 97.6%,导致执行时间几乎完全由通信决定。
  • 计算不平衡情况 (朴素 HDP): 朴素 HDP 通过按需分片减少了通信开销,但 rank 间存在严重不平衡。如图 18(b) 所示,最快的 rank 完成后需等待 171 秒,造成了大量空闲时间。
  • 平衡情况 (ByteScale): 平衡 HDP (ByteScale) 使得所有 rank 几乎同时完成执行。通过为执行快的 rank 分配更多任务,总步骤时间进一步缩短。
  • 总体比较: 朴素 HDP 将峰值执行时间减少了 1.7 倍,但 rank 间时间差异巨大。平衡 HDP 消除了这种差异,比朴素 HDP 进一步减少了 2.3 倍的执行时间。

图 18. 案例研究
图 18. 案例研究

8.4 消融研究

  • 动态通信的有效性: 如图 19 所示,应用朴素 HDP 后,峰值 RDMA 流量减半,张量核心利用率从 10% 提升到 40%,带来 1.59 倍的加速。但由于不平衡问题,硬件利用率不稳定。
  • 选择性卸载的有效性: 该技术减少了长序列所需的 rank 数量,使更多序列能被同时处理,将加速比从 1.59 倍提升到 2.01 倍。如图 21 所示,在 64K 上下文时,它能在不影响吞吐量的情况下节省 32.3% 的内存。
  • 平衡策略的有效性: 如图 19 所示,平衡策略稳定了 RDMA 流量,并使张量核心利用率持续保持在 40% 左右,实现了硬件的持续满负荷工作。这是提升最显著的部分,将加速比从 2.01 倍提升到 3.69 倍。
  • 远程数据加载器的有效性: 通过重叠数据读取和计算,将加速比从 3.69 倍进一步提升到 3.89 倍。

图 19. 网络流量和张量核心利用率
图 19. 网络流量和张量核心利用率

图 20. 消融研究
图 20. 消融研究

图 21. 激活值卸载的有效性
图 21. 激活值卸载的有效性

A5 结论

本文提出了 ByteScale,一个为大规模长短序列混合训练设计的高效、灵活且可扩展的分布式 LLM 训练框架。通过开发的通信优化器消除了冗余通信,并通过平衡调度器缓解了计算不平衡。在超过 12,000 个 GPU 的生产集群上,对从 7B 到 141B 的模型和从 256K 到 2M 的上下文长度进行了评估,实验结果显示 ByteScale 相较于 MegaScale 实现了高达 7.89 倍的性能提升。