LeanGRPO removes the update-time policy recomputation used by diffusion GRPO. It provides two execution modes alongside the unchanged native baseline:
- Retain keeps selected rollout graphs and applies the known policy loss after reward computation.
- Reweight backpropagates each selected rollout transition immediately, then applies group advantages to isolated gradients before synchronization.
In our experiments, LeanGRPO achieves up to 1.83x end-to-end training speedup over Native recomputation.
- Support DanceGRPO.
- Support FlowGRPO Fast.
- Support Wan2.1-T2V-1.3B.
- Support FLUX.1-dev.
- [2026.09.02] We released LeanGRPO training scripts for Wan2.1-T2V-1.3B and FLUX.1-dev🔥
| Integration | Model path | Native | Retain | Reweight |
|---|---|---|---|---|
| DanceGRPO | Wan2.1-T2V-1.3B | yes | yes | yes |
| FlowGRPO Fast | FLUX.1-dev | yes | yes | yes |
| MixGRPO / MixGRPO-Flash | FLUX.1-dev | planning | planning | planning |
The two upstream integrations have different dependency stacks. Install the common package and one integration environment at a time:
pip install -e '.[test]'
pip install -r integrations/dancegrpo/requirements.txt
# or, in a separate environment:
pip install -r integrations/flowgrpo/requirements.txtTraining requires Linux, a CUDA/NCCL environment supported by PyTorch 2.7, and
FFmpeg. The public configurations were designed for eight GPUs; multi-node
launches require a routable MASTER_ADDR and the same checkout on every node.
Model and reward weights are intentionally not bundled. Users are responsible for obtaining them under their respective licenses and configuring paths through the launchers.
The public launchers follow the upstream DanceGRPO style: the complete training
configuration is written directly in each torchrun command. Edit the model,
data, output paths, mode, and hyperparameters in the shell script, then run:
scripts/train_wan13b.sh
scripts/train_flux.shSet --mode to native, retain, or reweight. See
docs/data-layout.md for the prompt and reward-model files
expected by each launcher.
The launchers intentionally avoid an environment-variable configuration layer, so the exact resolution, sampling layout, optimizer, precision, and memory settings are visible in one place. The public Python entrypoints still accept integration-specific arguments for users building their own launchers.
The original upstream training programs remain as stable, thin entrypoints. Their implementations are split by responsibility:
integrations/dancegrpo/fastvideo/wan_training/
cli.py Wan arguments
trainer.py model construction and outer training loop
modes/ Separate Native, Retain, and Reweight update paths
sampling.py Wan rollout and transition construction
fsdp.py FSDP2, LoRA, and gradient-buffer operations
rewards.py video decoding and reward adapters
checkpoint.py resumable training state
runtime.py generation-group process state
integrations/flowgrpo/flow_grpo/training/
cli.py FLUX arguments and outer training loop
modes/ Separate Native, Retain, and Reweight epoch drivers
rollouts.py retained-graph and immediate-backward rollouts
sampling.py FlowGRPO rollout sampling
modeling.py model loading, log probability, and optimizer operations
policy.py advantage and PPO bookkeeping
data.py prompt data and text embeddings
checkpoint.py resumable training state and metrics
integrations/flowgrpo/flow_grpo/fsdp2_runtime.py
shared FSDP2, checkpoint, and gradient-buffer primitives
Read scripts/train_*.sh first to understand a complete experiment, then follow
the public Python wrapper into cli.py, modes/, and the selected rollout.
The hot Native, Retain, and Reweight loops intentionally remain separate so their
FSDP transactions and CUDA/NCCL ordering stay visible.
LeanGRPO builds on DanceGRPO, FastVideo, FlowGRPO, and Hugging Face Diffusers. We thank the authors and open-source communities behind these projects. See THIRD_PARTY_NOTICES.md for detailed provenance and license information.