从零开始掌握字节开源强化学习框架 veRL:原理、架构与实践
从零开始掌握字节开源强化学习框架 veRL:原理、架构与实践
1. 强化学习与 LLM 的奇妙结合:为什么要用 veRL?
1.1 强化学习的核心概念
强化学习(Reinforcement Learning, RL)的灵感来源于行为心理学。想象一下训练一只小狗:当它做出正确的动作(比如坐下)时,我们给它一块零食(正奖励);当它做错时,不予理睬或轻声呵斥(负奖励)。通过无数次这样的“尝试-反馈”,小狗逐渐明白做什么能获得更多零食。
在强化学习的术语中,做决策的“小狗”被称为智能体(Agent),它所处的环境是环境(Environment),它做出的行为是动作(Action),而环境反馈给它的“零食”或“惩罚”就是奖励(Reward)。智能体的目标很单纯:最大化它在整个过程中获得的累计奖励。
1.2 为什么大语言模型需要强化学习?
传统的语言模型训练(即监督微调,SFT)就像是“模仿学习”。模型看着海量的“标准答案”去学习如何组织语言。但这种方式有个局限:它只能学会“像”标准答案,却很难学会“推理”。
以数学题为例,模型可能记住了解题步骤的表面形式,但遇到没见过的变体就容易出错。强化学习的引入改变了这一点。我们不再只是给模型看答案,而是让它自己去“做”题,然后根据结果的好坏(答案是否正确、推理过程是否清晰)给予奖励。模型通过试错,逐渐学会如何“思考”才能获得更高奖励。这正是 DeepSeek-R1 等模型展现出强大推理能力的关键技术路径之一。
1.3 veRL:为 LLM 而生的强化学习库
veRL(volcano engine Reinforcement Learning)是由字节跳动开源的高性能强化学习训练库,专为大型语言模型设计。它的核心优势体现在三个层面:
- 灵活性强:它采用了一种巧妙的“混合控制器”架构,既能像灵活的脚本一样方便研究人员快速迭代算法逻辑,又能像分布式系统一样高效调度大规模计算资源。
- 性能卓越:veRL 集成了业界顶尖的训练框架(如 PyTorch FSDP、Megatron-LM)和推理引擎(如 vLLM、SGLang),并解决了训练与推理阶段模型切换时的内存和通信瓶颈,实现了业界领先的吞吐量。
- 生态开放:它与 HuggingFace 的模型生态无缝集成,你可以轻松加载 Qwen、Llama 等主流开源模型进行强化学习训练。
2. veRL 核心工作原理:不止是 GRPO
2.1 从 PPO 到 GRPO:简化与提效
近端策略优化(PPO)曾是强化学习微调语言模型的事实标准。但 PPO 需要额外训练一个与策略模型规模相当的价值网络(Critic),这几乎翻倍了显存和计算开销。
GRPO(Group Relative Policy Optimization)是 DeepSeek 提出的改进算法,它摒弃了价值网络,改为对每个输入问题,让模型生成一组(例如 N=5 个)不同的回答,然后在这组回答内部进行相互比较,评估每个回答的“相对优势”。这种组内比较的方式极大地降低了训练的资源门槛。veRL 是 GRPO 算法的高性能参考实现之一。
2.2 veRL 的 3D-HybridEngine 架构
veRL 高性能的秘密在于它解决了一个核心工程难题:LLM 的强化学习训练是一个多模型、多阶段的复杂流水线。这个过程至少涉及:
- Rollout 阶段:策略模型(Actor)作为推理引擎,为 prompts 生成回答。
- Ref 阶段:参考模型(Reference)计算旧策略的概率,用于约束策略更新幅度。
- 训练阶段:策略模型进行梯度计算和参数更新。
在早期的实现中,为了将庞大的模型参数从训练模式切换到推理模式,需要消耗大量的时间在 GPU 间复制数据。veRL 的3D-HybridEngine 实现了高效的 Actor 模型分片,消除了这种冗余,允许模型在训练和生成两种状态间快速切换,显著减少了通信开销。
2.3 技术蓝图:向模块化服务演进
veRL 的架构可以分为两层:
- verl-core:作为底层基础,它抽象出四个核心组件:模型引擎(负责训练)、Rollout 引擎(负责推理)、Checkpoint 引擎(负责权重同步)、传输队列(负责数据搬运)。这种抽象使得 veRL 可以轻松更换后端,比如用 Megatron-LM 替代 FSDP 来训练超大模型。
- verl-trainer:基于核心组件,构建了各种 RL 训练流水线,如同步策略、一步异步策略、全异步策略等,以适应不同规模的训练需求。
目前,veRL 正从“紧密耦合”的架构向 “服务化”架构演进。Rollout 引擎将成为一个独立的“推理服务”,通过 HTTP 接口与训练端通信。这样做的好处是训练和推理可以完全分离,部署在不同节点上,从而实现更高级的资源弹性和容错。
3. 实战演练:用 veRL 训练一个会解数学题的模型
接下来,我们通过一个实际例子来感受 veRL 的使用。我们将使用Qwen2.5-32B-Instruct 模型,在 GSM8K 数学数据集上进行 GRPO 强化学习微调。整个过程基于火山引擎机器学习平台(veMLP),但其核心概念和命令在本地环境同样适用。
3.1 准备工作:环境、数据与模型
-
环境配置:veRL 依赖 Ray 进行分布式管理。在开发机上,你需要配置好火山引擎的命令行工具
volc cli,并设置VOLC_ACCESS_KEY等环境变量,以便后续任务能访问存储和计算资源。 -
数据准备:veRL 要求数据为
.parquet格式。你可以使用 veRL 提供的预处理脚本生成 GSM8K 数据集:cd /root/code/verl/examples/data_preprocess/ python3 gsm8k.py --local_dir /root/code/data/gsm8k这会在指定目录下生成
train.parquet和test.parquet文件。每个数据条目至少需要包含一个prompt字段。 -
模型下载:从 HuggingFace 或其他镜像源将 Qwen2.5-32B-Instruct 模型下载到本地存储路径。veRL 可以直接读取本地路径或 HuggingFace 模型 ID。
3.2 核心概念:角色(Role)与资源池(Resource Pool)
在 veRL 的分布式编排中,两个核心概念是 角色(Role)和资源池(Resource Pool)。
- 角色:定义了不同功能模块,如
ActorRolloutRef(混合了策略、推理、参考模型功能)、Critic(价值网络)等。 - 资源池:是 GPU 资源的分组。你可以将模型的所有角色放置在同一个资源池中(共享 GPU),也可以将它们隔离到不同资源池中(独占 GPU),以实现更精细的资源管理。
3.3 定义奖励函数:RL 的灵魂
奖励函数决定了模型学习的方向。对于数学题,一个简单的规则是:答案正确得 1 分,错误得 0 分。
veRL 允许你轻松地注入自定义奖励函数。你需要创建一个 Python 文件,并实现一个特定签名的函数。
# custom_reward.py
import re
from typing import Optional
def extract_solution(solution_str: str) -> Optional[str]:
"""从模型回答中提取最终答案"""
# 假设答案被包含在 "####" 之后
solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str)
if solution:
return solution.group(0).split("#### ")[1].replace(",", "").replace("$", "")
return None
def compute_reward(data_source: str, solution_str: str, ground_truth: str, extra_info: dict) -> float:
"""
veRL 要求的奖励函数签名
"""
answer = extract_solution(solution_str)
if answer is None:
return 0.0
# 将答案与标准答案比较
return 1.0 if answer == ground_truth else 0.0
3.4 提交训练任务
一切就绪后,就可以提交任务了。veRL 使用 verl.trainer.main_ppo 作为入口,并通过命令行参数进行灵活配置。以下是关键参数的说明:
python3 -m verl.trainer.main_ppo \
algorithm.adv_estimator=grpo \ # 1. 指定使用 GRPO 算法
data.train_files=$HOME/data/gsm8k/train.parquet \ # 2. 数据路径
data.val_files=$HOME/data/gsm8k/test.parquet \
data.train_batch_size=1024 \ # 3. 每轮训练的样本总数
data.max_prompt_length=512 \ # 4. Prompt 最大长度
data.max_response_length=1024 \ # 5. 模型回答最大长度
actor_rollout_ref.model.path=Qwen/Qwen2.5-32B-Instruct \ # 6. 基础模型
actor_rollout_ref.actor.optim.lr=1e-6 \ # 7. 学习率
actor_rollout_ref.actor.use_kl_loss=True \ # 8. 启用 KL 散度惩罚
actor_rollout_ref.actor.kl_loss_coef=0.001 \ # 9. KL 惩罚系数,防止模型跑偏
actor_rollout_ref.rollout.name=vllm \ # 10. 指定 vLLM 作为推理引擎
actor_rollout_ref.rollout.n=5 \ # 11. GRPO 中每个 Prompt 生成的回答数
trainer.n_gpus_per_node=8 \ # 12. 每个节点的 GPU 数
trainer.nnodes=2 \ # 13. 节点总数
custom_reward_function.path=/path/to/custom_reward.py \ # 14. 自定义奖励函数路径
custom_reward_function.name=compute_reward # 15. 奖励函数名称
注:对于更大的模型,你可能需要开启
actor_rollout_ref.actor.fsdp_config.param_offload=True来将暂时不用的参数卸载到 CPU 内存,以节省显存。
训练启动后,你可以通过 Ray 的 Dashboard 或配置的 WandB 等工具监控训练指标,如 Loss 变化、KL 散度、平均奖励等。
4. 最小可行训练示例:让 0.5B 小模型学会解数学题
前面我们聊了很多原理,现在进入最实在的部分——亲手跑起来一个可工作的训练任务。这个示例的目标是:用最小的资源消耗,验证 veRL 的整个训练流程是否能跑通。
我们选择 Qwen2.5-0.5B-Instruct 模型和 GSM8K 数学数据集。这个组合的硬件门槛极低——单张 24GB 显存的 GPU 就够用。整个训练只需要 3 步。
4.1 准备数据集
veRL 要求训练数据为 .parquet 格式,框架提供了预处理脚本,帮你把 HuggingFace 上的原始 GSM8K 数据集转换成这种格式:
# 克隆 veRL 仓库(如果还没有)
git clone https://github.com/volcengine/verl.git
cd verl
# 预处理 GSM8K 数据集,生成 train.parquet 和 test.parquet
python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k
执行后,~/data/gsm8k/ 目录下会出现两个文件,分别用于训练和验证。
4.2 下载模型
我们使用 HuggingFace 上 Qwen2.5-0.5B-Instruct 作为基础模型。这个模型很小,下载快,训练也快:
python3 -c "import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-0.5B-Instruct')"
这条命令会自动从 HuggingFace 下载模型并缓存到本地。在国内网络环境下,可以设置环境变量 VERL_USE_MODELSCOPE=True 来使用 ModelScope 镜像源加速下载。
4.3 启动 GRPO 训练
这是最关键的一步。我们通过 verl.trainer.main_ppo 入口启动训练,并通过命令行参数完成所有配置。
PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \
data.train_files=$HOME/data/gsm8k/train.parquet \
data.val_files=$HOME/data/gsm8k/test.parquet \
data.train_batch_size=256 \
data.max_prompt_length=512 \
data.max_response_length=512 \
actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
actor_rollout_ref.actor.optim.lr=1e-6 \
actor_rollout_ref.actor.ppo_mini_batch_size=64 \
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
actor_rollout_ref.rollout.name=vllm \
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
algorithm.kl_ctrl.kl_coef=0.001 \
trainer.logger=console \
trainer.val_before_train=False \
trainer.n_gpus_per_node=1 \
trainer.nnodes=1 \
trainer.save_freq=10 \
trainer.test_freq=10 \
trainer.total_epochs=15
关键参数解读(按逻辑分组):
① 数据配置
data.train_files/data.val_files:指向刚刚生成的.parquet文件路径data.train_batch_size=256:每轮训练从数据集中采样 256 条 promptdata.max_prompt_length=512:输入 prompt 的最大 token 长度,超过会截断data.max_response_length=512:模型生成回答的最大长度
② 模型与训练配置
actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct:指定基础模型。框架会自动从 HuggingFace 加载这个模型,并同时实例化三个角色:Actor(策略模型,参与训练和生成)、Rollout(推理引擎,负责生成回答)、Reference(参考模型,用于计算 KL 散度约束)actor_rollout_ref.actor.optim.lr=1e-6:学习率,RL 训练通常比 SFT 用更小的学习率actor_rollout_ref.actor.ppo_mini_batch_size=64:一次梯度更新使用的样本数,受限于显存可以调小
③ 推理引擎配置
actor_rollout_ref.rollout.name=vllm:指定 vLLM 作为推理引擎。veRL 把生成回答(Rollout)这一步交给了 vLLM,利用其高吞吐能力加速actor_rollout_ref.rollout.gpu_memory_utilization=0.4:vLLM 使用 GPU 显存的比例,调低一点可以给训练留出空间actor_rollout_ref.rollout.tensor_model_parallel_size=1:张量并行度,单卡设为 1
④ 算法配置
algorithm.kl_ctrl.kl_coef=0.001:KL 散度惩罚系数。KL 散度衡量策略模型与参考模型输出分布的差异,这个惩罚防止模型为了“钻空子”获得高奖励而生成无意义的乱码
⑤ 训练运行配置
trainer.n_gpus_per_node=1和trainer.nnodes=1:单机单卡trainer.total_epochs=15:训练 15 轮trainer.logger=console:日志输出到控制台(也可以配 wandb 做可视化)
5.4 预期输出与验证
启动后,控制台会输出类似下面的训练日志:
step:0 - timing/gen:21.470 - actor/reward_kl_penalty:0.000 - response_length/mean:239.133 ...
step:1 - timing/gen:23.020 - actor/entropy_loss:0.440 - actor/pg_loss:0.000 ...
关键观测指标:
val/test_score/openai/gsm8k:验证集上的正确率(每trainer.test_freq步输出一次),值越高说明模型学得越好actor/entropy_loss:策略的熵,反映模型输出的多样性,太低可能意味着模型过于“自信”actor/pg_loss:策略梯度损失,应该在一个相对稳定的范围内波动
训练完成后,Checkpoint 会保存在 checkpoints/${trainer.project_name}/${trainer.experiment_name}/ 目录下。你可以用 verl.model_merger 模块将 FSDP 格式的 Checkpoint 合并成 HuggingFace 兼容的格式,方便后续加载使用。
5. 总结与展望
veRL 不仅仅是一个工具,它更是一套设计先进的 LLM 强化学习系统工程方案。它通过 GRPO 等高效算法降低了 RL 的门槛,通过 3D-HybridEngine 和模块化架构突破了性能瓶颈,并正朝着更加灵活、解耦的服务化方向演进。
正如其 GitHub Issue 和 v0.7 版本博客所展示的,veRL 的未来将包括更全面的异步训练支持、更丰富的模型引擎后端(如 JAX)以及与 Agent 框架的更深度整合。对于任何希望在大型语言模型上探索强化学习潜力的开发者或研究者来说,veRL 无疑是一个值得深入研究和投入的基石。