Optimus: Accelerating Large-Scale Multi-Modal LLM Training by Bubble Exploitation

作者/机构: Weiqi Feng1∗, Yangrui Chen2, Shaoyu Wang3∗, Yanghua Peng2, Haibin Lin2 and Minlan Yu1 (1Harvard University, 2Bytedance, 3University of Southern California)

A1 主要贡献

本文针对多模态大语言模型（MLLM）训练中存在的效率低下问题，提出了一个名为 Optimus 的分布式训练系统，旨在通过利用训练过程中的 GPU 空闲时间（“气泡”）来缩短端到端的训练时间。

核心问题: 现有的分布式训练系统在训练 MLLM 时效率低下，主要原因是 MLLM 中异构的模态模型和复杂的3D并行数据依赖导致了大量的 GPU 气泡。在大型 MLLM 训练任务中，超过 40% 的 GPU 周期处于空闲状态，具体分析发现，通信开销和流水线阶段不平衡是主要原因，导致 GPU 闲置时间占比高达 48%。

研究目标: 设计一个能够有效减少 MLLM 训练时间的分布式系统。核心思想是将计算量较小的编码器（encoder）计算任务调度到计算量巨大的语言模型（LLM）骨干网络训练过程中产生的气泡中。

创新点与主要贡献:
为了实现上述目标，Optimus 解决了三个核心挑战，并提出了相应的解决方案：

1. 解决 GPU 无法执行编码器计算的问题: 传统框架采用统一的并行策略，导致大多数 GPU 仅存储 LLM 的模型状态，无法在 LLM 气泡期间执行编码器计算。

   • Optimus 的方案: 采用分离的并行化方案，为编码器和 LLM 分别设计并行策略，并将两者共同部署（colocate）到每个 GPU 上。这确保了所有 GPU 都具备执行编码器计算的能力。Optimus 会系统性地搜索编码器的3D并行方案，并剔除超出 GPU 内存限制的方案。

2. 解决复杂的 MLLM 数据依赖约束: MLLM 训练中存在多层次的复杂数据依赖，包括同步训练迭代间的依赖、编码器内部流水线依赖，以及最复杂的编码器与LLM之间的微批次（microbatch）级别依赖。

   • Optimus 的方案: 采用两阶段依赖管理策略。第一阶段是本地调度（local scheduling），在单个编码器流水线内部处理迭代依赖和内部依赖，将编码器计算调度到可用的 LLM 气泡中。第二阶段是全局排序（global ordering），通过比较时间戳来处理编码器与 LLM 之间的微批次级别依赖，确保对于每个微批次，编码器的前向传播在 LLM 前向传播开始前完成，且编码器的反向传播在 LLM 反向传播结束后才开始。

3. 解决亚毫秒级气泡的利用难题: LLM 训练中的气泡时长差异巨大，从几百毫秒到亚毫秒级不等。现有的层级调度方法无法利用过短的亚毫秒级气泡（如张量并行通信产生的气泡）。

   • Optimus 的方案: 将编码器层的计算分解为一系列的核函数（kernels），从而能够将这些细粒度的计算任务调度到短暂的气泡中。此外，通过分析 LLM 气泡的常见模式，优化调度策略，将编码器的计算核与 LLM 的计算交错执行，最大化地减少了总训练时间。

实验成果: Optimus 基于 Megatron-LM 实现。在生产集群的实验中，使用 ViT-22B 和 GPT-175B 模型在 3072 个 GPU 上进行训练时，与基线系统相比，Optimus 将 MLLM 的训练速度提升了 20.5%-21.3%。

A3 背景知识/关键Observation/设计原则

2.1 多模态LLM的特点

多模态LLM的重要性与架构。多模态LLM是建立在LLM基础上的重要进展，它们集成了先进的自然语言处理方法，并将能力扩展到多种数据模态。像GPT-4【【索引21，Gpt-4v(ision) system card，2023】】这样的模型在处理图像和文本输入方面表现出人类水平的性能。一个典型的多模态大语言模型（MLLM）通常由三个主要部分组成：一个或多个模态编码器、输入投影器和一个大型语言模型骨干网络【【索引36，Mm-llms: Recent advances in multimodal large language models，2024】】。模态编码器处理非文本模态的输入，将其转换为特征表示；输入投影器则将这些特征与文本特征空间对齐；最后，LLM骨干网络利用对齐后的多模态和文本特征作为输入。图1展示了MLLM的架构。由于输入投影器的计算需求与编码器和LLM骨干网络相比微不足道，本文后续讨论中将其视为模态编码器的最后一层。

多模态LLM的独特性质。与同构的LLM架构不同，多模态LLM具有以下独有特征：

• LLM骨干网络在模型规模上的主导地位：在多模态LLM中，LLM骨干网络的参数数量远超编码器和投影器等其他组件。例如，Flamingo【【索引3，Flamingo: a visual language model for few-shot learning，2022，Advances in neural information processing systems】】模型总共有800亿参数，其中LLM骨干网络就占了700亿。
• 编码器与LLM骨干网络之间的依赖关系：在MLLM训练中，编码器和LLM之间存在两种数据依赖。在前向传播过程中，编码器必须先完成编码特征的生成，LLM骨干网络才能继续进行前向计算。相反，在反向传播过程中，LLM骨干网络先计算梯度，然后编码器才开始反向传播。

2.2 MLLM训练中的气泡

现有LLM流水线优化的局限性。现有的LLM流水线优化缺乏模型无关性，不足以应对MLLM的训练任务。在我们内部使用ViT编码器和GPT骨干网络（超过1000亿参数）的大规模MLLM训练任务中，尽管采用了包括MegaScale【【索引13，Megascale: Scaling large language model training to more than 10,000 gpus，2024，arXiv】】、零气泡流水线（Zero Bubble Pipeline）【【索引23，Zero bubble pipeline parallelism，2023】】以及细粒度通信-计算重叠【【索引32，Overlap communication with dependent computation via decomposition in large deep learning models，2022，ASPLOS】】等多项SOTA技术，但在超过3000个NVIDIA GPU上使用Megatron-LM时，我们观察到超过48%的GPU周期处于空闲状态。通过分析性能剖析的时间线，我们识别并研究了GPU空闲（即气泡）的发生情况。表1显示了不同类型气泡占用的总时长及其在平均训练步长时间（5.12秒）中的百分比。

