BigMac – 小红书开源的多模态大语言模型训练框架

BigMac是小红书dots infra团队开源的多模态大语言模型训练框架通过”准依赖安全嵌套流水线”设计,首次实现显存占用与训练速度的同步优化。传统方案中,多模态训练必须在”算力优先”(速度高但显存爆炸)和”显存优先”(显存低但速度慢)之间二选一,而BigMac将编码器与生成器计算嵌入LLM流水线间隙,既保持LLM主干执行节奏不受干扰,又避免激活值长期驻留显存,实测训练速度比传统方案快1.08~1.9倍,且显存占用随batch规模扩大保持稳定,已在小红书生产环境验证。

BigMac - 小红书开源的多模态大语言模型训练框架

BigMac核心特点

1. 显存与速度双突破

  • 彻底打破帕累托前沿:传统方案中显存与速度互斥,BigMac通过算法层面压缩模态激活显存至O(1)级别,使两者不再非此即彼。在Qwen3-30B-A3B+1.3B ViT编码器任务中,峰值显存随batch增大保持平稳,而”算力优先”方案因需保留编码器激活显存直接溢出。
  • 规模化训练稳定性:在生成型任务(含20B MMDiT生成器)中,所有batch规模下均避免OOM,而”算力优先”方案在任意batch下均失败。

2. LLM流水线主导设计

  • 复用成熟调度策略:不推翻现有LLM流水线(如1F1B、interleaved 1F1B),而是将编码器/生成器计算作为”搭便车”任务插入LLM执行间隙,确保主干节奏不受模态模块耗时波动影响。
  • 依赖安全嵌套机制:仅在输入数据就绪且不破坏LLM执行顺序的位置插入计算,避免传统方案中因模块耦合导致的”尾部气泡”问题。

3. 工程友好型工具链

  • 零感知流水并行接口:算法工程师只需定义模块输入输出关系,自动处理stage划分、激活传递与跨设备通信,单卡验证实验可直接扩展至流水并行。
  • 全流程诊断能力:提供操作符级profiler与simulator,可预估不同PP配置对吞吐/气泡率的影响,无需实际启动训练。

BigMac技术原理

1. 准依赖安全嵌套流水线

  • 动态间隙填充机制:LLM流水线执行过程中存在天然计算间隙(如数据加载等待期),BigMac将编码器/生成器计算任务精准插入这些间隙,输入就绪即启动计算,完成后立即释放显存。
  • 激活生命周期压缩:编码器激活无需保留至流水线结束,生成器激活在backward完成后即时释放,从根本上消除显存膨胀根源。

2. 全局调度架构

  • 三层调度解耦设计
    • Scheduler生成全局算子表:覆盖所有pipeline rank、microbatch及模块类型的完整执行序列。
    • Executor拆解本地指令:将全局序列分发至Megatron Core等后端,自动协调LLM、模态计算与通信任务
    • Simulator预验证配置:训练前模拟不同PP参数组合的性能表现,规避试错成本。
  • 跨模块依赖管理:通过显式声明数据消费/生产关系,自动处理激活值传递与梯度同步,避免人工干预。

3. 显存优化核心机制

  • O(1)显存占用保障:编码器激活仅保留至LLM处理对应microbatch所需时刻,不随batch规模线性增长
  • 生成器侧快速释放:在MLLM-Generation任务中,及时触发generator backward操作,避免20B参数生成器的激活长期驻留。

BigMac核心功能

1. 多模态训练加速

  • 理解型任务优化:在Qwen3-30B-A3B+ViT编码器任务中,比”算力优先”方案快1.08~1.1倍,比”显存优先”方案快1.6~1.9倍。
  • 生成型任务突破:支持含20B MMDiT生成器的复杂架构,比”显存优先”方案快1.5~1.9倍,且显存占用稳定。

2. 无缝工程集成

  • 单卡到集群平滑扩展:算法实验从单卡验证无需修改代码即可扩展至千卡集群,解决多模态训练中常见的扩展性断层问题。
  • 主流框架兼容:深度集成Megatron-Core等工业级训练栈,保留原有LLM调度逻辑,仅需新增模态模块描述。

3. 性能可视化诊断

  • 空泡率精准定位:可视化工具链可识别每个rank在每时刻的计算/空转状态,定位依赖阻塞点。
  • 吞吐预估能力:Simulator支持在启动训练前预测不同配置的吞吐量,指导硬件资源分配。

BigMac项目地址

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

BigMac应用场景

1. 多模态大模型研发

  • 统一架构训练:适用于图像/视频/音频等多模态输入输出模型,避免模态割裂导致的逻辑矛盾(如视频中物体运动违反物理规律)。
  • 生成式多模态任务:显著提升含MMDiT等大型生成器的训练效率,解决显存瓶颈制约模型规模的问题

2. 工业级AI生产部署

  • 高吞吐训练集群:在有限显存条件下支持更大batch规模训练,加速小红书场景中的图文理解、视频生成等业务模型迭代。
  • 资源受限环境适配:通过稳定显存占用特性,在中等规模GPU集群实现大规模多模态训练,降低硬件门槛。

3. 研究与工程提效

  • 快速实验验证:研究者可跳过多模态流水线工程调试,直接聚焦算法创新。
  • 生产环境稳定性保障:已通过小红书生产环境验证,避免实验室到落地的性能衰减,尤其适合需长期稳定运行的工业场景。

BigMac的本质是将多模态训练从”两难选择”转向”协同优化”:它不追求颠覆现有LLM训练体系,而是通过精准嵌入模态计算到LLM流水线间隙,在保留工业级调度优势的同时解决多模态特有瓶颈。对于正面临显存墙与速度瓶颈的团队,该框架将多模态训练的工程复杂度从”月级调优”降至”天级接入”,尤其适合需同时处理图像理解、视频生成等异构任务的场景。随着多模态模型向统一架构演进,能无缝整合异构计算单元的调度方案将成为大模型基础设施的关键一环。

© 版权声明
为这篇文章评分
10.0/ 10
1 人评价
点击⭐️进行评分

相关文章

暂无评论

none
暂无评论...