baidu-baige/LoongForge

A unified, high-performance framework for training LLMs, VLMs, diffusion, and embodied models on NVIDIA GPUs and Kunlun XPUs.

解決的問題

LoongForge 是一個高性能訓練框架，旨在統一大規模語言模型 (LLM)、視覺語言模型 (VLM)、擴散模型與具身智能模型 (VLA/WAM) 的訓練。它解決了在不同硬體上訓練不同類型模型時效率低下與碎片化的問題，提供了一個支援預訓練、持續預訓練與監督微調 (SFT) 的統一技術棧。

工作原理

LoongForge 根據模型類型利用多種分散式後端：LLM、VLM 與擴散模型運行在經過補丁修復的 Megatron-LM 版本上，而具身智能模型使用 torch 原生的 DDP/FSDP 棧。該框架同時支援 NVIDIA GPU 與昆侖 XPU。它採用了多種優化技術來提高吞吐量：

  • 異構並行 (Heterogeneous Parallelism)：允許為不同的模型組件（例如，分離 ViT 與 LLM）進行獨立的張量並行 (TP)、數據並行 (DP) 與重計算。
  • MoE 優化：包括重疊的 All2All、激活卸載 (activation offload) 以及拓撲感知專家負載平衡演算法。
  • 自適應 FP8 訓練：使用分塊 FP8，並具有可選的自適應模式，根據算子效率選擇精度。
  • DP 負載平衡：重新分配數據以緩解序列打包 (sequence-packing) 不平衡問題。
  • 自定義融合算子 (Custom Fused Operators)：實現高性能內核，如 FusedDSA，以加速特定的模型架構。

適用對象

LoongForge 面向從事大規模模型訓練的研究人員與工程師，特別是那些需要在異構硬體（NVIDIA 與昆侖 XPU）上訓練多模態或具身智能模型的人員。

亮點

  • 統一框架：在同一個地方支援 LLM、VLM、擴散模型與具身智能模型。
  • 廣泛的模型支援：相容包括 DeepSeek、Qwen、LLaMA 與 GR00T 在內的多種 SOTA 模型。
  • 硬體無關性：原生支援 NVIDIA GPU 與昆侖 XPU。
  • 顯著加速：展示了相對於基準測試的顯著訓練加速（例如，DeepSeek-V3.2 Lite 實現 5.04 倍加速）。
  • 靈活組合：透過可互換的組件進行配置驅動的 VLM 組裝。
  • 無縫檢查點：支援 Megatron 與 HuggingFace 格式之間的雙向轉換。