气泡的分类。根据其根本原因，这些气泡可以分为三类。

（1）数据并行（DP）中的通信。数据并行需要通过通信来聚合梯度，这导致了GPU在通信期间的空闲时间。具体来说，MegaScale【【索引13，Megascale: Scaling large language model training to more than 10,000 gpus，2024，arXiv】】和Megatron-LM【【索引25，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，CoRR】】采用了一种类似于ZeRO中Pos+g的分布式优化器【【索引24，Zero: Memory optimizations toward training trillion parameter models，2020，SC20】】，以节省训练大型模型时的内存，这会执行两种集体通信：all-gather和reduce-scatter。在每个训练步骤开始时，all-gather操作从所有DP ranks收集更新后的参数，产生一个DP all-gather气泡，占训练时间的3.3%。在训练步骤结束时，执行reduce-scatter来聚合梯度，产生一个DP reduce-scatter气泡，消耗了8.9%的训练时间。值得注意的是，尽管已经应用了Megascale【【索引13，Megascale: Scaling large language model training to more than 10,000 gpus，2024，arXiv】】中提出的数据并行重叠优化，但由于同步训练的性质，第一个模型块的上述DP通信无法被隐藏【【索引13，Megascale: Scaling large language model training to more than 10,000 gpus，2024，arXiv】】。图2展示了在每个训练步骤开始和结束时由all-gather和reduce-scatter操作产生的DP气泡。

（2）流水线并行（PP）中的依赖。尽管应用了Megascale【【索引13，Megascale: Scaling large language model training to more than 10,000 gpus，2024，arXiv】】的流水线发送-接收重叠优化，但由于前向和后向传播过程中各阶段之间固有的数据依赖性，流水线气泡仍然存在。重要的是，零气泡流水线方法无法消除MLLM训练中的流水线气泡，因为它需要对优化器进行更改【【索引23，Zero bubble pipeline parallelism，2023】】（详见第7节的讨论）。图2展示了MLLM训练流水线调度，包括三个阶段：预热（仅前向）、稳定（一前向一后向）和冷却（仅后向）。在整个流水线训练过程中，会产生三种类型的气泡：

（3）张量并行（TP）中的通信。张量并行涉及将单个层划分到多个GPU上，这需要在前向和后向传播过程中进行通信以同步GPU。在Megatron-LM中，一个transformer层的每次前向或后向传播需要两个all-gather和两个reduce-scatter核函数【【索引14，Reducing activation recomputation in large transformer models，2023，MLSys】】。图3详细展示了两个GPT-175B【【索引5，Language models are few-shot learners，2020，NeurIPS】】层前向传播过程中的CUDA计算和通信核函数。绿色核函数代表CUDA通信流中的all-gather通信，蓝色核函数代表reduce-scatter通信。在这些通信期间，计算流处于空闲状态。通常，这些TP气泡的持续时间为亚毫秒级，平均约为300微秒。然而，在MLLM训练期间，存在数千个TP气泡，它们共同占用了11.2%的训练时间。

2.3 挑战

核心观察与调度思想。为了最小化MLLM训练中的气泡，我们旨在利用MLLM独特的双组件结构，即编码器和LLM骨干网络。我们有两个关键观察：首先，MLLM训练期间的大多数气泡发生在LLM骨干网络的前向和后向传播过程中，根据生产数据，约90%的气泡源于LLM的通信。其次，由于编码器的参数规模较小，其所需的计算操作（FLOPs）远少于LLM骨干网络【【索引4，Qwen-vl: A versatile vision-language model for understanding, localization, text reading, and beyond，2023】，【索引7，Minigpt-v2: large language model as a unified interface for vision-language multi-task learning，2023】，【索引10，Palm-e: An embodied multimodal language model，2023】，【索引18，Visual instruction tuning，2023】】。基于这些见解，我们提出在LLM通信期间产生的LLM气泡中调度编码器计算，以最小化整个MLLM训练过程中的气泡。我们识别出将编码器计算调度到LLM气泡中的三个主要挑战。

挑战1：仅有少数GPU同时拥有编码器和LLM模型状态。当前的训练系统【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】，【索引39，Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning，2022，OSDI】】采用流水线并行将MLLM作为一个单一流水线分布在多个GPU上。由于编码器和LLM之间的依赖关系，编码器层被分配到初始的流水线阶段，而LLM层则被分配到后续的流水线阶段。因此，通常只有一个流水线阶段同时包含编码器和LLM层。如图4所示，使用3D并行（DP=1, PP=4, TP=2）在8个GPU上并行化MLLM时，只有流水线阶段1中的2个GPU同时持有编码器和LLM的模型状态。其余6个GPU由于缺乏编码器模型状态，无法在LLM气泡期间执行编码器计算。

挑战2：MLLM训练中的复杂依赖关系。MLLM训练中固有的复杂依赖关系在将编码器计算调度到LLM气泡中时带来了巨大挑战。首先，在同步训练中，LLM气泡的利用被限制为仅能在当前训练迭代内执行编码器计算（迭代依赖）。其次，编码器流水线内部的依赖要求，当前编码器流水线阶段i的前向计算只能在前序编码器阶段完成后进行调度，而后向计算只能在后续编码器阶段完成后进行调度。最后，编码器-LLM依赖施加了微批次级别的约束：对于微批次i，编码器必须在其前向传播完成后，LLM流水线才能开始该微批次的前向传播。类似地，只有在LLM流水线完成微批次i的后向传播后，编码器才能开始该微批次的反向传播。

挑战3：亚毫秒级的LLM气泡。现有的框架如MegaScale【【索引13，Megascale: Scaling large language model training to more than 10,000 gpus，2024，arXiv】】和Megatron-LM【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】通常在层级别进行调度。然而，LLM中的气泡持续时间范围很广，从亚毫秒（TP气泡）到数百毫秒（DP气泡）不等。例如，图3中所示的TP气泡在不同LLM层的前向和后向传播中平均持续约300微秒。这个时长不足以完成哪怕一个编码器层的前向或后向传播。举例来说，一个ViT-22B层通常需要约1.4毫秒来完成前向传播，2.0毫秒来完成反向传播。

