作者/机构: Zili Zhang∗, Yinmin Zhong∗, Ranchen Ming†, Hanpeng Hu†, Jianjian Sun†, Zheng Ge†, Yibo Zhu†, Xin Jin∗
(*北京大学 †StepFun)

A1 主要贡献

本文提出 DistTrain，一个高效且自适应的框架，旨在通过解决多模态异构性问题，革新多模态大语言模型（Multimodal LLMs）的训练方式。在大型集群上，DistTrain 实现了接近单模态LLM训练的SOTA级别的模型FLOPs利用率（MFU），并能有效扩展至数千个GPU。

  • 核心问题：训练多模态LLM面临两大挑战：模型异构性和数据异构性。

    • 模型异构性：多模态LLM通常包含模态编码器、LLM主干和模态生成器三个模块（如下图Fig 1所示），这些模块在模型大小和算子复杂度上差异巨大，导致流水线并行中出现严重的“气泡”（pipeline bubbles），降低GPU利用率。
    • 数据异构性：多模态输入数据本质上是复杂和非结构化的，不同模态数据的异构性会导致微批次间（inter-microbatch）和微批次内（intra-microbatch）的训练掉队者（stragglers），延长训练时间并加剧流水线气泡。这些挑战共同导致在生产级训练中MFU低至约20%。

  • 研究目标：开发一个高效且可扩展的训练框架，以最小化训练成本并加速模型训练。

  • 核心创新点 (Disaggregated Training - 分离式训练)：

    • 分离式模型编排 (Disaggregated Model Orchestration)：为了解决模型异构性，本文提出GPU训练分离，将模态编码器、LLM主干和模态生成器解耦。在此基础上，设计了分离式模型编排方法，通过搜索复杂的设计空间，为每个模块选择最优的资源和并行配置，从而最小化模型异构性导致的流水线气泡。
    • 分离式数据重排 (Disaggregated Data Reordering)：为了解决数据异构性，本文利用CPU预处理分离，将数据预处理（如解压和重排）与GPU训练解耦，实现无额外开销的数据重排。
      • 微批次间重排 (Inter-microbatch reordering)：对数据样本进行重排，以均衡不同数据并行组的负载。
      • 微批次内重排 (Intra-microbatch reordering)：针对（交错式）1F1B流水线方案对微批次进行重排，以最小化流水线气泡。
    • 系统优化：实现了一个自研的集合通信库StepCCL，用于将GPU通信与计算重叠，隐藏通信开销。

图1：多模态LLM的架构。
图1：多模态LLM的架构。

主要成果：在拥有数千GPU的生产集群上进行的实验表明，DistTrain 在训练一个72B参数的多模态LLM时，在1172个GPU上实现了54.7%的MFU，吞吐量最高比Megatron-LM高出2.2倍。

A3 背景知识与关键洞察

2.1 多模态大语言模型训练

大语言模型（LLM）。大语言模型（LLMs）【18，Attention is all you need，2017，Advances in Neural Information Processing Systems】、【19，Language models are unsupervised multitask learners，2019，OpenAI blog】、【20，Language models are few-shot learners，2020，Advances in Neural Information Processing Systems】通过在文本生成、翻译和摘要等广泛任务上取得SOTA性能，彻底改变了自然语言处理（NLP）。LLM的核心架构由一堆同构的Transformer层组成，利用自注意力机制捕捉文本中的上下文信息。LLM通常在大型文本语料库上进行无监督预训练，然后在特定任务的文本数据集上进行微调。训练过程可能需要在专用AI集群上花费数周甚至数月的时间，成本高昂，因此优化训练过程至关重要。

多模态LLM。单模态LLM仅限于处理文本数据，限制了其在多模态任务中的应用。因此，多模态LLM应运而生，通过将图像、音频和视频等多种模态集成到先进的LLM中来解决这一局限性。例如，GPT-4o【10，Hello gpt-4o，2024，https://openai.com/index/hello-gpt-4o/】通过视觉和听觉模态促进了与人类更自然的交互。此外，由于人类社会中以文本为主的数据是有限的【23，Position: Will we run out of data? limits of llm scaling based on human-generated data，2024，Forty-first International Conference on Machine Learning】，利用多模态数据是持续扩展和增强LLM能力的必然选择。

