A unified post-training framework for diffusion models

English | 中文

TeleBoost is a unified post-training framework for diffusion models, with DPO and GRPO supported. Used internally at TeleAI for diffusion-model alignment.

• 🎛️ Multi-paradigm post-training — DPO + GRPO
• 🔥 Memory-efficient DPO — ~40% memory cut, ~15× longer context on Wan 14B
• 🆕 Six GRPO algorithms — DanceGRPO, Flow-GRPO, GRPO-Guard, TempFlow-GRPO, BGPO, VIPO
• 🧩 Co-located + MPS reward heads — N rewards on actor GPU; wall ≈ max(model)
• 🎬 Ready-to-use sequence parallel — Ulysses SP for long-video training
• 🚀 Day-0 BGPO + VIPO (CVPR 2026) — TeleAI papers, implemented on release

_{Wan 14B DPO peak memory at 32× 80GB-Hopper GPUs — Decoupled DPO cuts ~40% peak memory on identical workload and scales to ~15× longer context. See the project page linked in the badges above.}

Methods

Method	Status	Use case	Path
DPO	✅ Ready	Preference alignment	recipe/dpo/
GRPO	✅ Ready	Reward-based optimization	recipe/grpo/
TempFlow-GRPO	✅ Ready	Noise-aware weighting + trajectory branching (arXiv 2508.04324)	recipe/tempflow/
GRPO-Guard	✅ Ready	Regulated clipping against implicit over-optimization (arXiv 2510.22319)	recipe/grpo_guard/
BGPO	✅ Ready	Bayesian-prior group optimization (details)	recipe/bgpo/
VIPO	✅ Ready	Pixel-weighted dense advantage (details)	recipe/vipo/
FSDP backend for DPO	🚧 Roadmap	Memory-efficient sharding without DeepSpeed-ZeRO	—

The DPO recipe ships with a precision anchor against the reference standalone-megatron implementation. See recipe/dpo/README.md.

Quickstart

Pick the recipe that matches your post-training need. The top-level README is the feature overview; the recipe docs contain the commands, environment variables, and dataset expectations needed for a real run.

• GRPO — see INSTALL.md. Launch with bash recipe/grpo/run_teleboost.sh after setting TRAIN_FILE / TEST_FILE / WAN_MODEL_PATH / REWARD_MODEL_PATH (TEST_FILE=$TRAIN_FILE works for a smoke run).
• DPO — see QUICKSTART.md and recipe/dpo/README.md. Launch with bash recipe/dpo/run_teleboost_dpo.sh after building the Dockerfile and exporting MEGATRON_LM_DIR.

Documentation map

Need	Start here
Understand the supported algorithms and system features	This README
Install and run GRPO training	INSTALL.md
Build the DPO image and run smoke / real DPO training	QUICKSTART.md
Understand DPO modes, precision alignment, and troubleshooting	recipe/dpo/README.md

---

🚀 What's new from TeleAI

Four TeleAI contributions ship in this repo: VIPO + BGPO (day-0 GRPO papers), co-located reward + MPS (GRPO systems), and Gradient Decoupled DPO (DPO systems).

VIPO — Visual Preference Policy Optimization  ·  GRPO  ·  CVPR 2026  ·  arXiv 2511.18719

Implementation scope. TeleBoost implements VIPO's structured advantage path: a DINOv2-derived allocation map M(p) forms dense advantages A^p = M(p) · A. The concrete map-construction defaults in this repo — PCA pooling, reverse normalization, Gaussian smoothing, and resize behavior — are TeleBoost implementation choices; treat them as this repository's reproducibility baseline, not as a claim of bit-for-bit official paper reproduction. See recipe/vipo/algorithm.py for the full scope statement.

Lifts scalar GRPO feedback into structured, pixel-level advantages via a perceptual structuring module that produces spatially-aware advantage maps. See arXiv:2511.18719.

_{Top — standard Group Relevant Policy Optimization: reward model output collapses into a scalar advantage before policy update. Bottom — VIPO: preference signals are allocated into a structured advantage map, redistributing optimization pressure toward perceptually important regions.}

BGPO — Bayesian-Prior Group Optimization  ·  GRPO  ·  CVPR 2026  ·  arXiv 2511.18919

Naming note. The paper names the method BPGO (Bayesian Prior-Guided Optimization). This repository keeps the existing BGPO identifier for backward compatibility with config keys (algorithm.bgpo.*), env vars (BGPO_*), and BGPOMixin. Use BPGO in external citations; use BGPO when referring to this repo's code/config surface.

Two levels of optimization grounded in a Bayesian prior: inter-group trust allocation (RAS) and intra-group prior-anchored renormalization (CRT). See arXiv:2511.18919.

_{BGPO operates at two levels grounded in a Bayesian prior. Left (RAS): group rewards + prior yield per-sample reliability weights → reliability-aware loss ℒ_RAS. Right (CRT): rewards are renormalized against the prior → recalibrated signal driving the next GRPO loss ℒ_CRT.}

Co-located reward + MPS-parallel multi-reward  ·  GRPO systems

Co-located reward (workers share actor GPUs) + MPS-parallel reward heads (N rewards on one GPU via CUDA MPS). Eliminates the idle reward-rank GPU and brings joint reward-head wall-time ≈ max(model) instead of sum. The trainer driver combines heads in advantage space. On by default in joint mode.