A2 方法细节

3.1 设计决策

设计决策1：通过分离的并行策略共置编码器和LLM。为确保每个GPU都同时拥有编码器和LLM的模型状态，我们提出为编码器和LLM在所有GPU上分配独立的并行方案。如图5所示，编码器使用并行方案（DP=2, PP=2, TP=2），LLM使用（DP=1, PP=4, TP=2）。每个GPU都保留了编码器和LLM的模型状态，从而使得所有GPU在LLM气泡期间执行编码器计算成为可能。需要注意的是，共置编码器和LLM的状态可能会需要更多的GPU内存，我们将在4.5节分析内存开销。

设计决策2：双阶段依赖管理。我们采用两个阶段来处理MLLM训练中的复杂依赖关系：本地调度和全局排序。每个编码器流水线经历本地调度，该调度将编码器计算与可用的LLM气泡进行安排，同时遵守迭代依赖和编码器内部依赖。全局排序通过对编码器前向传播的结束时间和后向传播的开始时间在所有微批次间进行排序，来确保编码器和LLM之间的微批次级别依赖。这涉及到比较时间戳以验证编码器-LLM依赖是否得到满足。如图6所示，本地调度独立应用于两个编码器流水线，维持了迭代依赖和编码器内部依赖。在全局排序中，检查所有微批次（共8个）的时间戳，以确认编码器-LLM依赖得到满足。

设计决策3：在核函数级别调度编码器计算。将编码器层分解为核函数，可以高效利用亚毫秒级的气泡。然而，编码器层中的TP通信核函数在LLM的TP气泡期间会竞争链路带宽，导致每次迭代的时间变长。为了解决这个问题，我们必须额外地在LLM计算期间调度编码器通信核函数（见图7）。

3.2 Optimus 概览

系统架构。Optimus是一个为MLLM设计的分布式训练系统，它通过在LLM气泡中调度编码器计算来改善端到端训练延迟。为应对3.1节中概述的挑战，Optimus由两个组件构成：模型规划器（Model Planner），用于解决挑战1，确保所有GPU都持有编码器和LLM的模型状态；以及气泡调度器（Bubble Scheduler），用于解决挑战2（MLLM训练中的复杂依赖）和挑战3（亚毫秒级LLM气泡）。

模型规划器（Model Planner）。模型规划器将编码器和LLM骨干网络分别划分到所有给定的GPU上（解决§3.1中的挑战1）。它为LLM骨干网络选择一个3D并行方案，并为编码器探索可能的3D并行方案，同时考虑部署LLM后可用的GPU内存。通过模型规划器，每个GPU都持有LLM和编码器的模型状态，从而能够在LLM气泡期间进行编码器计算。编码器和LLM的模型并行方案作为输入提供给气泡调度器，Optimus会根据执行时间最短的输出调度来选择并行方案。

Algorithm 1: Optimus workflow
1 Function Optimus(mllm):
2   encPlans, llmPlan = ModelPlanner(mllm)
3   bestLat, bestSchedule = +∞, None
4   for encPlan in encPlans do
5     schedule = BubbleScheduler(encPlan, llmPlan)
6     if schedule.lat < bestLat then
7       bestSchedule = schedule
8       bestLat = schedule.lat
9     end
10  end
11  return bestSchedule

气泡调度器（Bubble Scheduler）。气泡调度器负责将编码器计算调度到LLM气泡中。鉴于LLM训练流水线将数据划分为多个微批次，调度器以每个微批次为基础调度编码器计算，并满足微批次级别的编码器-LLM数据依赖（解决§3.1中的挑战2）。此外，调度器将编码器计算分解到核函数粒度，以便利用LLM训练期间的亚毫秒级气泡（TP气泡）（解决§3.1中的挑战3）。当前设计基于1F1B交错式流水线调度【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】，但并未与其紧密耦合，可以调整以支持其他流水线调度策略（如第6节所讨论）。

Optimus工作流程。Optimus使用模型规划器为编码器和LLM设计并行方案。随后，对于每个编码器并行方案，Optimus利用气泡调度器生成一个调度并估算延迟。延迟估算是基于模型规划器步骤中进行的离线性能分析，我们使用离线分析得到的编码器执行时间和LLM流水线训练时间线。最终，Optimus选择训练时间最短的调度方案，将编码器计算调度到LLM气泡中。Optimus的工作流程如算法1所示。

4.1 模型规划器

模型规划器工作流程。模型规划器的工作流程包括搜索编码器和LLM的并行方案、共置编码器和LLM、确保满足内存约束，以及为编码器和LLM流水线构建独立的微批次。

搜索独立的并行方案。首先，规划器根据Megatron-LM【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】中的见解，确定LLM骨干网络的3D并行方案($DP_{llm}$, $PP_{llm}$, $TP_{llm}$) 。随后，规划器枚举编码器可能的3D并行方案($DP_{enc}$, $PP_{enc}$, $TP_{enc}$) 。为确保多个编码器模型可以与每个LLM模型共置，我们强制要求$PP_{enc}$是$PP_{llm}$的因子，$TP_{enc}$是$TP_{llm}$的因子。

共置编码器和LLM。为确保每个GPU都能在LLM停机期间执行编码器计算，模型规划器将编码器和LLM的模型状态分配给每个GPU。如图5所示，所有GPU都包含编码器（绿色表示）和LLM（红色表示）的模型状态。如果没有这种共置，许多GPU将缺乏执行编码器计算所需的编码器模型状态。

根据内存约束剪枝并行方案。由于编码器和LLM阶段被共置在GPU上，我们根据选定的并行方案估算编码器和LLM模型状态以及LLM激活值的内存需求——这借鉴了【【索引14，Reducing activation recomputation in large transformer models，2023，MLSys】】中的内存分析。我们忽略了编码器激活值的估算，因为它们的内存占用可以忽略不计。任何超出GPU内存容量的方案都会被提前剪除。