多模态LLM架构。如图1所示，多模态LLM的架构包括三个模块：模态编码器、LLM主干和模态生成器【15，A survey on multimodal large language models，2023，arXiv】、【16，Mm-llms: Recent advances in multimodal large language models，2024，arXiv】、【24，Flamingo: a visual language model for few-shot learning，2022，Advances in Neural Information Processing Systems】。模态编码器（如用于图像的ViT【25，An image is worth 16x16 words: Transformers for image recognition at scale，2020，arXiv】和用于音频的Beats【26，Beats: Audio pre-training with acoustic tokenizers，2022，arXiv】）将不同模态的输入数据转换为中间表示（嵌入张量），然后通过输入投影层（如MLP和交叉注意力）投射到统一的嵌入空间。LLM主干（通常是Transformer模型，如GPT【19, 20】和Llama【22，Meta llama3，2024，https://llama.meta.com/】）处理多模态嵌入以辨别复杂的数据模式和模态间关系。LLM主干的输出数据随后由输出投影层进行精炼，为每个模态定制信息。最后，模态生成器（如用于图像的Diffusion【27，High-resolution image synthesis with latent diffusion models，2022，IEEE Conference on Computer Vision and Pattern Recognition】和用于音频的AudioLDM【28，Audioldm: Text-to-audio generation with latent diffusion models，2023，arXiv】）将LLM处理过的信息转换回各自的模态输出进行生成。

多模态LLM训练。多模态LLM的训练需要同时训练所有三个模块。在不同训练阶段，某些模块会被冻结以稳定训练损失并提高模型效果【15，A survey on multimodal large language models，2023，arXiv】。训练数据的输入序列包含文本、图像和音频令牌，这些来自不同模态的数据被分词成交错的子序列，形成固定长度的训练序列【12，Chameleon: Mixed-modal early-fusion foundation models，2024，arXiv】。如果输入序列不足，则用零令牌填充以进行批处理。不同的模态子序列由特殊令牌分隔，并通过特定模态的编码器处理。

现有训练框架。为了在大型集群上训练多模态LLM，事实上的解决方案是利用高效且强大的Transformer模型训练框架Megatron-LM。Megatron-LM对整个模型采用统一的并行策略，结合张量并行（TP）和流水线并行（PP）将模型参数分布在多个GPU上，并利用数据并行（DP）分发训练数据。对于多模态LLM，Megatron-LM通过将多模态模块视为LLM内的附加层来轻松扩展，并为模态编码器和生成器增加额外的PP阶段。如图2所示，LLM主干中使用的相同TP策略也应用于这两个多模态模块。如果模态编码器和生成器不够大，它们会在TP组内的GPU上复制以最大化资源利用。DP策略也与LLM主干相同。投影器与模态编码器和生成器共同部署，并在TP组的GPU上复制。

图2：使用Megatron-LM训练多模态LLM。
图2：使用Megatron-LM训练多模态LLM。

现有框架的局限性。这种训练框架因其僵化的模型编排方法（即多模态模块与LLM主干共享相同的DP和TP策略），引入了源于模型异构性和数据异构性的显著计算不平衡。它仅能实现约20%的MFU（§8.1），远低于训练单模态（纯文本）LLM时观察到的约50%的MFU【29，Megascale: Scaling large language model training to more than 10,000 gpus，2024，USENIX NSDI】。

2.2 模型异构性

模型异构性导致的计算不平衡。这是训练多模态LLM的第一个挑战。由于算子和输入的不同，多模态LLM中的每个模块（模态编码器、LLM主干、模态生成器）都有不同的计算需求。例如，作为模态编码器的ViT由较窄的Transformer层（即隐藏层尺寸较小）构成，而LLM主干则建立在较宽的Transformer层之上（即隐藏层尺寸较大）。同时，作为模态生成器的Diffusion模型则使用卷积和注意力层的组合（即U-Net）。这种架构多样性导致每个模块的计算时间各不相同。图3展示了在使用Megatron-LM时，不同输入配置下的前向传播时间差异显著。

图3：不同输入配置下的前向传播时间。
图3：不同输入配置下的前向传播时间。

