提升大模型训练 MFU:字节“拖后腿”现象分析和归因

发布于 2025-5-20 06:24
浏览
0收藏

一、背景

在之前的系列文章中,笔者已经系统性地介绍过大规模 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 的理想场景,并与实际情况对比。通过这种方式作者探讨了以下问题:

  1. Straggler 对于训练任务的影响频率及对训练任务的具体影响。
  2. Straggler 能否表现出时间或空间上的规律性。
  3. 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 的进度。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

四、方法

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 所示。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

为了降低追踪大量小型 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)操作组的依赖模型规定,任一单独操作的数据传输均需在所有操作启动后方能开始。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

模拟一个替代 Timeline。根据依赖模型和理想时长,可以通过以下规则模拟另一种执行 Timeline:

  • 模拟器在所有依赖操作完成后立即启动下一项操作。换言之,操作的开始时间为其依赖的上一个操作的最晚结束时间。
  • 计算操作一经启动即视为完成,其结束时间为开始时间 + OpDuration 中对应的操作时长。
  • 对于每个通信操作,模拟器会等待同一集合通信组(或 P2P)中的所有对等操作都启动。操作的结束时间为该组中最晚启动时间 + OpDuration 中存储的相应“传输时长”。

4.3 与 Straggler 相关的降速和 GPU 浪费的指标

模拟器可以估算在没有 Straggler 干扰的替代 Timeline 的 JCT,那么应该计算哪些指标来量化 Straggler 的影响呢?作者用 Tideal 表示无 Straggler 的 JCT。为了考虑模拟过程中引入的误差,也使用未修改的操作时长对原始 Timeline 进行模拟,并将得到的 JCT 记为 T。模拟误差相对较小。此外,部分追踪数据存在较大的模拟误差,为了确保分析的准确性,作者剔除了模拟误差 >= 5% 的追踪数据。为了量化与 Straggle 任务相关的性能下降程度,作者计算慢速指标 S,即比例:

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

除了与整体 Straggle 相关的性能下降外,还希望量化由不同类型操作中的 Straggle 导致的性能下降(如 Table  1)。为此,作者首先为每种操作类型 t 计算其操作类型降速值 St,具体方法如下:

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

其中,Tideal 是按上述方法计算出的理想 JCT,而 T−tideal 表示当操作类型 t 的 OpDuration 中的元素未被固定时的 JCT。

在作者的集群中,一个作业在其整个运行期间独占 GPU 资源。因此,JCT 的增加(或下降)可以直接转化为该作业所浪费的 GPU 小时数。具体来说,通过以下公式估算作业的资源浪费比例,即浪费的 GPU 小时百分比:

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

同样地,可以通过计算 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。在大多数情况下,速度缓慢的操作集中在计算环节,而非通信过程。基于此,作者认为服务器问题(可能是硬件故障或配置错误)通常是导致这些问题的罪魁祸首。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

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 任务就足以实现成本效益最大化。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

5.3 哪些操作导致了 Straggle

如下图 Figure 5 展示了追踪数据中所有作业按操作类型划分的资源浪费比例。与 FALCON 的研究结论不同,作者发现多数减速源于计算操作而非通信环节。这得益于作者集群中充足的网络带宽、专用 LLM 训练集群的使用以及多项内部网络优化与调参措施。

Figure 5 同时显示,PP-level 通信的影响略高于 DP-level 通信操作。与预期相符,因为 DP-level 通信存在大量 Overlap,相比 PP-level 通信 Kernel(多发生于训练 Step 关键路径上的预热和冷却阶段)能容忍更多降速。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

5.4 作业规模与 Straggle 有何关联

令人惊讶的是,并未观察到任务降速与任务规模之间存在明显的正相关关系。这表明,任务规模并非导致滞后的决定性因素,而模型类型或人为因素等其他因素可能扮演着更为关键的角色。例如,超大规模任务通常由值班团队精心维护,并能比其他任务获得更好的优化,因此它们的降速情况未必比小型任务更糟。再比如,长上下文任务更容易观察到受到 Straggle 影响,但这些任务通常规模较小,当与其他模型混合时,这种偏差反而使结果呈现出任务规模与降速之间的反向关联。

六、根本原因