_{Left: co-located reward shares the actor GPUs, eliminating the reward-rank idle gaps. Right: CUDA MPS — N reward models concurrent on one GPU, wall-time ≈ max(model) instead of sum.}

Gradient Decoupled DPO  ·  DPO systems

Per-branch backward + immediate reduce-scatter — frees each branch's full-shape gradient before the next backward starts. Mathematically equivalent to single-backward; on Wan 14B DPO at 32× 80GB-Hopper GPUs: ~40% peak memory cut and ~15× longer context.

_{Per-branch backward + immediate reduce-scatter. Decoupled DPO frees each branch's full-shape tensor before the next backward starts — visibly cutting the peak. (Result figure is at the top of this README.)}

---

Layout

teleboost/          shared lib
  teletron/         TeleTron fork: megatron-side runtime (parallel state, CP/TP layers, adaptor)
  train/            algorithm-agnostic trainer + shared math (algos/) + reward orchestration (rewarding/)
  workers/          verl worker plug-ins (diffusion rollout, sharding manager, reward managers)
  reward/           reward functions, adapters, and reward model implementations
  models/           Wan adapters + model-family facts
  datasets/         DPO dataset base + transforms + video IO
  patches/          runtime monkey-patches over pinned upstream (applied on import)
  runtime/          cross-process transport (transfer queue)
recipe/
  dpo/              DPO recipe (verl plug-in: MegatronEngineWanVideo + split-DPO multi-backward)
  grpo/             GRPO recipe backbone (verl FSDP worker + diffusion actor + backend wiring)
  bgpo/ vipo/ tempflow/ grpo_guard/   per-algorithm verticals (trainer + math + launcher)
docs/               wan DPO workflow + figures referenced by READMEs
tools/              one-shot helpers (e.g. convert_wan_to_teletron.py)
tests/              unit + special-distributed tests
data_preprocess/    GRPO data prep

Dockerfile / makefile / pyproject.toml / requirements.txt   build + deps
LICENSE / NOTICE / CITATION.cff                              upstream attributions
.github/                                                     CI + CODEOWNERS

Both recipes share the Dockerfile / requirements.txt / teleboost/ lib — choose your recipe at launch via recipe/{dpo,grpo}/run_*.sh.

License

TeleBoost is released under Apache 2.0 — see LICENSE.

Acknowledgments

This project builds on the following upstreams. Full per-package attributions (license texts + redistribution terms) are in NOTICE.

RL training stacks

• volcengine/verl — Apache 2.0. Bytedance's RL training framework; TeleBoost is built as a recipe layer on top.
• DanceGRPO — Apache 2.0. TeleBoost's GRPO-family algorithms implement the DanceGRPO method described in arXiv 2505.07818.
• Tele-AI/TeleTron — Apache 2.0. TeleAI's long-context multi-modal training framework; TeleBoost is built on top, with Gradient Decoupled DPO added.

Generation models

• Wan-Video/Wan2.1 — Apache 2.0. Alibaba's Wan2.1 / Wan2.2 video diffusion models, vendored under third_party/wan/ with the upstream LICENSE retained.
• modelscope/DiffSynth-Studio — Apache 2.0. ModelScope's diffusion toolkit; the FlowMatchScheduler (teleboost/models/flow_match.py) and the Wan VAE adapter are ported from it.

Reward models

• tgxs002/HPSv2 — Apache 2.0. Human Preference Score v2.
• alibaba-pai/VideoCLIP-XL — CC-BY-NC-SA-4.0 (non-commercial). Alibaba's video-text alignment model. Not redistributed; the videoclip reward loads a user-supplied copy placed under third_party/VideoCLIP_XL/.
• Hritikbansal/videophy — MIT. UCLA's video physical-plausibility model.
• LAION-AI/aesthetic-predictor — MIT. LAION's CLIP + linear-head aesthetic predictor.
• mlfoundations/open_clip — MIT. The ViT-L-14 vision encoder behind the aesthetic reward; teleboost/models/offline_clip.py is a standalone reimplementation (dual OpenCLIP-MIT / TeleAI-Apache header).
• princeton-vl/RAFT — BSD-3-Clause. Princeton's optical-flow model, used as a temporal-consistency reward.
• TencentARC/VideoAlign — Apache 2.0. Referenced for reward-model design; not vendored in this repo (pull upstream if you want to train it).

Citation

@article{teleboost2026,
  title  = {TeleBoost: A Systematic Alignment Framework for High-Fidelity,
            Controllable, and Robust Video Generation},
  author = {Liang, Yuanzhi and Wu, Xuan'er and Liu, Yirui and Fang, Yijie and
            Fan, Yizhen and Hao, Ke and Li, Rui and Liu, Ruiying and Ni, Ziqi and
            Yu, Peng and Wang, Yanbo and Huang, Haibin and Weng, Qizhen and
            Zhang, Chi and Li, Xuelong},
  year   = {2026},
}

For per-algorithm citations:

• GRPO algorithms (DanceGRPO, Flow-GRPO, GRPO-Guard, BGPO, VIPO) — see CITATION.cff.
• Gradient Decoupled DPO — cite the TeleBoost paper above.