
提升大模型训练 MFU:字节“拖后腿”现象分析和归因
一、背景
在之前的系列文章中,笔者已经系统性地介绍过大规模 LLM 训练面临的各种挑战以及可能涉及的问题和解决方案。在对大规模任务进行 Profiling 分析的时候,面对成千上万的 kernel 也经常苦不堪言,想要通过统计分析来诊断相应的问题,并为优化提供更多的可能性。碰巧看到了字节跳动 Seed 的这篇文章,虽然社区内没有看到太多讨论,不过其确实与我们的一些思路不谋而合,这里进行简单介绍。
其实文章中的大部分结论性内容笔者在之前的文章中都已总结过,比如:
- 大规模 GPU 集群运维实践:假装万卡 GPU 集群经验中,我们提到过硬件(GPU 降频,PCIe 链路降级)导致的训练任务降速;以及软件(Python 垃圾回收)导致的周期性降速等。
- 大模型训练中的负载均衡:挑战与解决方案解析中,我们也提到了一些因为负载不均衡导致无法充分发挥算力的问题,包括:
流水线并行(PP) Stage 的负载不均衡导致的问题和解决方案。
MoE 中专家负载不均衡导致的问题和方案。
长序列负载不均衡问题和优化方案。
多模态和长序列场景中样本序列长度不均衡导致的问题和改进方案。
一些碎碎念:最近停更了将近半个月,一方面是确实太忙,没有太多精力;另一方面是技术迭代太快,看了很多东西都没来得及记录。比如,笔者不久之前还推测“不久的未来一定会有 30B 规模 Dense 模型 Reasoning 能力接近或超越 DeepSeek R1”,没想到这个“不久”这么快到来;在介绍 Reasoning 相关工作时也提到,混合 Reasoning 和非 Reasoning 模型必定会出现,没想到现在基本已是标配;最后,LLaMA-4、Qwen3 等模型相继推出,Google Gemini 也屡创新高,期待 DeepSeek R2 能再创辉煌。
二、摘要
LLM 训练是当前对分布式计算要求最高的任务之一,通常需要成千上万的 GPU,并且需要频繁进行跨机器同步。这种工作模式容易受到 Straggler 的影响(即,训练可能因少数慢 Worker 而变慢)。在字节跳动的场景中,作者发现 Straggler 并非总是简单的硬件故障引起,可能源于多种复杂因素。
本文研究旨在通过对字节跳动 LLM 训练集群 5 个月的追踪数据,全面探讨 LLM 训练中的 Straggler 问题。核心方法是采用假设分析,模拟无 Straggler 的理想场景,并与实际情况对比。通过这种方式作者探讨了以下问题:
- Straggler 对于训练任务的影响频率及对训练任务的具体影响。
- Straggler 能否表现出时间或空间上的规律性。
- Straggler 现在的潜在根源是什么。
三、引言
3.1 问题背景
LLM 训练中,Straggler 的影响与集群中训练任务采用的分布式并行策略密切相关。典型的 LLM 训练会采用混合并行策略,包括 DP(Data Parallelism)、PP(Pipeline Parallelism)、TP(Tensor Parallelism);在 MoE 模型中,也通常会采用 EP(Expert Parallelism);在长序列场景还会采用 SP(Sequence Parallelism)或 CP(Context Parallelism)。所有这些策略都需要频繁协调以在 GPU 和服务器间同步结果,因而容易受到 Straggler 的影响。
为了缓解 Straggler 问题,常见的手段如下所示,通常比较少采用:
- 使用冗余的备份节点,资源浪费较严重。
- 异步 SGD,数学不等价,精度有损。
- 丢弃慢 Worker 的更新,数学不等价,精度有损。
鉴于 LLM 训练任务容易受 Straggler 影响,作者分析了字节跳动 LLM 训练集群 2024 年 1 月到 5 月期间收集的追踪数据,以便进行分析。
3.2 混合并行训练和 Straggler
Straggler 定义:若所有 Worker 完成分配任务所需时间相同,则采用混合并行技术的 LLM 训练任务被视为无 Straggler 现象。这可以最大限度减少了同步所需时间,实现理想的无 Straggler 场景,达到最佳性能。基于这一定义,任何导致 Worker 进度滞后的问题均被视为 Straggle 现象。不仅包括针对单一 Worker 的硬件故障,也包括影响所有 Worker 的不可预测停滞(比如垃圾回收),以及数据不均等问题导致 Workload 不均。
DP/ZeRO/FSDP:将训练数据分割至多个 Worker,每个 Worker 均持有整个模型的副本。
- 采用 DP 时,每个 Step 各 Worker 被分配一个训练 Batch。该 Worker 执行其 Batch 的 Forward 与 Backward 计算,随后所有 Worker 在进入下一 Step 训练前进行梯度 AllReduce 操作。此梯度 AllReduce 步骤要求 Worker 间同步,任一 Worker 的延迟都可能导致整体停滞,形成拖尾效应。
- ZeRO 和 FSDP 对 DP 进行了扩展,通过跨 Worker 切分优化器状态、参数和/或梯度来降低单 GPU 内存需求。与梯度 AllReduce 不同,ZeRO 和 FSDP 需执行梯度计算的 ReduceScatter 步骤、Device-Local 参数更新步骤,以及参数 AllGather 步骤来完成每次训练迭代。其中 ReduceScatter 与 AllGather 操作均需跨 Worker 同步,因此同样易受 Straggler Worker 影响。
PP:将模型切分至多个 Worker,每个 Worker 持有模型连续层的不相交子集,称为 PP Stage。PP 降低了模型权重和激活值对单块 GPU 内存的需求。训练过程中,一个 Batch 的数据被划分为若干 Micro-Batch,通过 PP Stage 进行流水线式训练。已有多种 Micro-Batch 调度策略,如 GPipe、1F1B 及 VPP(Virtual PP)。这些调度策略均假设计算在各流水线阶段间均匀分配,旨在最小化 PP Bubble —— 即某 Stage 因等待前一 Stage 数据而空闲的时间。若 PP Stage 分配不均,最慢 Stage 将阻碍其他阶段,形成性能瓶颈。因此,PP 易因 Stage 间计算分配不均而受拖累。
TP 和 CP:除 PP 外,还可采用 TP 和 CP,进一步降低单 GPU 内存需求。TP 在各 Worker 间划分每层权重,CP 则在节点间分割序列 Token。二者均需在每层 Transformer 后执行同步步骤,以聚合 TP 或 CP 组内所有 Worker 的部分计算结果。由于同步时慢速设备会拖累整体进度,TP 与 CP 同样易受 Straggler Worker 影响。本文不分析 TP 和 CP 组内的系统性落后 Worker 问题。
在实际训练中,通常会采用混合并行策略,其性能优于任何单一策略。采用混合并行策略时,Worker 可组织为一个超立方结构,每个维度对应一种并行策略。如下图 Figure 1 所示,展示了一种 DP-PP-TP 的混合并行策略。
- 一个速度较慢的 TP Worker 会拖慢所属的 PP Rank,进而产生 PP 性能瓶颈。
- 这个 PP 瓶颈又会进一步减缓该 Worker 所属的整个 DP Rank,延后梯度同步,并阻碍其他 DP Rank 的进度。
四、方法
4.1 LLM 训练任务追踪
集群配置:收集追踪数据的集群专用于训练,并由多个团队内部共享。集群中的机器采用与 NVIDIA DGX 服务器类似的硬件配置:每台服务器 8 个 GPU,通过 NVLink 或 PCIe 链路互连,包括 4 或 8 个百 Gbps NIC、一个专用于存储和管理的独立 NIC、数百个 CPU Core 及数 TB 内存。服务器间通过高性能交换机以三层 CLOS 拓扑结构互联。网络无收敛,并且经过精心调优,确保不会因网络拥塞导致性能下降。
任务调度:在集群上可同时调度多个任务,每个任务在其执行期间独占分配的 GPU 资源。调度器确保每项任务使用相同类型的 GPU,以实现硬件配置的一致性。此外,调度器会以最优方式分配 GPU,保证任务所需 GPU 在网络拓扑中位置相邻。由于大型任务以 8 的倍数申请 GPU,不同大规模任务不会共用同一台服务器。加之网络无拥塞,这意味着即便任务运行于同一集群,也不会因资源争抢出现 Straggler 现象。
数据采集:本研究所用追踪数据采集自 2024 年 1 月 1 日至 2024 年 5 月 31 日期间提交的 LLM 预训练任务。鉴于研究聚焦大规模任务,仅分析至少使用 128 块 GPU 的任务,并根据规则剔除无效或不适合分析的数据,最终得到 3079 个有效任务样本。这些任务包含采用 Dense 架构和 MoE 架构的模型,分别配置了短序列或长序列上下文训练。其中 31.7% 的任务使用 ≥256 块GPU,18.3% 使用 ≥ 512 块 GPU,3.6% 使用 ≥ 5000 块 GPU。总体而言,这些分析样本覆盖了 LLM 训练任务总 GPU 时长的一半左右。
其中所有任务均采用开源 Megatron-LM 的定制版本完成。使用的 Megatron-LM 框架已经集成自研的性能分析工具 NDTimeline(GitHub - volcengine/veScale: A PyTorch Native LLM Training Framework [3]),默认情况下,该工具会对任务中 10% 的训练 Step 进行采样分析。对于每个被分析的 Step,工具会记录一系列重要操作的开始和结束时间,这些操作包括 Forward 和 Backward 计算以及通信操作。详细如下图 Table 1 所示。
为了降低追踪大量小型 Kernel 的成本,分析中的 Forward 和 Backward 操作通常包括多个 GPU Kernel。此外,NDTimeline 会定期同步任务中所有机器的时钟,这使得能够对齐不同机器上的相关操作,以便进行假设分析。
在追踪记录中的每个操作,其日志包括操作类型、开始和结束时间戳,以及一组元数据,如训练 Step ID、Micro-Batch ID、PP Rank 和 DP Rank。这些元数据可以帮助重建操作间的依赖关系,对于模拟无 Straggler 情况下的替代 Timeline 至关重要。
4.2 用于假设分析的模拟器
假设分析的目标是通过回答以下问题来评估 Straggler 的影响:
- 如果所有 Straggler 都不存在,任务会耗时多久?
- 如果除特定一组 Straggler 外其他都不存在,任务又会耗时多久?
为了解答这些问题,作者模拟了一个不存在 Straggler 的替代 Timeline。核心观点在于,在没有 Straggler 的情况下,相似操作的耗时应当是相同的。基于这一观点,作者首先尝试估算在无 Straggler 的理想情境下每个操作的持续时间。然后,模拟这样一个替代时间线:操作按其依赖关系启动,并在估算的理想时长内完成。通过比较模拟作业完成时间(JCT)与实际追踪记录中的时间,便能评估 Straggler 现象对整体效率的影响。
估计无 Straggle 的理想操作时长。从概念上讲,可以将追踪到的操作组织成一个四维张量,维度分别为:训练 Step、Micro-Batch、PP Rank 和 DP Rank。作者将此张量称为 OpDuration 张量。如上 Table 1 中的每种操作类型都有一个这样的张量。
- 对于计算操作,张量中的元素直接对应追踪到的操作时长。
- 对于通信操作,计算追踪时长中的一个子部分,称为传输时长。由于通信是作为集合操作(或 P2P)的一部分进行的,单个操作的追踪时长受两个因素影响:(1)数据传输到另一 Rank 所需的时间,即“传输时长”;(2)等待同一集合操作(或 P2P)中其他操作开始的时间,即“阻塞时长”。在这两个因素中,“传输时长”是集合操作固有的,而“阻塞时长”则由操作调度决定。因此,在 OpDuration 张量中,作者仅为通信操作类型(如 Params-Sync、Forward-Send 等)存储“传输时长”。为了估算一个操作的“传输时长”,作者取同一集合操作(或 P2P 操作)中所有对等操作的最大开始时间,并从该操作的结束时间中减去这一最大开始时间。
提取操作依赖关系。模拟器需要两项输入:理想化操作时长与操作依赖模型。接下来,阐述模拟器的依赖模型,该模型源自当前使用的基于 Megatron-LM 的训练系统。
在此依赖模型中,每个 Worker 运行若干“Stream”以执行其操作。同一 Stream 上安排的所有操作按序执行,而跨 Stream 的操作只要满足依赖关系便可并发执行。具体而言,每个 Worker 拥有:
- 一个执行全部 Forward 与 Backward 操作的 Stream
- 一个执行所有 DP 专用通信操作的 Stream
- 四个各执行不同类型 PP 专用通信操作的流:
Forward-recv(RF)
Backward-recv(RB)
Forward-send(SF)
Backward-send(SB)
各操作间的依赖关系如下图 Figure 2 所示:
- 相同 Stream 的依赖关系:Stream 内的操作根据其在追踪记录中的启动时间进行排序。假设相邻操作之间存在隐式依赖。
- DP 通信与计算依赖关系:在每个 PP Stage:
首个 Micro-Batch 的 Forward 操作需在完成对应的 params-sync 集合通信(以获取该 Stage 参数)之后进行,如下图 Figure 2 所示(Syncparams → CF, mid=1)。
这些参数会被本地缓存,供后续 Micro-Batch 使用。不同 Micro-Batch 计算出的梯度会在本地累积,然后在 DP Rank 间进行聚合。因此,最后一个 Micro-Batch 的 Backward 应在执行 grads-sync 集体通信之前完成,以便跨 DP Rank 聚合该 PP Stage 的梯度,如下图 Figure 2 所示(CB,mid=8 → Syncgrads)。
- PP 通信与计算依赖关系:
除首个 PP Rank 外,任何一个 Micro-Batch 在 PP Rank p 上的 Forward 与 Backward 操作,必须等待同一 GPU 上该 Micro-Batch 的 Forward-receive(RF) 与 Backward-receive(RB) 通信操作完成后才能启动,如下图 Figure 2 所示(例如,RF,mid=1 → CF,mid=1,RB,mid=7 → CB,mid=7)。
同理,除最后一个 PP Rank 外,任何 Micro-Batch 在 PP Rank p 上的 Forward-send(SF) 与 Backward-send(SB) 操作,必须待同一 GPU 上该 Micro-Batch 的 Forward 与 Backward 操作结束后方可开始,如下图 Figure 2 所示(例如,CF,mid=1 → SF,mid=1,CB,mid=7 → SB,mid=7)。
- 跨 Rank 通信依赖关系:对于给定的 Micro-Batch,其 DP 通信操作(即 params-sync 和 grads-sync)在所有具有相同 PP Rank 的 DP 进程间形成一个集合通信组。同理,该 Micro-Batch 的 PP send 与 receive 操作则在相邻的、具有相同 DP Rank 的 PP 进程之间形成配对关系。此类集合(或 P2P)操作组的依赖模型规定,任一单独操作的数据传输均需在所有操作启动后方能开始。
模拟一个替代 Timeline。根据依赖模型和理想时长,可以通过以下规则模拟另一种执行 Timeline:
- 模拟器在所有依赖操作完成后立即启动下一项操作。换言之,操作的开始时间为其依赖的上一个操作的最晚结束时间。
- 计算操作一经启动即视为完成,其结束时间为开始时间 + OpDuration 中对应的操作时长。
- 对于每个通信操作,模拟器会等待同一集合通信组(或 P2P)中的所有对等操作都启动。操作的结束时间为该组中最晚启动时间 + OpDuration 中存储的相应“传输时长”。
4.3 与 Straggler 相关的降速和 GPU 浪费的指标
模拟器可以估算在没有 Straggler 干扰的替代 Timeline 的 JCT,那么应该计算哪些指标来量化 Straggler 的影响呢?作者用 Tideal 表示无 Straggler 的 JCT。为了考虑模拟过程中引入的误差,也使用未修改的操作时长对原始 Timeline 进行模拟,并将得到的 JCT 记为 T。模拟误差相对较小。此外,部分追踪数据存在较大的模拟误差,为了确保分析的准确性,作者剔除了模拟误差 >= 5% 的追踪数据。为了量化与 Straggle 任务相关的性能下降程度,作者计算慢速指标 S,即比例:
除了与整体 Straggle 相关的性能下降外,还希望量化由不同类型操作中的 Straggle 导致的性能下降(如 Table 1)。为此,作者首先为每种操作类型 t 计算其操作类型降速值 St,具体方法如下:
其中,Tideal 是按上述方法计算出的理想 JCT,而 T−tideal 表示当操作类型 t 的 OpDuration 中的元素未被固定时的 JCT。
在作者的集群中,一个作业在其整个运行期间独占 GPU 资源。因此,JCT 的增加(或下降)可以直接转化为该作业所浪费的 GPU 小时数。具体来说,通过以下公式估算作业的资源浪费比例,即浪费的 GPU 小时百分比:
同样地,可以通过计算 1 − 1/St 来估算因不同操作类型导致的资源浪费情况。
五、Straggler 的影响
5.1 Straggler 普遍存在,并导致不可忽视的资源浪费
如下图 Figure 3 展示了所有作业中资源浪费百分比的累积分布函数(CDF)。可以看出,追踪的训练任务中,42.5% 存在 Straggler。更为严重的是,由于 Straggler 的存在,超过 10% 的作业至少浪费了分配到的 GPU 时间的 21.3%,而约 1% 的作业浪费了分配的 GPU 资源的 45% 以上。整体数据显示,由于 Straggler,追踪的所有作业中有 10.4% 的 GPU 时间被浪费。
此外,作者还研究了那些执行速度大幅下降(S > 3)的作业,发现这些均为大型作业,且问题往往源于不到 3% 的 Worker。在大多数情况下,速度缓慢的操作集中在计算环节,而非通信过程。基于此,作者认为服务器问题(可能是硬件故障或配置错误)通常是导致这些问题的罪魁祸首。
5.2 Step 在落后任务中表现出类似的降速现象
接下来,作者探究了任务降速究竟是少数极慢步骤还是多数步骤共同作用。作者将显著降速比 S > 1.1 的任务定义为 Straggle 任务。为此,作者将单个 Step 的降速比定义为该 Step 执行时间与理想 Step 执行时间(例如,Tideal/n 对应含 n 个训练步骤的任务)之比。
如下图 Figure 4 展示了经过任务整体降速比归一化后的单 Step 降速比 CDF。数据显示,中位数 Step 的归一化降速比为 1.0,即使 90 分位 Step 的归一化降速比也仅为 1.06。表明大多数步骤的降速程度与任务整体降速比相当,说明 Straggle 现象并非由临时环境因素导致,而是源于更持续性的问题。这一结果同时意味着,仅需采样少量训练 Step 来分析 Straggle 任务就足以实现成本效益最大化。
5.3 哪些操作导致了 Straggle
如下图 Figure 5 展示了追踪数据中所有作业按操作类型划分的资源浪费比例。与 FALCON 的研究结论不同,作者发现多数减速源于计算操作而非通信环节。这得益于作者集群中充足的网络带宽、专用 LLM 训练集群的使用以及多项内部网络优化与调参措施。
Figure 5 同时显示,PP-level 通信的影响略高于 DP-level 通信操作。与预期相符,因为 DP-level 通信存在大量 Overlap,相比 PP-level 通信 Kernel(多发生于训练 Step 关键路径上的预热和冷却阶段)能容忍更多降速。
5.4 作业规模与 Straggle 有何关联
令人惊讶的是,并未观察到任务降速与任务规模之间存在明显的正相关关系。这表明,任务规模并非导致滞后的决定性因素,而模型类型或人为因素等其他因素可能扮演着更为关键的角色。例如,超大规模任务通常由值班团队精心维护,并能比其他任务获得更好的优化,因此它们的降速情况未必比小型任务更糟。再比如,长上下文任务更容易观察到受到 Straggle 影响,但这些任务通常规模较小,当与其他模型混合时,这种偏差反而使结果呈现出任务规模与降速之间的反向关联。
六、根本原因
作者重点针对那些执行速度明显减慢的作业进行研究,即速度降幅 S ≥ 1.1 的情况。由于确定作业降速的根本原因需手动检查,因此作者并没有探究所有可能的成因,而是聚焦于常见诱因。
6.1 是否应该归因到个别 Worker
作者首先分析了有多少 Straggle 是由于 Worker 的硬件或软件问题所致。由于作者执行了健康检查,预计仅有少数节点容易出现问题。因此,若某项任务因少数 Worker 的问题导致降速,则修正这些问题 Worker 上的操作的执行时间(将其设置为理想执行时间),足以优化整个任务的完成时间。
作者基于这一观察来衡量问题 Worker 对落后任务的影响。采用与之前相同的方法,通过估算操作降速百分比来计算每个 Worker 造成的降速,具体而言,定义 Worker w 的降速 Sw 为:
其中,T−wideal 是仅修复非 Worker w 上执行的操作时的 JCT(由模拟器得出),而 Tideal 是所有操作均被修复时的 JCT。
接下来,针对每项任务,筛选出速度减慢程度 Sw 位居该任务前 3% 的 Worker 集合 W。若少数 Worker 因硬件(或软件)问题表现不佳,那么 W 将包含这些 Worker。随后,计算这些 Worker 对任务整体降速所贡献的比例 MW。
其中,T 代表模拟的原始 Step 时长(未修正任何 Straggle),TWideal 为仅修正选定 Worker 上运行的操作时的模拟 Step 时长,而 Tideal 则是所有 Straggle 被修正后的理想模拟 Step 时长。
对于大规模作业(涉及数千个 Worker),计算 MW 成本很高,每个 Worker w 需要运行数千次模拟以计算 Sw。因此,作者采用一种近似方法来扩展分析:不单独计算各 Worker 的降速程度,而是测量整个 DP Rank 和 PP Rank 的降速情况。随后,为每个 Worker 赋予其所属 DP 和 PP Rank 中较小的(即最小的)降速值(任何 Worker 必属于一个 DP Rank 和一个 PP Rank)。这一方法将所需模拟次数从 DP-Degree × PP-Degree 减少到 DP-Degree + PP-Degree,从而大幅降低计算复杂度。
如下图 Figure 6 展示了 Straggle 作业的 MW CDF。可以看出,仅有 1.7% 的 Straggle 作业中,Worker 问题导致了超过 50% 的显著时延。据此可以得出结论:在追踪的数据中,问题 Worker 并非大多数 Straggle 作业的主导因素。作者进一步调查了少数由问题 Worker 主导作业降速的情况,发现这些情况造成的降速程度显著更大:有问题 Worker 的作业降速比为 3.04,而平均降速比仅为 1.28。
6.2 PP Stage 不均衡
在作者的追踪分析中,发现 PP 最后一个 Stage 相比其他 Stage 的计算不均衡问题是导致 Straggle 现象的常见原因。最后一个 Stage 包含 LM Head,还有 Loss 计算,因此在模型切分时经常无意间造成最后一个 Stage 计算量突增的情况,进而引发性能瓶颈。
为了验证损失层的耗时问题,作者运行了一个包含 4 个 PP Stage 的任务:前 3 个 Stage 各 9 个 Transformer Layer,最后一个 Stage 额外多一个损失层(包括 LM Head)。最后一层的逻辑运算耗时是普通 Transformer 层的 9 倍以上,导致最后一个 Stage 的 Forward(Backward)速度比平均水平慢 2.07x(1.41x)。
采用与之前类似的量化方法,作者获得一些观测结果,如下图 Figure 7 所示,展示了不同任务中 MS(最后一个 PP Stage 的影响系数)的 CDF,可以看出,39.3% 的任务中主要降速(MS >= 0.5)源自最后一个 PP Stage。
接下来,作者进一步探讨了能否缓解这一问题。采用与 LLaMA 3 类似的方法:给最后一个 PP Stage 少分配 ε 个层数。然而,手动调整 ε 值存在多重挑战:
- 首先,划分模型时,必须将完整的 Transformer 层分配给 PP Stage,限制了各 Stage 间的计算分配灵活性;
- 其次,当词表增大或最大序列长度/隐藏层尺寸减小时,损失层耗时(相对于 Transformer 层)占比会上升。随着词表增加,这种现象将愈发普遍,这会导致需要在前置阶段执行更多 Transformer 层以平衡最后一个 PP Stage 的耗时。如此一来,可用 PP Stage 数量受限,每个 Virtual Stage 包含的层数超出理想值,制约了模型并行效率。因此,即便 ε 取值合理,仍可能导致性能未达最优。
针对前述包含 9 个 Transformer 层的任务,作者也尝试手动调整各 PP Stage 的层数分配,发现通过人工切分可获得 9.9% 的加速效果。但即便经过手动调优,各 Stage 计算负载仍不完全均衡 —— 例如调优后最后一个 PP Stage 的 Forward 计算耗时仍为其他 Stage 的 1.55x。
6.3 序列长度不均衡
作者又进一步分析了 Worker 速度慢、PP Stage 负载不均无法解释的任务。深入研究发现,对于长上下文任务,训练数据间序列长度的差异是导致任务降速的重要原因。如下图 Figure 10 展示了某个最大序列长度为 32K 的训练任务的序列长度分布情况,可以看出有比较明显的长尾现象:
序列长度的差异之所以成为问题,主要是 SelfAttention 的算法复杂度呈二次方关系。在训练中通常会采用 Sample Packing 的方式尽可能降低序列不均导致的问题,但是 SelfAttention 部分的不均衡依然无法避免,比如一个 32K 序列的 SelfAttention 计算量是 32 个 1K 序列计算量的 32 倍(预训练阶段通常也会采用 Sample Packing,但不会添加 Attention Mask,可以保障负载比较均衡,而长度扩展或 SFT 等阶段会添加 Attention Mask,因此会存在这个问题)。
不同 Micro-Batch 间计算时间的差异会导致 PP Stage 出现 Bubble,并使得 DP 的各个 Worker 完成时间不同步,如下图 Figure 8 所示,这两种情况共同造成了 Straggle 问题。
上述分析基于计算时间与 O(Σsi²) 成正比的观察。如下图 Figure 9 中通过实验验证了这一假设:针对典型任务的前几个训练 Step,绘制了每个 Micro-Batch 的处理时长与 Σsi² 的关系图,结果证实二者确实存在正比关系。
接下来,作者量化了因序列长度不平衡而导致作业降速的比例。由于追踪数据未能提供足够信息来修正序列长度不平衡的问题,故无法沿用之前的度量标准。因此采用一种 Forward-Backward 相关性指标来衡量序列长度失衡引发的降速现象。该指标基于以下观察:若某 Micro-Batch 的 Forward 计算因序列长度不均而变慢,则其 Backward 同样会以相近程度放缓,且 Forward-Backward 计算时间应存在关联性。如下图 Figure 11 展示了每个 Straggle 作业(即 S ≥ 1.1 的作业)中某 PP Stage 的皮尔逊相关系数。经验表明,相关系数 ≥ 0.9 的作业极可能因序列长度失衡而降速。以此阈值为界,发现 21.4% 的作业受到序列长度不平衡影响,其平均降速程度达 1.34 倍。
此外,如下图 Figure 12 中,作者分析了最大序列长度的变化如何影响作业降速。随着最大序列长度的增长,序列长度不平衡的影响更为显著。然而,上下文长度也在不断增加,因此解决这一可扩展性挑战变得愈发重要。
解决这一问题的思路也比较简单,因为通过序列长度可以准确预测 Micro-Batch 的计算时间(如 Figure 9),因此可以在一个 Global Batch 进行数据的重排(数学等价),让不同的 DP Rank 中的计算尽可能均衡,这一目标可以通过贪心算法实现。在一个最大序列长度为 32K 的典型任务上测试此方法后,观察到吞吐量提升了 23.9%。
6.4 Python 自动垃圾回收
Python 的垃圾回收(GC)是导致降速的另一个主要原因。集群中的任务都使用 Python 运行,其 Runtime 在认为必要时会触发 GC。一旦触发,GC 可能耗费数百毫秒,在此期间,用户程序会被暂停,新 Kernel 无法启动,进而阻塞 Forward 计算。需要注意的是,Backward 不受影响,因为它们由 C++ 发起。
由于不同 Python 进程在不同时间触发 GC,该问题进一步恶化,如下图 Figure 13 所示,当 GC 暂停单个 Worker 进程时,整个训练作业都会阻塞,随着 Worker 进程数量的增加,此类暂停次数也随之增加。因此,随着模型规模的扩大,由 GC 引发的 Straggle 愈发严重。
因此,作者工程团队实施了一项有计划的 GC 优化方案,旨在减少因 GC 导致的 Straggle 问题。该优化措施关闭了 Python 的自动垃圾回收机制,改为根据用户设定的训练 Step 间隔手动触发 GC,确保所有 Worker 同步执行 GC。此措施显著降低了 GC 相关的 Straggle:在使用 128 个 DP Rank 的任务中,通过设定每 500 Step 执行一次 GC,性能提升了 12.6%。
然而如何确定合适的 GC 间隔也很有挑战,间隔过长可能导致内存耗尽而任务崩溃,间隔过短可能影响性能。由于不同任务的内存分配速率不同,合适的 GC 与具体任务密切相关,因此作者默认不开启,而让用户自主选择。
此外,作者还观察到,随着任务推进,GC 导致的暂停时间会逐渐延长,造成训练吞吐持续下降。推测原因是内存泄露导致堆内存持续增长,进而延长了 GC 暂停时间。需要说明的是,计划性 GC 能有效掩盖此类泄露的影响,维持训练吞吐的稳定。
6.5 其他根本原因
除了上述问题外,作者还探讨了其他可能得原因:
- CUDA 内存碎片化。偶发情况下,发现内存碎片化会显著拖慢 PyTorch 的 CUDA 内存分配性能,致使 cudaFree 与 cudaMalloc 调用频次激增,进而造成 Forward 或 Backward 计算操作异常缓慢。在数据追踪中,表现为同一 TP Group 内不同 TP Rank 上运行的通信 Kernel 启动时间参差不齐但却近乎同时完成,暗示部分通信 Kernel 存在启动延迟。通过调取 PyTorch 运行日志确认相关任务确实频繁调用内存管理接口后,作者启用了内存分配器追踪功能进行复现测试。重跑过程中观察到大量 segment_alloc 与 segment_free 调用,证实 PyTorch 内存分配器为症结所在。
- 虚假 Kernel 依赖问题。在初期训练大型 MoE 模型时,发现用于梯度同步的 ReduceScatter Kernel 会阻塞其他无依赖关系的 Kernel 启动,导致任务严重降速。推测成因在于无关 Kernel 共享同一 CUDA 硬件队列引发的虚假依赖。实践证明调高 CUDA_DEVICE_MAX_CONNECTIONS 参数可缓解此现象。但值得注意的是,该问题会随模型框架迭代时隐时现,作者仍在持续探究其深层机制。
6.6 总结:观察和启发
基于之前的根因分析可以得到如下结论:
- 极少数的 Straggle 任务是机器问题(无论硬件还是软件)引起,这意味着传统的健康指标不太可能有助于检测或预防大多数 Straggle 问题。
- 导致 Straggle 的最常见原因包括:PP Stage 分配不均、每个 Micro-Batch 中序列长度的不均衡,以及 Python GC 引起的暂停。
- 作者还发现两种不太常见的 Straggle 原因:PyTorch 内存碎片化和虚假内存依赖。
七、在线检测 Straggler
作者还开发了一个名为 SMon 的 Online 服务。该服务在每次 NdTimeline 性能分析会话(记录数十个训练 Step)后自动运行。SMon 能够提供热力图来呈现 Worker 节点的降速情况:
- 每个单元格表示一个 Worker,其 x 和 y 轴对应 DP 和 PP 的 Rank,颜色深浅表示降速程度。
- 可以简化识别 Straggle Worker 的过长;也可以帮助识别降速的根源。
- 如下图 Figure 14 所示,Worker 问题、PP Stage 划分不均衡以及序列长度不均导致的降速展现出了各自独特的模式。
八、参考链接
