AAllHere AI
登录注册

小红书开源 BigMac：打破多模态训练显存与吞吐帕累托前沿

小红书 dots infra 团队开源 BigMac：打破多模态训练“帕累托困境”

近日，小红书 dots infra 团队正式开源了多模态大模型训练系统 BigMac，并同步发布了相关技术论文《BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training》。该系统旨在解决多模态大语言模型（MLLM）训练中计算效率与显存占用难以兼顾的行业痛点，目前已成为支撑小红书 dots 系列大模型训练的核心组件之一。

  • 论文地址：https://arxiv.org/pdf/2605.25451
  • 开源地址：https://github.com/Dots-Infra/BigMac/

困境：多模态训练的“帕累托前沿”

随着头部厂商纷纷探索覆盖视觉、音频和文本的统一模型体系，底层计算的异构性成为了新的瓶颈。一个典型的 MLLM 包含模态编码器、LLM 主干和模态生成器三个模块，它们在 GPU 集群上遵循不同的计算节奏。

将这三个模块接入同一套训练流水线后，业界通常面临两难选择（即材料中提到的帕累托前沿）：

1. 计算高效但显存昂贵

策略：将编码器和生成器从 LLM pipeline 中分离（类似 LongCat-Flash-Omni 的 modality-decoupled parallelism）。
缺点：编码器 activation 需保留至梯度返回，LLM activation 甚至需保留至生成器结束。随着 microbatch 数量增加，显存占用激增，容易触碰上限。

2. 显存高效但存在空泡（Bubble）

策略：将所有模块整合进同一条 pipeline（如 Megatron Core 的多模态示例）。
缺点：模块间被 pipeline 依赖耦合。若 encoder 或 generator 耗时波动，LLM pipeline 必须等待，导致尾部产生大量空泡，牺牲了计算吞吐。

方案：依赖安全的嵌套流水线 (Nested Pipeline)

BigMac 的核心创新在于提出了一种原生多模场景下的流水并行训练新范式。它不试图替换成熟的 LLM 调度策略，而是以 LLM 流水线为主干，在不打乱 LLM 执行顺序的前提下，有序嵌入编码器和生成器计算。

关键机制与优势

  • 保持 LLM 主干稳定：LLM pipeline 继续运行在优化过的 schedule（如 1F1B）上，不受模态模块耗时波动的干扰。
  • 激活显存有界 (O(1))：BigMac 在梯度就绪后尽快运行 encoder 和 generator 的 backward 过程，尽早释放 activation。这使得模态 activation 显存降低到 O(1)，同时保持 LLM activation 行为不变。
  • 解耦调度与执行：工程上解耦了全局调度与运行时执行，降低了模型接入成本。

落地工程能力

为了让理论落地，BigMac 提供了三大关键工程能力：

  1. 全局计划可视化：Scheduler 生成覆盖所有 rank 的全局 operator 表，Executor 分发任务。调度策略变得可见、可检查，且后端兼容 Megatron Core 等现有优化。
  2. PP 透明接口：算法工程师只需定义模块的生产/消费关系，BigMac 自动处理 stage 划分、activation handoff 和通信，无需修改底层 PP 逻辑。
  3. Schedule-aware 工具链：提供 Profiler 和 Simulator。Profiler 导出 per-rank timeline 以定位 bubble 来源（如 compute imbalance 或 communication wait）；Simulator 可在训练前预估不同配置对吞吐的影响，减少试错成本。

实验结果

团队在两个代表性负载上评估了 BigMac 的性能，对比基线包括计算高效的 Optimus 和显存高效的 Megatron-DistTrain。

1. MLLM-Understanding（理解任务）

  • 配置：Qwen3-30B-A3B 作为 Backbone，1.3B 参数的 ViT encoder。
  • 性能提升：
    • 相比 Optimus：训练速度提升 1.08x - 1.1x。
    • 相比 Megatron-DistTrain：训练速度提升 1.6x - 1.9x。
  • 显存表现：随着 batch size 增加，BigMac 的峰值显存保持稳定；而 Optimus 因保留大量 activation，在大 batch size 下会 OOM。

2. MLLM-Generation（生成任务）

  • 配置：同上 Backbone 和 Encoder，新增一个 20B 参数的 MMDiT generator。
  • 性能提升：
    • 相比 Megatron-DistTrain：训练速度提升 1.5x - 1.9x。
  • 关键突破：在此负载下，Optimus 在所有测试 batch size 上均发生 OOM（因为 LLM activation 需保留至巨大的 Generator 计算结束）。BigMac 通过及时运行 generator backward 避免了这一问题，实现了速度与显存的兼顾。

背景与应用

BigMac 已应用于小红书内部生产环境，作为 dots 多模态模型 训练的核心组件。dots infra 团队负责支撑 dots 系列大模型的预训练、后训练及推理服务。

据材料显示，依托该工程体系，小红书「dots-note-3.0」内部版本在第 67 届国际数学奥林匹克竞赛（IMO 2026）中，于全部六道题中获得满分 7 分，取得 42/42 的满分金牌成绩。

作者与团队

  • 杨程旭：小红书 dots infra 组工程师，北京大学计算机学院博士毕业。
  • 章梓立：小红书 dots infra 组实习生，北京大学计算机学院在读博士生。
  • 铭曦：小红书 dots infra 组工程师，卡内基梅隆大学硕士毕业。

参考来源：机器之心，原文链接 https://mp.weixin.qq.com/s/Ost0idlB6oxOZqPkkC5L3A