构建独立的微批次。由于编码器和LLM的并行方案不同，对于给定的GPU集合，编码器流水线的数量是LLM流水线的 $m = DP_{enc}DP_{llm}$ 倍（例如，在图5中m=2）。对于属于同一个LLM流水线的GPU，会共置m个编码器流水线。根据LLM流水线训练中使用的微批次数$N_{mb}$，这些$N_{mb}$个微批次的数据需要被分配到这m个编码器流水线中，其中每个编码器流水线i处理$N_{enc}$个微批次数据的前向和后向计算。模型规划器枚举了将这$N_{mb}$个微批次在m个编码器流水线之间划分的可能方式。例如，如果LLM训练中有8个微批次，且m=2个编码器流水线，那么总共有7种可能的划分选项，如[1, 7], [2, 6], ..., [7, 1]。

Algorithm 2: BubbleScheduler
1 Function BubbleScheduler(encPlan, llmPlan):
2   schedules = InitSchedule(encPlan, llmPlan)
3   dep = GetEncLLMDep(llmPlan)
4   bestLat, bestSchedule = +∞, None
5   for schedule in schedules do
6     schedule = OptimizeSchedule(schedule, dep, FWD)
7     schedule = OptimizeSchedule(schedule, dep, BWD)
8     if schedule.lat < bestLat then
9       bestSchedule = schedule
10      bestLat = schedule.lat
11    end
12  end
13  return bestSchedule
14
15 Function OptimizeSchedule(schedule, dep, mode):
16  while True do
17    encPPID = findCritical(schedule, mode)
18    newSchedule, success = ScheduleKernels(encPPID, schedule, mode)
19    if success and checkEncLLMDep(schedule, dep) then
20      schedule = newSchedule
21    else
22      return schedule
23    end
24  end

4.2 气泡调度

LLM气泡的通用模式。尽管不同GPU中的LLM气泡具有不同的开始时间和持续时长，但存在一种通用的LLM气泡模式，如图8所示。在任何LLM计算开始之前，有一个大的单一气泡（DP all-gather气泡和PP-warmup气泡之和）；在所有LLM计算结束后，也有一个大的单一气泡（PP-cooldown气泡和reduce-scatter气泡之和）。此外，还有许多小气泡（PP气泡和TP气泡）【【索引14，Reducing activation recomputation in large transformer models，2023，MLSys】，【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】，【索引25，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，CoRR】】与LLM计算交错出现。

气泡调度器算法。如算法2所述，气泡调度器首先进行粗粒度气泡利用，通过创建初始调度，将编码器计算纳入LLM计算之前和之后的气泡中（第2行）。然而，这两个气泡可能没有足够的时间来完成所有的编码器计算，导致一些编码器计算未被调度到气泡内。为了减少总训练时间，气泡调度器接着执行细粒度气泡利用。这包括通过将编码器前向计算分配到与LLM计算交替出现的气泡中来优化调度（第7行），然后将编码器后向计算调度到这些相同的气泡中（第8行）。气泡调度器的最终输出是实现最短可能运行时间的调度。

粗粒度气泡利用。对于每种可能的数据划分方法，气泡调度器通过将编码器前向操作安排在LLM计算之前，将编码器后向操作安排在LLM计算之后来初始化调度。图9展示了当有m=2个编码器流水线且数据划分方式为[3, 5]时的初始化调度，即3个微批次分配给第一个编码器流水线，5个分配给第二个。

细粒度气泡利用。OptimizeSchedule函数（算法2第15行）通过迭代方法优化初始调度。首先，气泡调度器使用findCritical来识别其计算位于端到端MLLM训练关键路径上的编码器流水线（第17行）。随后，它利用ScheduleKernels将该编码器计算的一个微批次分配到与LLM计算交错的气泡中（第18行）。如果找到足够的气泡来容纳编码器计算并且满足编码器-LLM数据依赖（如§4.3所述），气泡调度器会继续调度。否则，它将终止并返回迄今为止找到的最佳调度。

关键路径识别。当优化编码器前向计算的调度时（算法2第7行），findCritical会识别出前向计算是关键路径的编码器流水线。如图10左侧所示，在初始调度中，编码器流水线2的前向计算（微批次8的前向）最初位于关键路径上。在成功将该微批次的前向计算调度到后续的气泡后，编码器流水线1占据了关键路径位置。这个迭代过程在每一步之后都会减少端到端的MLLM训练时间。类似地，图10右侧展示了后向计算是关键路径的编码器流水线。每一步之后，气泡调度器必须在进行下一步之前验证是否仍然满足编码器-LLM数据依赖。

核函数级别的调度。当将编码器计算调度到与LLM计算交错的气泡中时（第17行的ScheduleKernels），气泡调度器将编码器计算分解为核函数粒度，并根据气泡的持续时间调度这些核函数。对于每个气泡，气泡调度器调度多个核函数，同时确保这些核函数的总执行时间在气泡持续时间内。此外，气泡调度器必须满足编码器的内部数据依赖。如图11所示，设备1持有编码器的前两层，而设备2持有接下来的两层。在调度前向传播过程中的核函数时，设备2只能利用在设备1完成其前向传播之后出现的气泡来执行编码器计算。对于前向计算，气泡调度器从上游编码器流水线阶段向下游调度。相反，对于后向计算，调度顺序相反。虽然每个编码器层也包含通信核函数，但调度器确保这些核函数不会被分配到LLM通信期间出现的TP气泡中。相反，调度器识别LLM层内持续时间较长的计算核函数，并用它们与编码器通信核函数重叠。由于LLM和编码器层交替执行计算和通信任务，它们高效地利用了GPU带宽和流式多处理器（SMs）。这种设计策略有助于最小化资源争用并提高整体GPU利用率【【索引15，Pytorch distributed: Experiences on accelerating data parallel training，2020】】。