流水线气泡问题。模块间的计算不平衡导致了流水线并行中的两种类型的“气泡”。第一种出现在模态编码器和生成器阶段（如图4(a)所示），原因是它们未能充分利用分配给它们的GPU资源。第二种出现在LLM主干的阶段（如图4(b)所示），这是因为编码器和生成器密集的计算需求延长了它们的阶段持续时间，由于流水线依赖性，LLM阶段被迫等待多模态阶段完成，从而产生流水线气泡。后一个问题在大型多模态LLM训练中尤其突出，因为大部分GPU资源都分配给了LLM主干。这些源于模型异构性的流水线气泡，显著降低了训练期间的MFU。

图4：由模型异构性引起的两类流水线气泡。
图4：由模型异构性引起的两类流水线气泡。

2.3 数据异构性

数据异构性导致的计算不平衡。这是第二个挑战。多模态LLM训练的每个输入序列（即训练样本）由交错的模态子序列组成，这些子序列呈现出高度倾斜的分布。通过对LAION-400M数据集【30，Laion-400m: Open dataset of clipfiltered 400 million image-text pairs，2021，arXiv】进行数据特征分析，发现文本和图像子序列的大小（图5(a)和5(b)）以及每个训练样本的图像子序列数量（图5(c)）都显示出倾斜分布。不同的样本大小导致模态编码器和生成器阶段的计算时间变化。

图5：多模态LLM训练中的数据异构性。
图5：多模态LLM训练中的数据异构性。

图5(b)：图像子序列大小的分布。 图5(c)：图像子序列数量的分布。
图5(b)：图像子序列大小的分布。 图5(c)：图像子序列数量的分布。

数据异构性导致的掉队者（stragglers）。这种数据异构性在模态编码器和生成器的流水线并行（PP）阶段导致了微批次内（intra-microbatch）和微批次间（inter-microbatch）的掉队者，加剧了计算不平衡并进一步降低GPU利用率。需要注意的是，LLM主干内的所有微批次计算时间相同，因为序列长度是固定的。我们不考虑全局批次之间的数据异构性，因为每个全局批次包含大量随机打乱的训练样本，这有效地平滑了数据异构性。

微批次内掉队者（Intra-microbatch straggler）。当特别大的训练样本减慢训练速度时，就会出现微批次内掉队者，因为不同的数据并行（DP）组处理大小可变的训练样本。如图6所示，第一个DP组（DP1）在两个微批次内处理两个大的训练样本，而第二个DP组（DP2）处理两个较小的样本，完成得更快。因此，DP1落后并成为掉队者，从而延迟了整个训练过程。

图6：微批次内掉队者（DP组之间）。
图6：微批次内掉队者（DP组之间）。

微批次间掉队者（Inter-microbatch straggler）。微批次间掉队者源于微批次之间的流水线不平衡。如图7所示，第一个流水线阶段是模态编码器，其后是一个LLM主干阶段。图7(a)展示了没有数据异构性的流水线，模态编码器处理每个微批次的时间一致。相比之下，图7(b)描绘了存在数据异构性的流水线，模态编码器的前向传播时间在不同微批次之间差异显著。掉队者（即微批次a）显著延迟了后续PP阶段的训练过程，导致了大的流水线气泡。

图7：微批次间掉队者。
图7：微批次间掉队者。

A2 方法细节

3 DistTrain 概览

DistTrain是一个高效且自适应的框架，通过解决多模态异构性来重塑多模态LLM的训练。它通过分离式模型编排（§5）消除模型异构性，并通过分离式数据重排（§6）利用数据异构性，此外还采用了针对多模态LLM训练的系统优化（§7）。

图8：DistTrain 概览。
图8：DistTrain 概览。

DistTrain管理器 (DistTrain manager)。训练开始前，DistTrain使用一个训练管理器来确定多模态LLM中每个模块的资源分配和并行策略。该调度器首先从用户处收集模型架构和训练配置，并随机抽样一部分训练数据以分析数据分布。利用这些信息，它运行一系列基准测试训练试验，并构建一个带有线性插值的性能分析器来估计每个模块的计算和通信时间。基于分析结果，训练管理器通过分离式模型编排为特定训练任务决定最优的资源分配和并行策略。

DistTrain初始化器 (DistTrain initializer)。DistTrain接着为模态编码器、LLM主干和模态生成器分别初始化GPU训练分离。DistTrain为每个并行单元分配不同数量的GPU，每个单元随后建立特定的通信组。单元从分布式文件系统加载模型检查点，并对模型参数和优化器状态进行分片。最后，DistTrain进行几次通信试验以预热系统并测试连通性。

