最新消息:关注人工智能 AI赋能新媒体运营

小红书开源BigMac:多模态训练不再纠结显存与速度二选一

科技资讯 admin 浏览

小红书开源 BigMac:打破多模态大模型训练的显存与速度两难

多模态大语言模型(MLLM)正成为新一代 AI 的基础设施,但其训练系统长期面临一个经典困境:追求计算速度,显存便难以承受;压缩了显存,训练速度又大幅滑落。小红书 dots infra 团队将这一矛盾称为"多模态流水训练的帕累托前沿",而他们最新推出的开源方案 BigMac,正在这道前沿上撕开缺口。

7 月 22 日,团队正式开源了这款面向原生多模态场景的流水并行训练新范式,代码已托管于 GitHub 的 Dots-Infra 仓库。

多模态训练的结构性难题

与标准 Transformer 不同,多模态模型的架构呈明显的三段式结构:负责将图像、音频转为嵌入向量的模态编码器,承担核心推理的 LLM 主干,以及将 LLM 输出映射回图像、语音等模态的生成器。这三部分计算特性差异显著,强行塞进同一条训练流水线,必然引发效率损耗。

业界此前主要两条路径。计算高效型方案将编码器和生成器从 LLM 流水线中剥离、独立运行,避免了模态模块的耗时波动在 LLM 流水线中产生空泡,但代价是激活显存随 microbatch 数量线性膨胀,生产规模下成本高昂。显存高效型方案则将全部模块塞入同一条流水线做首尾阶段,激活生命周期缩短、显存占用降低,然而一旦某个编码器或生成器延迟,整条 LLM 流水线只能干等,尾部空泡随之产生。规模扩大后,两种设计的瓶颈都会暴露。

准依赖安全的嵌套流水线

BigMac 的核心思路是:不以替换成熟的 LLM 流水调度为代价。团队将 1F1B、interleaved 1F1B 等已高度优化的 LLM 调度作为底层时间线,在输入就绪、且不扰乱 LLM 执行顺序的位置,插入编码器和生成器的计算。这种设计被命名为"准依赖安全的嵌套流水线"。

该架构带来两个关键性质。一是编码器与生成器的耗时波动不再向 LLM 流水线传导,LLM 按自身节奏推进;二是模态激活显存被压缩至 O(1),编码器无需为所有 microbatch 保留激活直至流水线结束,生成器也不会拖曳长尾激活。BigMac 并非在显存与空泡之间做权衡交换,而是通过重构调度结构,使两个目标得以同时实现。

从算法到系统的三件工具

BigMac 面向工程落地提供了三项核心能力。

首先是全局调度的透明化。运行时 Scheduler 生成覆盖所有 pipeline rank、microbatch 和模块类型的全局算子表,Executor 将其拆分为各 rank 的本地执行序列,分发给 Megatron Core 等 LLM 后端、模态 runtime 及通信后端。调度策略由此变得可见、可检查,同时保持对后端的兼容性。

其次是对算法工程师透明的流水并行接口。工程师只需声明各模块的产出与消费关系,stage 划分、激活与梯度交接、跨设备通信等细节均由 BigMac 底层处理,使单卡验证过的多模态实验能平滑扩展至流水并行。

第三是一套与调度结构深度绑定的分析工具链,包含 profiler、simulator 和可视化组件。训练迭代可被拆解至算子级别,清晰呈现各 rank 在各时间点的计算内容、空转位置及依赖阻塞;simulator 支持在正式训练前模拟不同 PP 配置与 microbatch 组合,预估对空泡和吞吐的影响。

实验验证:速度提升与显存稳定兼得

团队在两类典型负载上验证了 BigMac 的效果。

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,BigMac 仍取得 1.5–1.9 倍加速,且显存保持稳定。这正是嵌套流水线的价值所在——系统必须同时协调编码器与生成器的复杂依赖,又无法承受大量激活驻留或严重空泡。

生产验证与开源进展

目前,BigMac 已作为 dots 多模态模型训练的核心组件,部署于小红书生产环境。相关论文《BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training》已发布于 arXiv,团队同时放出了交互式 PP Profiler 的 trace 示例。对于正被多模态训练显存与速度双重压力困扰的研究者和工程师而言,这份经过生产验证的开源方案,有望减少大量重复造轮子的工作。