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 格式的工具。 |
典型工作流程
- 数据准备 – 使用提供的脚本将数据集 (例如
openai/gsm8k) 转换为仓库的 Parquet 格式。 - Rollout –
Strategy启动 Rollout 引擎 (SGLang 或 vLLM) 来为每个 Prompt 生成多个响应。 - 奖励计算 – 通过基于规则的标签或训练好的奖励模型来计算奖励。
- 优势估计 | 选择一个估计器 (
group_norm,reinforce++,cpgd, 等) 将奖励转化为 Token 级别的优势值 (advantages)。 | - 策略更新 | 运行所选的优化器 (GRPO, DAPO, PPO/GAE, 等) 在 Actor 模型上,并可选择性地混合 Teacher Logits 以进行 OPD。
- 同步与检查点 | 权重在分布式进程中同步,并保存检查点以供后续使用。
快速上手
# 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 与文档; 没有添加任何外部假设。
相关
- 项目
- 项目
- 项目
- 项目
- 项目