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

Optimus：通过气泡利用加速大规模多模态LLM训练
作者/机构: Weiqi Feng1∗, Yangrui Chen2, Shaoyu Wang3∗, Yanghua Peng2, Haibin Lin2 and Minlan Yu1 1Harvard University, 2Bytedance, 3University of Southern California

A1 主要贡献

本文旨在解决训练大规模多模态语言模型（MLLM）时的效率低下问题。现有系统在训练MLLM时，由于异构的模态模型和3D并行中复杂的数据依赖性，会产生大量的GPU空闲时间（即“气泡”）。

核心问题与研究目标
- 核心问题: 在使用3D并行（数据并行、流水线并行、张量并行）训练大型MLLM（如包含数百亿参数的模型）时，超过40%的GPU周期处于空闲状态。这些空闲“气泡”主要由两大原因造成：
1. 3D并行中广泛且频繁的通信开销导致GPU空闲。
2. MLLM的流水线阶段不平衡，以及阶段间的依赖关系导致了显著的数据等待时间。
- 研究目标: 提出一个名为Optimus的分布式MLLM训练系统，通过有效利用这些GPU空闲时间来减少端到端的训练时间。核心思想是在语言模型（LLM）计算过程中产生的“气泡”内调度编码器（encoder）的计算任务。

创新点与主要贡献
Optimus系统通过以下三个核心设计解决了在LLM气泡中调度编码器计算的挑战：
1. 分离式并行策略: 传统的训练框架对整个MLLM模型（编码器和LLM）采用统一的并行策略，导致大多数GPU上只有LLM的模型状态，无法在LLM气泡期间执行编码器计算。Optimus为编码器和LLM搜索并采用独立的并行计划，从而将编码器和LLM的模型状态共置（colocate）在每个GPU上，使得所有GPU都能在空闲时执行编码器计算。
2. 两阶段依赖管理策略: MLLM中存在复杂的依赖关系，包括训练迭代同步依赖、编码器内部依赖以及编码器与LLM之间的微批次（microbatch）级别依赖。Optimus采用两阶段策略来处理这些依赖：
- 本地调度（Local Scheduling）: 在每个编码器流水线内部，根据迭代同步依赖和编码器内部依赖，将编码器计算调度到可用的LLM气泡中。
- 全局排序（Global Ordering）: 通过比较时间戳，确保编码器和LLM之间在微批次级别的依赖关系得到满足（即编码器的前向传播在LLM前向传播之前完成，编码器的反向传播在LLM反向传播之后开始）。
3. 核函数级别的细粒度调度: LLM气泡的持续时间差异巨大，从亚毫秒到数百毫秒不等。传统的层级别调度无法利用短暂的亚毫秒级气泡。Optimus将编码器层的计算分解为一系列核函数（kernels），从而能够有效利用这些短暂的气泡。此外，系统分析了3D并行中常见的“气泡”模式，通过将编码器核函数计算与LLM计算交错执行来优化调度，最小化整体训练时间。

实验成果
- Optimus基于Megatron-LM实现，并在生产集群上进行了全面评估。
- 实验结果表明，在使用ViT-22B编码器和GPT-175B语言模型，并扩展到3072个GPU时，Optimus相比基线系统能够将MLLM的训练速度提升20.5%-21.3%。

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

2.1 多模态LLM的特点

多模态LLM的重要性与架构。多模态LLM（MLLM）在大型语言模型（LLM）的基础上，通过集成先进的自然语言处理方法，并将其范围扩展到涵盖多种数据模态，变得越来越重要。例如，GPT-4【21】就是这样一种进步，它增强了其前代模型的能力，包括了多模态的理解和生成，并在涉及图像和文本输入的各种基准测试中展现出人类水平的性能。一个MLLM通常包含三个主要部分：一个或多个模态编码器、输入投影器和一个大型语言模型骨干【36】。模态编码器处理非文本模态的输入，将其转化为各自的特征表示；输入投影器则将这些特征与文本特征空间对齐。随后，LLM骨干将对齐后的多模态和文本特征作为输入。图1展示了MLLM的架构。由于输入投影器的计算需求相比编码器和LLM骨干要小得多，因此在后续讨论中不予考虑，根据Llava【18】的细节，本文将输入投影器视为模态编码器的最后一层。

