Baidu's internal trainer chases every modality at once

What it does

LoongForge is a training framework that wraps Megatron-LM for language, vision, and diffusion models, and a separate torch-native DDP/FSDP stack for embodied AI. It targets pre-training, continued pre-training, and supervised fine-tuning across both NVIDIA GPUs and Kunlun XPUs. Before it was open-sourced as LoongForge, it ran inside Baidu as AIAK-Training-LLM, reportedly delivering 30%–50% speedups over customer baselines on production runs scaling past 5,000 XPUs.

The interesting bit

The framework treats heterogeneity as a first-class feature: you can assign independent tensor-parallel, data-parallel, and recomputation strategies to different components of a VLM—say, the ViT encoder versus the LLM backbone—rather than forcing one global strategy. For MoE models, it even implements topology-aware dynamic expert replica placement (TAOT) to fix Expert Parallel load imbalance, which is the kind of optimization usually locked inside proprietary clouds.

Key highlights

• Dual-backend architecture: Megatron-LM for LLM/VLM/diffusion, torch-native DDP/FSDP for VLA and world-action models, each heavily optimized rather than abstracted into a least-common-denominator API.
• Heterogeneous parallelism: Independent parallelization and memory strategies per model component, plus decoupled encoder-decoder training to kill pipeline bubbles.
• MoE-native optimizations: Overlapped All2All, activation offload, and adaptive FP8 training with per-operator precision selection; the TAOT paper claims low-communication expert rebalancing.
• Hardware coverage: Native support for NVIDIA GPUs and Kunlun XPUs via a plugin design, with pre-built Docker images and dedicated XPU tutorials.
• Checkpoint diplomacy: Offline bidirectional conversion between Megatron and HuggingFace formats, plus native online HF load/save.

Caveats

• The headline 5× speedup for DeepSeek-V3.2 was measured on a reduced-layer “Lite” configuration because the test bed lacked the scale to run the full model.
• The highest-performance fused DSA kernels are only available on the Baidu Baige platform; the open-source release provides a TileLang version.
• Embodied training is explicitly decoupled from the Megatron core, so “unified” here means coexistence under one umbrella rather than a single execution engine.

Verdict

Worth a look if you are training large open-source models at scale and want Baidu’s collection of memory and parallelism tricks without writing your own kernels. Skip it if you need a lightweight, single-GPU fine-tuning script; this is a heavy industrial stack.