复杂度分析。我们的气泡调度算法复杂度较低。给定n个GPU，n的素因子数量为$n_p$，并行方案的搜索空间为$C^{np+1}_{2}$。微批次划分的数量为$O(N^{m-1}_{mb})$。因此，调度气泡的复杂度为 $O(C^{np+1}_{2} * N^{m-1}_{mb} * (F + B))$。在我们的实验中，对于一个大型MLLM模型，该算法在单个CPU核心上运行不到一分钟就能找到最优调度（见§5.3.2），这也是一次性成本。

4.3 解决编码器-LLM依赖

依赖问题的复杂性。模型规划器为编码器和LLM骨干网络提供不同的并行策略，包括微批次的数量，这导致了编码器和LLM之间以及内部都存在复杂的数据依赖。同时，编码器和LLM的通信和计算是交错执行的，如果协调不当，可能会引入额外的流水线气泡，加剧了系统依赖的复杂性。

微批次级别的依赖检查。气泡调度器通过检查每个微批次i的编码器-LLM前向和后向依赖点来解决依赖问题。这些依赖点分别表示为$F_i$和$B_i$，代表LLM需要相应激活值$A_i$（由编码器输出）进行前向传播的时间，以及LLM在后向传播期间生成相应梯度$G_i$（编码器的输入）的时间。为确保满足编码器-LLM依赖，气泡调度器使用两个函数：GetEncLLMDep（算法2第3行）和CheckEncLLMDep（算法2第19行）。

获取依赖点 (GetEncLLMDep)。鉴于交错式1F1B调度【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】是LLM训练最高效的流水线调度之一，我们深入探讨该调度中数据依赖点$F_i$和$B_i$的具体细节。图12的上图展示了一个包含两个模型块的交错式1F1B调度实例。在这里，前向依赖点表示第一流水线阶段（设备1）为第一个模型块开始前向执行的时刻（深蓝色表示），而后向依赖点则表示第一流水线阶段（设备1）完成第一个模型块后向执行的时刻（深绿色表示）。

调整依赖点以创造调度机会。我们观察到，推迟最后四个微批次（F5至F8）的前向数据依赖点是可行的，而不会对整体流水线延迟产生任何不利影响。为实现这一点，我们可以调整每个流水线阶段的预热微批次数，如图12的下图所示。这种调整使气泡调度器能够在优化初始调度时，利用从预热阶段到1F1B稳定阶段的过渡期间的气泡来调度编码器前向计算。GetEncLLMDep函数会产出调整后的1F1B交错式调度前向和后向数据依赖点。

验证依赖 (CheckEncLLMDep)。通过考虑将编码器计算调度到气泡中，气泡调度器估算出编码器完成分布在不同编码器流水线上的微批次前向传播的时间。调度器将这些完成时间按升序排列为$EF_i$（全局排序），表示编码器为参与LLM流水线训练的微批次i完成前向操作的时间。如果对于所有微批次（i=1...$N_{mb}$），编码器在其指定的时间点$F_i$之前完成前向操作（$EF_i \leq F_i$），则认为前向依赖得到满足。类似地，如果编码器的后向操作不早于$B_i$时间点开始（$EB_i \geq B_i$），则后向依赖得到满足。当CheckEncLLMDep确认前向和后向依赖都成功满足时，返回true。图13提供了一个评估编码器-LLM依赖的例子，其中有两个编码器流水线，每个处理四个微批次。编码器完成其前向传播的顺序决定了激活值在LLM流水线中的使用方式：来自编码器流水线1的激活值被指定为第1、3、7、8个微批次，而来自编码器流水线2的激活值被用作第2、4、5、6个微批次。然后，气泡调度器通过确保每个编码器的前向操作在相应LLM前向传播开始前结束，并且每个编码器的后向操作在LLM结束后才开始，来验证微批次级别的依赖。

插入P2P通信。当依赖关系得到满足时，气泡调度器会在编码器流水线的最后阶段和LLM流水线的第一阶段之间，将必要的点对点（P2P）通信集成到训练调度中。例如，如果编码器流水线j完成了微批次i的前向传播，调度器将在编码器流水线j的最后阶段插入一个P2P发送（发送激活值），并在LLM流水线的第一阶段插入一个P2P接收（接收激活值）。同样，当LLM流水线完成微批次i的反向传播时，调度器会在LLM流水线的第一阶段添加一个P2P发送（发送梯度），并在编码器流水线j的最后阶段添加一个P2P接收（接收梯度）。在图13所示的场景中，训练流水线处理8个微批次，调度器在设备1和2之间插入8对P2P发送-接收操作来管理编码器流水线1和LLM流水线之间的依赖。这包括4对用于前向激活值发送/接收，和4对用于后向梯度发送/接收。同样，另外8对P2P发送-接收操作被插入到设备3和4之间，以处理编码器流水线2和LLM流水线之间的依赖。

4.4 多分支编码器调度

并行化多分支编码器。为支持具有多个编码器的MLLM【【索引6，Mm-vit: Multi-modal video transformer for compressed video action recognition，2021】，【索引35，Dual-encoder transformers with cross-modal alignment for multimodal aspect-based sentiment analysis，2022，AACL-IJCNLP】】，模型规划器为每个编码器独立地应用一个编码器并行方案($DP_{enc}$, $PP_{enc}$, $TP_{enc}$)。对于流水线并行，每个编码器内的层被划分为$PP_{enc}$个阶段（如图14所示）。每个编码器的每一层随后根据$TP_{enc}$进行并行化。

调度多分支编码器核函数。气泡调度器将不同编码器的层分解为核函数级别的粒度，并将它们的调度安排得好像这些核函数是单个编码器的一部分。这是因为MLLM内的编码器是独立运行的，它们之间没有任何数据依赖。

4.5 内存分析

模型状态内存计算。当使用$n_{gpu}$个GPU进行MLLM训练时，模型规划器根据并行方案需要$DP_{enc}$份复制的编码器模型状态和$DP_{llm}$份复制的LLM模型状态。假设编码器的参数数量为$\phi_{enc}$，LLM的参数数量为$\phi_{llm}$，每个参数需要k字节的内存。存储模型状态的平均GPU内存使用量$MEM_{model}$计算如下：

