BigMac – 小红书开源的多模态大模型流水并行训练框架

1. BigMac是什么

BigMac是小红书dots infra团队开源的多模态大模型流水并行训练框架，旨在解决多模态大模型（MLLM）训练中计算效率与显存占用之间的根本矛盾。该框架提出依赖安全的嵌套流水线范式，以成熟的大语言模型（LLM）流水线为稳定主干，在不打乱执行顺序的前提下嵌入编码器和生成器的前向与反向计算，从而实现训练速度提升1.08～1.9倍的同时保持显存占用稳定。BigMac已作为dots多模态训练核心组件投入生产，验证了其在实际场景中的有效性和可靠性。

技术定位与领域： BigMac属于大规模分布式训练框架领域，专注于多模态大模型的流水并行训练。与传统的Megatron-LM等框架不同，BigMac专门针对多模态模型中编码器、语言模型和生成器之间的复杂依赖关系进行了优化，通过嵌套流水线技术打破了计算效率与显存占用之间的帕累托权衡，为多模态训练提供了一种高效且显存可控的新方案。

研发背景： 该框架由小红书dots infra团队开发，源于内部多模态训练的实际需求。在训练大规模多模态模型（如视觉-语言模型）时，团队发现传统流水并行方法在引入编码器和生成器计算后会导致显著的流水线气泡和显存压力，因此开发了BigMac以解决这一痛点。其设计理念体现了工业界对训练效率与资源利用率的追求。

核心价值： BigMac的核心价值在于打破了计算效率与显存占用之间的帕累托前沿。传统方案中，提高计算效率往往意味着更高的显存占用（如减少流水线气泡但保留更多激活），反之亦然。BigMac通过依赖安全的嵌套流水线调度，在保持LLM流水线高效执行的同时，将编码器和生成器的激活显存降至O(1)级别，使得训练速度提升的同时显存占用可控，支持更大批次大小的训练，从而提升了模型训练的可扩展性。

技术特点： BigMac的技术特点包括依赖安全的嵌套流水线调度、显存有界的模态激活、全局调度与执行解耦、流水线无感接口以及Schedule-aware工具链。其中，依赖安全的嵌套流水线是核心创新，通过分析模块间的数据依赖关系，在LLM流水线的“空档”插入模态计算，既保证了执行顺序的正确性，又避免了跨模块的气泡。此外，全局调度与执行解耦的设计使得调度策略的调整不影响执行层，提高了系统的灵活性。

2. 主要功能

• 依赖安全的嵌套流水线： 以LLM流水线为稳定主干，通过分析模块间的数据依赖，在输入已就绪且不会打乱LLM执行顺序的“空档”插入编码器和生成器的计算。这种设计避免了将模态模块作为独立stage所带来的跨模块气泡，同时保证了LLM的执行连续性，从而在整体上提升了训练吞吐。
• 显存有界的模态激活： 编码器和生成器的前向激活显存被限制在O(1)级别，即不随序列长度或批次大小线性增长。这意味着在训练过程中，模态计算部分的显存占用保持恒定，而LLM的激活行为保持不变，从而支持更大的batch size而不导致内存溢出（OOM），显著提高了显存利用率。
• 全局调度与执行解耦： Scheduler生成覆盖所有rank、microbatch和模块的全局operator表，这张表显式描述了整个训练过程的调度策略。Executor则负责在本地解释对应的operator序列，并分发给LLM后端（Megatron Core）或模态后端（数据并行/FSDP）。这种解耦使得调度策略的调整无需修改执行层代码，便于快速迭代。
• 流水线无感接口： 算法工程师只需通过PP-transparent接口描述各模块的输入输出边界，BigMac自动处理pipeline stage划分、activation handoff、gradient handoff以及跨设备通信。模型代码保持模块化，无需感知底层的流水线runtime，降低了开发复杂度，使得单卡验证的模型可以无缝扩展至流水并行。
• Schedule-aware工具链： 内置profiler、simulator和可视化工具，支持per-operator级别的trace导出，并兼容Perfetto UI。工程师可以利用这些工具在训练前模拟不同并行策略的性能，在训练中定位流水线气泡、慢算子以及通信等待，从而加速性能优化过程。
• 灵活的并行策略配置： 支持设置流水并行（PP）度数、虚拟流水并行（VPP）度数、microbatch数量以及模块在设备上的放置策略。用户可以根据具体的模型结构、集群拓扑和显存预算灵活调整，以找到最优的训练配置。