MLLM的独有特性。与同构的LLM架构不同，多模态LLM具有以下独有特性：
- LLM骨干占据主导的模型规模：在多模态LLM中，LLM骨干的参数数量远超其他组件，如编码器和投影器。例如，Flamingo【3】总共有800亿参数，其中其LLM骨干就占了700亿。
- 编码器与LLM骨干之间的依赖关系：在MLLM训练中，编码器和LLM之间存在两种类型的数据依赖。在前向传播过程中，编码器必须完成编码特征的生成，LLM骨干才能继续进行前向计算。相反，在反向传播过程中，LLM骨干先计算梯度，然后编码器才开始反向传播。

2.2 MLLM训练中的气泡

现有LLM流水线优化的局限性。现有的LLM流水线优化缺乏模型无关性，不足以应对MLLM训练任务。在我们内部使用ViT编码器和GPT骨干（超过1000亿参数）的大规模MLLM训练任务中，我们在超过3000个NVIDIA GPU上使用了Megatron-LM，并观察到超过48%的GPU周期处于空闲状态，尽管已经应用了多种SOTA技术，包括MegaScale【13】，Zero Bubble Pipeline【23】和细粒度的通信-计算重叠【32】。我们分析了性能剖析的时间线，以识别和调查GPU空闲（即气泡）的发生情况。表1显示了不同类型气泡占用的总时长和平均训练步长时间（5.12秒）的百分比。

气泡的分类。这些气泡可以根据其根本原因分为三类：
1. 数据并行（DP）中的通信。数据并行需要通信来聚合梯度，导致在这些通信期间GPU空闲。具体来说，MegaScale【13】和Megatron-LM【25】采用了一种分布式优化器，类似于ZeRO【24】中的Pos+g，以节省大型模型训练的内存，这需要执行两次集体通信：all-gather和reduce-scatter。在每个训练步骤开始时，一次all-gather操作从所有DP ranks收集更新后的参数，导致一个DP all-gather气泡，占训练时间的3.3%。在训练步骤结束时，执行reduce-scatter来聚合梯度，产生一个DP reduce-scatter气泡，消耗了8.9%的训练时间。值得注意的是，尽管已经应用了Megascale【13】中提出的数据并行重叠优化，但由于同步训练的性质【13】，上述针对第一个模型块的DP通信无法被隐藏。图2展示了在每个训练步骤开始和结束时由all-gather和reduce-scatter操作产生的DP气泡。
2. 流水线并行（PP）中的依赖。尽管应用了来自Megascale【13】的流水线发送-接收重叠优化，但由于前向和反向传播过程中各阶段之间固有的数据依赖，流水线气泡依然存在。重要的是，由于需要更改优化器【23】，Zero Bubble Pipeline方法无法消除MLLM训练中的流水线气泡（详见第7节讨论）。图2展示了MLLM训练流水线调度，包括三个阶段：预热（仅前向）、稳定（一前向一反向）和冷却（仅反向）。在整个流水线训练过程中，会产生三种类型的气泡：
- PP预热气泡：由于第一次前向传播的依赖性，除了第一阶段外的所有阶段都会出现，平均占训练时间的5.0%。
- PP冷却气泡：由于最后一次反向传播的依赖性，除了第一阶段外的所有阶段都会出现，平均占训练时间的9.2%。
- 其他PP气泡：由于其他前向和反向传播的依赖性，除了最后阶段外的所有阶段都会出现，占训练时间的8.7%。例如，在PP预热阶段之后，由于初始反向传播的依赖性，会立即出现PP气泡。此外，在MLLM模型异构性导致流水线阶段不平衡的情况下，还会出现图2中未描绘的额外流水线气泡。
3. 张量并行（TP）中的通信。张量并行涉及将单个层划分到多个GPU上，这需要在前向和反向传播期间进行通信以同步GPU。在Megatron-LM中，一个transformer层的每次前向或反向传播需要两个all-gather和两个reduce-scatter核函数【14】。图3详细展示了两个GPT-175B【5】层前向传播过程中的CUDA计算和通信核函数。绿色核函数代表CUDA通信流中的all-gather通信，而蓝色核函数代表reduce-scatter通信。计算流在这些通信期间处于空闲状态。通常，这些TP气泡持续时间为亚毫秒级，平均约为300微秒。然而，在MLLM训练期间，存在数千个TP气泡，它们共同占用了11.2%的训练时间。

