Announcing LoongForge-Embodied: a high-performance training subsystem for embodied models
Today we are releasing LoongForge-Embodied inside LoongForge — a torch-native training subsystem purpose-built for VLA (Vision-Language-Action) and WAM (World-Action Model) models. It covers mainstream embodied models such as Pi0.5, GR00T N1.6, X-VLA, FastWAM, LingBot-VA, Cosmos3, and DreamZero in one shot, and delivers 1.6× ~ 2.67× training speedups over each model's official implementation on the same hardware.
📚 GitHub: https://github.com/baidu-baige/LoongForge/tree/master/loongforge/embodied
1. What it is: a torch-native embodied training subsystem
LoongForge is a unified training framework spanning LLM / VLM / Diffusion / Embodied. The LLM / VLM / Diffusion stacks are built on Megatron-LM, targeting large-scale, model-parallel pre-training and SFT. Embodied models have a fundamentally different shape, so we built a dedicated subsystem for them.
The reason is that embodied models differ from typical large models in a fundamental way: they have far fewer parameters (usually under 10B), and a typical VLA is often just a VLM plus an action head. This scale does not need the TP / PP / EP model parallelism designed for giant models — forcing it on only adds complexity and overhead. LoongForge-Embodied is therefore built directly on native PyTorch DDP / FSDP, focusing on training efficiency and model-onboarding experience in small-to-medium, data-parallel scenarios.
The two stacks share the same repository, release process, and toolchain, but do not share the Megatron core engine, nor args / parser / core, so each can evolve independently:
| Dimension | LoongForge core (LLM / VLM / Diffusion) | LoongForge-Embodied |
|---|---|---|
| Compute / distributed | Megatron-LM — TP / PP / EP / CP / FSDP | torch-native DDP / FSDP |
| Workload profile | Large-scale, model-parallel pre-training / SFT | Small-to-medium, data-parallel SFT |
| Model scale | Billions to hundreds of billions of parameters | Usually under 10B |
2. Supported models: mainstream VLA and WAM, all at once
The first release already covers mainstream VLA and WAM models, with a unified onboarding path:
| Category | Models |
|---|---|
| Pi | pi0.5 |
| GR00T | groot_n1_6, groot_n1_7 |
| XVLA | xvla |
| FastWAM | fastwam |
| LingBot-VA (WAM) | lingbot_va |
| Cosmos3 | cosmos3_nano |
| DreamZero | Series variants (Wan2.1 14B / Wan2.2 5B, LoRA and full fine-tuning, covering DROID / LIBERO / AgiBot / YAM and other datasets) |
Full fine-tuning is a standard capability; on the data side it supports formats such as LeRobot / HDF5 (supported formats vary by model), and some models (e.g. DreamZero) also provide LoRA variants.
3. Performance: measured speedups on the same hardware
The following are measured speedups for LoongForge-Embodied on the same hardware configuration as each model's official implementation:
| Model | Type | Baseline | Speedup | Measured version |
|---|---|---|---|---|
| Pi0.5 | VLA | OpenPI | 2.23× | main · 2026-07 |
| GR00T N1.6 | VLA | LeRobot | 2.31× | main · 2026-07 |
| X-VLA | VLA | X-VLA | 1.6× | main · 2026-07 |
| DreamZero (DROID Wan2.2-5B full) | WAM | DreamZero | 2.67× | main · 2026-07 |
| LingBot VA | WAM | LingBot-VA | 1.80× | main · 2026-07 |
This speedup spans both VLA and WAM categories and multiple mainstream implementations, rather than holding for a single model only.
The numbers above reflect the performance of each baseline and LoongForge version at measurement time (see the "Measured version" column) and may change as implementations evolve; the exact speedup also varies with model, data, and configuration.
4. Architecture: share common capabilities, push model differences down
The design philosophy of this subsystem can be summed up in one line: share common capabilities, push model differences down — make the common layer deep and the model layer thin, so adding a model requires no changes to the training main loop.
① Model layer (model/) — one directory per model. Each model registers into
a unified entry point via @register_model: modeling_<name>.py handles
the network / forward / loss, and model_configuration_<name>.py is the
architecture-hyperparameter dataclass. A unified interface is exposed upward, so onboarding a new model
does not touch the training loop.
② Data pipeline (data/) — common backends + per-model differences.
Dataset-reading backends (lerobot / hdf5 / dummy), a stateful distributed sampler, and a
composable transform framework are consolidated as common capabilities; model-specific data reading
(such as fastwam's multi-frame geometry), action / image transforms, and batch assembly are
pushed down to datasets/<name>/; each model declares image size, action dimensions,
normalization statistics, etc. via DataConfig.
③ Three-layer config parsing (train/parser.py) — frozen into immutable
objects. The YAML model: section → ModelConfig, the YAML
data: section → DataConfig, the command line → TrainingArgs;
--model-name is routed via config_map.py to the corresponding YAML and type,
and the command line can also override fields via dotlist (e.g. model.action_horizon=64),
all finally frozen into a global singleton.
④ Distributed trainer (train/trainers/, distributed/) — flexible
strategies. Standard SFT uses FinetuneTrainer directly; special paradigms such
as multi-stream and CUDA Graph inherit BaseTrainer and register in
trainer_builder.py, selected via --trainer-type; distribution can switch on
demand among ddp (data parallel), ddp + --zero-optimizer (ZeRO
Stage-1), fsdp (full sharding), and hsdp (hybrid sharding), matching the scale
and memory characteristics of embodied models. BaseTrainer also consolidates optimizer / LR
scheduling, gradient clipping and NaN cleanup, checkpoint resume, and determinism control into the
common layer, letting every model reuse the same implementation.
Adding a new model takes just 5 steps: ① add the network and config under
model/<name>/ and @register_model; ② add
data/datasets/<name>/ (DataConfig + transform + collator); ③ add a YAML under
configs/models/embodied/ and register it in config_map.py; ④ inherit
BaseTrainer if the paradigm differs, otherwise reuse FinetuneTrainer; ⑤ add a
launch script under examples/.
5. Quick start
Training is launched with torchrun running loongforge/embodied/train.py:
export LOONGFORGE_PATH=/workspace/LoongForge
PYTHONPATH=$LOONGFORGE_PATH:$PYTHONPATH \
torchrun --nproc_per_node 8 --nnodes 1 \
$LOONGFORGE_PATH/loongforge/embodied/train.py \
--model-name pi05 \
--trainer-type FinetuneTrainer \
--dataset-format lerobot_datasets \
--distributed-strategy fsdp \
--train-iters 30000 \
...Different models differ in data format, processing pipeline, and performance-optimization config, so the
ready-to-run command varies by model. Out-of-the-box example scripts live under
examples/embodied/:
bash examples/embodied/pi05/run_pi05_fsdp_finetune.sh # ddp / ddp_zero1 variants also available
bash examples/embodied/groot_n1_6/run_groot_n1_6_ddp_finetune.sh6. Evaluation: offline benchmarks, process-decoupled
Beyond training, LoongForge-Embodied ships with an offline evaluation module covering LIBERO / CALVIN / SimplerEnv / RoboTwin / ManiSkill. It splits the benchmark client and the model policy server into separate processes, connected via a WebSocket / msgpack-numpy RPC protocol, with YAML as the sole user entry point:
cd /path/to/LoongForge
examples/embodied/pi05/eval/run_libero_eval.shThe model side only needs to implement a unified predict_action interface; the generic
policy layer GenericPredictActionPolicy handles image-view selection, action shape
validation, chunk caching, latency statistics, and normalization-statistics loading. Onboarding a new
model only requires adding a lightweight factory under eval/factories/, with no changes to
the server main flow.
The evaluation module is still under active development; more benchmarks and capabilities will be added.
7. Closing thoughts
The release of LoongForge-Embodied marks the point where LoongForge's training capabilities for embodied models graduate from scattered VLA support inside the main framework into an independent, decoupled subsystem: mainstream VLA and WAM in one shot with unified onboarding, delivering 1.6× ~ 2.67× training speedups over official implementations on the same hardware.
We believe the R&D efficiency of embodied models depends on whether you can train them with infrastructure that best fits their shape. We welcome you to try it out, and to join the community via GitHub Issues, WeChat, or Slack to help us make it better together.