opendilab/LightRFT

LightRFT: Light, Efficient, Omni-modal & Reward-model Driven Reinforcement Fine-Tuning Framework

LightRFT – LLM 与多模态生成模型的强化学习微调框架

什么是它

  • 一个开源库,让您可以使用强化学习目标来训练大型语言模型 (LLM) 以及视觉语言 (VLM) 或音频语言模型。它支持经典的 RLHF 风格的流水线,奖励模型驱动的策略优化,以及 on-policy 蒸馏。

为什么它很重要

  • 强化学习微调是使 LLM 与人类偏好对齐的常用方法,但现有工具通常与单一后端(例如 DeepSpeed)或单一模态绑定。LightRFT 统一了 DeepSpeed, PyTorch FSDP v2, SGLang, 和 vLLM 的训练与 Rollout,使得使用相同的代码库即可在文本、图像、视频和音频上运行 RL。

关键设计亮点

特性 功能描述
Unified Strategy abstraction 单一接口 (Strategy) 将分布式训练后端 (DeepSpeed 或 FSDP) 与 Rollout 后端 (SGLang 或 vLLM) 连接起来。这消除了后端特定的样板代码,并允许相同的 GPU 进程同时处理 Rollout 与权重更新(“Colocate Anything”模式)。
Distributed & parameter-efficient training 支持 FSDP v2 和 DeepSpeed ZeRO stages 1-3, BF16, 梯度检查点 (gradient checkpointing), CPU offload, LoRA, 视觉前缀冻结 (visual-prefix freezing), FlashAttention 2, 以及融合的对数概率路径 (fused log-probability paths)。
Reward-model driven optimization 接受基于规则的奖励、自定义 Python 函数、本地神经网络奖励模型或远程奖励服务。每个任务可以聚合多个奖励来源。
Multimodal actors 为文本 (ActorLanguage), 视觉语言 (ActorVL), 和音频语言 (ActorAL) 提供独立的 Actor 类。示例涵盖 GSM8K (文本推理), Geo3K (视觉几何), 视频奖励 RL, 以及音频 Q&A。
Algorithm matrix 实现了一套策略优化和优势估计算法:GRPO, Dr.GRPO, REINFORCE++, DAPO, CPGD, FIRE sampling, 以及 on-policy 蒸馏 (OPD)。目前已为 GSPO 和 GMPO 预留了实验性 Hook。
Experiment tooling 内置于 Weights & Biases 或 TensorBoard 的日志记录、轨迹保存、高熵 Token 标注、分布式检查点 (checkpointing), 以及将检查点转换为 Hugging Face 格式的工具。

典型工作流程

  1. 数据准备 – 使用提供的脚本将数据集 (例如 openai/gsm8k) 转换为仓库的 Parquet 格式。
  2. Rollout – Strategy 启动 Rollout 引擎 (SGLang 或 vLLM) 来为每个 Prompt 生成多个响应。
  3. 奖励计算 – 通过基于规则的标签或训练好的奖励模型来计算奖励。
  4. 优势估计 | 选择一个估计器 (group_norm, reinforce++, cpgd, 等) 将奖励转化为 Token 级别的优势值 (advantages)。 |
  5. 策略更新 | 运行所选的优化器 (GRPO, DAPO, PPO/GAE, 等) 在 Actor 模型上,并可选择性地混合 Teacher Logits 以进行 OPD。
  6. 同步与检查点 | 权重在分布式进程中同步,并保存检查点以供后续使用。

快速上手

# 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 并使用 --gpus all 运行。
  • 运行示例 (使用 Qwen2.5-0.5B 模型在 8 个 GPU 上通过 GRPO 进行 GSM8K 任务):
ENGINE_TYPE=sglang \
  bash examples/gsm8k_geo3k/run_grpo_gsm8k_qwen2.5_0.5b.sh

需要根据需要调整 PATH_TO_YOUR_BASE_MODEL, 数据集路径, 以及 GPU 数量。

谁应该使用它

  • 研究人员,正在实验基于 RL 的 LLM 或多模态生成模型的对齐算法。
  • 需要使用单一代码库在文本 and 视觉/音频数据上运行 RL 的团队。
  • 工程人员,正在寻找一个生产就绪、分布式训练循环,且能同时支援 DeepSpeed 和 PyTorch FSDP。

成熟度 & 社群

  • 当前稳定版本: v0.1.1 (2024 年 6 月)。除了标记为 WIP (GSPO, GMPO) 的两个算法外,列出的所有算法皆可使用。
  • 文档非常详尽 (Sphinx site, GSM8K/Geo3K 教程, 算法指南, 故障排除 FAQ)。
  • 在 Apache 2.0 许可下开源; 欢迎通过标准的 fork-PR 流程进行贡献。

在哪里学习更多

  • 完整文档: docs/ (安装, 算法指南, 策略设计, 运行时架构).

  • 快速上手教程: docs/source/quick_start/grpo_gsm8k_geo3k_tutorial.md.

  • Issues & roadmap: GitHub issues 链接于 README.


以上所有信息皆直接取自于仓库的 README 与文档; 没有添加任何外部假设。

相关

  • 项目
  • 项目
  • 项目
  • 项目
  • 项目