内存开销分析。与现有的3D并行训练方案（其中$DP_{enc} = DP_{llm}$）相比，估计的内存开销$MEM_{overhead}$可以表示为：

开销与调度的权衡。随着$DP_{enc}$值的增大，由于需要更多复制的编码器模型状态，内存开销会更高。然而，这会导致调度时编码器内部依赖关系变得不那么复杂（表现为较小的$PP_{enc}$）。模型规划器根据估计的内存使用量$MEM_{model}$过滤编码器并行方案，确保遵守GPU内存限制。在实践中，内存开销通常低于12%（见§5.3.1），因为$\phi_{enc}$很小（例如，最大的视觉编码器有220亿参数【【索引9，Scaling vision transformers to 22 billion parameters，2023，ICML】】），且k很小（例如，当使用bf16参数和fp32梯度以及分布式优化器时，k=6【【索引1，GitHub - NVIDIA/Megatron-LM: Ongoing research training transformer models at scale】】）。

A4 实验环境

• 硬件配置:
  • GPU: 生产训练集群，包含数千个NVIDIA Hopper GPU，每个GPU拥有80GB内存和989 TFLOPS计算性能。
  • 互联: 服务器内为NVLink，服务器间为高带宽RDMA网络。
  • （注：与Alpa和FSDP的对比实验使用了8个NVIDIA A100 GPU）
• 模型架构:
  • 图像编码器: ViT-22B, ViT-11B, 和 ViT-5B (ViT-22B的缩减版)。
  • 语言模型: LLAMA-70B 和 GPT-175B。
  • 所有实验中的序列长度均为2048。详细模型配置见附录A。
• 软件配置:
  • 代码实现: 基于开源的Megatron-LM框架【【索引1，GitHub - NVIDIA/Megatron-LM: Ongoing research training transformer models at scale】】开发。
  • 基线系统:
    1. PyTorch FSDP: 一种分布式数据并行训练模块。
    2. Alpa: 一个自动生成并行执行计划的编译器系统。
    3. Megatron-LM: 一个集成了3D并行技术的SOTA LLM训练框架。
    4. Megatron-LM balanced: 一个“稻草人”基线，通过动态规划算法平衡MLLM中不同子模块在流水线阶段的层划分，以实现计算负载均衡。
• 数据集与用途:
  • 论文未明确指定具体数据集，实验主要关注不同规模模型和硬件配置下的系统性能（迭代时间和模型FLOPs利用率MFU），而非模型在特定任务上的准确率。实验数据在300次训练迭代（预热10次后）上取平均值。

A4 实验结果

5.2 端到端性能

5.2.1 弱扩展性实验

• 实验内容: 随GPU数量增加而扩展模型规模，评估Optimus和基线系统的训练性能。具体配置见表3。
• 实验结果:
  • Optimus相比Megatron-LM最高提速1.22倍，相比Megatron-LM balanced最高提速1.18倍。Alpa和FSDP在这些模型上均出现内存溢出（OOM）问题。（见图15）
  • 在一个较小的模型（ViT-3B + GPT-11B）上，Optimus相比Alpa提速3.09倍，相比FSDP提速15.1%。（见表4）
    

    

    
    

    

    
    

    

5.2.2 强扩展性实验

• 实验内容: 使用ViT-22B+GPT-175B模型，保持全局批量大小（1536）不变，增加GPU数量（1536, 2048, 3172）。
• 实验结果:
  • Optimus相比Megatron-LM最高减少21.3%的迭代时间，相比Megatron-LM balanced最高减少20.5%。
  • 随着GPU数量增加，Optimus的优势更加明显，因为保持批量大小不变会增加气泡比例，为Optimus提供了更多调度机会。
  • Optimus的MFU（模型FLOPs利用率）保持稳定，而基线系统的MFU在更大规模下有所下降。（见表5）
    

    

5.2.3 多编码器MLLM实验

• 实验内容: 在512个GPU上评估Optimus和Megatron-LM在含多个编码器（2个或3个）的MLLM上的训练性能。具体配置见表6。
• 实验结果:
  • Optimus在三种多编码器MLLM上分别取得了高达1.25倍、1.26倍和1.27倍的加速比。
  • 加速比更高是因为Megatron-LM将所有编码器置于第一个流水线阶段，导致更严重的流水线不平衡。（见图16）
    

    

    
    

    

5.3 微基准测试

5.3.1 Optimus内存消耗

• 实验内容: 测量Optimus和基线系统在训练不同规模MLLM（表3配置）时的GPU内存消耗。
• 实验结果:
  • 与内存效率最高的基线相比，Optimus的最大GPU内存开销为12%。
  • 在某些情况下，Optimus的内存使用甚至低于基线，因为基线策略可能因编码器和LLM层的隐藏大小不同而导致内存不平衡。（见图17）
    

    

5.3.2 气泡调度器算法性能

• 实验内容: 在单核CPU上运行气泡调度算法，计算强扩展性实验设置下的调度方案。评估指标为“调度效率”（Eff_coarse和Eff_fine）和算法运行时间。
• 实验结果:
  • 随着GPU数量增加，调度效率（Eff_coarse和Eff_fine）更高，因为LLM流水线中的气泡比例增加。
  • 启用细粒度气泡利用后，效率相比仅用粗粒度最高可提升1.67倍。
  • 算法运行时间很短，并且随着LLM流水线中微批次数量的减少而减少。（见表7）
    

    

A7 补充细节

6 讨论

冻结参数的MLLM训练。Optimus可以轻松支持实际中常见的多阶段训练工作流，例如LLaVA【【索引18，Visual instruction tuning，2023】】所采用的方法。虽然本文主要关注所有参数（包括编码器和LLM）都更新的一般情况，但Optimus可以自然地适应只训练适配器（adapter）的阶段。在这种情况下，Optimus会将编码器+适配器的前向传播和适配器的反向传播调度到LLM流水线气泡中，同时由于参数被冻结而跳过编码器的反向计算。这既保持了正确的数据依赖关系，又继续有效地利用了气泡。

