If you train models across more than one modality, you've probably felt this: Megatron-LM is battle-tested for LLMs but wasn't built for vision-language-action robots. LeRobot and OpenPI are built for research iteration on embodied models, not production throughput. Diffusion training lives in yet another silo. You end up maintaining three or four codebases — each with its own parallelism strategy, checkpoint format, and performance ceiling.
We just open-sourced LoongForge to collapse that into one framework. It's the open-source release of AIAK-Training-LLM, Baidu Baige's training acceleration stack, which has run production training for enterprise customers in education, computer vision, and embodied AI — typically delivering 30–50% speedups over customer baselines, with the largest runs reaching 5,000+ accelerators.
Repo: https://github.com/baidu-baige/LoongForge (Apache-2.0)
What makes it different
1. Embodied models training faster.
This is the part most training frameworks don't touch. LoongForge ships a dedicated torch-native DDP/FSDP subsystem for VLA and world-action (WAM) models — Pi0.5, GR00T N1.6, X-VLA, FastWAM, DreamZero, LingBot-VA — decoupled from the Megatron core, with flexible DDP / ZeRO-1 / FSDP / HSDP strategies. If you train robot policies today, production-grade throughput engineering is rare. Here it's built in.
2. One framework, four modalities.
LLMs, VLMs, diffusion, and embodied — under a unified model abstraction. VLMs are assembled config-driven from interchangeable ViT + LLM components, so a custom ViT + LLM backbone is a YAML change, not a fork.
3. Two hardware backends, natively.
NVIDIA GPUs and Kunlun XPUs via a minimally-intrusive plugin design — the same training code runs on both.
Performance
Tuned to beat the specific open-source baseline each model type is usually trained on.
The engineering underneath
- Flexible Multi-Modal Composition — Configuration-driven assembly of VLMs from interchangeable ViT and LLM components.
- Heterogeneous Parallelism — Independent TP / DP / recompute per model component (e.g., ViT vs. LLM) for optimal throughput and memory. [blog]
- Decoupled Encoder-Decoder Training — Separates ViT and LLM into independent tasks, eliminating encoder-induced pipeline bubbles.
- DP Load Balancing — Load-aware data redistribution mitigates sequence-packing imbalance, improving multi-node scaling efficiency. [blog]
- MoE-Native Optimization — Overlapped All2All / activation offload / compute, with further memory reduction beyond upstream Megatron-LM on DeepSeek-V3, Qwen3-MoE, etc.
- MoE Expert Load Balancing — Dynamically replicates hot experts using a topology-aware algorithm to balance Expert Parallel (EP) workloads and improve training efficiency.
- Adaptive FP8 Training — End-to-end FP8 for LLMs and VLMs with standard blockwise FP8; optional adaptive mode picks per-operator precision by GEMM shape and efficiency.
- Custom Fused Operators — Fused kernels like FusedDSA for DSA-style models.
- Flexible Checkpointing — Offline bidirectional Megatron ↔ HuggingFace conversion plus native online HF load/save — no format barriers across your workflow.
- Versatile Pipelines & Data Tools — Out-of-the-box Pretrain / MidTrain / SFT / LoRA, with built-in dataset format conversion and sequence packing.
- Embodied Model Training — A dedicated torch-native DDP/FSDP subsystem for VLA and world-action (WAM) models (e.g. Pi0.5, GR00T N1.6, FastWAM), decoupled from the Megatron core, with flexible DDP / ZeRO-1 / FSDP / HSDP strategies.
- Heterogeneous Hardware — Native support for NVIDIA GPUs and Kunlun XPUs via a minimally-intrusive plugin design.
Who it's for
It fits teams pre-training or fine-tuning at scale who are tired of maintaining separate stacks for language, vision-language, diffusion, and robotics, and who care about throughput. Apache-2.0, covers Pretrain / MidTrain / SFT / LoRA out of the box, runs on NVIDIA and Kunlun.
Try it / contribute
If you're training in this space, I'd genuinely value your feedback — especially on the embodied subsystem. What models would you want supported next?
Top comments (0)