MegaScale-Omni: A Hyper-Scale, Workload-Resilient System for MultiModal LLM Training in Production

发表时间: 2026-05 · arXiv:2605.08962 (ByteDance Seed)

作者/机构: Chunyu Xue(1,2), Yangrui Chen(1), Jianyu Jiang(1), Ningxin Zheng(1), Junda Feng(1), Jingji Chen(1), Shixiong Zhao(1), Shen Yan(1), Yi Lin(1), Lei Shi(1), Zanbo Wang(1), Lishu Luo(1), Faming Wu(1), Haibin Lin(1), Xin Liu(1), Yanghua Peng(1,†), Quan Chen(2,†)

(1)ByteDance Seed, (2)Shanghai Jiao Tong University


A1 主要贡献

本文的核心问题在于,现有的多模态大语言模型(MLLM)训练系统在处理动态工作负载时效率低下。这种动态性主要来源于两个因素:训练过程中不同模态数据集的混合比例动态变化,以及多模态输入样本的长度分布差异巨大。现有系统通常静态地耦合编码器(encoder)和LLM主干的资源分配与模型并行化策略,导致在编码器工作负载增加时出现计算瓶颈、内存溢出(OOM)问题,或在工作负载减少时导致设备闲置。

为了解决这些问题,本文的研究目标是设计一个能够适应动态工作负载并支持超大规模部署的工业级MLLM训练系统。该系统名为MegaScale-Omni,它基于一种名为编码器-LLM多路复用(encoder-LLM multiplexing)的训练方案,其核心思想是解耦编码器和LLM的并行化策略,并将它们并置在共享资源上进行协同调度。

本文的主要创新点和贡献如下:

  1. 提出解耦的并行化策略以适应超大规模和可变长度工作负载

    • 针对编码器,提出了长短序列并行(long-short sequence parallelism, LSSP),该策略统一了数据并行(DP)、Ulysses序列并行(SP)和ZeRO2/3,以高效处理可变长度的样本。
    • 针对LLM主干,集成了成熟的5D并行化技术,并设计了通信高效的并行化布局,以减轻执行干扰。
  2. 提出统一的编码器-LLM表示和工作负载弹性的联合流水线

    • 为了实现灵活、可扩展的编码器与LLM的并置,设计了统一的编码器-LLM表示,将编码器抽象为LLM流水线的“锚点”,封装了复杂的并行化和通信细节,避免了对LLM底层代码的侵入式修改。
    • 基于此表示,提出了一种新的编码器-LLM联合流水线范式,其特点是均匀、按需地插入编码器计算,从而在动态工作负载下保持流水线结构的稳定性。
  3. 提出针对多模态数据和超大规模并行化的工作负载均衡技术

    • 在数据加载阶段,通过在去中心化的数据加载器中采用分组重排序(grouped reordering)技术,来平衡各个编码器的工作负载。
    • 在将多模态嵌入(embeddings)从编码器重分发(resharding)到LLM计算单元时,通过自适应样本分片(adaptive sample sharding)对称分发(symmetric dispatching)技术,来平衡LLM各计算单元的计算负载并减少通信瓶颈。

MegaScale-Omni已在公司内部署,作为数千张GPU规模的大型MLLM训练任务的基础设施。实验证明,在生产级的动态工作负载下,与四种最先进的系统相比,MegaScale-Omni的吞吐量提升了1.27倍至7.57倍。

图1 多模态LLM架构和数据处理。
图1 多模态LLM架构和数据处理。
图2 超大规模MLLM训练中的动态工作负载及现有系统导致的吞吐量下降。
图2 超大规模MLLM训练中的动态工作负载及现有系统导致的吞吐量下降。

A3 背景知识与设计动机

2.1 工业级多模态训练

图3 MLLM混合训练与混合打包。
图3 MLLM混合训练与混合打包。

2.2 生产环境中动态工作负载的特征

图4 训练方案中跨模态和前5大任务的动态比例。-L/-S表示长/短上下文。
图4 训练方案中跨模态和前5大任务的动态比例。-L/-S表示长/短上下文。
图5 图像(OpenImages, RefCOCOg)、音频(LibriSpeech, GigaSpeech)和文本(BytedLong, BytedOCR)模态的数据分布,以编码后的样本长度衡量。
图5 图像(OpenImages, RefCOCOg)、音频(LibriSpeech, GigaSpeech)和文本(BytedLong, BytedOCR)模态的数据分布,以编码后的样本长度衡量。

2.3 动态工作负载下的编码器-LLM多路复用

2.3.1 为什么现有系统效率低下?
2.3.2 机遇

A2 方法细节

3 系统概览

MegaScale-Omni是一个为动态工作负载适应和数千GPU规模部署而设计的生产级MLLM训练系统。它基于§2.3.2中提出的编码器-LLM多路复用方案,解决了实现高效系统设计的三个主要挑战:

系统架构:如图6所示,MegaScale-Omni架构包含两个主要组件:模型多路复用器(Model Multiplexer)工作负载均衡器(Workload Balancer)
* 模型多路复用器:作为核心组件,它通过为编码器采用长短序列并行(§4.1)和为LLM主干采用全面的5D并行来解耦编码器-LLM的并行化。为了灵活地在GPU上并置编码器和LLM,它通过将编码器抽象为LLM流水线代码的“锚点”来提供统一的模型表示(§4.2)。利用这些表示,多路复用器通过一个具有均匀、按需编码器插入的工作负载弹性流水线来联合编排编码器-LLM的执行(§4.3)。
* 工作负载均衡器:为缓解动态多模态工作负载引起的不平衡,均衡器通过分组重排序和零冗余过滤技术增强了去中心化的数据加载器(§5.1)。在编码器和LLM之间传输数据时,均衡器通过自适应样本分片和对称分发来重分片嵌入和梯度(§5.2)。

