MegaScale-Omni: A Hyper-Scale, Workload-Resilient System for MultiModal LLM Training in Production
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的并行化策略,并将它们并置在共享资源上进行协同调度。
本文的主要创新点和贡献如下:
-
提出解耦的并行化策略以适应超大规模和可变长度工作负载:
- 针对编码器,提出了长短序列并行(long-short sequence parallelism, LSSP),该策略统一了数据并行(DP)、Ulysses序列并行(SP)和ZeRO2/3,以高效处理可变长度的样本。
- 针对LLM主干,集成了成熟的5D并行化技术,并设计了通信高效的并行化布局,以减轻执行干扰。
-
提出统一的编码器-LLM表示和工作负载弹性的联合流水线:
- 为了实现灵活、可扩展的编码器与LLM的并置,设计了统一的编码器-LLM表示,将编码器抽象为LLM流水线的“锚点”,封装了复杂的并行化和通信细节,避免了对LLM底层代码的侵入式修改。
- 基于此表示,提出了一种新的编码器-LLM联合流水线范式,其特点是均匀、按需地插入编码器计算,从而在动态工作负载下保持流水线结构的稳定性。
-
提出针对多模态数据和超大规模并行化的工作负载均衡技术:
- 在数据加载阶段,通过在去中心化的数据加载器中采用分组重排序(grouped reordering)技术,来平衡各个编码器的工作负载。
- 在将多模态嵌入(embeddings)从编码器重分发(resharding)到LLM计算单元时,通过自适应样本分片(adaptive sample sharding)和对称分发(symmetric dispatching)技术,来平衡LLM各计算单元的计算负载并减少通信瓶颈。
MegaScale-Omni已在公司内部署,作为数千张GPU规模的大型MLLM训练任务的基础设施。实验证明,在生产级的动态工作负载下,与四种最先进的系统相比,MegaScale-Omni的吞吐量提升了1.27倍至7.57倍。
A3 背景知识与设计动机
2.1 工业级多模态训练
-
多模态大语言模型(MLLM)的应用场景:MLLM已广泛应用于多种AI应用,如GPT-4o【34, OpenAI, 2024】, Qwen2.5-VL【5, Qwen2.5-vl technical report, 2025】和Seedream-3.0【14, Seedream 3.0 technical report, 2025】。在我们的生产环境中,典型的训练场景包括:(1) 视觉-语言(Vision-Language)【7, An introduction to visionlanguage modeling, 2024】,集成视觉编码器(如用于图像的ViT【12, An image is worth 16x16 words: Transformers for image recognition at scale, 2021】和用于视频的VideoSwin【31, Video swin transformer, 2021】)与LLM主干;(2) 视觉理解与生成【29, Mogao: An omni foundation model for interleaved multi-modal generation, 2025】,通过变分自编码器(VAE)【24, Auto-encoding variational bayes, 2022】理解视觉信息并根据文本提示合成图像;(3) 三模态【54, Google usm: Scaling automatic speech recognition beyond 100 languages, 2023】,使用通用语音模型(USM)【54】和离散语音编码器【44, A comparative study of discrete speech tokens for semanticrelated tasks with large language models, 2024】统一处理视觉、听觉和文本模态。
-
并行化策略:大规模MLLM训练涉及多种并行策略。(1) 数据并行(DP)将全局批次数据划分到多个GPU上,每个GPU复制模型参数和优化器状态。ZeRO【37, Zero: Memory optimizations toward training trillion parameter models, 2020】【57, Pytorch fsdp: experiences on scaling fully sharded data parallel, 2023】通过在GPU间分发模型状态来扩展DP。(2) 张量并行(TP)【40, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019】沿隐藏维度切分层参数。(3) 流水线并行(PP)【32, Pipedream: Generalized pipeline parallelism for dnn training, 2019】【33, Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】将层分组为阶段,并将全局批次拆分为多个微批次进行流水线执行。(4) 序列并行(SP)沿序列维度切分输入和激活值,其主要变体包括Megatron SP【25, Reducing activation recomputation in large transformer models, 2022】、Ulysses SP【20, Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, 2023】和上下文并行(CP)【3, Nvidia context parallel package, 2025】。(5) 专家并行(EP)【38, Deepspeedmoe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022】专为混合专家(MoE)模型设计。
-
混合训练与混合打包:现代大规模MLLM训练通过混合来自不同模态和任务领域的多个数据集来构建数据批次(图3(a))。单模态训练中的一个直观做法是将相同模态的样本打包成序列,然后将不同模态的序列批处理(图3(b))。然而,由于不同模态间工作负载的动态变化,这种方法会引入严重的序列间不平衡或大量填充。因此,工业实践采用混合打包(hybrid packing),允许跨模态打包样本以实现更好的对齐(图3(c))。不同模态的样本被打包成长度大致均匀的序列,然后组合成一个全局批次。
2.2 生产环境中动态工作负载的特征
- 动态的模态间比例:研究人员通常设计多阶段训练方案【5, Qwen2.5-vl technical report, 2025】【8, Internlm2 technical report, 2024】【16, Seed1.5-vl technical report, 2025】来调度数据集的处理顺序,并在每个阶段调整模态比例以实现不同目标。如图4所示,一个VLM训练任务为了增强长上下文能力,可能使用35%的512K长度文本、55%的图像和10%的视频样本。即使在同一模态内,不同任务领域的比例也不同(例如,47%的书籍文本 vs. 13%的代码文本),这会因样本长度可变而进一步改变模态间的相对工作负载。此外,一些方法在训练阶段中平滑地调整模态比例(例如每一步或几步调整一次),这使得在单个阶段内进行静态工作负载优化变得不可能。例如,一个三模态训练任务在最初的100亿个token中将图文混合比例设为1:1,然后逐渐增加音频模态的比例,最终达到13:74:13的图-音-文混合比例。
- 动态的数据分布:混合不同模态和任务领域的数据集加剧了样本长度分布的偏斜,如图5所示。例如,OpenImages【26, The open images dataset v4, 2020】数据集的平均样本长度为3.8K,而LibriSpeech【35, Librispeech: An asr corpus based on public domain audio books, 2015】和BytedLong(一个生产数据集)的平均长度分别为0.34K和6K。即使在同一模态内,RefCOCOg【53, Modeling context in referring expressions, 2016】的平均长度为1.4K,比OpenImages短2.71倍。这种样本长度的差异导致编码器在不同模态之间以及同一模态内部的工作负载不平衡。数据加载器通常按固定的样本数量比例(如1:1的图文样本比例)检索数据,忽略了偏斜的数据分布,导致一些编码器处理过长的样本,而另一些只处理短样本。
- 模态计算的动态性:不同模态间的计算异构性进一步放大了动态多模态数据集的影响。以一个三模态训练为例,序列长度为16K,批大小为32,图-音-文比例为4:4:2。ViT-1B【12】、USM-2B【54】和LLaMA-12B【41, Llama: Open and efficient foundation language models, 2023】分别消耗14.7K、4.54K和29.0K TFLOPs,因为它们包含不同的计算操作(例如USM中的卷积层)。考虑到音频样本通常比图像样本短,这种异构性进一步加剧了不同模态计算时间的方差。此外,一些特定于模态的操作也引入了工作负载的动态性。例如,ViT处理动态分辨率的图像【11, Patch n’ pack: Navit, a vision transformer for any aspect ratio and resolution, 2023】,通过固定的下采样策略(例如每token 3x14x14像素)生成可变长度的token。
2.3 动态工作负载下的编码器-LLM多路复用
2.3.1 为什么现有系统效率低下?
- 现有系统的局限性:当前的MLLM训练系统主要分为两类。第一类是类单模态(unimodal-like)系统(如Megatron-LM【40】, Transformers【49】),它们将编码器前置于第一个LLM流水线阶段,以替代文本嵌入层。这种方式下,整个MLLM像单模态LLM一样进行并行化,即编码器和第一阶段共享相同的资源和并行策略。当编码器工作负载增加时(例如,占MLLM总TFLOPs的30%),第一阶段会成为严重的计算和内存瓶颈,导致高达2.84倍的吞吐量下降和OOM问题(图2(c))。第二类是解聚(disaggregation)系统(如DistMM【18, DISTMM: Accelerating distributed multimodal model training, 2024】, DistTrain【55, Disttrain: Addressing model and data heterogeneity with disaggregated training for multimodal large language models, 2024】),它们静态地将编码器和LLM主干解耦为两个独立的模型,各自拥有固定、不同的资源和并行策略。当编码器和LLM之间的相对工作负载变化时,这种静态解聚会面临编码瓶颈(当编码器负载超出其分配资源时)或长时间的设备空闲(当编码器样本不足时)。另一个近期系统Optimus【13, Optimus: Accelerating large-scale multi-modal llm training by bubble exploitation, 2024】,虽然解耦了资源分配和并行化决策,但它是基于静态、合成的工作负载进行并行化,无法处理生产环境中的动态变化。
2.3.2 机遇
- 编码器-LLM多路复用方案:针对现有系统的局限性,本文提出了一种新的MLLM训练方案——编码器-LLM多路复用(encoder-LLM multiplexing),以实现具有工作负载弹性的工业级MLLM训练。其关键思想包括:
- 解耦MLLM并行化:为编码器和LLM主干采用不同的并行策略。具体来说,编码器应以高并发度和足够的资源进行并行化,例如在所有GPU上使用数据并行(DP),以避免成为LLM流水线的瓶颈。对于LLM,理想方式是继承单模态训练中所有成熟的并行技术,如虚拟流水线【33】和高效算子【9, Flux: Fast software-based communication overlap on gpus through kernel fusion, 2024】【10, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】。
- 并置编码器和LLM主干:在共享资源上并置编码器和LLM,以实现工作负载弹性的联合编排。通过在所有可用的GPU上部署编码器和LLM模型状态,可以时间上调度编码器和LLM计算的细粒度执行。当相对工作负载发生变化时,它们各自的执行时间比例可以被动地动态调整,从而避免因编码器样本不足或LLM流水线气泡造成的设备空闲。
A2 方法细节
3 系统概览
MegaScale-Omni是一个为动态工作负载适应和数千GPU规模部署而设计的生产级MLLM训练系统。它基于§2.3.2中提出的编码器-LLM多路复用方案,解决了实现高效系统设计的三个主要挑战:
- 长上下文和超大规模下高效的解耦编码器-LLM并行化:MegaScale-Omni确保编码器的并行化不会干扰LLM主干的全面并行化,避免了LLM流水线瓶颈或因复杂数据依赖而使流水线调度复杂化。同时,它解决了静态编码器并行化在处理可变长度样本时的低效问题,例如,序列并行(SP)对长上下文样本有利,但对短上下文样本则可能导致GPU利用率不足。
- 灵活的编码器-LLM并置和具工作负载弹性的联合编排:MegaScale-Omni解决了在生产环境中并置编码器和LLM的实际困难。由于并行策略和数据依赖的差异,直观地将编码器包装为LLM层或嵌套地将编码器操作插入到LLM流水线代码中是不可行的。MegaScale-Omni提出了一种原则性的编排方法,在动态工作负载下保持结构稳定性。
- 数据加载和重分片期间高效的编码器-LLM工作负载均衡:MegaScale-Omni在加载多模态数据时增强了编码器间的工作负载平衡,解决了因数据分布偏斜而加剧的负载不平衡问题。此外,在将数据从编码器重分片到LLM时,它平衡了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)。
4 模型并行化
本节详细描述基于编码器-LLM多路复用方案构建的MegaScale-Omni的系统设计。
4.1 解耦的并行化
4.1.1 编码器的长短序列并行
- 长短序列并行(LSSP):为了高效处理可变长度的样本,MegaScale-Omni提出了长短序列并行(Long-Short Sequence Parallelism, LSSP),它统一了数据并行(DP)、Ulysses序列并行(SP)和ZeRO2/3策略。如图7(a)所示,一种直观的统一策略是空间上划分DP和SP节点,但这会在工作负载变化时导致严重的负载不平衡。相反,MegaScale-Omni设计的LSSP能在DP和Ulysses SP策略之间进行时间上的节点状态切换。如图7(b)所示,对于每个微批次,每个节点本地按长度对样本进行重排序,并用一个长度阈值η进行划分。长度不超过η的样本在DP状态下进行编码,而超过η的样本则在SP状态下编码,并伴有额外的all-to-all通信。由于DP和Ulysses SP都在各计算单元间复制编码器参数(或使用ZeRO2/3进行分片),状态切换不会引入额外的模型重分片开销。图7(c)展示了LSSP的执行分解。DP和SP状态的执行时间会根据动态工作负载被动地调整,从而在无需跨节点样本重排序的情况下高效平衡各GPU的工作负载。
- 并行策略选择的理由:编码器并行化的原则是高并发且不干扰LLM的执行。因此,数据并行(DP)是首选,因为它能提升编码并发性且通信密集度低。对于长上下文样本,我们结合DP采用Ulysses序列并行(SP)【20】,在每个DP组内将patch后的token沿序列维度分片。这主要基于两个原因:首先,对于长上下文样本和中等大小的编码器,激活值是内存消耗的主要部分,Ulysses SP可以将其分片,实现SP倍的内存缩减。其次,Ulysses SP相比其他SP变体具有更好的兼容性和工作负载平衡性。它不像Megatron SP那样与TP紧密耦合,可以灵活地与其他并行策略(如ZeRO3)集成。与因果注意力掩码而易于负载不平衡的CP【48, Wlb-llm: Workload-balanced 4d parallelism for large language model training, 2025】相比,Ulysses SP通过沿头维度分片注意力来保持每个计算单元的完整注意力,实现最优平衡。当扩展到更大的编码器时,我们采用ZeRO2/3来分发编码器状态和参数。为避免超大规模部署中的P2P死锁风险并支持复杂的LLM流水线【33】【36, Zero bubble pipeline parallelism, 2023】,我们不为编码器采用PP,因为它会使编码器-LLM的数据依赖和联合流水线调度变得复杂【13】。
4.1.2 编码器-LLM并行化布局
-
LLM的全功能5D并行化:为了利用高度优化的训练技术【9】【33】,MegaScale-Omni为LLM主干继承了全功能的5D并行化,包括DP(ZeRO1)、TP、PP(支持多种调度)、SP和EP(用于MoE)。我们没有选择全分片数据并行(FSDP)【57】作为TP或PP的替代方案,因为它在缩减本地批次大小的同时,引入了模型参数的恒定通信开销,在扩展到数千GPU时容易成为通信瓶颈。对于序列并行,我们主要使用Ulysses以实现工作负载平衡和易用性;对于超长上下文(如>512K),则选择CP,因为Ulysses SP的扩展受限于注意力头的数量。
-
通信高效的布局:为了消除编码器和LLM主干之间的通信拥塞,对于每个模态,MegaScale-Omni将一个编码器专门与每个LLM流水线阶段并置,如图8所示。在这种布局中,通信密集的组(如编码器SP和LLM TP)被限制在节点内以利用高带宽链路,并通过时间调度来避免拥塞。节点间的通信(如编码器和LLM主干之间的P2P操作)则根据它们的数据依赖和流水线调度进行联合调度。值得注意的是,LLM各阶段共享同质的编码器分配(例如,所有阶段都有模态1和模态2的编码器),而不是将模态1分配给PP0和PP1,模态2分配给PP2和PP3。这是为了在不同模odalitities计算异构性的情况下,平衡各LLM阶段上编码器的工作负载。
4.2 并置表示
- 稻草人方案:嵌套代码插入:一种自然的并置方式是嵌套地将编码器操作插入到LLM代码的指定位置。然而,LLM流水线遵循“单程序,多数据”(SPMD)范式,具有复杂的调度【39, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020】【40】,如图9(a)所示。为了精确定位插入点,需要额外的(PP rank, microbatch index)信息(如图9(b)中的
condition_1(pp, i)),这使得编码器-LLM的代码难以开发、维护和优化。更糟糕的是,工业代码库在多个项目间共享训练基础设施,阻碍了对LLM代码的侵入式修改。
- 统一编码器表示:为了解决上述问题,我们提出了一种统一的编码器表示,它在表达所有编码器插入操作的同时,保持对LLM代码的非侵入性。关键思想是将编码器在LLM代码中的定位与其数据流解耦。如图9(c)所示,MegaScale-Omni将编码器表示为LLM流水线调度的一个锚点(EncoderAnchor)。
EncoderAnchor通过指定以下信息来实例化:(1) 在此PP rank上并置的模态编码器类;(2) 并行化配置,如是否启用ZeRO2/3或激活卸载;(3) 数据配置,如patch大小和最大序列长度。以下是一个启用两个编码器和ZeRO3的示例:
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的格式中定义:
4.3 联合编排
- 具有工作负载弹性的统一流水线:尽管上述表示带来了灵活性,但激进地将编码器插入LLM流水线会给后面的LLM阶段分配更多的微批次,如图10(a)所示。当编码器工作负载增加时,这种非均匀插入会破坏原始LLM流水线的结构。具体来说,当第一阶段的编码时间增加Δt时,最后阶段会经历(Nm-1/Nm0)Δt的延迟,导致原本紧密配合的LLM微批次严重错位(图10(a)中的红色箭头),从而引入1.63倍的气泡和P2P死锁风险【22, Dynapipe: Optimizing multitask training through dynamic pipelines, 2024】【50, Adaptra: Straggler-resilient hybrid-parallel training with pipeline adaptation, 2025】。为了实现工作负载弹性,我们提出在LLM各阶段均匀地插入编码器,如图10(b)所示。我们发现,由于混合打包中使用了统一的序列长度,LLM阶段的延迟在工作负载变化下保持相对稳定。相反,由于编码器之间的工作负载是平衡的,它们的延迟会随着工作负载的变化而均匀地增减,从而保持结构稳定性,避免干扰LLM的执行。
- 按需编码器插入:为了均匀插入编码器,一种直观的方法是在LLM流水线开始前计算所有编码器的前向微批次,并在其完成后计算所有反向微批次。虽然这不干扰LLM流水线执行,但会对GPU或主机存储造成巨大的内存压力,并与多源数据加载器等内存密集型操作产生竞争。此外,这种时间上的“隔离”放弃了细粒度重叠的机会。因此,我们提出在需要前按需插入编码器,如图10(b)所示。例如,PP0恰好在LLM计算需要其输出之前,为Encoder1和Encoder2计算第5个微批次。其他阶段(PP1-3)并发计算后续微批次并异步发送给PP0。这通过提早进行反向计算并回收内存来减少编码器激活值的峰值内存占用,同时将编码器通信与LLM计算重叠。值得注意的是,按需插入对数据加载器中的样本重排序有影响(§5.1)。
5 工作负载均衡
MegaScale-Omni的效率与工作负载均衡性能密切相关。因此,我们引入了专门的技术来解决动态工作负载引起的不平衡问题。
5.1 均衡编码器数据加载
- 去中心化分组重排序:先前的工作如DistTrain【55】和DynaPipe【22】依赖于全局重排序来平衡微批次间的训练样本,因为它们在第一个流水线阶段采用中心化数据加载器。相比之下,MegaScale-Omni将编码器分布在所有流水线阶段(即所有GPU),使得中心化数据加载器易于受限于存储读取并发性,并因向所有GPU大规模分发而产生高昂开销。为解决此问题,我们采用去中心化数据加载器来增强数据加载并发性。然而,全局重排序因昂贵的跨节点乃至跨机架通信和all-to-all挂起风险而变得不可行。为了协调工作负载均衡与开销,我们提出了去中心化分组重排序,该方法根据网络局部性将GPU计算单元划分为重排序组,如图11所示。在每个组内,我们首先进行元数据all-gather操作交换数据大小信息,然后使用Karp算法【6, Analysis of the karmarkar-karp differencing algorithm, 2008】重排序模态样本,并通过all-to-all操作重新分发数据。该技术不会对收敛性产生负面影响,因为编码器数据加载遵循i.i.d.采样原则,且梯度同步后,跨DP副本的样本重排序等效于不重排序。
- 零冗余过滤:为了缓解远程存储的I/O带宽拥塞,工业实践通常只将数据加载到PP0的第一个TP rank,然后本地广播到其他TP rank【56】。然而,在MegaScale-Omni中,每个阶段都应先加载整个批次,然后根据PP rank进行分片,这导致了冗余的数据传输和内存占用。为了减少数据冗余,我们在远程数据加载器中集成了一个过滤器,利用元数据(如PP world size, batch size)来预取每个阶段将要消耗的数据。虽然有效,但零冗余加载器无法像通常那样进行检查点和恢复。我们通过在DP0上维护一个数据缓冲区来跟踪过滤前的数据来解决这个问题,在检查点操作中保存完整的预过滤数据,并在恢复时重新加载和过滤,以确保与原始执行流的一致性。
5.2 均衡编码器-LLM数据重分片
-
发送后重分片策略:由于编码器和LLM主干采用不同的并行策略,需要对嵌入进行重分片以进行通信。一种直观的方法是先在每个编码器rank上分片嵌入,然后将每个分片路由到LLM PP0上的目标rank。这种方法需要复杂的全局all-to-all rank映射。相反,我们提出了“先发送后重分片”(send-then-reshard)策略以获得更好的代码可维护性。每个PP阶段的编码器输出被异步发送到PP0,然后通过自适应分片和对称分发进行重分片。
-
自适应样本分片:为了在可变长度样本下平衡数据重分片,我们根据LLM并行化所用的SP变体采用自适应分片策略。(1) 对于Ulysses SP(图12(a)),我们沿序列维度均匀分片数据。(2) 对于上下文并行(CP,图12(b)),沿序列维度分片数据会导致LLM CP rank间的注意力不平衡。我们因此只对长样本(如图中的图像和文本)进行分片,并集成混合数据并行【15, Bytescale: Efficient scaling of llm training with a 2048k context length on more than 12,000 gpus, 2025】来处理分片样本(使用CP)和未分片样本(使用DP)。
- 对称分发:分片后,我们遵循对称原则将分片后的嵌入分发到LLM rank。如图12(a)所示,对于Ulysses SP,我们使用对称的all-to-all操作来减少内存占用。然而,对于CP,由于样本内分片,all-to-all操作变得不对称,导致复杂的rank映射和通信瓶颈。因此,我们使用带有循环内存缓冲区的all-reduce操作在每个LLM CP rank上聚合样本,然后照常进行分片和合并token(图12(b))。
6 实现与优化
MegaScale-Omni基于Megatron-LM【40】构建,并已集成到我们公司内部的代码库中,作为大多数大规模MLLM训练任务的基础。
-
并行度调优:在扩展到数千个GPU时,直接确定最优并行配置非常困难。我们通过在一个按比例缩减的“代理”集群上进行并行度调优来高效地找到最优配置。该代理集群按比例减少DP度、GPU数量和全局批处理大小,但保留了与原始集群相同的最优配置。
-
选择性激活卸载:MegaScale-Omni为LLM和编码器部分采用了不同的激活卸载和重计算策略。对于LLM,我们在虚拟流水线中实现了计算和通信的块级重叠。对于编码器,由于其非流水线并行化,我们以层级方式重叠卸载通信和层计算。卸载和重计算的选择在算子粒度上进行,我们选择性地重计算内存密集型算子(如核心注意力)并仅卸载计算密集型算子(如MLP投影),并进一步卸载重计算所需的持久输入激活,以实现GPU上零激活驻留。
-
高效算子:MegaScale-Omni集成了先进的算子级优化。对于注意力层,我们采用FlashAttention2【10】来改善长上下文输入的工作分区。对于密集层,我们使用Flux【9】来减少通信开销,它将GEMM操作与all-gather和reduce-scatter集合通信融合到大型内核中进行细粒度重叠。此外,MegaScale-Omni实现了编码器和LLM微批次之间的通信重叠,例如,LLM梯度同步与编码器计算重叠,编码器数据加载和参数all-gather与LLM计算重叠。
A4 实验环境
-
硬件配置:
- Cluster-A: 一个拥有512个GPU的集群。每个节点包含8个GPU(通过NVLink连接),1.8TB内存,88个vCPU。所有节点通过8个400Gbps的Infiniband (RDMA) 网络连接。
- Cluster-B: 一个拥有数千个GPU的超大规模生产集群(具体规格因商业机密未披露)。
-
模型与工作负载:
- 主要使用视觉-语言模型(VLM),包括不同尺寸的ViT作为图像编码器和LLaMA/GPT作为LLM主干(详见表1)。
- 同时评估了包含USM作为音频编码器的三模态工作负载。
- 最大序列长度设为16K,部分工作负载(C/D)为8K以避免基线系统大量OOM。
- 数据集包括开源视觉数据集(OpenImages, RefCOCOg)和一个生产文本数据集(BytedLong,文档长度达512K)。
- 使用工业级多源数据加载器以流式方式加载数据。
-
软件配置与基线系统:
- 实现: MegaScale-Omni基于Megatron-LM构建。
- 基线系统:
- Megatron-LM: 将编码器视为LLM第一个流水线阶段的嵌入层。
- Megatron-Dist: 基于Megatron的解聚系统,将编码器和LLM作为具有不同资源和并行策略的独立阶段。
- AutoParallel: 基于Alpa的3D并行系统,可自动构建流水线以解决不平衡问题。
- Optimus: 一种先进的MLLM训练系统,为静态工作负载并置编码器和LLM,并采用多层气泡调度。
| 工作负载 | 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所示,在64到512个GPU的不同规模下,MegaScale-Omni在各种图文混合比例下均显著优于所有基线系统。与最强的基线相比,其吞吐量分别在64、128、256和512 GPU规模下最高提升了1.54倍、3.30倍、7.57倍和1.98倍。随着图像样本比例的增加,MegaScale-Omni的优势愈发明显,最高提升可达5.08倍,这证明了其核心设计——编码器-LLM多路复用方案的有效性。
- 序列长度扩展性: 如图14所示,在处理更长的序列(16K和32K)时,MegaScale-Omni依然保持稳定性能,吞吐量分别比基线高出6.88倍和4.72倍。这得益于其长短序列并行策略和工作负载均衡技术,能够高效处理长上下文样本并避免计算掉队。
- 内存占用: 如图15所示,MegaScale-Omni在第一个流水线阶段的内存消耗显著低于其他系统(最高减少2.21倍),因为它将编码器激活值分布在所有GPU上并采用了选择性卸载/重计算。虽然其在最后阶段的内存消耗略高(因并置了编码器),但整体内存效率更高。
- 模型浮点运算利用率 (MFU): 如图16所示,在动态工作负载下,MegaScale-Omni的MFU比基线高出23.0%,并且随着编码器工作负载和序列长度的变化保持稳定和良好的扩展性,从17.0%提升至38.2%。
- 三模态实验: 如图17所示,在动态的三模态工作负载下,MegaScale-Omni的吞吐量比基线高出2.01倍,并对编码器工作负载的增加表现出强大的弹性。
7.3 消融实验
- 性能分解: 如图18所示,消融实验表明编码器-LLM多路复用方案是性能提升的最关键因素,禁用它会导致吞吐量下降超过50%。其次是工作负载均衡技术,禁用后吞吐量下降约45-50%。长短序列并行、并行度调优、激活卸载重叠和高效算子重叠等优化也分别贡献了约16-19%的性能提升。
- 多路复用方案的鲁棒性: 如图19所示,即使禁用了MegaScale-Omni的其他优化,其核心的多路复用方案在不同的并行策略组合下(如改变虚拟流水线层数、PP度、禁用激活卸载等)依然能稳定地带来高达1.64倍的吞吐量提升,证明了其架构设计的优越性。
- 工作负载均衡与开销分析: 如图20所示,对分组数据重排序的分析表明,更大的重排序组可以提升吞吐量(最高1.33倍),但会带来更高的通信开销,并且存在收益递减的效应。
7.4 超大规模训练经验
- 性能分析: 在数千GPU的生产环境中,MegaScale-Omni表现稳定,平均训练吞吐量约为每秒800万token(图21)。训练过程中经历了59次重启,大部分由软硬件故障导致并能自动恢复。通过优化张量检查策略,在保证收敛稳定性的同时减少了对吞吐量的影响。
-
MFU优化: 通过将非训练通信(如多源token计数器)与模型计算重叠,解决了网络干扰问题,避免了2.2%的吞吐量下降。对于512K长序列训练,采用“forward-then-backward”流水线【19, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019】并结合更激进的内存优化,显著提高了MFU,节省了约2/3的GPU小时。
-
数据加载器问题: 解决了主机内存卸载与多进程数据加载之间的内存竞争问题。通过采用异步加载器状态快照技术,将状态保存的开销从6-20分钟减少到约800毫秒,解决了掉队者问题。
-
通信与检查点异常: 解决了大规模(>512 GPUs)下all-to-all和P2P通信原语不稳定的问题,采用非P2P操作来保存检查点,避免了挂起。通过为检查点持久化保存计划缓存,将重启后的首次检查点开销从15分钟大幅缩短。
A7 补充细节
-
大型模型训练框架:许多工作致力于提升大模型训练效率,如Megatron-LM【40】、DeepSpeed【39, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020】和PyTorch FSDP【57】等是最流行的开源训练框架,它们提供了多种并行策略。一些工作【1, Nvidia transformer engine, 2022】【9, Flux: Fast software-based communication overlap on gpus through kernel fusion, 2024】【46, Overlap communication with dependent computation via decomposition in large deep learning models, 2022】探索了通信与计算的重叠。Alpa【58, Alpa: Automating inter-and intra-operator parallelism for distributed deep learning, 2022】、nnScaler【30, {nnScaler}:{ConstraintGuided} parallelization plan generation for deep learning training, 2024】和Unity【42, Unity: Accelerating {DNN} training through joint optimization of algebraic transformations and parallelization, 2022】则自动寻找最优的混合并行策略。然而,这些框架主要关注通用大模型,缺乏针对MLLM架构、编码器-LLM编排和动态多模态工作负载的专门设计。
-
多模态训练系统:随着多模态任务的兴起,近期有工作开始研究为MLLM构建专门的训练系统。Megatron-LM【40】和Transformers【49】以单模态的方式支持MLLM训练。DistMM【18】和DistTrain【55】提出将编码器和LLM主干解聚为两个独立模型,并探索了异构感知的模型划分、资源分配和数据均衡。Optimus【13】、GraphPipe【21, Graphpipe: Improving performance and scalability of dnn training with graph pipeline parallelism, 2024】、Spindle【47, Spindle: Efficient distributed training of multi-task large models via wavefront scheduling, 2025】和PipeWeaver【51, Pipeweaver: Addressing data dynamicity in large multimodal model training with dynamic interleaved pipeline, 2025】提出了细粒度的多层气泡利用或动态流水线调度算法来减少执行
💬 评论讨论
欢迎在这里分享您的想法和见解!