复杂计算图。Optimus专注于典型MLLM模型架构上的气泡调度，该架构由多模态编码器后接一个LLM组成。我们未来可能会探索对复杂MLLM计算图的气泡调度。这需要一种新的划分算法，将计算图划分为骨干流水线调度和用于填充气泡的工作负载。Optimus的气泡调度算法可以很容易地扩展到划分后的计算图上。

其他流水线调度。我们在MLLM训练中使用了广泛应用的Megatron-LM交错式1F1B流水线调度。然而，还存在其他流水线调度（例如Chimera【【索引16，Chimera: efficiently training large-scale neural networks with bidirectional pipelines，2021，SC】】和零气泡流水线【【索引23，Zero bubble pipeline parallelism，2023】】），它们在某些场景下可能性能更优。Optimus的气泡调度与这些流水线调度优化是正交的，只要分析并解决了特定的编码器-LLM依赖关系，Optimus就可以应用于其他流水线调度。

在线调度。我们的气泡调度算法为简化执行，没有考虑CUDA核函数运行时间的波动。我们通过收集性能统计数据（如核函数执行时间）来检测一个训练步骤中的气泡，并假设未来步骤的行为保持一致。然而，与预测执行时间的偏差可能导致次优调度，产生更大或改变了的流水线气泡。一个可能的解决方案是进行实时性能监控，以动态调整调度。

排除DiffusionPipe和DistTrain作为基线。我们明确排除了DiffusionPipe【【索引30，Diffusionpipe: Training large diffusion models with efficient pipelines，2024，MLSys】】和DistTrain【【索引37，Disttrain: Addressing model and data heterogeneity with disaggregated training for multimodal large language models，2024，arXiv】】作为我们的基线，原因如下。DiffusionPipe专为扩散模型设计，并且仅在小规模集群（≤64个GPU）上进行了评估。其关注的模型家族和规模不同，使其不适合与Optimus进行比较，后者旨在跨越数千个GPU进行大规模多模态LLM（MLLM）训练。DistTrain依赖于一种简化的模型划分策略，当应用于MLLM的异构结构时，可能导致严重的流水线不平衡。此外，DistTrain并非开源，无法进行直接的实证比较。

7 相关工作

多模态训练。Pytorch FSDP训练【【索引38，Pytorch fsdp: experiences on scaling fully sharded data parallel，2023，arXiv】】仅支持数据并行，效率低于混合并行策略。Alpa【【索引39，Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning，2022，OSDI】】能为多种模型自动化并行，但其缺点是不支持最先进的1F1B交错式流水线并行【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】，且比优化的Megatron-LM框架【【索引25，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，CoRR】】需要更多内存，并且由于其对编码器和解码器的统一视角，错失了流水线优化的机会。DistMM【【索引12，Distmm: Accelerating distributed multimodal model training，2024，NSDI】】提供了协调多个并行编码器的解决方案，但它是为对比学习设计的，忽略了解码器，在综合训练效率上留下了空白。DiffusionPipe【【索引30，Diffusionpipe: Training large diffusion models with efficient pipelines，2024，MLSys】】和DistTrain【【索引37，Disttrain: Addressing model and data heterogeneity with disaggregated training for multimodal large language models，2024，arXiv】】是另外两项关于多模态训练的工作，各自的局限性已在前一节中概述。

气泡减少。先前减少“气泡”的努力从不同角度解决了这个问题。1F1B交错式流水线【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】技术通过将模型分块并在不同阶段交错这些块来最小化气泡，而零气泡流水线【【索引23，Zero bubble pipeline parallelism，2023】】方法则进一步细化反向传播计算以消除气泡。然而，在实践中，零气泡流水线调度无法完全消除所有流水线气泡，因为它需要对优化器进行更改，这引发了对端到端模型收敛性的担忧。另一方面，异步张量并行【【索引26，Communication-minimizing asynchronous tensor parallelism，2023】】和谷歌的重叠技术【【索引32，Overlap communication with dependent computation via decomposition in large deep learning models，2022，ASPLOS】】旨在将张量并行通信与计算重叠，但受限于特定的硬件配置，并且随着计算能力的提升难以维持完全重叠。

气泡利用。Pipefisher【【索引22，Pipefisher: Efficient training of large language models using pipelining and fisher information matrices，2023，MLSys】】利用跨多个训练步骤的流水线气泡来完成K-FAC计算，而我们的方法在单个同步训练步骤内操作，专注于即时优化。Hydro的Bubble Squeezer【【索引11，Hydro:{Surrogate-Based} hyperparameter tuning service in datacenters，2023，OSDI】】利用GPT模型的气泡执行独立任务，如超参数调优，这无法提升训练步骤本身的性能。Bamboo【【索引29，Bamboo: Making preemptible instances resilient for affordable training of large {DNNs}，2023，NSDI】】利用流水线气泡进行冗余计算，以减轻在易失性实例上训练时抢占的影响，其假设是后续流水线阶段承载更多层，这在大型语言模型（LLM）训练场景中通常不成立。

A5 结论

本文提出了Optimus，一个分布式的MLLM训练系统，它通过将编码器的计算调度到LLM的气泡中，从而减少了端到端的MLLM训练时间。为了减少MLLM训练期间的GPU气泡，Optimus将多模态编码器和LLM骨干网络分开划分，并将编码器计算调度到LLM的气泡中。我们考虑了内存和计算资源的约束，为编码器搜索最优的并行方案，从而平衡了用于填充气泡的编码器计算。Optimus进一步采用了一种气泡调度算法来处理编码器-LLM的依赖关系，并选择最优的调度方案，将核函数级别的编码器计算填充到亚毫秒级的LLM气泡中。我们的大量实验表明，与基线相比，Optimus在使用ViT-22B和GPT-175B模型、超过3072个GPU的配置下，可将MLLM训练加速20.5%-21.3%，并且平均性能比现有MLLM训练系统高出20.3%。

A6 附录

A MLLM模型配置

这里我们列出了Optimus评估实验中使用的所有MLLM配置。ViT编码器配置见表8。LLM骨干网络配置见表9。在所有实验中，我们使用的序列长度为2048。

B Megatron-LM balanced DP算法

