BigMac – 小紅書開源的多模態大模型流水並行訓練框架

0 0 0

BigMac是什麼

BigMac 是小紅書 dots infra 團隊開源的多模態大模型流水並行訓練框架，針對 MLLM 訓練中計算效率與顯存佔用難以兼顧的問題，提出依賴安全的嵌套流水線範式。框架用成熟 LLM 流水線爲主幹，在不打亂執行順序的前提下嵌入編碼器和生成器計算，實現訓練速度提升 1.08～1.9 倍的同時保持顯存穩定，已作爲 dots 多模態訓練核心組件投入生產。

BigMac的主要功能

依賴安全的嵌套流水線：以 LLM pipeline 爲穩定主幹，在依賴滿足的位置嵌入 encoder / generator 計算，避免跨模塊 bubble 干擾。
顯存有界的模態激活：編碼器和生成器激活顯存降至 O(1)，LLM 激活行爲保持不變，支持更大 batch size 而不 OOM。
全局調度與執行解耦：Scheduler 生成覆蓋所有 rank、microbatch 和模塊的全局 operator 表；Executor 解釋本地序列並分發給對應後端。
流水線無感接口：算法工程師只需描述模塊輸入輸出邊界，BigMac 自動處理 schedule、handoff 和跨設備通信，模型代碼保持模塊化。
Schedule-aware 工具鏈：內置 profiler、simulator 和可視化工具，支持 per-operator 級別 trace 導出（兼容 Perfetto UI），快速定位 bubble、慢算子與通信等待。

BigMac的技術原理

依賴安全的嵌套流水線調度：BigMac將其作爲基礎時間線。調度器分析 encoder forward / LLM forward / generator forward 及對應 backward 的數據依賴關係，僅在輸入已就緒且不會打亂 LLM 執行順序的”空檔”插入模態計算。 LLM 按原 schedule 持續推進，encoder 和 generator 的激活也能在梯度就緒後儘快釋放，backward 及時執行，打破”計算高效必高顯存、顯存高效必多 bubble”的帕累託權衡。
全局 Operator 表與後端解耦：系統將調度策略顯式化爲一張全局 operator 表，包含所有 pipeline rank 上 LLM / encoder / generator 的 forward/backward operator 及模塊邊界通信。執行層僅解釋本地 operator 序列，LLM operator 複用 Megatron Core，模態 operator 保留自身 data-parallel / FSDP 實現，schedule 調整無需重寫模型 runtime。

BigMac – 小紅書開源的多模態大模型流水並行訓練框架

微信關注回覆“開源”，加入AI開源項目交流羣

如何使用BigMac

定義模塊邊界：按模塊化方式編寫 encoder_forward、llm_forward、generator_forward 函數，聲明各模塊的輸入輸出。
接入 BigMac 接口：通過 PP-transparent 接口註冊模塊，BigMac 自動處理 pipeline stage 劃分、activation handoff、gradient handoff 和跨 rank 通信。
配置並行策略：設置 PP / VPP、microbatch 數量、模塊放置策略等參數。
運行訓練：啓動 Scheduler 生成全局計劃，Executor 在各 rank 上執行 operator 序列。
性能診斷：用內置 profiler 採集 per-operator 時間，導出 rank-aligned timeline 至 Perfetto UI 分析 bubble 與瓶頸；或使用 simulator 在啓動訓練前快速迭代並行策略。

BigMac的核心優勢

打破帕累託前沿：同時實現高計算效率與低顯存佔用，無需二選一。
LLM 流水線零干擾：保留成熟 LLM schedule，編碼器/生成器耗時波動不會沿 pipeline 傳播。
顯存可擴展：encoder / generator 激活顯存 O(1)，支持生產級大 batch 訓練。
算法友好：模型代碼無需感知 pipeline runtime，單卡驗證的實驗可自然擴展至流水並行。
可調試可部署：全局 schedule 可見可檢查，配套 trace / simulation 工具鏈降低性能優化成本。

BigMac的項目地址

項目官網：https://dots-infra.github.io/BigMac/
GitHub倉庫：https://github.com/Dots-Infra/BigMac/
arXiv技術論文：https://arxiv.org/pdf/2605.25451

BigMac的同類競品對比

維度	BigMac	Megatron-LM
調度範式	依賴安全的嵌套流水線，LLM 爲主幹	顯存高效的單條流水線，模態模塊作爲獨立 stage
計算效率	高，無跨模塊 bubble	中，encoder/generator 波動會引入 pipeline bubble
顯存佔用	穩定，模態激活 O(1)	較低，但隨模塊耦合顯存優化空間有限
擴展性	支持理解+生成雙路徑大規模訓練	生成路徑擴展時 bubble 與顯存壓力顯著
接口友好度	PP-transparent，算法代碼零侵入	需理解並適配 pipeline runtime 與 stage 劃分
工具鏈	內置 schedule-aware profiler + simulator	依賴傳統 nsys / torch trace 手動關聯分析

BigMac的應用場景

多模態理解模型預訓練：如圖文理解、視頻理解等 MLLM 的端到端預訓練，需高效處理 ViT / Whisper 等編碼器與 LLM 的流水協同。
多模態生成模型訓練：同時包含理解與生成路徑（如圖像生成、語音合成），需協調 LLM 與 MMDiT / LDM 等生成器的雙向依賴。
大規模 batch 訓練：在顯存受限集羣上追求更大 global batch size，避免傳統計算高效設計因 activation 累積導致 OOM。
快速實驗迭代：算法團隊可在單卡或 DP 環境驗證模型結構，通過 BigMac 接口無縫擴展至 pipeline 並行，無需重寫代碼。
訓練性能調優：工程團隊用 schedule-aware profiler 和 simulator 定位 pipeline bubble，優化並行配置與模塊放置策略。