作者/机构: 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)
本文针对大规模语言模型(LLM)长短序列混合训练场景,提出了一个名为 ByteScale 的高效、灵活且可扩展的训练框架。现有框架通常采用静态的通信网格(如二维网格)来组织设备,将数据并行(DP)和上下文并行(CP)视为正交的技术,但这导致在处理长度可变的序列时,出现冗余通信和计算不平衡的问题,从而降低了训练效率。
核心问题:
1. 冗余通信: 静态并行策略强制所有序列(无论长短)都使用为最长序列配置的上下文并行(CP)组进行分区和通信,即使短序列并不需要跨设备分区,也产生了不必要的通信开销。
2. 计算不平衡: 尽管通过打包(packing)技术可以在设备间均匀分配令牌数量,但由于自注意力机制的计算复杂度为 $O(L^2)$($L$为序列长度),不同打包序列的实际计算负载(FLOPs)差异巨大,导致设备间出现计算不平衡,产生空闲等待时间(气泡)。
研究目标与创新点:
为了解决上述挑战,ByteScale 提出了以下核心贡献:
Transformer 架构【40, Attention is All you Need, 2017, NeurIPS 2017】是当前大语言模型(LLM)的主流基础架构。它由一系列 Transformer 层堆叠而成,每个层包含一个注意力模块和一个前馈网络(FFN)模块。如图 1 所示,自注意力机制需要在整个序列的所有令牌间进行计算以捕获上下文信息,而其他操作(如归一化、线性投影和激活函数)则是逐令牌计算的,每个令牌可以独立处理。
为了在静态并行策略中支持可变长度序列,需要使用填充(padding)和打包(packing)技术。如图 2 所示,填充将同一批次中的序列补齐到相同长度,但这会造成计算浪费。打包【22, Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance, 2021】则将多个序列拼接成一个长序列,不使用填充令牌,并采用一种特殊的分段注意力掩码(segmented attention mask)来确保每个序列被独立处理。
自注意力的时间和内存复杂度均为 $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 上。
观察 1:真实世界数据集中的序列长度呈偏态分布。如图 4 所示,我们分析了开源数据集 GitHub 和一个用于长上下文训练的生产数据集 Byted。两者都显示出序列长度的偏态分布。例如,在 Byted 数据集中,随机抽样的全局批次中近 80% 的样本长度小于等于 4K,而只有 0.05% 的样本能达到 2M。然而,从令牌分布来看,这 0.05% 的长样本(>=2M)贡献了全局批次中 12.1% 的令牌。GitHub 数据集中,超过 128K 的序列贡献了 16.2% 的令牌,显示出显著的数据异构性。
* 观察 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】也展示了在强化学习过程中,逐步增加和多样化的响应长度有助于提升模型性能。
* 挑战:数据异构性导致效率下降。虽然混合训练有益,但现有系统的静态并行策略不适应动态工作负载,导致了冗余通信和计算不平衡问题。
不平衡的 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 存在巨大差异。
* 不平衡数据与流水线并行。微批次间执行时间的不平衡进一步降低了数据并行和流水线并行的效率。在数据并行中,所有 DP rank 必须执行相同数量的微批次,然后同步梯度。如图 3(c) 所示,计算量较小的 rank-2 会提前完成,进入空闲等待(DP 气泡)。在流水线并行中,除了固有的 PP 气泡外,微批次间的 FLOPs 不平衡会导致各阶段执行时间无法理想重叠,产生额外的 PP 气泡(级间等待),如图 5 所示。更严重的是,每个微批次需在 $P_{pp}$ 个阶段上顺序执行,任何 DP 气泡都会被放大 $P_{pp}$ 倍。如图 5 所示,流水线(b)在执行完 8 个微批次后,会长时间空闲等待流水线(a),导致 DP 气泡占总执行时间的 30% 以上。
为了应对上述挑战,我们提出了 ByteScale。如图 7 所示,它包含三个主要组件:
1. Profiler (分析器): 用于分析环境、模型配置、数据分布,并为其他组件构建成本模型。
2. Communication Optimizer (通信优化器): 通过数据感知分片、动态通信和选择性卸载,来提高长短序列的通信效率。
3. Balance Scheduler (平衡调度器): 通过并行感知的数据分配,来解决计算不平衡问题。
本节描述 ByteScale 如何优化通信开销。首先,通过动态序列分片和通信减少短序列的冗余通信。其次,通过选择性卸载进一步压缩长序列的通信成本。
混合数据并行 (Hybrid Data Parallelism, HDP)。我们引入了一种新的并行策略 HDP,以实现对不同长度序列的高效训练。DP 和 CP 都将训练数据划分到设备上,但 DP 是数据间划分,CP 是数据内划分。HDP 统一了这两种划分方式,其目标是均匀地将令牌分布到设备上。HDP 可以替代传统的 DP 和 CP,其并行度 $D_{hdp} = D_{dp} \times D_{cp}$。与 DP 和 CP 不同,HDP 允许其 ranks 之间存在异构行为,具有两个关键特性:
* 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}$ 可表示为:
由于参数是复制的,令牌被均匀分布到 HDP ranks (R) 上,每个 rank $r$ 上的局部累积梯度对应于分配给它的令牌子集 $B_r$ 的梯度之和。因此,与 DP 类似,在所有 HDP ranks 间执行一次全局集合通信(如 All-Reduce 或 Reduce-Scatter)来聚合部分梯度,同样可以得到来自所有令牌的梯度 $G_{\theta}$:
公式(2)等价于公式(1),确保了 HDP 中的梯度累积结果与标准 DP 在数学上等价。我们通过令牌级损失(token-level loss)来实现这一点,即用总令牌数而不是样本数来缩放损失。
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()
offload_ratio。这种方法有效地将长序列所需的 rank 数量从 $L_i/T$ 压缩到 $C(S_i)$,如图 11(a) 所示,显著减少了通信开销,并使更多 HDP rank 可用于处理数据,从而提高效率。ByteScale 的整体流程如算法 1 所示。简而言之,算法遍历全局批次中的每个序列 $S_i$。对于长序列,它推导出卸载比例 $\alpha$ 和所需的 rank 数量 $C(S_i)$(1-6行)。对于短序列,它将它们打包以填满每个 rank 的容量 $T$(7-9行)。处理后的序列被分配给 $D_{hdp}$ 个 rank,算法返回每个 rank 的微批次和 offload_ratio 用于执行(10-12行)。
本节介绍平衡调度器如何解决 DP 和 PP 的不平衡问题。通过精心设计数据分配(替代算法 1 的第 10 行),它在保持最小通信开销的同时缓解了这些不平衡。
梯度累积要求不同 DP rank 执行相同数量的微批次,这是基于所有微批次计算负载相同的假设。然而,实际执行时间差异很大。在 ByteScale 中,我们重新定义了一个更灵活的策略,允许不同的 HDP rank 处理不同数量的微批次(大小相同但工作负载不同),以缓解不平衡问题。如图 13 所示,这使得所有 rank 能在同一时间完成计算。更重要的是,该策略不影响模型收敛,因为我们最终计算的是全局批次中所有令牌的梯度总和,保证了数学等价性。
为执行时间较短的流水线分配更多微批次。如图 12(a)-(b) 所示,pipeline-0 处理平均执行时间较长的微批次,因此只分配了 8 个微批次。相比之下,pipeline-1 被分配了 18 个微批次,以与 pipeline-0 同步完成。此外,由于微批次更多,其气泡率也进一步降低。
算法 2 描述了平衡策略。
1. 首先,按长度降序对全局批次 B 中的序列进行排序。然后将这些序列划分为 FLOPs 总和近似相等的桶(buckets),因此平均长度较长的桶包含的序列较少(3-5行)。
2. 其次,确定哪些 rank 的执行时间较短,以便后续分配(7-9行)。
3. 第三,如果使用 DP-Balance 策略,则从同一个桶中选择序列;如果使用 PP-Balance 策略,则从所有桶中顺序选择序列。实际上,执行时间较短的 rank 会被分配更多序列(12-15行)。
4. 最后,重复第二和第三步,直到所有桶都为空。
ByteScale 基于 Python, C++ 和 CUDA 实现,代码约 16K 行,并已集成到高性能 LLM 训练框架 MegaScale【18, MegaScale: scaling large language model training to more than 10,000 GPUs, 2024, NSDI’24】中。
带打包的 Dist-attn: 为避免打包序列在 CP 组内造成异构计算和通信,我们优化了 dist-attn。如图 14 所示,我们将打包序列的每个子序列均匀划分为 $2C$ 份,并对称地分配给 $C$ 个设备。这确保了每个设备持有所有子序列的 $1/C$ 部分,并覆盖 $1/C$ 的注意力掩码区域,所有设备参与相同的环形 P2P 通信,数据交换量也相同。
* 远程数据加载器 (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 进程根据计划从服务器读取部分数据)。
* 融合的 SoftmaxCrossEntropy: 现代 LLM 的词汇表很大,计算 SoftmaxCrossEntropyLoss 前将 logits 从 BF16 转换为 FP32 会消耗大量内存。如图 16 所示,我们开发了 FusedSoftmaxCrossEntropy,它将多个操作融合成一个内核,接收 BF16 输入,但在内部以 FP32 精度进行在线计算,从而节省了时间和内存。
模型架构:
* 数据集:
* GitHub: 开源数据集,长序列比例较低。
* Byted: 生产数据集,长序列比例较高。
* 数据分布如图 4 所示。
* 基线:
1. MegaScale + 静态并行: 使用 DP, TP, PP, CP 策略,并对打包序列应用了优化的 dist-attn(图 14)。
2. MegaScale + 朴素 HDP: 只应用通信优化(算法 1)。
3. MegaScale + 平衡 HDP (ByteScale): 同时应用通信和平衡优化(算法 2)。
本文提出了 ByteScale,一个为大规模长短序列混合训练设计的高效、灵活且可扩展的分布式 LLM 训练框架。通过开发的通信优化器消除了冗余通信,并通过平衡调度器缓解了计算不平衡。在超过 12,000 个 GPU 的生产集群上,对从 7B 到 141B 的模型和从 256K 到 2M 的上下文长度进行了评估,实验结果显示 ByteScale 相较于 MegaScale 实现了高达 7.89 倍的性能提升。