2.3 挑战

核心思想与关键观察。为了最小化MLLM训练中的气泡，我们的目标是利用MLLM独特的双组件结构，即由编码器和LLM骨干组成。我们有两个关键观察：首先，MLLM训练过程中的大多数气泡发生在LLM骨干的前向和反向传播期间，根据生产数据，大约90%的气泡源于LLM的通信。其次，由于编码器的参数规模较小，其所需的计算操作（FLOPs）少于LLM骨干【4, 7, 10, 17, 18】。基于这些见解，我们提出在LLM通信期间出现的LLM气泡中调度编码器计算，以最小化整个MLLM训练过程中的气泡。我们识别出将编码器计算调度到LLM气泡中的三个主要挑战。

挑战1：只有少数GPU同时拥有编码器和LLM的模型状态。当前的训练系统【20, 39】采用流水线并行，将MLLM作为一个单一流水线分布在多个GPU上。由于编码器和LLM之间的依赖关系，编码器层被分配到初始的流水线阶段，而LLM层则被分配到后续的流水线阶段。因此，通常只有一个流水线阶段同时包含编码器和LLM层。如图4所示，使用3D并行（DP=1, PP=4, TP=2）将MLLM并行化到8个GPU上，只有流水线阶段1中的2个GPU同时持有编码器和LLM的模型状态。其余6个GPU因为缺少编码器模型状态，无法在LLM气泡期间执行编码器计算。

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

挑战3：亚毫秒级的LLM气泡。现有的框架如MegaScale【13】和Megatron-LM【20】通常在层级别进行调度。然而，LLM中的气泡持续时间范围很广，从亚毫秒（TP气泡）到数百毫秒（DP气泡）不等。例如，图3中所示的TP气泡在不同LLM层的前向和反向传播中平均约为300微秒。这个时长甚至不足以完成单个编码器层的前向或反向传播。举例来说，一个ViT-22B层通常需要约1.4毫秒来完成前向传播，2.0毫秒来完成反向传播。

3 设计决策与系统概览

我们讨论驱动Optimus设计的核心决策，并提供Optimus的概览。下一节将讨论详细设计。

3.1 设计决策

设计决策1：通过分离式并行共置编码器和LLM。为确保每个GPU都拥有编码器和LLM的模型状态，我们提议为所有GPU上的编码器和LLM分配独立的并行计划。如图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），它通过确保所有GPU都持有编码器和LLM的模型状态来解决挑战1；以及气泡调度器（bubble scheduler），它解决了挑战2（MLLM训练中的复杂依赖）和挑战3（亚毫秒级LLM气泡）。

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

# 算法 1: Optimus 工作流程
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

气泡调度器。气泡调度器负责将编码器计算调度到LLM气泡中。鉴于LLM训练流水线将数据划分为多个微批次，调度器按微批次调度编码器计算，并在微批次级别满足编码器-LLM的数据依赖（解决了§3.1中的挑战2）。此外，调度器将编码器计算分解到核函数粒度，以利用LLM训练期间的亚毫秒级气泡（TP气泡）（解决了§3.1中的挑战3）。当前的设计基于1F1B交错式流水线调度【20】，但它并非与此紧密耦合，可以适应其他流水线调度策略（如第6节所讨论）。

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

A2 方法细节