图6 MegaScale-Omni的系统架构。
图6 MegaScale-Omni的系统架构。

4 模型并行化

本节详细描述基于编码器-LLM多路复用方案构建的MegaScale-Omni的系统设计。

4.1 解耦的并行化
4.1.1 编码器的长短序列并行
图7 编码器的长短序列并行。每个节点维护一个并行化编码器的实例。
图7 编码器的长短序列并行。每个节点维护一个并行化编码器的实例。
4.1.2 编码器-LLM并行化布局
图8 编码器-LLM并行化布局。
图8 编码器-LLM并行化布局。
4.2 并置表示
图9 原生LLM流水线代码、嵌套编码器插入和统一编码器表示的图示。
图9 原生LLM流水线代码、嵌套编码器插入和统一编码器表示的图示。
EncoderAnchor ([ ViT , USM ] , cfg ( zero3 = True ) )
工程师只需在其自定义的`step_func()`中将锚点挂钩到LLM阶段,而无需侵入式地插入代码。例如,要将第5-8和9-12个编码器微批次插入到PP0-3上第5和第6个LLM前向微批次之前,相应的锚点注册如下:
anchor.hook(llm_stage, True) # step_func()
然后,编码器和LLM微批次的数据流在一个类似JSON的格式中定义: 数据流定义 其中,键表示编码器微批次索引,值`(pp, [left, right])`表示在流水线rank `pp`上,插入到第`right`个LLM微批次之前(在第`left`个之后)。负值表示反向微批次。在运行时,`pp_schedule`会暴露给锚点以指示下一个编码器微批次。数据流必须满足编码器和LLM之间的数据依赖关系。
4.3 联合编排
图10 动态工作负载下具有工作负载弹性的编码器-LLM流水线调度图示。
图10 动态工作负载下具有工作负载弹性的编码器-LLM流水线调度图示。

5 工作负载均衡

MegaScale-Omni的效率与工作负载均衡性能密切相关。因此,我们引入了专门的技术来解决动态工作负载引起的不平衡问题。

5.1 均衡编码器数据加载
图11 去中心化分组重排序工作流程。8个GPU,组大小为4,微批次大小为4。
图11 去中心化分组重排序工作流程。8个GPU,组大小为4,微批次大小为4。
5.2 均衡编码器-LLM数据重分片
图12 为三模态均衡编码器-LLM重分片。编码器和LLM rank并置在2个GPU上。
图12 为三模态均衡编码器-LLM重分片。编码器和LLM rank并置在2个GPU上。

6 实现与优化

MegaScale-Omni基于Megatron-LM【40】构建,并已集成到我们公司内部的代码库中,作为大多数大规模MLLM训练任务的基础。


A4 实验环境

工作负载 LLM Backbone ViT Encoder #GPUs 全局批大小
A LLaMA-7B ViT-L/14 64 256
B LLaMA-13B ViT-g/14 128 512
C GPT-45B ViT-G/14 256 1024
D GPT-175B ViT-H/14 512 2048


表1 评估中主要使用的模型和工作负载。


A4 实验结果

7.2 端到端性能

图13 在Workload-A/B/C/D和不同图文混合比例下的训练吞吐量(每秒处理的token数),固定序列长度为16K/8K。缺失的条形表示在指定配置下发生OOM问题。
图13 在Workload-A/B/C/D和不同图文混合比例下的训练吞吐量(每秒处理的token数),固定序列长度为16K/8K。缺失的条形表示在指定配置下发生OOM问题。
图14 Workload-C和256个GPU在序列长度扩展到16K和32K时的训练吞吐量。
图14 Workload-C和256个GPU在序列长度扩展到16K和32K时的训练吞吐量。
图15 Workload-B和128个GPU在不同混合比例下的内存占用,序列长度为16K,PP度为4。
图15 Workload-B和128个GPU在不同混合比例下的内存占用,序列长度为16K,PP度为4。
图16 Workload-B和128个GPU的MFU,(a)在不同混合比例下(序列长度16K),(b)在不同序列长度下(图文混合比例7:3)。
图16 Workload-B和128个GPU的MFU,(a)在不同混合比例下(序列长度16K),(b)在不同序列长度下(图文混合比例7:3)。
图17 Workload-B和128个GPU在三种扩展策略下的三模态实验。
图17 Workload-B和128个GPU在三种扩展策略下的三模态实验。

7.3 消融实验

图18 在Workload-A/B、图文比例7:3、序列长度16K下的性能分解。
图18 在Workload-A/B、图文比例7:3、序列长度16K下的性能分解。
图19 Workload-B和128个GPU在不同并行策略下encoder-LLM多路复用的性能,图文比例7:3,序列长度16K。
图19 Workload-B和128个GPU在不同并行策略下encoder-LLM多路复用的性能,图文比例7:3,序列长度16K。
图20 Workload-B和128个GPU在不同重排序组大小下的吞吐量和开销权衡。
图20 Workload-B和128个GPU在不同重排序组大小下的吞吐量和开销权衡。

7.4 超大规模训练经验

图21 在我们的生产环境中,使用数千个GPU进行超大规模MLLM训练的性能。
图21 在我们的生产环境中,使用数千个GPU进行超大规模MLLM训练的性能。

A7 补充细节