首页
看点啥
插画图片
首页 看点啥 小红书 dots infra BigMac 开源 : 突破多模态大模型训练的帕累托前沿

小红书 dots infra BigMac 开源 : 突破多模态大模型训练的帕累托前沿

2026-07-23 0

小红书技术REDtech 2026-07-22 18:06 上海

导读:

BigMac是原生多模场景下的流水并行训练新范式。它针对多模态大模型训练中计算效率与显存占用难以兼顾的问题,提出了依赖安全的嵌套流水线:以成熟的 LLM 流水线为主干,在不打乱LLM执行顺序的前提下,有序嵌入编码器和生成器计算,从而在不增加 LLM 流水线空泡、保持激活显存有界的同时,高效实现多模态流水训练。相比传统的流水并行实现方式,BigMac 在工程实现上进一步解耦了全局调度与运行时执行,并设计了一系列接口和工具链降低了模型接入与系统调优的成本。实验中,其训练速度较基线提升 1.08~1.9 倍,并在global batch size增大时保持的稳定显存占用。目前,BigMac已作为dots多模态模型训练的核心组件之一应用于生产中。

开源地址:

https://github.com/Dots-Infra/BigMac/

多模态大语言模型(Multimodal large language models, MLLM)正在成为新一代 AI 系统的核心基础设施。相比只处理文本的 LLM,多模态模型需要把语言、图像、音频等不同类型的信息放到同一个模型里理解;在很多场景下,它们还要进一步生成图像、语音等内容。能力边界被打开了,训练系统的复杂度也随之上升。

难点在于,MLLM 通常不是一整块同构的 transformer,而是由多个形态差异很大的模块拼接起来。一个典型的多模态模型通常包含三个部分:模态编码器、LLM 主干,以及模态生成器。编码器把图像或音频等原始输入转换成 embedding。LLM 在得到的多模态 token 序列上进行推理。生成器则把 LLM 的输出映射回目标模态,例如用图像扩散模型生成像素。

图 1:一个典型的 MLLM 架构,包含模态编码器、LLM 主干和模态生成器。

这种模块化架构带来了很强的表达能力,也把训练系统推向了普通 LLM 流水线很少遇到的复杂局面。现有 MLLM 训练系统通常会选择两种设计之一。两种设计都有其合理性,但在规模变大后都会暴露明显瓶颈。

第一种设计是计算高效的。它把编码器和生成器计算从 LLM pipeline 中分离出来,类似 LongCat-Flash-Omni 中用于应对大规模多模态训练异构性的modality-decoupled parallelism。例如,系统可能先对所有 microbatch 运行 encoder forward,保留这些 activation,然后再启动 LLM pipeline。这样可以让 LLM pipeline 不被模态模块的耗时波动打乱:编码器和生成器慢一点,不会直接在 LLM pipeline 内部制造 bubble。

然而,由于编码器 activation 必须一直保留到对应的梯度返回,activation 显存会随着 microbatch 数量增长。如果模型中还有生成器,情况会更糟:LLM activation 也可能需要一直保留到生成器计算结束。在生产规模模型中,这尤其昂贵,因为 LLM activation 往往远大于 encoder activation。

图 2:计算高效的 pipeline。该设计将模态计算从 LLM pipeline 中分离出来,减少 bubble,但会保留大量 activation。

第二种设计是显存高效的。它把编码器和生成器整合进同一条 pipeline,把它们看作围绕 LLM 的首尾阶段;例如 Megatron Core 的多模态训练示例中,vision-language 模型作为一个独立的 pipeline stage,加入到了 LLM 的流水线中。这样可以缩短 activation 生命周期,因此显存占用更低。但是,所有模块现在都被 pipeline 依赖耦合在一起。如果某个 encoder 或 generator microbatch 很慢,LLM pipeline 就必须等待。如果生成器运行时间变化较大,pipeline 尾部就会产生 bubble。系统通过牺牲计算效率来节省显存。

图 3:显存高效的 pipeline。该设计将所有模块整合到一条 pipeline 中,减少显存占用,但会产生跨模块 bubble。

这正是 BigMac 想要突破的 Pareto frontier。计算高效的系统速度快但显存开销大;显存高效的系统节省显存但容易产生 bubble。对于小规模训练任务,这种二选一可能还可以接受。但对于大规模多模态训练,尤其是包含模态生成器的训练,这种两难会成为瓶颈。