本节详细介绍Optimus的设计。4.1节描述模型规划器如何搜索编码器的并行计划；4.2节详细说明气泡调度器如何通过本地调度利用粗粒度和细粒度的气泡；4.3节讨论气泡调度器如何通过全局排序处理编码器-LLM的数据依赖；4.4节设计了多分支编码器模型中的气泡调度；4.5节分析了气泡调度算法的内存消耗。

4.1 模型规划器

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

搜索独立的并行计划。首先，规划器根据Megatron-LM【20】中的见解，确定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 = \frac{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]。

4.2 气泡调度

LLM气泡的通用模式。尽管不同GPU中的LLM气泡具有不同的开始时间和持续时长，但存在一个如图8所示的通用模式。在任何LLM计算开始之前，有一个大的单一气泡（DP all-gather气泡和PP-warmup气泡之和），在所有LLM计算结束后，也有一个大的单一气泡（PP-cooldown气泡和reduce-scatter气泡之和）。并且，在LLM计算之间穿插着许多小气泡（PP气泡和TP气泡）【14, 20, 25】。

# 算法 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

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

算法复杂度。我们的气泡调度算法复杂度很低。给定n个GPU，n的素因子数量为$np$，并行计划的搜索空间为$C_{np}^{2}+1$。微批次划分的数量为$O(N_{mb}^{m-1})$。因此，调度气泡的复杂度为$O(C_{np}^{2}+1 * N_{mb}^{m-1} * (F + B))$。在我们的实验中（见§5.3.2），找到最优调度所需的时间不到2分钟，这也是一次性的成本。

4.3 解决编码器-LLM依赖问题

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

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

获取依赖点 (GetEncLLMDep)。该函数获取编码器-LLM的前向和反向依赖点。鉴于交错式1F1B调度【20】是LLM训练中最有效的流水线调度之一，我们深入研究了该调度中数据依赖点$F_i$和$B_i$的具体细节。图12的上图展示了一个具有两个模型块的交错式1F1B调度实例。在这里，前向依赖点表示第一个流水线阶段（设备1）开始为第一个模型块执行前向计算的时刻（深蓝色表示），而反向依赖点则表示第一个流水线阶段（设备1）完成第一个模型块反向计算的时刻（深绿色表示）。

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

验证依赖 (CheckEncLLMDep)。该函数验证微批次级别的编码器-LLM依赖是否得到满足。考虑到已将编码器计算调度到气泡中，气泡调度器估算出编码器完成分布在不同编码器流水线上的微批次前向传播的时间。调度器将这些完成时间按升序排序为$EF_i$（全局排序），表示编码器为参与LLM流水线训练的微批次i完成前向操作的时间。如果对于所有微批次（$i=1...N_{mb}$），编码器在其指定的时间点$F_i$之前完成前向操作（$EF_i \le F_i$），则认为前向依赖得到满足。同样，如果对于每个微批次（$i=1...N_{mb}$），编码器的反向操作不早于时间点$B_i$开始（$EB_i \ge 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 多分支编码器调度