DistTrain运行时 (DistTrain runtime)。在运行时，专用的CPU节点（即CPU预处理分离）从分布式文件系统检索训练样本进行预处理。它执行分离式数据重排，在一个全局批次内对训练样本进行重排，而不破坏同步训练语义【31，Techniques and systems to train and serve bigger models，2022，https://icml.cc/virtual/2022/tutorial/18440】。这种重排有效地消除了微批次间和微批次内的数据异构性。在每次迭代中，主训练过程从CPU节点异步接收预处理后的数据。数据随后在训练流水线中依次通过模态编码器、LLM主干和模态生成器。最后，它通过all-gather操作同步梯度和模型参数，采用ZERO-1优化【32，Zero: Memory optimizations toward training trillion parameter models，2020，International Conference for High Performance Computing, Networking, Storage and Analysis】和混合精度训练【33，Mixed precision training，2017，arXiv】。此外，DistTrain采用一个专用进程定期异步地将模型检查点保存到分布式文件系统以实现容错。

4 分离式训练 (Disaggregated Training)

分离式训练的核心原则。为了解决多模态LLM训练中的模型和数据异构性，我们引入了DistTrain的核心原则：分离式训练。这包括GPU训练分离和CPU预处理分离。GPU训练分离为跨三个模块的自适应模型编排提供了机会，以解决模型异构性。CPU预处理分离则促进了数据预处理，以微不足道的运行时开销解决数据异构性。

4.1 GPU训练分离

GPU训练分离的拓扑结构。图9展示了DistTrain中分离式GPU训练的训练拓扑。与Megatron-LM中僵化的模型编排（图2）不同，DistTrain能够自适应地调整资源分配和并行策略。例如，DistTrain可以为模态编码器分配4个GPU（DP=2，TP=2），为每个PP阶段的LLM主干分配12个GPU（DP=3，TP=4），为模态生成器分配4个GPU（DP=1，TP=4）。此外，投影器层与模态编码器或生成器共同部署，其副本数量根据需要进行调整。我们通过一个专用模块，即并行单元（parallelism unit），来实现GPU训练分离。

图9：DistTrain中的GPU训练分离。
图9：DistTrain中的GPU训练分离。

并行单元 (Parallelism unit)。在训练初始化时，需要根据资源分配和并行策略建立通信组。DistTrain引入了一个由一个或多个PP阶段组成的模块，即并行单元。每个单元可以采用自己的DP和TP策略，并形成一个特定的通信组。单元间的连接由一个通信代理（communication broker）促成，该代理在并行单元之间桥接PP通信。用户只需为每个并行单元指定DP和TP配置，DistTrain会自动设置通信组和通信代理。DistTrain将模态编码器、LLM主干和模态生成器视为三个独立的并行单元。

4.2 CPU预处理分离

CPU预处理分离的动机与实现。在训练多模态LLM时，训练样本通常结合了轻量级的文本和重量级的多模态数据，后者显著增加了数据预处理时间。例如，一个典型的训练样本可能包含256个词的文本序列和十张1024×1024的RGB图像，文本只有几KB，而图像总共有120MB。预处理（如解压、调整大小和重排）这些样本可能需要几秒钟，并干扰共同部署的训练过程。DistTrain通过生产者-消费者模型将CPU数据预处理与GPU训练过程分离。生产者在专用CPU节点上运行，从分布式文件系统获取数据并与GPU训练过程异步地预处理训练数据。消费者，即主训练过程，接收这些预处理后的数据进行训练。生产者和消费者通过RPC调用进行通信，并在可用时使用RDMA网络以降低延迟。这种分离保证了CPU预处理不会干扰GPU训练过程，并实现了可忽略不计的数据预处理开销。

5 解决模型异构性

分离式训练使得可以根据训练工作负载在不同模块间进行分离式模型编排。我们首先将分离式模型编排问题形式化，以最小化每次迭代的训练时间。然后，我们提出了详细的算法来最优地解决由模型异构性引起的问题。

5.1 问题形式化

