Light, Efficient, Omni-modal & Reward-model Driven Reinforcement Fine-Tuning Framework
English | 简体中文
LightRFT (Light Reinforcement Fine-Tuning) is a reinforcement fine-tuning framework for large language models (LLMs) and vision-language or multimodal generative models (VLMs). It provides a structured and extensible workflow for reinforcement learning with verifiable rewards (RLVR), reinforcement learning from human feedback (RLHF), and model-reward-driven policy optimization, covering policy sampling, reward computation, advantage estimation, and policy updates. The repository also includes reward-model training and on-policy distillation workflows.
LightRFT uses torchrun and PyTorch distributed communication as its runtime foundation. A unified Strategy interface connects FSDP v2 and DeepSpeed ZeRO training backends with SGLang and vLLM rollout backends. Current code paths and examples include text, image, video, and audio tasks. “Omni-modal” means that the repository contains dedicated model, data, or example paths for these modalities; it does not imply that every model and modality combination works without adaptation.
This document describes repository version
0.1.1. Source code and runnable examples define the implemented feature boundary; roadmap items are not treated as released features.
- Design highlights
- Supported algorithm matrix
- Runtime architecture
- Installation
- Quick start
- Configuration
- Examples and applications
- Monitoring, trajectories, and checkpoints
- Repository layout
- Documentation and troubleshooting
- Roadmap
- Contributing
- Citation, license, and acknowledgements
- A common interface connects DeepSpeed/FSDP v2 training backends with SGLang/vLLM rollout backends, reducing backend-specific coupling in the upper-level training workflow.
- Distributed jobs follow a single-program, multiple-data (SPMD) topology launched with
torchrun; Ray is not required for scheduling, and standard PyTorch distributed tools remain applicable for debugging. - Training and rollout reuse the same GPU process set by phase. The rollout engine can sleep during policy updates and receives refreshed Actor weights before the next rollout.
LightRFT calls this logical colocation and phase-oriented resource sharing model Colocate Anything. See Runtime Architecture and Resource Reuse for Strategy boundaries, evaluation flow, model placement, and weight synchronization.
- FSDP v2 and DeepSpeed ZeRO stages 1/2/3; DeepSpeed is selected when
--fsdpis absent. - BF16, gradient checkpointing, Adam offload, and FSDP CPU offload.
- LoRA, visual-prefix freezing, and sample packing.
- Optional FlashAttention 2 and a fused log-probability path.
- Rule rewards, custom reward functions, local reward models, and remote reward services.
- Multiple reward sources with task-specific aggregation.
- Training entry points for vision scalar reward models (SRM), vision generative reward models (GRM), and audio SRMs.
- On-policy distillation (OPD) with teacher log-probabilities, either as a distillation-only objective or combined with task rewards.
- Text (
ActorLanguage), vision-language (ActorVL), and audio-language (ActorAL) policy paths. - Image inputs are handled by the vision-language path; the experience-generation path also handles video fields.
- Examples cover GSM8K text reasoning, Geo3K visual geometry reasoning, video reward-model RL, and audio question answering.
- Weights & Biases and TensorBoard logging.
- Trajectory saving and analysis for repetition, reflection patterns, and policy entropy.
- High-entropy-token annotation and local visualization.
- Distributed checkpoints, optional Hugging Face checkpoints, and conversion utilities.
LightRFT organizes policy optimization, advantage estimation, sampling, and knowledge distillation as composable modules. See the algorithm guide for principles and detailed configuration.
| Algorithm | Type | Main improvement | Current implementation and entry point | Reference |
|---|---|---|---|---|
| GRPO | Policy Optimization | Group-normalized advantage estimation | Supported: use --advantage_estimator group_norm; requires multiple responses per prompt |
arXiv:2402.03300 |
| GSPO (WIP) | Policy Optimization | Group sequence policy optimization | Experimental interface: --use_gspo and related options are available while integration is in progress |
arXiv:2507.18071 |
| GMPO (WIP) | Policy Optimization | Geometric-mean policy optimization | In development: the end-to-end training path is being completed | arXiv:2507.20673 |
| Dr.GRPO | Policy Optimization | Mitigation of length bias | Supported: unbiased group-relative optimization reduces length bias and improves token efficiency | arXiv:2503.20783 |
| REINFORCE++ | Advantage Estimation | Improved baseline estimation | Supported: use --advantage_estimator reinforce++ for return and advantage estimation |
arXiv:2501.03262 |
| DAPO | Policy Optimization | Decoupled clipping and dynamic sampling | Supported: includes --dynamic_sampling, --overlong_buffer, and related training mechanisms |
arXiv:2503.14476 |
| CPGD | Advantage Estimation | KL-drift constraint | Supported: use --advantage_estimator cpgd; --use_cpg_loss enables asymmetric clipping |
arXiv:2505.12504 |
| FIRE Sampling | Sampling Strategy | High-temperature first-token sampling for greater diversity | Supported: configure with --use_fire and --first_token_temperature |
arXiv:2410.21236 |
| OPD | Knowledge Distillation | On-policy teacher–student token-level distillation | Supported: reads teacher log-probabilities from --teacher_model_url and supports pure or task-reward-hybrid distillation |
Blog |
The main training entry point, examples/gsm8k_geo3k/train_colocate.py, also provides the following foundational training paths:
| Method | --advantage_estimator |
Critic | Description |
|---|---|---|---|
| PPO / GAE | gae |
Required | Computes GAE from value estimates and trains with a value loss |
| REINFORCE | reinforce |
No | Builds token-level returns from sequence rewards |
| RLOO | rloo |
No | Uses a leave-one-out group baseline and requires multiple responses per prompt |
| REINFORCE with baseline | reinforce_baseline |
No | Uses the group mean as the baseline without standard-deviation scaling |
These training paths can be combined with the following stability and efficiency mechanisms:
- Sample filtering and length control:
--dynamic_samplingmasks groups with no reward variation, while--overlong_bufferadds a length-dependent penalty to overlong responses. - Token-level updates:
--high_entropy_token_ratiorestricts policy-gradient updates to a selected fraction of high-entropy tokens;0.0disables filtering. - Numerical stability:
--reward_running_norm,--reward_clip,--advantages_norm, and--advantage_clipcontrol reward normalization, reward clipping, advantage whitening, and advantage clipping.
Implementation status: All algorithms in the matrix are supported except GSPO and GMPO, which remain WIP. WIP entries expose their corresponding designs or experimental interfaces but are not complete training paths in the current release.
A typical LightRFT training cycle is:
data preparation → rollout generation → reward and experience construction
→ advantage estimation and policy update → weight synchronization
The Trainer organizes the iteration, Strategy provides the distributed training and rollout interfaces, and reward components evaluate generated responses. The relationships among Actor, Reference Model, Critic, and rollout policy—and the exact sequence of engine sleep/wake, model reload/offload, and weight synchronization—are documented in Runtime Architecture and Resource Reuse.
| Component | Source installation requirement or note |
|---|---|
| Python | >= 3.12 |
| PyTorch | >= 2.9.1 in pyproject.toml |
| GPU | Distributed training requires a CUDA-capable NVIDIA GPU environment |
| Default rollout backend | SGLang >= 0.5.6.post2 |
| Optional rollout backend | vLLM >= 0.18.1 |
| Training backend | DeepSpeed >= 0.18.3, or PyTorch FSDP v2 |
CUDA, PyTorch, FlashAttention, SGLang, and vLLM have binary compatibility constraints. Select versions compatible with the installed driver and CUDA runtime; the repository Dockerfile is one pinned reference environment.
SGLang is included in the default dependency set:
git clone https://github.com/opendilab/LightRFT.git
cd LightRFT
pip install -e .Install the optional vLLM backend with:
pip install -e ".[vllm]"Alternatively, install a compatible vLLM release after the default installation:
pip install "vllm>=0.18.1"Running a GPU container requires Docker and NVIDIA Container Toolkit. The published example image is version v0.1.0:
docker pull opendilab/lightrft:v0.1.0
docker run --gpus all -it --rm \
--ipc=host \
-v /path/to/data:/app/data \
-v /path/to/checkpoints:/app/checkpoints \
opendilab/lightrft:v0.1.0 /bin/bashBuild the repository Dockerfile with:
make dbuild
make dbuild IMAGE_NAME=your-custom-tag:latestThe current Dockerfile starts from nvcr.io/nvidia/pytorch:25.01-py3 and explicitly installs a PyTorch 2.9.0 CUDA 12.8 wheel, DeepSpeed 0.18.3, vLLM 0.18.1, FlashAttention 2.8.3, and SGLang 0.5.6.post2. Its PyTorch version is lower than the >=2.9.1 source-package declaration. Verify the intended version set before treating the Dockerfile as a release reference.
If a FlashAttention source build fails, select a wheel that exactly matches Python, PyTorch, CUDA, and the C++ ABI. For example, the repository Docker environment uses:
pip install flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_x86_64.whlWhen no matching wheel is available, build from source in an environment with the required compiler toolchain. See the installation guide and troubleshooting guide.
The launchers in this repository are training templates. Before running them, review model and dataset paths, GPU count, rollout tensor parallelism, sequence lengths, batch sizes, and logging configuration.
The example uses Qwen2.5-0.5B-Instruct, GSM8K, group-normalized advantages, and rule rewards.
python examples/gsm8k_geo3k/data_preprocess/gsm8k.py \
--local_save_dir /path/to/data/gsm8kThe preprocessing script reads openai/gsm8k and writes training and test Parquet files. Each example contains the prompt, reference answer, and the gsm8k_rule reward label; the training recipe uses answer-correctness and output-format rules rather than a neural reward model.
Edit examples/gsm8k_geo3k/run_grpo_gsm8k_qwen2.5_0.5b.sh and verify at least:
PATH_TO_YOUR_BASE_MODEL="Qwen/Qwen2.5-0.5B-Instruct"
PATH_TO_YOUR_GSM8K_DATASET="/path/to/data/gsm8k"
export NNODES=1
export GPUS_PER_NODE=8
ENGINE_TP=2Also review W&B credentials, master address/port, batch sizes, and sequence lengths. ENGINE_TP must divide the total process count.
If W&B is not required, leave WANDB_API_KEY empty and remove or adjust the corresponding launcher options. Multi-node execution also requires correct NODE_RANK, MASTER_ADDR, and MASTER_PORT values.
# Default SGLang backend
ENGINE_TYPE=sglang \
bash examples/gsm8k_geo3k/run_grpo_gsm8k_qwen2.5_0.5b.sh
# Optional vLLM backend
ENGINE_TYPE=vllm \
bash examples/gsm8k_geo3k/run_grpo_gsm8k_qwen2.5_0.5b.shThe launcher uses torchrun to start train_colocate.py; its default recipe enables FSDP, BF16, FlashAttention, engine sleep/wake, rule rewards, and group-normalized advantages. It is an eight-GPU template. When reducing GPU count, also adjust tensor parallelism, global and micro batch sizes, sequence lengths, or model size.
Prepare Geo3K, review the model and dataset paths in the launcher, and then run:
python examples/gsm8k_geo3k/data_preprocess/geo3k.py \
--local_save_dir /path/to/data/geo3k
ENGINE_TYPE=sglang \
bash examples/gsm8k_geo3k/run_grpo_geo3k_qwen2.5_vl_7b.shSee the GSM8K/Geo3K tutorial for the complete workflow.
Model, data, algorithm, distributed-backend, rollout-engine, logging, and checkpoint options are organized in the configuration guide. Entry points do not necessarily expose identical arguments, so also inspect the selected launcher and its command-line help:
python examples/gsm8k_geo3k/train_colocate.py --helpFor reproducible experiments, use launchers, documentation, and argument parsers from the same repository revision, and treat the selected entry point's --help and source implementation as authoritative.
| Directory | Modality or task | Purpose |
|---|---|---|
examples/gsm8k_geo3k/ |
Text and image | GRPO, PPO, LoRA, and rule-reward training |
examples/orm_rl_demo/ |
Image | Combined format, general model, and accuracy rewards |
examples/grm_training/ |
Image/video reward | Vision GRM training |
examples/grm_vl_rl/ |
Video | Policy optimization with a vision reward model |
examples/srm_training/ |
Image and audio | Vision/audio SRM training |
examples/r1_aqa/ |
Audio | Audio-question-answering GRPO |
examples/on_policy_distillation/ |
Text | Teacher service and OPD training |
examples/math_benchmarks/ |
Text evaluation | Math500, AIME, GPQA, and related benchmarks |
examples/entropy_viz/ |
Analysis | Local visualization of high-entropy tokens |
examples/chat/ |
Interactive inference | Check exported model generation |
Example shell files contain cluster-specific paths, ports, and GPU settings and should be treated as templates.
LightRFT supports Weights & Biases, TensorBoard, trajectory recording and analysis, high-entropy-token visualization, distributed training-state recovery, and Hugging Face-format checkpoints. See the configuration guide for the relevant options, lightrft/utils/ckpt_scripts/README.md for checkpoint conversion, and examples/entropy_viz/render_trajectories.html for local trajectory visualization.
LightRFT/
├── lightrft/
│ ├── datasets/ # Text and multimodal datasets
│ ├── evaluation/ # Evaluation and reward functions
│ ├── models/ # Text, vision, and audio Actors and reward models
│ ├── strategy/
│ │ ├── deepspeed/ # DeepSpeed strategy
│ │ ├── fsdp/ # FSDP v2 strategy
│ │ ├── sglang_utils/ # SGLang engines and weight synchronization
│ │ └── vllm_utils/ # vLLM engines and weight synchronization
│ ├── trainer/ # Advantage computation, experience generation, and trainers
│ └── utils/ # Logging, trajectory, and checkpoint utilities
├── examples/ # Training, distillation, evaluation, and analysis examples
├── docs/ # Sphinx documentation
├── tools/ # Version and Docker helper tools
├── README.md
└── README_zh.md
- Installation
- GSM8K/Geo3K tutorial
- Algorithms
- Configuration
- Strategy guide
- Strategy design philosophy
- Runtime architecture
- Reward models
- FAQ
- Troubleshooting
- Contributing
For rollout-backend, GPU-memory, distributed-initialization, multimodal-data, and training-stability issues, consult the FAQ and troubleshooting guide.
pip install -r requirements-doc.txt
make docsThe HTML output is written to docs/build/html/index.html. For live preview:
make docs-live
# Open http://localhost:8000 in a browserRoadmap entries describe proposed work and are not guarantees of current functionality.
Issues and pull requests are welcome. The recommended workflow is:
- Fork the repository and create a feature or documentation branch from
main. - Keep the change scoped and add the necessary tests or documentation checks.
- Use a Conventional Commits style commit message.
- Push the branch and open a pull request describing the motivation, changes, and validation.
Common repository commit types include feature, fix, polish, docs, style, and refactor. Documentation branch names should contain doc when the documentation deployment workflow is required.
Run the development checks with:
pip install -r requirements-dev.txt
make format # YAPF
make fcheck # Flake8See the contribution guide for the repository workflow.
If LightRFT supports your research or application, please cite:
@misc{lightrft,
title={LightRFT: Light, Efficient, Omni-modal & Reward-model Driven Reinforcement Fine-Tuning Framework},
author={Niu, Yazhe and Pu, Yuan and Shi, Dongxing and Lu, Yudong and Xiong, Yingtong and Ge, Ruijun and Sun, Jiaxuan and Wan, Zunian and Zhang, Shaoang},
publisher={GitHub},
howpublished={\url{https://github.com/opendilab/LightRFT}},
year={2025},
}LightRFT is licensed under the Apache License 2.0.
LightRFT is based on OpenRLHF, with some files and implementations adapted or reused. The project also builds on or learns from verl, SGLang, vLLM, DeepSpeed, and PyTorch FSDP.
The project is developed in collaboration with colleagues from the System Platform Center and the AI Safety and Trustworthiness Center at Shanghai AI Laboratory.
- Issues: opendilab/LightRFT
- Email: [email protected]