多模态大语言模型是新一代AI的地基,这个判断应该没人反对。但它的训练系统,如今被一个老问题卡得死死的——算得快的,显存撑不住;显存省着的,速度又掉下去。小红书Dots Infra团队把这个两难局面叫做多模态流水训练的帕累托前沿,而他们现在把这个前沿撕开了一道口子。7月22日,这个团队正式开源了BigMac,一个专门针对原生多模态场景的流水并行训练新范式,开源地址就挂在GitHub的Dots-Infra下面。
问题在于,多模态模型从来不是一块规整的transformer。一个典型的MLLM,是由三块拼起来的:把图像、音频转成embedding的模态编码器,在其上做推理的LLM主干,以及把LLM输出映射回图像、语音等目标模态的生成器。这三块形态差异巨大,把它们塞进同一条流水线,麻烦自然就来了。
从能设计出schedule到真正把模型训起来,中间还隔着系统集成、接口定义和性能排查这一整套工程关卡。BigMac正是冲着这些环节去的,它给了三件东西。
第一,把全局schedule摆到明面上。运行时的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示例。对正在被多模态训练显存与速度两头夹击的研究者和工程师来说,这份带着生产验证的开源,或许能省下不少重新造轮子的时间。