并行计划应用。为了支持具有多个编码器【6, 35】的MLLM，模型规划器为每个编码器独立地应用一个编码器并行计划($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内存限制。在实践中，由于$\phi_{enc}$很小（例如，最大的视觉编码器有220亿参数【9】）且k很小（例如，当使用bf16参数和fp32梯度以及分布式优化器时，k=6【1】），内存开销通常不到12%（见我们的评估§5.3.1）。

A4 实验环境

• 硬件配置:
  • 平台: 生产训练集群，包含数千个NVIDIA Hopper GPU。
  • GPU: 每个GPU拥有80GB内存和989 TFLOPS的计算性能。
  • 连接: 服务器内通过NVLink连接，服务器间通过高带宽RDMA网络连接。
• 模型架构:
  • 图像编码器: ViT-22B【9】、ViT-11B和ViT-5B（ViT-22B的缩减版本）。
  • 语言模型: LLAMA-70B【31】和GPT-175B【5】。
  • 详细模型配置见附录A。
• 软件配置与基线:
  • 实现: Optimus基于开源的Megatron-LM框架【1】开发。
  • 基线系统:
    1. PyTorch FSDP【38】: 一种分布式数据并行训练模块。
    2. Alpa【39】: 一种自动生成并行执行计划的编译器系统。
    3. Megatron-LM【20】: SOTA的LLM训练框架，将多模态编码器置于第一个流水线阶段进行适配。
    4. Megatron-LM balanced: 一个我们实现的基线，使用动态规划算法来平衡MLLM各子模块在不同流水线阶段的层划分，以实现大致相等的计算量。
• 性能指标:
  • 迭代时间 (Iteration time): 完成一次训练迭代所需的时间。
  • 模型FLOPS利用率 (MFU)【8】: 衡量硬件计算效率的指标。
  • 所有报告的性能数据都是在10次预热迭代后，对300次训练迭代取平均值。

A4 实验结果

5.2 端到端性能

5.2.1 弱扩展性实验

• 实验内容: 遵循常见的机器学习实践，随GPU数量增加而扩展模型规模（具体配置见表3），评估Optimus和基线方法的弱扩展训练性能。
• 实验结果:
  • 如图15所示，Optimus相比Megatron-LM实现了高达1.22倍的加速，相比Megatron-LM balanced实现了1.18倍的加速。
  • Alpa和FSDP在处理这些大型模型时遇到了GPU内存不足（OOM）的问题。
  • 在一个较小的MLLM模型（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的基线方法的强扩展性。
• 实验结果:
  • 如表5所示，Optimus相比Megatron-LM最多可将迭代时间减少21.3%，相比Megatron-LM balanced最多减少20.5%。
  • 随着GPU数量增加，Optimus相对于基线方案的加速效果更明显。Optimus的MFU保持稳定，而基线的MFU在更大规模下有所下降。这是因为固定的批次大小在更多GPU上导致了更高的气泡比例，使Optimus能调度更多编码器计算到LLM气泡中。

5.2.3 多编码器MLLM实验

• 实验内容: 在512个GPU上，使用批次大小256，评估Optimus和Megatron-LM在多编码器MLLM（配置见表6）上的训练性能。
• 实验结果:
  • 如图16所示，Optimus在这些MLLM上分别实现了高达1.25倍、1.26倍和1.27倍的加速。
  • 加速效果更显著的原因是，Megatron-LM将所有编码器都放在第一个流水线阶段，由于编码器总参数量更大，导致了更严重的流水线不平衡。

5.3 微基准测试

5.3.1 Optimus内存消耗

• 实验内容: 测量Optimus和基线在训练不同规模MLLM（见表3）时的GPU内存消耗。
• 实验结果:
  • 如图17所示，与内存效率最高的基线相比，Optimus在不同模型上的最大GPU内存开销为12%。
  • 值得注意的是，对于模型C和模型D，Optimus使用的GPU内存少于部分或全部基线。这是因为基线的计算负载划分策略可能因编码器和LLM层中隐藏大小的不同而导致内存不平衡。

5.3.2 气泡调度器算法

• 实验内容: 在单核CPU上运行气泡调度器算法，为强扩展性实验（5.2.2节）中的场景计算调度方案。引入“调度效率”指标（$Eff_{coarse}$和$Eff_{fine}$），量化可被有效调度到LLM气泡中的编码器计算百分比。
• 实验结果:
  • 如表7所示，随着GPU数量增加，调度效率更高。这是因为固定的批次大小导致每个LLM流水线的微批次数量减少，而DP和PP气泡的固定时长使得气泡比例增加。
  • 启用细粒度气泡利用相比仅使用粗粒度利用，效率最多可提升1.67倍。
  • 算法的运行时间随着LLM流水线中微批次数量的减少而减少，这与算法复杂度分析一致。

A7 补充细节

6 讨论

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

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

其他流水线调度。我们使用广泛应用的Megatron-LM交错式1F1B流水线调度进行MLLM训练。然而，存在其他可能在特定场景下性能更优的流水线调度（例如，Chimera【16】和zero-bubble pipeline【23】）。Optimus的气泡调度与这些流水线调度优化是正交的，当分析并解决了特定的编码器-LLM依赖关系后，Optimus可以应用于其他流水线调度。

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

排除DiffusionPipe和DistTrain作为基线的原因。我们明确排除了DiffusionPipe【30】和DistTrain【37】作为我们的基线，原因如下。DiffusionPipe专为扩散模型设计，并且仅在小规模集群（≤64个GPU）上进行过评估。其关注的模型家族和规模不同，使其不适合与Optimus进行比较，后者目标是在数千个GPU上进行大规模多模态LLM（MLLM）训练。DistTrain依赖于一种简化的模型划分策略，当应用于MLLM的异构结构时，可能导致严重的流水线不平衡。此外，DistTrain不是开源的，无法进行直接的实证比较。

7 相关工作

多模态训练。Pytorch FSDP训练【38】仅支持数据并行，效率低于混合并行策略。Alpa【39】能为各种模型自动化并行，但缺点是不支持最先进的1F1B交错式流水线并行【20】，并且比优化的Megatron-LM框架【25】需要更多内存，同时由于其对编码器和解码器的统一视角，也错失了流水线优化的机会。DistMM【12】为协调多个并行编码器提供了解决方案，但它是为对比学习设计的，忽略了解码器，在综合训练效率上留下了空白。DiffusionPipe【30】和DistTrain【37】是多模态训练领域的另外两项工作，各自的局限性已在前一节中概述。

气泡减少。以往减少“气泡”的努力从不同角度解决了这个问题。1F1B交错式流水线【20】技术通过将模型分块并在不同阶段交替处理这些块来最小化气泡，而零气泡流水线【23】方法则进一步将反向传播计算细粒化以消除气泡。然而，在实践中，零气泡流水线调度无法完全消除所有流水线气泡，因为它需要修改优化器，这引发了对端到端模型收敛性的担忧。另一方面，异步张量并行【26】和谷歌的重叠技术【32】旨在将张量并行通信与计算重叠，但受限于特定的硬件配置，并且随着计算能力的提升，难以保持完全重叠。

气泡利用。Pipefisher【22】利用跨多个训练步骤的流水线气泡来完成K-FAC计算，而我们的方法在单个同步训练步骤内操作，专注于即时优化。Hydro的Bubble Squeezer【11】利用GPT模型的气泡来执行独立任务，如超参数调整，这无法提升训练步骤本身的性能。Bamboo【29】利用流水线气泡进行冗余计算，以减轻在易失性实例上训练时抢占的影响，其基于的假设是后续流水线阶段承载更多层，这在大型语言模型（LLM）训练场景中通常不成立。

A5 结论

我们提出了Optimus，一个分布式的MLLM训练系统，它通过在LLM气泡中调度编码器计算来减少端到端的MLLM训练时间。为了减少MLLM训练期间的GPU气泡，Optimus分别对多模态编码器和LLM骨干进行划分，并在LLM气泡中调度编码器计算。我们考虑内存和计算资源约束，为编码器搜索最优的并行计划，从而在GPU之间平衡编码器计算以填充气泡。Optimus进一步采用了一种气泡调度算法，以解决编码器-LLM的依赖关系，并选择最优的调度方案，将核函数级别的编码器计算填充到亚毫秒级的LLM气泡中。我们广泛的实验表明，在使用ViT-22B和GPT-175B模型，并在3072个GPU上进行训练时，Optimus相比基线可将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】将层分配到不同的虚拟阶段。遵循Alpa【39】的方法，该DP算法旨在最小化最慢阶段的延迟，以减少流水线调度的端到端延迟。给定流水线并行大小PP和配置的V个模型块，DP算法寻求最小化最慢虚拟阶段的延迟。它确定了将层分布在这V × PP个虚拟阶段上的最优层划分策略。

公式定义。我们定义函数$F(l,m)$来表示当l层被分配到前m个虚拟阶段时，单个虚拟阶段的最大延迟。计算从$F(l, 1) = \sum_{i=1}^{l} 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。