模型编排问题的定义。通过分离式训练，我们能够自适应地编排三个模块。现在的问题在于确定最优的资源分配和并行策略，以最小化每次迭代的训练时间。由于搜索空间巨大，尤其是在大型集群中，穷举搜索是不可行的。一个简单的解决方案是按每个模块的模型FLOPs比例分配资源，但这种方法因忽略了并行训练中的复杂模式而效果不佳。在深入探讨分离式模型编排之前，我们首先将优化问题形式化。

LLM主干的公式化。我们首先对LLM主干进行公式化，重点关注前向传播，因为后向传播时间与此类似。假设一次迭代的全局批量大小为$BS$，LLM主干的TP大小为$TP_{lm}$，PP和DP大小分别为$PP_{lm}$和$DP_{lm}$。分配给LLM主干的GPU数量为$y = TP_{lm} \times DP_{lm} \times PP_{lm}$。设整个LLM的前向传播时间（包括通信时间）为$C_{lm}(TP_{lm})$。因此，一个微批次在一个PP阶段的前向传播时间为$T_{lm} = \frac{C_{lm}(TP_{lm})}{PP_{lm}}$。此外，每次迭代的微批次数量为$\frac{BS}{DP_{lm}}$。

模态编码器和生成器的公式化。在DistTrain中，模态编码器被视为一个PP大小为$PP_{me}$的并行单元。设TP大小为$TP_{me}$，DP大小为$DP_{me}$。分配给模态编码器的GPU数量为$x = TP_{me} \times DP_{me} \times PP_{me}$。微批次大小由LLM主干决定，为$\frac{DP_{lm}}{DP_{me}}$。设整个模态编码器的前向传播时间（包括通信时间）为$C_{me}(TP_{me})$。因此，模态编码器中一个微批次在一个PP阶段的前向传播时间为$T_{me} = \frac{DP_{lm}}{DP_{me}} \times \frac{C_{me}(TP_{me})}{PP_{me}} = \frac{DP_{lm} \times TP_{me}}{x} \times C_{me}(TP_{me})$。类似地，模态生成器中一个PP阶段的前向传播时间为$T_{mg} = \frac{DP_{lm} \times TP_{mg}}{z} \times C_{mg}(TP_{mg})$，其中$z$是分配给模态生成器的GPU数量。

目标函数。基于上述分析，我们定义优化问题的目标函数，即一次迭代的训练时间。如图10所示，多模态LLM训练的流水线分为预热（warm-up）阶段和稳定（steady）阶段。预热阶段的持续时间为$T_{warmup} = T_{lm} \times PP_{lm} + T_{me} \times PP_{me} + T_{mg} \times PP_{mg}$，其公式如下：

Warm-up阶段时间公式
Warm-up阶段时间公式

稳定阶段的持续时间由PP阶段中的最大计算时间决定，为$T_{steady} = max(T_{lm}, T_{me}, T_{mg}) \times (\frac{BS}{DP_{lm}} - 1)$，其公式如下：

稳定阶段时间公式
稳定阶段时间公式

因此，目标函数是最小化$T_{iter} = T_{warmup} + T_{steady}$。对于后向传播，目标函数类似，只需将$C_{lm}, C_{me}, C_{mg}$从前向传播时间函数更改为前向和后向时间之和的函数。此公式适用于GPipe和1F1B流水线方案。

图10：多模态LLM训练流水线。
图10：多模态LLM训练流水线。

约束条件。除了目标函数，还必须考虑约束条件以确保训练的可行性。
1. 资源约束：分配给每个模块的GPU总数$x+y+z \le N$，其中$N$是集群中的GPU总数。
2. 内存约束：以LLM主干为例，内存分配包括模型参数、梯度、优化器状态和激活状态四部分。一个GPU上的模型参数和梯度内存为$\frac{P}{PP_{lm} \times TP_{lm}}$，优化器状态内存（使用ZeRO-1）为$\frac{S \times DP_{lm}}{y}$，激活状态的峰值内存为$\frac{A \times DPlm}{L \times PPlm}$。在1F1B中，第一个PP阶段需要存储$PP_{lm}$个微批次的激活状态。内存约束确保单个GPU上的总内存消耗不超过其容量。模态编码器和生成器的内存约束公式类似。

内存约束公式
内存约束公式

5.2 分离式模型编排

