小红书推出BigMac：打破多模态训练中的内存-速度权衡

多模态大型语言模型正逐渐成为下一代人工智能的基石，但其训练基础设施长期以来一直受制于一个难以调和的矛盾：针对速度优化的系统往往面临内存限制，而优先考虑内存效率的系统则通常运行缓慢。 小红书的 Dots Infra 团队将这一瓶颈视为多模态管道训练的帕累托前沿，并已成功突破。 7月22日，该团队开源了BigMac——一种专为原生多模态场景设计的创新管道并行训练范式。该项目现已在GitHub的Dots-Infra仓库中发布。

与标准Transformer模型不同，多模态模型本身就具有复杂性。典型的多模态大型语言模型（MLLM）由三个独立组件构成：将图像和音频转换为嵌入向量的模态编码器、用于推理的大型语言模型（LLM）骨干网络，以及将LLM输出映射回图像或语音等目标模态的生成器。 这些组件之间的显著结构差异，在将其整合到单一训练管道时带来了巨大的挑战。

当前的行业解决方案通常分为两类。第一种方法优先考虑计算效率，将编码器和生成器从 LLM 管道中解耦，并分别运行。这虽然可以防止模态模块持续时间的波动在 LLM 管道中引发“气泡”，但会导致激活内存随微批次数量的增加而增长，从而在规模化时产生高成本。 第二种方法侧重于内存效率，将所有模块保留在同一管道中，这缩短了激活生命周期并减少了内存占用。然而，如果任何编码器或生成器运行缓慢，整个 LLM 管道就会停滞，从而导致尾部气泡。随着模型规模的扩大，这两种设计都会暴露出关键瓶颈。

BigMac 通过一个简单却至关重要的原则来应对这一挑战：核心 LLM 管道仍是首要关注点。大规模 LLM 训练已经依赖于 1F1B 或交错 1F1B 等成熟的调度策略，这些策略已深度集成到生产级训练堆栈中。 BigMac 并未取代这些成熟方法，而是将 LLM 调度作为其底层时间线，在输入准备就绪时策略性地插入编码器和生成器的计算，同时不干扰 LLM 的执行顺序。研究团队将这种架构称为“准依赖安全嵌套管道”。

该设计带来了两大关键优势。首先，编码器和生成器耗时的波动不再向 LLM 管道中传播，从而使 LLM 能够保持最佳运行节奏。其次，在算法层面上，模态激活的内存需求被降低至 O(1)。 编码器不再需要保留所有微批次的激活值直至管道结束，生成器也避免了延长冗长的激活尾部。本质上，BigMac 并未以牺牲内存为代价来换取性能提升；相反，它重构了调度逻辑，使这两个目标相互兼容而非相互排斥。

弥合调度设计与实际执行之间的差距涉及重大的工程障碍，包括系统集成、接口定义和性能调试。BigMac 通过提供三项核心能力来应对这些具体挑战。 首先，它使全局调度透明化：调度器在运行时生成一个全面的运算符表，涵盖所有管道秩、微批次和模块类型。 随后，执行器将这些内容分配给每个 rank 的本地序列，并与 Megatron Core 等 LLM 后端、模态运行时以及通信层无缝集成。这种可视性使得检查和后端优化变得更加容易。

其次，BigMac 提供了一套管道并行接口，该接口对算法工程师而言完全不可见。用户只需定义每个模块的输入和输出；系统会在内部处理阶段划分、激活与梯度传递以及跨设备通信。这使得在单张 GPU 上验证过的多模态实验能够自然地扩展到管道并行架构。 第三，它提供了一套性能分析、模拟和可视化工具。这些工具将训练迭代分解到运算符级别，精确揭示每个时间步中每个排的运行状况，识别空闲时段，并突出显示依赖关系瓶颈。 该模拟器还允许团队在进行全规模训练之前，测试各种 PP 配置和微批量组合，并估算其对气泡和吞吐量的影响。

BigMac的有效性已在两个具有代表性的工作负载上得到验证。在MLLM-Understanding任务中，以Qwen3-30B-A3B作为骨干模型并采用13亿维的ViT编码器， BigMac 相较于计算效率高的基线模型 Optimus，实现了 1.08 倍至 1.1 倍的加速；相较于内存效率高的基线模型 Megatron-DistTrain，则实现了 1.6 倍至 1.9 倍的加速。 关键在于，随着每张 GPU 的批量大小增加，内存使用量保持稳定，而 Optimus 则会出现内存激增，最终在较大批量下导致内存不足（OOM）错误。 MLLM-Generation 任务更为复杂，涉及一个 20B 的 MMDiT 生成器，此时性能差距进一步拉大：Optimus 因 OOM 问题在所有测试批量大小下均失败，而 BigMac 通过高效地反向执行生成器并快速释放激活值，成功避免了这一问题。 与 Megatron-DistTrain 相比，BigMac 仍能提供 1.5 倍至 1.9 倍的加速效果，且内存使用量保持稳定。这凸显了嵌套管道的主要价值：既能满足编码器和生成器的依赖关系，又不会导致过度的激活保留或显著的空闲时间。

BigMac 现已成为 Dots 多模态模型训练基础设施的核心组件，并在小红书的生产环境中运行。 相关论文《BigMac：突破多模态大型语言模型训练中计算与内存的帕累托前沿》已发布在 arXiv 上，并附有交互式 PP Profiler 跟踪示例。 对于在多模态训练中苦于应对内存与速度权衡问题的研究人员和工程师而言，这个经过生产环境验证的开源项目提供了一种切实可行的解决方案，可节省大量时间和精力。