opendilab/LightRFT

LightRFT – Reinforcement‑Fine‑Tuning Framework for LLMs & Multimodal Generative Models

What it is

• An open‑source library that lets you train large language models (LLMs) and vision‑language (VLM) or audio‑language models using reinforcement‑learning objectives. It supports classic RLHF‑style pipelines, reward‑model‑driven policy optimization, and on‑policy distillation.

Why it matters

• Reinforcement fine‑tuning is the go‑to method for aligning LLMs to human preferences, but existing tools are often tied to a single backend (e.g., DeepSpeed) or a single modality. LightRFT unifies training and rollout across DeepSpeed, PyTorch FSDP v2, SGLang, and vLLM, making it possible to run RL on text, images, video, and audio with the same code base.

Key design highlights

Feature	What it does
Unified Strategy abstraction	A single interface (Strategy) connects the distributed training backend (DeepSpeed or FSDP) with the rollout backend (SGLang or vLLM). This removes backend‑specific boiler‑plate and lets the same GPU processes handle both rollout and weight updates (the “Colocate Anything” model).
Distributed & parameter‑efficient training	Supports FSDP v2 and DeepSpeed ZeRO stages 1‑3, BF16, gradient checkpointing, CPU offload, LoRA, visual‑prefix freezing, FlashAttention 2, and fused log‑probability paths.
Reward‑model driven optimization	Accepts rule‑based rewards, custom Python functions, local neural reward models, or remote reward services. Multiple reward sources can be aggregated per task.
Multimodal actors	Separate actor classes for text (ActorLanguage), vision‑language (ActorVL), and audio‑language (ActorAL). Examples cover GSM8K (text reasoning), Geo3K (visual geometry), video‑reward RL, and audio Q&A.
Algorithm matrix	Implements a suite of policy‑optimization and advantage‑estimation algorithms: GRPO, Dr.GRPO, REINFORCE++, DAPO, CPGD, FIRE sampling, and on‑policy distillation (OPD). Experimental hooks for GSPO and GMPO are present.
Experiment tooling	Built‑in logging to Weights & Biases or TensorBoard, trajectory saving, high‑entropy token annotation, distributed checkpointing, and utilities for converting checkpoints to Hugging Face format.

Typical workflow

1. Data prep – Convert a dataset (e.g., openai/gsm8k) to the repository’s Parquet format with the provided scripts.
2. Rollout – The Strategy launches a rollout engine (SGLang or vLLM) that generates multiple responses per prompt.
3. Reward – Compute rewards via rule‑based labels or a trained reward model.
4. Advantage estimation – Choose an estimator (group_norm, reinforce++, cpgd, etc.) to turn rewards into token‑level advantages.
5. Policy update – Run the selected optimizer (GRPO, DAPO, PPO/GAE, etc.) on the actor model, optionally mixing in teacher logits for OPD.
6. Sync & checkpoint – Weights are synchronized across the distributed processes and saved for later use.

Getting started

# Clone and install (editable mode)
git clone https://github.com/opendilab/LightRFT.git
cd LightRFT
pip install -e .          # core deps (torch, deepspeed, sglang, …)
pip install -e "[vllm]"   # optional vLLM rollout backend

• Docker: docker pull opendilab/lightrft:v0.1.0 and run with --gpus all.
• Run an example (GSM8K with GRPO on a Qwen2.5‑0.5B model, 8 GPUs):

ENGINE_TYPE=sglang \
  bash examples/gsm8k_geo3k/run_grpo_gsm8k_qwen2.5_0.5b.sh

Adjust PATH_TO_YOUR_BASE_MODEL, dataset path, and GPU counts as needed.

Who should use it

• Researchers experimenting with new RL‑based alignment algorithms for LLMs or multimodal generative models.
• Teams that need a single code base to run RL on text and vision/audio data.
• Engineers looking for a production‑ready, distributed training loop that works with both DeepSpeed and PyTorch FSDP.

Maturity & community

• Current stable release: v0.1.1 (June 2024). All listed algorithms are functional except the two marked WIP (GSPO, GMPO).
• Documentation is extensive (Sphinx site, tutorials for GSM8K/Geo3K, algorithm guide, troubleshooting FAQ).
• Open‑source under the Apache 2.0 license; contributions are welcomed via the usual fork‑PR workflow.

Where to learn more

• Full docs: docs/ (installation, algorithm guide, strategy design, runtime architecture).
• Quick‑start tutorial: docs/source/quick_start/grpo_gsm8k_geo3k_tutorial.md.
• Issues & roadmap: GitHub issues linked from the README.

---

All information above is taken directly from the repository’s README and documentation; no external assumptions have been added.