算法描述。我们采用一种动态规划（DP）算法，为Megatron 1F1B交错式调度【【索引20，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，SC】】将层分配到不同的虚拟阶段。遵循Alpa【【索引39，Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning，2022，OSDI】】的方法，该DP算法旨在最小化最慢阶段的延迟，以减少流水线调度的端到端延迟。给定一个流水线并行大小为PP和V个模型块的配置，DP算法试图最小化最慢虚拟阶段的延迟。它确定了最优的层划分策略，将层分布在这V × PP个虚拟阶段上。

公式定义。我们定义函数$F(l,m)$来表示将前l层分配到前m个虚拟阶段时，单个虚拟阶段的最大延迟。计算从$F(l, 1) = \sum_{i \le l, i=1} t_i$开始，其中$t_i$表示第i层的执行时间（基于FLOPs估算）。F的最优结构为：

求解与局限性。对于一个有L层的MLLM模型，通过计算$F(L, V \times PP)$并记录划分结果来确定层划分策略，以找到最优解。这确保了在1F1B交错式流水线调度中，所有虚拟阶段中最长虚拟阶段的延迟$F(L, V \times PP)$被最小化。上述动态规划算法适用于具有单个编码器的MLLM配置，其中编码器层和LLM层遵循线性结构。然而，该DP算法不适用于具有多个编码器的MLLM模型，因为这些编码器之间没有数据依赖关系。

C Optimus、Alpa和FSDP的训练性能比较

实验设置。为了便于与Alpa和FSDP进行比较，我们构建了一个小规模的MLLM，包括ViT-3B和GPT-11B，具体配置在附录A中提供。我们在8个NVIDIA A100 GPU上评估了训练性能，因为在NVIDIA Hopper GPU上运行Alpa时遇到了CUDA库的问题。全局批量大小设置为16，序列长度为2048。

结果。根据表10，Optimus比Alpa快3.09倍，比FSDP快15.1%。

D 基于Megatron-LM的基线详细配置

D.1 弱扩展性实验

表11显示了弱扩展性实验中基于Megatron-LM的基线的详细配置。

D.2 强扩展性实验

表12显示了强扩展性实验中基于Megatron-LM的基线的详细配置。

D.3 多编码器MLLM实验

在多编码器MLLM实验中，对于所有MLLM模型，我们为Megatron-LM使用(DP=8, TP=8, PP=8)的配置，并将微批次大小配置为2。

方法细节中的引用汇总

1. [20] Narayanan, D., et al. (2021). Efficient large-scale language model training on gpu clusters using megatron-lm. SC.

   • 引用位置: Introduction, Background, Design Decisions, Optimus Design, Related works, Appendix B
   • 引用描述: 本文多次引用Megatron-LM作为SOTA的LLM训练框架和实现基础。具体引用了其1F1B交错式流水线调度方案作为分析和优化的基础，并将其作为主要的性能比较基线。

2. [23] Qi, P., et al. (2023). Zero bubble pipeline parallelism.

   • 引用位置: Introduction, Background, Discussion, Related works
   • 引用描述: 引用了"Zero Bubble Pipeline"作为一种减少流水线气泡的技术。文中指出，由于需要修改优化器，该技术在MLLM训练中无法完全消除气泡，并可能影响模型收敛性。

3. [13] Jiang, Z., et al. (2024). Megascale: Scaling large language model training to more than 10,000 gpus. arXiv.

   • 引用位置: Introduction, Background, Challenges
   • 引用描述: 引用MegaScale作为现有的SOTA训练系统之一，并指出尽管应用了其优化技术（如通信计算重叠），MLLM训练中仍存在大量气泡。

4. [14] Korthikanti, V. A., et al. (2023). Reducing activation recomputation in large transformer models. MLSys.

   • 引用位置: Background, Optimus Design
   • 引用描述: 引用了其关于Transformer层中通信核（all-gather, reduce-scatter）的分析，以及模型训练内存消耗的分析方法。

5. [39] Zheng, L., et al. (2022). Alpa: Automating inter-and {Intra-Operator} parallelism for distributed deep learning. OSDI.

   • 引用位置: Challenges, Related works, Appendix B
   • 引用描述: 引用Alpa作为一个自动并行化的编译器系统。文中指出，Alpa等系统采用统一的并行策略，导致了挑战1（大部分GPU无编码器状态）。实验部分将其作为基线进行比较。附录中提到了其DP算法思想。

6. [9] Dehghani, M., et al. (2023). Scaling vision transformers to 22 billion parameters. ICML.

   • 引用位置: Memory Analysis
   • 引用描述: 在分析内存开销时，引用这篇论文来说明最大的视觉编码器（ViT-22B）的参数规模，以论证编码器相对于LLM的尺寸较小，从而内存开销可控。

7. [1] NVIDIA/Megatron-LM. GitHub.

   • 引用位置: Memory Analysis, Evaluation
   • 引用描述: 引用Megatron-LM的开源实现作为Optimus的开发基础，并引用其在使用bf16参数和fp32梯度时的内存占用数据（k=6）。

8. [16] Li, S., & Hoefler, T. (2021). Chimera: efficiently training large-scale neural networks with bidirectional pipelines. SC.

   • 引用位置: Introduction, Discussion
   • 引用描述: 引用Chimera作为一种不同于1F1B的流水线调度策略，并指出Optimus的设计是正交的，可以适用于这类调度。

9. [22] Osawa, K., et al. (2023). Pipefisher: Efficient training of large language models using pipelining and fisher information matrices. MLSys.

   • 引用位置: Related works
   • 引用描述: 在相关工作中，将Optimus与Pipefisher进行比较，指出Pipefisher利用气泡执行跨训练步骤的任务，而Optimus专注于优化单步训练的性能。

10. [11] Hu, Q., et al. (2023). Hydro:{Surrogate-Based} hyperparameter tuning service in datacenters. OSDI.

    • 引用位置: Related works
    • 引用描述: 引用Hydro的Bubble Squeezer作为利用气泡的另一种方法，但指出其用于超参数调整等独立任务，不能提升训练步骤本身的性能。