SGLang & veRL: 开创性的端到端多轮 RLHF 实现

作者:SGLang Team

https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/verl-multiturn-rollout-Release.md
>> 加入青稞AI技术交流群,与青年AI研究员/开发者交流最新AI技术

图片


我们很高兴地宣布发布第一个功能齐全、经过收敛验证、端到端开源多轮强化学习与人类反馈 (RLHF) 框架,该框架由 SGLang 提供支持并与 veRL 集成。

经过两个月的紧张开发和最后五天的冲刺,我们的团队终于交付了一套强大的解决方案,能够在 Agentic RL 中实现异步多轮对话和工具调用。此版本标志着面向大型语言模型的可扩展 RLHF 向前迈出了重要一步。

拉取请求:

https://github.com/volcengine/verl/pull/1037

训练性能(在 8×H100 GPU 上,使用 Qwen2.5-3B-Instruct、GRPO 策略、FSDP+TP 混合并行、启用工具调用、256 批次大小进行 GSM8K 任务的多轮训练):

https://wandb.ai/swordfaith/gsm8k_async_rl/runs/0pv3qwcy?nw=nwuserswordfaith

我们解决了什么问题?

多轮 RLHF 对于训练语言模型处理复杂交互式对话(例如编码任务)至关重要,因为生成的脚本的输出需要作为下一代脚本的反馈。然而,现有框架缺乏多轮部署支持,缺乏标准化的工具调用方法。

我们的目标是:

  • 支持异步多轮交互,效率高。

  • 以易于扩展的方式将工具调用集成到 RLHF 工作流中。

  • 证明我们在复杂、多回合适用任务上的训练表现。

  • 在大规模分布式环境中可靠运行。

我们的解决方案

从批处理级别到请求级别的推出:我们扩展了 veRL 的推出接口以支持异步、每个请求的多轮交互,其中每个对话可以独立地跨越多轮和工具调用,与批处理结构分离。

具有统一模式的通用工具调用:我们支持OpenAIFunctionToolSchema,可与 veRL 的模型上下文协议 / Agent2Agent 协议进行转换。这允许在训练和推理工作流程中实现工具的无缝集成。

参数注入机制:用户可以在采样过程中动态选择工具(例如,通过need_tools_kwargs)并注入调用参数(tool_kwargs),从而简化工具集成。

GRPO 策略梯度:我们采用 GRPO 策略(adv_estimator = grpo)在多轮序列上实现稳定的奖励传播。

如何使用

环境设置

创建一个新的 Docker 容器
docker run \
    -it \
    --shm-size 32g \
    --gpus all \
    -v /models/shared/.cache:/root/.cache \
    --ipc=host \
    --network=host \
    --privileged \
    --name sglang_{your-name} \
    lmsysorg/sglang:dev \
    /bin/zsh
退出容器后重新启动容器:
docker start -i sglang_{your-name}
更新 Python
apt update
apt install -y python3.10 python3.10-venv
python3 -m ensurepip --upgrade
使用 Python 虚拟环境
# Create a virtual environment
python3 -m venv ~/.python/veRL-multiturn-rollout

# Activate the virtual environment
source ~/.python/veRL-multiturn-rollout/bin/activate

# Install uv
python3 -m pip install uv
克隆 veRL 主存储库
cd ~
git clone https://github.com/volcengine/verl.git
cd verl
设置Python环境
# Install SGLang
python3 -m uv pip install -e ".[sglang]"

# Manually install flash-attn
python3 -m uv pip install wheel
python3 -m uv pip install packaging
python3 -m uv pip install flash-attn --no-build-isolation --no-deps

# Install veRL
python3 -m uv pip install .
python3 -m uv pip install -r ./requirements.txt
8-GPU SGLang 测试

跑 WANDB_API_KEY 前设置

如果您不确定在哪里找到您的 API 令牌,请参阅本指南[1]。

export WANDB_API_KEY={YOUR_WANDB_API_KEY}

# Define a timestamp function
function now() {
    date '+%Y-%m-%d-%H-%M'
}
下载数据集
python3 ./examples/data_preprocess/gsm8k_multiturn_w_tool.py
运行
# First make sure the now() function is available in current shell
# Create logs directory if it doesn't exist
mkdir -p logs

# Set GPUs and run with better log organization
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

nohup bash examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn.sh trainer.experiment_name=qwen2.5-3b_rm-gsm8k-sgl-multiturn-$(now) > logs/gsm8k-$(now).log 2>&1 &

配置

基本配置

要启用多轮推出,请确保在推出配置中配置以下字段:

actor_rollout_ref:
  rollout:
    name: "sglang_async"
    multi_turn:
      enable: True

此配置激活 AsyncSGLangRollout 引擎,以便在推出期间进行多轮交互。

自定义工具配置

工具是我们框架的关键组件,它支持环境交互,例如执行脚本、查询 API 或计算奖励。要集成自定义工具,您可以在单独的 YAML 文件中定义其行为,并在部署配置中引用该文件。要集成工具,请按以下步骤操作:

1.用Python定义工具逻辑:

每个工具必须为子类 BaseTool,并实现:

  • create(instance_id, ...):每次推出时初始化工具状态。

  • execute(instance_id, parameters, ...):执行工具的核心功能(例如,评估输出)。

  • calc_reward(instance_id):根据工具状态和交互计算奖励。

  • release(instance_id):清理所有分配的资源。

2. 使用 YAML 描述该工具:

您必须使用 OpenAI 的调用格式为每个工具提供一个架构function,包括名称、描述、参数和类型。然后在部署配置中引用该 YAML 文件。

