====================================================================================== LightRFT 项目概览 ====================================================================================== LightRFT 是一个专为 RLHF（基于人类反馈的强化学习）和 RLVR（基于可验证奖励的强化学习）设计的框架，能够实现大语言模型（LLM）和视觉语言模型（VLM）的高效微调。该项目的结构旨在支持跨多个 GPU 和节点的分布式训练，同时保持简洁的模块化架构。 高层架构 ====================================================================================== .. image:: ./lightrft_system.png :alt: LightRFT 系统 :width: 100% :align: center LightRFT 框架采用分层架构，旨在将推理大模型应用到不同领域。系统分为四个主要层级，从基础硬件基础设施逐步演进到高级推理应用： **异构硬件层 (Heterogeneous Hardware Layer)**：基础层支持多种 GPU 架构，包括 NVIDIA GPU（A100/A800/H100）、华为 GPU (Work in Progress) 以及其他硬件平台，实现在不同计算环境下的灵活部署。 **分布式执行层 (Distributed Execution Layer)**：该层通过多种训练引擎（DeepSpeed ZeRO、PyTorch FSDP）和高性能推理引擎（vLLM、SGLang）提供分布式执行能力，确保高效的资源利用和可扩展性。 **RLVR 框架层 (RLVR Framework Layer)**：作为系统的核心，该层集成了各种模型组件——Actor 模型、Critic 模型、参考模型（Reference Models）、可验证规则（Verifiable Rules）和奖励模型（Reward Models），并结合了包括 GRPO、DAPO、TTRL、CPGD 在内的高级优化算法。模型和算法通过 **colocate** 机制进行交互，该机制能够在同一计算资源内无缝优化并共同部署不同的组件。 **推理大模型应用层 (Reasoning Large Models Application Layer)**：最顶层专注于提供跨多个领域的推理能力，包括形式数学推理（AIME/GSM8K）、代码生成（LiveCodeBench/SWE）、基于搜索的高效推理以及多模态推理（图像/视频/音频/音乐）。(Work in Progress) 该架构在分布式执行框架和 RLVR 框架内运行，协调所有组件，为大语言模型、视觉语言模型、音频语言模型和扩散模型提供可扩展、高效且可信的推理应用。 训练流程图 ====================================================================================== .. image:: ./lightrft_pipeline.png :alt: LightRFT 流水线 :width: 100% :align: center 与传统的强化学习任务（如游戏或机器人技术，它们严重依赖 CPU 密集型环境）不同，大语言模型主要通过神经网络推理在线生成数据，产生大量的 token 序列。这使得在不同计算资源之间划分不同角色（如生产者和消费者）的解耦系统架构变得不再适用。 在这种“模型即环境”的设置下，**colocate** 架构将训练流水线的所有阶段整合在相同的计算资源中。每台机器通过数据并行承担相同的角色，不同的阶段通过底层引擎的切换进行转换。 这种设计实现了简洁的可扩展性，并确保了计算资源的高利用率。 在 LightRFT 中，我们将这一概念扩展到训练中使用的所有模型，建立了我们的 **通用 Colocate 机制 (Universal Colocate Mechanism)**。 - **多模式 Colocate (Multi-Mode Colocate)**：多回复生成（rollout）、奖励判定（辅助计算）和优化（训练）在统一的硬件系统中运行——就像厨师、骑手和质检员在单个调度中心内无缝协作。这消除了交接延迟，提升了端到端效率。 - **即插即用模块 (Plug-and-Play Modules)**：无需重建系统即可添加新机制——如奖励、约束或教师模型（teachers）。这就像在飞行中更换乐高车轮或转向模块，能够快速适应不断变化的需求。 - **高效模式切换 (Efficient Mode Switching)**：在训练、生成（rollout）和校验评分模式之间无缝切换，开销微乎其微，类似于一级方程式赛车的快速进站。结合灵活的数据路由和模型共享，这最大限度地减少了冗余的资源分配和切换开销。系统在保持峰值效率的同时减少了总训练时间。 核心模块 ====================================================================================== 数据集模块 (Datasets Module) -------------------------------------------------------------------------------------- Datasets 模块为各种训练场景（包括 prompt、有监督微调 (SFT) 和奖励建模）提供数据加载和预处理能力。 核心组件： - 用于 RL 训练的 Prompt 数据集（文本和视觉语言） - 用于有监督微调的 SFT 数据集（文本和视觉语言） - 奖励模型数据集（生成式和标量奖励模型） - 音频语言数据集支持 - 数据预处理和分词 (tokenization) 工具 模型模块 (Models Module) -------------------------------------------------------------------------------------- Models 模块定义了语言模型和视觉语言模型进行强化学习所需的神经架构和适配。它专注于通过 monkey patch 方法对现有模型架构进行最小限度的修改。 核心组件： - Actor 网络实现（语言、视觉语言、音频语言） - 奖励模型实现（生成式和标量奖励模型） - 策略优化的损失函数定义 - 通过 monkey patching 进行的非侵入式模型修改 - 支持多种模型架构（LLaMA、Qwen 等） - 针对 RL 训练优化的自定义生成方法 策略模块 (Strategy Module) -------------------------------------------------------------------------------------- Strategy 模块实现了不同的分布式训练方法，使框架能够跨多个 GPU 和节点高效扩展。它为各种并行技术和优化策略提供了抽象。 核心组件： - 用于基于 ZeRO 优化的 DeepSpeed 集成 - 完全分片数据并行 (FSDP) 实现 - 高效推理引擎 (vLLM, SGLang) - 训练引擎与推理引擎间的模型参数高效同步 - 分布式张量操作工具 - 检查点 (Checkpoint) 管理和广播 - 策略选择与配置 训练器模块 (Trainer Module) -------------------------------------------------------------------------------------- Trainer 模块实现了强化学习算法和训练循环。它协调模型、策略和优化过程之间的交互。 核心组件： - 策略优化：PPO/GRPO、CPGD、DAPO 组件与高熵 token 筛选；GSPO、GMPO、Dr.GRPO 仍属于未完整接通的接口 - 优势估计与回报计算：GAE、REINFORCE、RLOO、REINFORCE with baseline、组归一化和 CPGD；REINFORCE++ 尚未接入优势计算器工厂 - 经验生成与采样策略：FIRE Sampling - 奖励处理与稳定性增强：Reward Norm/Clip, High Entropy Token Selection - 策略与价值函数更新 - 重放缓存 (Replay buffer) 管理 - KL 散度控制 - SPMD (单程序多数据) 训练支持 工具模块 (Utils Module) -------------------------------------------------------------------------------------- Utils 模块提供了支持整个框架的通用工具函数，包括 CLI 参数解析、日志记录、分布式采样和数据处理。 核心组件： - 命令行界面 (CLI) 参数解析 - 用于数据加载的分布式采样器 (Distributed sampler) - 日志记录与监控工具 - 多模态输入的处理器 - 远程奖励模型工具 - 计时器与性能分析 (profiling) 工具 - 用于训练效果分析的轨迹 (Trajectory) 保存 系统工作流 ====================================================================================== LightRFT 框架通过协调的工作流运行： 1. **初始化**：系统首先通过 Utils 模块设置分布式 environment 并加载配置。 2. **数据加载**：Datasets 模块加载并预处理训练数据，为 prompt、SFT 或奖励建模创建 dataloader。 3. **模型准备**：加载预训练模型，并通过 Models 模块中的策略性 monkey patching 将其适配用于 RL。 4. **策略选择**：根据配置选择并初始化适当的分布式训练策略。 5. **经验生成**：系统利用 Strategy 模块集成的推理引擎（如 vLLM/SGLang）进行高效的大规模并行采样，生成用于训练的经验数据 (Rollout)。 6. **模型优化**：Trainer 模块基于生成的经验计算奖励和优势，并执行策略更新。更新后的模型参数随后被同步回推理引擎，为下一轮迭代做准备。 7. **检查点与评估**：系统定期保存模型状态并评估性能。 扩展点 ====================================================================================== LightRFT 在设计时充分考虑了扩展性： - **新数据集**：可以通过扩展 Datasets 模块中的基础数据集类来添加额外的数据集类型。 - **新模型**：可以通过 monkey patch 系统或创建新的 actor/reward 模型类来增加对其他模型架构的支持。 - **新的分布式训练策略**：可以通过扩展 strategy 基类来实现新的分布式训练方法。 - **自定义 RL 算法**：不同的强化学习算法可以实现为新的 trainer 类。 - **推理优化**：可以通过 strategy 模块集成额外的推理引擎，并持续集成 vLLM 和 SGLang 等推理引擎的最新版本特性。 - **自定义工具**：可以添加新的工具函数和处理器以支持特定用例。 这种模块化设计使 LightRFT 能够适应新的研究方向、模型架构和硬件配置，同时为用户保持一致的接口。