优化算法。该优化问题是非凸的，穷举搜索所有可能的变量值（$x, y, z$以及DP, TP大小）是不切实际的。我们的关键洞察是将这个非凸优化问题分解为一系列简化的凸问题，其中变量为资源分配$x, y, z$。我们将TP大小限制在一个小的离散集合（如[1,2,4,8]），并将DP大小设置为$BS$的因子，以平衡DP组间的计算。LLM主干的PP大小则根据$y, DP_{lm}, TP_{lm}$计算得出。这样，可能的并行策略集合是可管理的有限集。

凸优化转换。通过枚举所有可行的TP和DP大小，我们将原始优化问题转化为一组简化问题。在每个简化问题中，目标函数是关于$1/x, 1/y, 1/z$的线性函数和最大值函数，因此是凸函数。类似地，约束函数也是凸的。因此，简化后的优化问题是凸的，可以通过现有的求解器【34，Cvx: Matlab software for disciplined convex programming, version 2.1，2014】、【35，Scipy 1.0: fundamental algorithms for scientific computing in python，2020，Nature methods】高效地求解以获得最优解。算法1详细描述了这一过程，能有效地找到最优的资源分配和并行策略。

虚拟流水线并行 (VPP) 的适配。虚拟流水线并行（即交错式1F1B）通过将模型划分为更细粒度的虚拟PP（VPP）阶段来减少预热时间。在预热阶段，每个PP阶段启动一个VPP阶段的计算，预热时间被VPP大小除。为了使我们的公式与VPP对齐，我们根据VPP大小按比例减少预热时间。

Algorithm 1: 分离式模型编排算法
1. 初始化：设置最优迭代时间 $T_{iter}^*$ 为无穷大。
2. 枚举并行策略：遍历所有可能的TP大小（$TP_{me}, TP_{lm}, TP_{mg}$）和DP大小（$DP_{me}, DP_{lm}, DP_{mg}$）组合。
3. 求解凸优化问题：对于每种并行策略组合，将资源分配（$x, y, z$）作为变量，构建一个凸优化问题。
4. 使用凸优化求解器找到当前策略下的最优资源分配 $x^*, y^*, z^*$ 以及对应的最小迭代时间 $T_{iter}$。
5. 更新最优解：如果当前 $T_{iter}$ 小于 $T_{iter}^*$，则更新最优解。
6. 返回全局最优的资源分配和并行策略。

6 解决数据异构性

分离式训练使得数据预处理可以在几乎没有运行时开销的情况下进行。基于此，我们引入了无缝集成到数据预处理中的分离式数据重排，以解决数据异构性问题，且无额外开销。

6.1 微批次内重排 (Intra-microbatch Reordering)

核心洞察。为了解决微批次内掉队者问题，我们通过识别拥有最大训练样本的DP组来定位掉队者。如图6所示，第一个DP组因包含两个最大的训练样本而成为掉队者。为消除这种不平衡，我们提议按大小对全局批次内的训练样本进行重排。具体来说，如图11所示，我们将训练样本重排为序列[1, 3, 2, 4]。这一策略能更均匀地分配计算负载并提高可扩展性，且由于重排仅在单个全局批次内进行，因此不会破坏同步训练语义。

图11：微批次内重排。
图11：微批次内重排。

重排算法。我们采用贪婪数分区算法，该算法能保证近似比≤4/3【37，Approximation algorithms for maximin fair division，2020，ACM Transactions on Economics and Computation (TEAC)】。算法2总结了详细过程。该算法首先按样本大小升序对训练样本进行排序。然后，遍历所有训练样本，并将每个样本分配给当前计算负载最低的DP组。最后返回重排后的样本序列。该算法的时间复杂度为$O(n \log n + m \times n)$。

算法2：批内重排。
算法2：批内重排。

6.2 微批次间重排 (Inter-microbatch Reordering)

问题分析。第二个子问题是微批次间掉队者。数据异构性导致模态编码器和生成器中不同微批次的计算时间不同，掉队者微批次会通过产生大的流水线气泡来延长训练时间。在1F1B流水线方案（如图12所示）中，总迭代时间主要由流水线气泡和模态编码器第一个PP阶段的计算时间决定。我们抽象出一个创新概念——流水线间隔（pipeline intervals）。如图12所示，这些间隔通常由前向传播填充，但最后的$p-1$个间隔除外。掉队者微批次会延长这些间隔或增加未填充区域（即气泡）。