小紅書推出 BigMac：突破多模態訓練中的記憶體與速度權衡

多模態大型語言模型正逐漸成為下一代人工智慧的基石，然而其訓練基礎設施長期以來一直受制於一個難以調和的權衡：針對速度進行優化的系統往往面臨記憶體限制的挑戰，而優先考量記憶體效率的系統則通常運行緩慢。 小紅書的 Dots Infra 團隊將此瓶頸視為多模態管道訓練的帕累托極限，並已成功突破此限制。 7月22日，該團隊開源了 BigMac——這是一種專為原生多模態場景設計的創新管道並行訓練範式。該專案現已上架 GitHub，位於 Dots-Infra 儲存庫中。

與標準的 Transformer 模型不同，多模態模型本質上相當複雜。典型的多模態大型語言模型（MLLM）由三個截然不同的組件構成：將圖像和音訊轉換為嵌入向量的模態編碼器、用於推論的大型語言模型（LLM）骨幹，以及將 LLM 輸出映射回目標模態（如圖像或語音）的生成器。 這些組件之間的顯著結構差異，在將其整合至單一訓練管道時，會造成巨大的挑戰。

當前業界的解決方案大致可分為兩類。第一種方法優先考量運算效率，將編碼器和生成器從 LLM 管道中解耦，並分別執行。此舉雖能防止模態模組執行時間的波動在 LLM 管道中產生「氣泡」，但會導致激活記憶體隨微批次數量增加而擴增，在規模擴大時造成高昂成本。 第二種方法則著重記憶體效率，將所有模組保留在同一條管道中，藉此縮短活化週期並降低記憶體使用量。然而，若任何編碼器或生成器運作緩慢，整個大型語言模型管道便會停滯，進而產生尾部氣泡。隨著模型規模擴大，這兩種設計都會暴露出關鍵瓶頸。

BigMac 透過一個簡單卻至關重要的原則來解決此挑戰：將核心 LLM 管線作為首要焦點。大規模 LLM 訓練本就依賴 1F1B 或交錯式 1F1B 等成熟的調度策略，這些策略已深度整合至生產級訓練堆疊中。 BigMac 並非要取代這些既定方法，而是將 LLM 排程作為其底層時間軸，在輸入資料準備就緒時，策略性地插入編碼器和生成器的運算，同時不干擾 LLM 的執行順序。研究團隊將此架構稱為「準依賴安全嵌套管線」。

此設計帶來兩大關鍵優勢。首先，編碼器與生成器的執行時間波動不再傳遞至 LLM 管線，使 LLM 能維持最佳執行節奏；其次，模態活化值的記憶體需求在演算法層級上降至 O(1)。 編碼器無需再保留所有微批次（microbatches）的激發值直至管線結束，而生成器則可避免延長冗長的激發值尾部。本質上，BigMac 並未以記憶體為代價來換取效能提升；相反地，它透過重構排程邏輯，使這兩項目標得以相容，而非相互排斥。

彌合排程設計與實際執行之間的差距涉及重大的工程障礙，包括系統整合、介面定義以及效能除錯。BigMac 透過提供三項核心功能，專門針對這些具體挑戰進行設計。 首先，它使全局排程透明化：排程器會在執行時生成一份涵蓋所有管線 rank、微批次及模組類型的完整運算子表。 執行器隨後將其分配至各級的本地序列中，並與 Megatron Core 等大型語言模型（LLM）後端、模態執行環境及通訊層無縫整合。這種可視性使檢查與後端優化更加容易。

其次，BigMac 提供了一種對演算法工程師而言完全不可見的管線並行介面。使用者只需定義每個模組的輸出與輸入；系統會在內部處理階段劃分、激發函數與梯度傳遞，以及跨裝置通訊。這使得在單一 GPU 上驗證的多模態實驗，能夠自然地擴展至管線並行處理。 第三，它提供了一套效能分析、模擬與視覺化工具。這些工具將訓練迭代分解至運算子層級，精確揭示每個 rank 在每個時間步所執行的操作，識別閒置時段，並突顯依賴性瓶頸。 該模擬器還允許團隊在進行全面訓練之前，測試各種 PP 配置與微批次組合，並評估其對氣泡和吞吐量的影響。

BigMac 的有效性已透過兩項具代表性的工作負載獲得驗證。在 MLLM-Understanding 任務中，採用 Qwen3-30B-A3B 作為骨幹模型，並搭配 13 億維度的 ViT 編碼器， BigMac 與計算效率高的基準模型 Optimus 相比，速度提升達 1.08 倍至 1.1 倍；相較於記憶體效率高的基準模型 Megatron-DistTrain，則提升 1.6 倍至 1.9 倍。 關鍵在於，隨著每張 GPU 的批次大小增加，BigMac 的記憶體使用量保持穩定；反觀 Optimus 則會出現記憶體用量的尖峰飆升，最終在較大的批次下導致記憶體不足 (OOM) 錯誤。 MLLM-Generation 任務更為複雜，涉及一個 200 億維度的 MMDiT 生成器，此時性能差距進一步擴大：Optimus 因 OOM 問題在所有測試批次大小下均失敗，而 BigMac 則透過高效地反向執行生成器並快速釋放活化值來避免此問題。 相較於 Megatron-DistTrain，BigMac 仍能在維持穩定記憶體使用量的同時，提供 1.5 倍至 1.9 倍的加速效果。這凸顯了嵌套管線的主要價值：既能滿足編碼器與生成器的依賴關係，又不會造成過度的活化值保留或顯著的閒置時間。

BigMac 現已成為 Dots 多模態模型訓練基礎架構的核心組件，並在小紅書的生產環境中運行。 相關論文《BigMac：突破多模態大型語言模型訓練中運算與記憶體的帕累托前沿》（BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training）已發布於 arXiv，並附有互動式 PP Profiler 追蹤範例。 對於在多模態訓練中苦於記憶體與速度權衡問題的研究人員和工程師而言，這個經過生產環境驗證的開源專案提供了一種實用的解決方案，能節省大量時間與精力。