作者重点针对那些执行速度明显减慢的作业进行研究,即速度降幅 S ≥ 1.1 的情况。由于确定作业降速的根本原因需手动检查,因此作者并没有探究所有可能的成因,而是聚焦于常见诱因。

6.1 是否应该归因到个别 Worker

作者首先分析了有多少 Straggle 是由于 Worker 的硬件或软件问题所致。由于作者执行了健康检查,预计仅有少数节点容易出现问题。因此,若某项任务因少数 Worker 的问题导致降速,则修正这些问题 Worker 上的操作的执行时间(将其设置为理想执行时间),足以优化整个任务的完成时间。

作者基于这一观察来衡量问题 Worker 对落后任务的影响。采用与之前相同的方法,通过估算操作降速百分比来计算每个 Worker 造成的降速,具体而言,定义 Worker w 的降速 Sw 为:

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

其中,T−wideal 是仅修复非 Worker w 上执行的操作时的 JCT(由模拟器得出),而 Tideal 是所有操作均被修复时的 JCT。

接下来,针对每项任务,筛选出速度减慢程度 Sw 位居该任务前 3% 的 Worker 集合 W。若少数 Worker 因硬件(或软件)问题表现不佳,那么 W 将包含这些 Worker。随后,计算这些 Worker 对任务整体降速所贡献的比例 MW。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

其中,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。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

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。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

接下来,作者进一步探讨了能否缓解这一问题。采用与 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 的训练任务的序列长度分布情况,可以看出有比较明显的长尾现象:

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

序列长度的差异之所以成为问题,主要是 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 问题。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

上述分析基于计算时间与 O(Σsi²) 成正比的观察。如下图 Figure 9 中通过实验验证了这一假设:针对典型任务的前几个训练 Step,绘制了每个 Micro-Batch 的处理时长与 Σsi² 的关系图,结果证实二者确实存在正比关系。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

接下来,作者量化了因序列长度不平衡而导致作业降速的比例。由于追踪数据未能提供足够信息来修正序列长度不平衡的问题,故无法沿用之前的度量标准。因此采用一种 Forward-Backward 相关性指标来衡量序列长度失衡引发的降速现象。该指标基于以下观察:若某 Micro-Batch 的 Forward 计算因序列长度不均而变慢,则其 Backward 同样会以相近程度放缓,且 Forward-Backward 计算时间应存在关联性。如下图 Figure 11 展示了每个 Straggle 作业(即 S ≥ 1.1 的作业)中某 PP Stage 的皮尔逊相关系数。经验表明,相关系数 ≥ 0.9 的作业极可能因序列长度失衡而降速。以此阈值为界,发现 21.4% 的作业受到序列长度不平衡影响,其平均降速程度达 1.34 倍。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

此外,如下图 Figure 12 中,作者分析了最大序列长度的变化如何影响作业降速。随着最大序列长度的增长,序列长度不平衡的影响更为显著。然而,上下文长度也在不断增加,因此解决这一可扩展性挑战变得愈发重要。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

解决这一问题的思路也比较简单,因为通过序列长度可以准确预测 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 愈发严重。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

因此,作者工程团队实施了一项有计划的 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 内存碎片化和虚假内存依赖。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

七、在线检测 Straggler

作者还开发了一个名为 SMon 的 Online 服务。该服务在每次 NdTimeline 性能分析会话(记录数十个训练 Step)后自动运行。SMon 能够提供热力图来呈现 Worker 节点的降速情况:

  • 每个单元格表示一个 Worker,其 x 和 y 轴对应 DP 和 PP 的 Rank,颜色深浅表示降速程度。
  • 可以简化识别 Straggle Worker 的过长;也可以帮助识别降速的根源。
  • 如下图 Figure 14 所示,Worker 问题、PP Stage 划分不均衡以及序列长度不均导致的降速展现出了各自独特的模式。

提升大模型训练 MFU:字节“拖后腿”现象分析和归因-AI.x社区

八、参考链接

  1. ​https://www.arxiv.org/abs/2505.05713​
  2. ​https://github.com/ByteDance-Seed/StragglerAnalysis​
  3. ​https://github.com/volcengine/veScale​

本文转载自​​​​​​AI闲谈​​​​​​,作者:AI闲谈

收藏
回复
举报
回复
相关推荐