BigMac 的出发点很简单:调度的主干仍然应该是 LLM pipeline。大规模 LLM 已经依赖 1F1B 或 interleaved 1F1B 等精心优化过的 pipeline schedule。这些 schedule 成熟、高效,并且与生产级 LLM 训练栈深度绑定。BigMac 不试图替换它们,而是把 LLM schedule 作为基础时间线。

BigMac 的思路是先稳住最重要的 LLM pipeline,再把编码器和生成器的计算插入到合适的位置。这里的“合适”,指的是输入已经准备好,而且插入这些计算不会打乱原本的 LLM 执行顺序。这样一来,LLM pipeline 可以持续推进,编码器和生成器的激活也能更早释放。我们称这种设计为依赖安全的嵌套流水线 (nested pipeline)。

图 4:BigMac 保留原始 LLM pipeline schedule,并把 encoder/generator 工作嵌入到依赖安全的位置。

这个设计给 BigMac 带来了两个重要性质。

第一,它保留了计算效率。由于编码器和生成器不再被强行塞进 LLM 流水线,作为独立的 pipeline stage 逐级推进,它们的耗时波动也不会沿着 LLM 流水线一路传下去。LLM 仍然按照它在计算高效系统中本该遵循的 schedule 运行。

第二,它限制了模态 activation 显存。BigMac 会在梯度就绪后尽快运行 encoder backward 和 generator backward。编码器不需要一直为所有 microbatch 保留 activation,直到 LLM pipeline 结束。生成器也不需要保留很长的 activation 尾部。从算法层面看,BigMac 将编码器和生成器 activation 显存降低到O(1),同时保持 LLM activation 行为不变。

这就是整个流水线的核心:BigMac 不是在显存和 bubble 之间做交换,也不是用 bubble 换显存。它改变了 schedule 结构,让这两个目标不再必须相互冲突。

在真实的多模态大模型训练中,pipeline parallelism 的难点往往不只是“如何排 microbatch”,框架还必须解决一系列系统集成、接口定义、性能排查等工程问题。

BigMac 关注的正是这些从“能设计 schedule”到“能真正训起来”的系统环节。它提供了三个关键能力:把 schedule 变成可见、可检查、可执行的全局计划;流水线无感的接口让模型开发者仍然以模块化方式编写模型代码;schedule-aware 工具链则把 profiling、simulation 和可视化接入到性能诊断闭环中。下面我们分别展开这三点。

1. 把全局 schedule 摆到明面上

先看系统集成层。BigMac运行时的核心组件是 Scheduler 和 Executor,Scheduler 会生成一张全局 operator 表,覆盖所有 pipeline rank、microbatch 和模块类型。这张表包含 LLM forward/backward operator、encoder forward/backward operator、generator forward/backward operator,以及模块边界上所需的通信。Executor 随后解释每个 rank 上的本地 operator 序列,并把每个 operator 分发给对应后端:LLM pipeline operator 交给 Megatron Core,模态 operator 交给 encoder 或 generator runtime,数据搬运交给通信后端。

这种拆分给 BigMac 带来了两个实际优势。第一,调度策略变得可见、可检查。我们可以直接看到 encoder、LLM 和 generator 如何在所有 rank 上交错执行,而不需要从一段很长的 rank-local runtime 函数中反推出全局行为。第二,执行仍然对后端友好。LLM operator 可以继续复用 Megatron Core 等优化过的 pipeline runtime,而模态模块可以保留自己的 data-parallel、FSDP 或 optimizer 实现。Schedule 也可以在不重写模型 runtime 的情况下加入通信 operator、进行死锁检查。

更多细节建议通过 BigMac 官方的可视化工具查看。

2. 让模型代码自然接入 PP

在多模态训练里,算法工程师会不断调整 encoder、LLM、generator 的连接方式、loss 设计和数据形态,问题在于,传统 PP 系统一旦扩展到多模态,这条协作边界很容易变得模糊:一次模型结构改动可能要求重写 stage 划分和 send/recv 逻辑,系统侧的优化也可能反过来侵入模型代码。如果算法工程师每次 scale up 都要理解 PP runtime,BigMac 仍然不是一个好用的训练系统。BigMac 的第二个关键能力,就是对算法工程师无感的流水并行接口。

