告别“选速度还是选显存”？小红书开源 BigMac，把多模态训练的两难一刀切开

作者提米哥

训练多模态大模型，就像在高速公路上既要跑得快又要省油——但现实是，你往往只能二选一：要么计算快但显存爆掉，要么显存省下来但速度慢得让人心焦。小红书 dots infra 团队把这个困境称为“多模态流水训练的帕累托前沿”，意思是，在现有技术框架下，想不牺牲任何一方几乎不可能。

但就在 7 月 22 日，他们扔出了一把“手术刀”——BigMac，一个专门为原生多模态场景设计的流水并行训练新范式。代码已经开源在 GitHub 的 Dots-Infra 仓库下。简单说，它让“既要又要”不再是空话。

为什么多模态训练这么拧巴？

一个典型的多模态大语言模型（MLLM）不是一块整整齐齐的 Transformer。它由三块拼成：

  • 模态编码器：把图像、音频转成 embedding；
  • LLM 主干：做推理；
  • 生成器：把 LLM 的输出映射回图像、语音等目标模态。

这三块形态差异巨大，想塞进同一条流水线，麻烦就来了。过去业界只有两条路：

  • 计算高效路线：把编码器和生成器从 LLM 流水线里单独拎出来跑。好处是各模块的耗时波动不会互相传染，坏处是激活显存会随着 microbatch 数量膨胀，生产规模下贵得吓人。
  • 显存高效路线：把所有模块塞进同一条流水线，让激活生命周期缩短，显存降下来。但只要某个编码器或生成器慢一拍，整条 LLM 流水线就得干等，产生大量“空泡”。

规模一大，两条路都走到尽头。你只能在显存和速度之间做痛苦的交换。

BigMac 的解法：不换 LLM 调度，而是“见缝插针”

BigMac 的出发点朴素但关键：调度主干仍然是那条已经高度优化的 LLM 流水线。大规模 LLM 训练早已依赖成熟的 schedule（如 1F1B、interleaved 1F1B），并与生产级训练栈深度绑定。BigMac 不替换它们，而是把 LLM schedule 当作底层时间线，然后把编码器和生成器的计算，插入到输入已就绪、且不会打乱 LLM 执行顺序的位置。

团队把这种设计称为 “准依赖安全的嵌套流水线” 。它带来了两个关键性质：

  1. 编码器与生成器的耗时波动不再沿 LLM 流水线传导，LLM 按自己本该有的节奏推进；
  2. 模态激活显存被算法层面压到 O(1)——编码器不必为所有 microbatch 保留激活直到流水线结束，生成器也不用拖着长长的激活尾巴。

换句话说，BigMac 不是在显存和空泡之间做交换，而是改写了 schedule 的结构，让这两个目标不再非此即彼。

从设计到落地：三件“工程硬货”

能设计 schedule 是一回事，真正跑起来是另一回事。BigMac 给了三件实实在在的东西：

第一件：全局调度可视化
运行时的 Scheduler 会生成一张覆盖所有 pipeline rank、microbatch 和模块类型的全局 operator 表。Executor 再把它拆成每个 rank 上的本地序列，分发给 Megatron Core 等 LLM 后端、模态 runtime 和通信后端。调度策略从此可见、可检查，对后端也很友好。

第二件：对算法工程师无感的流水并行接口
工程师只需要描述每个模块“生产什么、消费什么”，剩下的 stage 划分、activation 与 gradient 交接、跨设备通信，全由 BigMac 在底层料理。这意味着，一个在单卡上验证过的多模态实验，可以更自然地扩展到流水并行，而不需要改代码结构。

第三件：理解 schedule 的工具链
团队提供了一套 profiler、simulator 和可视化工具，能把一次训练迭代拆回 operator 级别，让你看清每个 rank 在每个时间点究竟在算什么、在哪里空转、哪条依赖卡住了后续。simulator 还能在正式花钱启动训练前，先试不同的 PP 配置和 microbatch 组合，预估对空泡和吞吐的影响。

效果如何？两类负载验证

团队在两类代表性负载上做了测试：

  • MLLM-Understanding（理解类）：主干用 Qwen3-30B-A3B，配一个 1.3B 的 ViT 编码器。相比计算高效基线 Optimus，BigMac 提速 1.08~1.1 倍；相比显存高效基线 Megatron-DistTrain，提速 1.6~1.9 倍。更关键的是显存：随着 per-GPU batch size 增大，BigMac 的峰值显存保持平稳，而 Optimus 因为要保留编码器激活，显存一路飙升，最终在更大 batch 下直接 OOM。

  • MLLM-Generation（生成类）：更复杂，加上了一个 20B 的 MMDiT 生成器。差异被放大：Optimus 在所有测试 batch size 上全部 OOM，而 BigMac 靠及时跑 generator backward、快速释放生成器侧激活，顺利躲过。相比 Megatron-DistTrain 仍取得 1.5~1.9 倍加速，显存依旧稳定。

这正是嵌套流水线最值钱的场景——系统必须同时伺候编码器和生成器的依赖，又没法承受大量激活驻留或严重空泡。

写在最后

目前，BigMac 已经作为 dots 多模态模型训练的核心组件之一，跑在了小红书的生产环境里。相关论文《BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training》也已挂在 arXiv 上，团队同步放出了交互式 PP Profiler 的 trace 示例。

对于正在被多模态训练“显存与速度”两头夹击的研究者和工程师来说，这份带着生产验证的开源，或许能省下不少重新造轮子的时间。毕竟，在 AI 基础设施里，最值钱的事不是做一个新轮子，而是让已有的轮子不再互相打架。

微信联系：timidamen

类似文章