actor_rollout_ref:
  rollout:
    multi_turn:
      tool_config_path: <path_to_tool_yaml_file>
      format: chatml
  • tool_config_path:定义所有工具模式和相关 Python 类的 YAML 文件的路径。

  • format:工具交互消息的格式(目前chatml仅支持)。

示例:GSM8K 工具配置

为了说明如何配置工具,我们提供了一个基于 Gsm8kTool 的示例,该工具旨在评估 GSM8K 数学问题的答案以计算奖励。工具配置定义如下。

配置细分

  • class_name:指定实现工具逻辑的 Python 类。该类必须可在代码库中访问。

  • config:用于附加工具特定配置(例如 API 密钥、模型参数)的可选字段。

  • tool_schema:使用与 OpenAIFunctionToolSchema 或 veRL 协议兼容的模式定义工具的接口。

  • 类型:“function”:表示该工具是基于功能的工具。

  • function.name:工具的标识符(calc_gsm8k_reward),在工具选择期间使用。

  • function.description:该工具用途的人类可读的描述。

  • function.parameters:描述该工具所需的输入参数。必填字段定义了必需的参数。

工具交互如何工作
在推出期间:

  • LLM 根据模式生成包含结构化函数调用(工具调用)的响应。

  • AsyncSGLangRollout使用FunctionCallParser从输出中检测并提取工具调用。

  • 推出引擎tool.execute()使用解析的参数进行调用。

  • 该工具可以:

  • 返回文本响应

  • 返还步数奖励

  • 更新内部工具状态

  • 引擎将工具响应附加到消息历史记录并继续对话。

  • 全面推出后,calc_reward()将调用该工具来计算最终奖励。

通过这种插件式架构,工具可以在任务和代理之间灵活地重复使用,而无需改变核心引擎或推出循环。

挑战与方法

填充策略不匹配:veRL 对提示采用左填充策略,对响应采用右填充策略,这与我们初始版本的代码存在填充不一致的问题。为了解决这个问题,我们的实现明确跟踪了各个段的 token 位置,并应用了定制的掩码来维护正确性。

多轮损失掩蔽:大多数现有的 RLHF 框架都假设单轮生成模式,缺乏对跨多轮对话进行细粒度、token 级损失掩蔽的支持。然而,在多轮对话设置中(尤其是在工具交互的情况下),并非所有生成的 token 都应计入学习信号。例如,某些工具生成的响应应该被排除在优化过程之外。我们通过设计自定义多轮损失掩蔽机制来解决这个问题,该机制允许对哪些 token 包含在策略梯度更新中进行细粒度控制,从而确保准确的奖励计算。

通用工具 API 设计:RLHF 训练场景中的环境可能非常复杂,智能体需要定制工具才能与外界交互。为了支持灵活且可重用的工具集成,我们设计了一个通用工具接口。此设计使用户能够将具有自定义功能的工具注册到部署流程中。通过以模式驱动的格式统一工具定义,我们使框架具有高度的可扩展性,可以轻松添加工具并在任务和模型之间重用它们,而无需进行大量修改。

工具调用中的 SPMD 冲突:在张量并行 (TP) 环境中,必须控制对外部工具(例如 API 调用、脚本评估器或奖励计算器)的调用,以避免并发问题。简单的实现可能会导致同一工具在不同的 TP 等级中被并行调用多次,从而造成不一致或死锁。为了避免这种情况,所有工具调用都限制在 TP 等级 0 上,并将结果广播到其他等级。这可以避免因冗余调用而导致的性能瓶颈。

多轮交互的异步部署:同步部署经常受到长尾问题的影响,即批次中最慢的样本会延迟整个流程。这个问题在涉及可变长度对话和工具调用的多轮任务中尤为突出。为了解决这个问题,我们在请求级别实现了异步部署,允许每个对话独立进行。

异步执行中的事件循环冲突:在测试过程中,我们遇到了一个问题:在启用 enable_memory_saver 的情况下,async_generate 挂起了。经过深入调查,我们发现根本原因是存在多个并发事件循环,这违反了 Python 的 asyncio 设计。SGLang 内部管理着自己的异步事件循环,以协调令牌流传输、多轮交互和高效的内存生成。我们错误地添加了第二个事件循环,导致程序永远卡住了。我们的修复方案是通过运行现有循环而不是在 async_generate 中调用 asyncio.run() 来确保所有异步执行都发生在 SGLang 自己的循环中。

我们的工作计划

  • 在训练过程中,我们观察到了 NaN 损失。我们分析认为,这可能是由于参与者模型和参考模型对数概率之间的巨大差异导致重要性采样率不稳定造成的。我们目前正在努力解决这个问题。

  • 集成多回合 RLHF 适用任务(例如 R1-Searcher、沙盒融合)

  • 支持除 GRPO 之外的更多策略梯度估计器

  • 改进阶梯式奖励整合

  • 模拟多轮 RLHF 的人机交互

  • 引入微批处理和代理推出循环

我们欢迎社会各界与我们合作,共同推动RLHF研究和应用的前沿。

引用链接

[1] 本指南: https://community.wandb.ai/t/where-can-i-find-the-api-token-for-my-project/7914

往期推荐

SGLang Team:在 96 个 H100 GPU 上部署具有 PD 分解和大规模专家并行性的 DeepSeek

强化学习算法梳理:从 PPO 到 GRPO 再到 DAPO

veRL:All in RL元年的必修课

思考 RLHF 的本质以及主流算法思路

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值