图 5:BigMac 的 pipeline-parallelism-transparent 接口示意。算法工程师描述模块边界,BigMac 在底层处理 schedule、handoff 和通信。

在这个接口下,算法工程师只需要说明每个模块生产什么、消费什么;BigMac 负责把这些模块接入全局 schedule,并在底层处理 pipeline stage、activation handoff、gradient handoff 和跨设备通信。换句话说,BigMac 让一个已经在单卡或 data-parallel 环境中验证过的多模态实验,可以更自然地扩展到 pipeline-parallel training。

3. 用 schedule-aware 工具链定位瓶颈

BigMac 的最后一个关键能力,是一套理解 schedule 结构的 profiler、simulator 和可视化工具链。对于大规模多模态流水并行训练来说,性能问题很少能只靠 iteration time 或几行日志定位:bubble 可能来自某个慢 microbatch,可能来自 encoder/generator 的运行时间波动,也可能来自一次 activation handoff、gradient handoff 或跨 rank 通信等待。工程同学真正需要的是把一次训练 iteration 拆回 operator 级别,看清楚每个 rank 在每个时间点到底在执行什么、哪里空转、哪个依赖正在阻塞后续计算。

图 6:传统性能排查通常需要下载并手动关联多个 rank 的重型 trace;BigMac PP profiler 将这些信息整理成 schedule-aware 的 pipeline 诊断流程。

BigMac executor 会记录每个 operator 的执行时间,并导出 per-rank timeline / trace,如图 7 所示。通过这个 trace,工程同学可以直接观察 encoder forward/backward、LLM forward/backward、generator forward/backward 和通信操作在不同 rank 上如何交错执行,从而判断 pipeline bubble 是由 compute imbalance、communication wait、模块 handoff,还是不合适的 microbatch grouping 引起的。

图 7:由 per-operator 执行时间生成的 pipeline trace 示例。