3. 如何使用

1. 环境要求与前置条件： BigMac基于PyTorch和Megatron Core构建，因此需要安装PyTorch 2.0或更高版本，以及Megatron Core库。推荐使用Linux操作系统，配备NVIDIA GPU（如A100或H100），并配置好CUDA和NCCL通信库。此外，需要准备多模态模型的定义，包括编码器、LLM和生成器的模块代码，并确保这些模块能够在前向和反向传播中正确执行。
2. 定义模块边界： 按照模块化方式编写encoder_forward、llm_forward和generator_forward函数，每个函数需明确声明输入输出的张量形状和数据类型。例如，编码器可能接受图像张量并输出视觉特征，LLM接受文本嵌入并输出隐藏状态，生成器可能从LLM输出生成最终结果。这些边界信息将作为BigMac调度器分析依赖关系的基础。
3. 接入BigMac接口： 通过PP-transparent接口将上述模块注册到BigMac框架中。具体操作包括调用bigmac.register_module()函数，传入模块实例、输入规范（input_spec）和输出规范（output_spec）。注册后，BigMac会自动完成pipeline stage划分、activation handoff和gradient handoff的设置，无需手动干预。
4. 配置并行策略： 设置流水并行（PP）度数、虚拟流水并行（VPP）度数、microbatch数量以及模块在设备上的放置策略。这些参数可以通过配置文件（如YAML）或代码中的配置字典传入。BigMac提供了默认配置，但建议根据具体模型大小和集群拓扑进行调整，以达到最佳性能。
5. 运行训练与性能诊断： 启动Scheduler生成全局operator计划，然后在各rank上通过Executor执行。训练过程中，可以使用内置profiler采集每个operator的执行时间，导出rank-aligned timeline到Perfetto UI进行可视化分析。也可以使用simulator在训练前模拟不同并行策略的性能，以选择最优配置，避免实际训练中的性能损失。

4. 优缺点分析

5. 同类工具对比

选型建议： 对于专注于多模态大模型（如视觉-语言模型、图像生成模型）训练的团队，BigMac提供了最优的调度效率和显存控制，尤其适合理解与生成双路径的场景。其流水线无感接口和内置工具链能够显著降低开发与调优成本。如果团队已经在使用Megatron-LM生态系统，并且主要进行纯文本或简单多模态训练，Megatron-LM仍然是成熟稳定的选择，但需要手动处理模态模块的流水线集成。对于希望利用自动并行减少手动调优的团队，ColossalAI提供了灵活的并行组合，但在多模态场景下需要额外配置才能达到类似效率。DeepSpeed则适合已经深度使用DeepSpeed、需要ZeRO优化的大规模训练，但流水线方面需要额外配置。

6. 编辑总结

BigMac在流水并行调度上提出了依赖安全的嵌套流水线范式，这是一个具有创新性的技术突破。它通过将编码器和生成器计算嵌入LLM流水线的“空档”，有效解决了传统流水并行中计算效率与显存占用之间的权衡问题，使得多模态大模型训练能够在保持高吞吐的同时控制显存开销。从技术创新角度看，BigMac的全局调度与执行解耦设计、显存有界激活机制以及流水线无感接口，都体现了对分布式训练系统设计的深入思考。其实用价值已经在小红书内部生产环境中得到验证，训练速度提升1.08~1.9倍的数据表明其在实际场景中的有效性。BigMac主要面向从事多模态大模型研究和开发的算法工程师和分布式系统工程师，尤其适合需要高效训练大规模多模态模型的团队。当前其社区生态和文档完善程度是影响广泛采用的关键因素，未来在这些方面的投入将决定其普及速度。总体而言，BigMac是一个技术扎实、具有实际应用价值的开源项目，值得多模态训练领域关注和尝试。

