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 格式之间的双向转换。