我们也提供了一个可交互查看的 PP Profiler trace example(https://github.com/Dots-Infra/BigMac/blob/main/assets/pp_profiler_example.json)。下载后可以直接在Perfetto UI(https://ui.perfetto.dev/)中打开,观察不同 rank、不同 operator 之间的时间线和依赖关系。

在 profiler 之外,BigMac 还提供 simulator 来支持更快的并行策略迭代。Profiler 告诉我们当前训练为什么慢;simulator 则让我们在启动下一次昂贵训练任务之前,先尝试不同的 PP/VPP 配置、microbatch 数量、模块放置和 scheduling policy,预估它们对 bubble 和吞吐的影响。这样,并行策略优化就不再完全依赖大规模训练试错,而可以先在 schedule 层面完成快速探索。

我们在两个代表性多模训练负载(MLLM-Understanding 和 MLLM-Generation)上评估 BigMac。

MLLM-Understanding 代表多模态理解训练,这个负载使用 Qwen3-30B-A3B 作为 LLM backbone,并使用一个 1.3B 参数的 ViT encoder。相比计算高效 baseline Optimus,BigMac 实现了 1.08x-1.1x 的训练加速;相比显存高效 baseline Megatron-DistTrain,实现了 1.6x-1.9x 的训练加速。相比 Megatron-DistTrain 的加速来自消除跨模块 pipeline 干扰。相比 Optimus 的优势更微妙:BigMac 避免了 activation 显存压力,而这种压力会随着 global batch size 增大让计算高效设计越来越昂贵。

图 8:MLLM-Understanding 训练负载上的整体性能。

显存趋势和速度提升同样重要。随着 per-GPU batch size 增加,BigMac 的 peak memory 保持稳定,类似显存高效 baseline。相比之下,Optimus 会在 LLM pipeline 期间保留 encoder activation;它的显存使用快速增长,并最终在更大的 batch size 下 OOM。

MLLM-Generation 代表同时包含理解和生成路径的训练,我们使用相同的 LLM 和 encoder,但增加了一个 20B 参数的 MMDiT generator。在这个负载上,差异变得更加明显。Generator 会让计算高效设计更难扩展,因为 LLM activation 可能需要一直保留到 generator computation 完成。在论文实验中,Optimus 在 MLLM-Generation 负载的所有测试 batch size 上都会 OOM。BigMac 通过及时运行 generator backward 并快速释放 generator 侧 activation 避免了这个问题。

图 9:MLLM-Generation 训练负载上的整体性能。

相比 Megatron-DistTrain,BigMac 在 MLLM-Generation 负载上实现了 1.5x-1.9x 的训练加速,同时保持稳定显存使用。这正是 nested pipeline 最有价值的场景:系统必须同时处理 encoder 和 generator 依赖,但又不能承受大量 activation retention 或严重 pipeline bubble。

多模态大模型训练的难点,并不只是模型更大或数据更复杂,而是不同模态模块把流水线训练限制在了一个尴尬的帕累托前沿内。BigMac 的核心是打破这个前沿。它保留 LLM pipeline 作为稳定主干,只在依赖满足且不打乱 LLM 执行顺序的位置嵌入 encoder 和 generator 计算。这样,系统既能维持 LLM 流水的计算效率,又能尽早释放模态模块的 activation。更重要的是,BigMac 不只给出一个 schedule,还把 schedule 变成可见、可检查、可执行的全局计划,并配套 PP-transparent 接口和 schedule-aware 工具链,让这一设计真正进入可开发、可调试、可部署的训练系统。

论文:BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training(https://arxiv.org/abs/2605.25451)

如果 BigMac 对你的研究或系统有帮助,欢迎引用:

@article{zhang2026bigmac,

title={BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training},

author={Zili Zhang and Chengxu Yang and Shenglong Zhang and Chenyu Wang and Yufan Zhang and Tuo Dai and Zhouyang Li and Yuhong Ge and Chao Jin and Xin Jin and Yuliang Liu},

journal={arXiv preprint arXiv:2605.25451},

year={2026},

doi={10.48550/arXiv.2605.25451}

}

杨程旭

小红书dots infra组工程师,博士毕业于北京大学计算机学院,曾在TOSEM、TMC、FSE、WWW等顶级会议/期刊发表论文4篇。

主要研究方向:多模训练、分布式训练框架

章梓立

小红书dots infra组实习生,北京大学计算机学院在读博士生,在OSDI、NSDI、sigcomm等系统领域顶级会议/期刊发表9篇论文。

主要研究方向:分布式系统、多模态模型

铭曦

小红书dots infra组工程师,硕士毕业于卡内基梅隆大学,主要从事多模态模型,生成模型端到端训练性能分析和优化工作。

主要研究方向:多模/生成训练,分布式训练框架

我们是负责支撑dots系列大模型的工程团队 —— dots infra。团队业务覆盖大模型infra的核心场景,包括文本/多模理解大模型的预训练、多模生成模型的预训练、通用/应用等方向后训练、dots 系列模型的推理服务、data infra、模型codesign等多个方向。这里有硬核且纯粹的工程师团队,成员拥有顶级技术背景(多篇系统顶会、OI 选手/ICPC world final),日常讨论技术氛围浓厚;还有行业内最顶级的施展空间,包括世界级算力规模、全栈自研技术栈、千万级 DAU 战场、Deep Co-design 机会。

热招岗位速览

训练框架

负责LLM/多模态训练框架的研究与开发,优化分布式 RL 训练吞吐,处理长序列/推理/CLI Agent 任务

算子优化

深入 CUDA/CuteDSL/TileLang 底层,攻克性能瓶颈(如 FlashMLA、Mega Kernel、低精度算子、通信计算重叠优化)

推理引擎

打造高并发、低延迟的推理架构,优化 PD/AF 分离 schedule、解耦架构及动态资源调度

数据工程

详负责支撑 dots(LLM / 多模态 / Agentic)训练与对齐所需的超大规模数据工程体系,覆盖从 数据生产、清洗、存储、调度、分布式读取到质量闭环的全链路。

招聘官网

https://job.xiaohongshu.com/social/position/14633

喜欢(0)

上一篇

纠偏率 98.15%!方向感觉醒:第一视角认路的国产 AI 开挂了

纠偏率 98.15%!方向感觉醒:第一视角认路的国产 AI 开挂了

下一篇

真实工作流:正在成为下一代训练数据

真实工作流:正在成为下一代训练数据
猜你喜欢