7. 应用场景

• 多模态理解模型预训练： 如图文理解、视频理解等MLLM的端到端预训练，需要高效处理ViT、Whisper等编码器与LLM的流水协同。BigMac通过依赖安全的嵌套流水线，确保编码器计算不会干扰LLM的执行顺序，同时将编码器激活显存降至O(1)，使得预训练速度提升1.08~1.9倍，显存占用稳定，支持更大批次的数据处理。
• 多模态生成模型训练： 同时包含理解与生成路径（如图像生成、语音合成）的模型训练，需要协调LLM与MMDiT、LDM等生成器的双向依赖。BigMac的全局调度器能够处理这种复杂依赖，在生成器计算就绪时无缝插入，避免气泡，同时保持生成器激活显存恒定，提升生成模型训练效率。
• 大规模batch训练： 在显存受限集群上追求更大global batch size，传统计算高效设计常因activation累积导致OOM。BigMac的显存有界激活机制使得编码器和生成器显存占用恒定，从而支持更大的batch size，提高训练吞吐，加速模型收敛。
• 训练性能调优： 工程团队可以使用BigMac内置的schedule-aware profiler和simulator，在训练前模拟不同并行策略的性能，在训练中定位pipeline bubble和慢算子，从而优化并行配置和模块放置策略，提升训练效率。这种工具链支持per-operator级别的分析，使得性能调优更加精确。

8. 常见问题FAQ

Q：BigMac与Megatron-LM的主要区别是什么？
A：BigMac针对多模态模型训练进行了专门优化，通过依赖安全的嵌套流水线将编码器和生成器计算嵌入LLM流水线，避免了传统Megatron-LM中将模态模块作为独立stage引入的流水线气泡和显存开销。同时，BigMac提供了更友好的PP-transparent接口和内置的schedule-aware工具链，降低了使用门槛。

Q：BigMac支持哪些硬件平台？
A：目前BigMac基于Megatron Core，主要支持NVIDIA GPU集群，推荐使用A100或H100等高端GPU。对于其他硬件平台（如AMD GPU、华为昇腾），官方尚未公布支持计划，但未来可能扩展。

Q：BigMac是否支持单卡训练？
A：BigMac主要用于流水并行训练，但可以通过设置PP=1在单卡上运行，此时其调度优化效果有限。算法团队可以在单卡或数据并行环境下验证模型结构，再通过BigMac的流水线无感接口无缝扩展至流水并行，无需重写代码。

Q：如何获取BigMac并开始使用？
A：BigMac的源代码已在GitHub上开源（https://github.com/Dots-Infra/BigMac），用户可按照README中的指引安装依赖并运行示例。同时提供了技术论文和项目官网文档，建议用户先阅读技术论文了解设计原理，再参考官方示例进行实践。

Q：BigMac的性能提升具体能达到多少？
A：根据官方数据，BigMac在典型多模态训练任务中可实现1.08～1.9倍的训练速度提升，同时保持显存占用稳定。具体提升幅度取决于模型结构、硬件配置和并行策略，建议用户使用内置simulator在自己的配置下进行模拟。

Q：BigMac是否支持自定义算子或模块？
A：是的，BigMac的调度器支持任意模块，只要模块定义了清晰的输入输出边界。用户可以通过PP-transparent接口注册自定义的编码器或生成器算子，BigMac会自动处理依赖分析和调度。

9. 项目地址

• 项目官网：https://dots-infra.github.io/BigMac/
• GitHub仓库：https://github.com/Dots-Infra/BigMac/
• arXiv技术论文：https://arxiv.org/pdf/2605.25451