This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,360 @@
|
||||
# GYM环境训练
|
||||
|
||||
GYM 风格的环境训练把"模型 → 环境 → 奖励"这条链路封装成一个抽象接口，让 LLM 像 Agent 一样与环境进行多轮交互，每一步的奖励直接由环境给出，无需再单独写 reward 函数从轨迹里反推。本文先介绍接口，再用一个完整的自定义示例（FrozenLake）说明如何接入训练。
|
||||
|
||||
## Gym 接口
|
||||
|
||||
GYM 源自 [Gymnasium库](https://github.com/Farama-Foundation/Gymnasium)。在 ms-swift 中我们定义了如下接口：
|
||||
|
||||
```python
|
||||
class Env(ABC):
|
||||
|
||||
def __init__(self, env_config):
|
||||
"""env_config 来自数据集每行的 env_config 列，可承载初始化参数"""
|
||||
self.env_config = env_config
|
||||
|
||||
@abstractmethod
|
||||
async def reset(self, config: RolloutInferRequest) -> Tuple[str, Dict[str, Any], str]:
|
||||
"""
|
||||
Returns:
|
||||
- observation: 作为首轮 user 消息发送给模型
|
||||
- info: 调试/日志信息，记录到 completions.jsonl
|
||||
- system_message: 本条轨迹的 system prompt
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def step(self, action: Messages) -> Tuple[str, float, bool, Dict[str, Any]]:
|
||||
"""
|
||||
Args:
|
||||
action: 截止当前的完整对话消息，最后一条即模型最新回复
|
||||
Returns:
|
||||
- next_observation: 下一轮 user 消息
|
||||
- reward: 当前 step 奖励
|
||||
- done: 轨迹是否结束
|
||||
- info: 调试/日志信息
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def close(self):
|
||||
"""释放资源"""
|
||||
pass
|
||||
```
|
||||
|
||||
`reset` 接收到的 `RolloutInferRequest` 包含数据集行的 `messages`、`data_dict`（额外列，包括 `env_config`）等。完整示例参见 [入参示例](./multi_turn.md#多轮规划器-multiturnscheduler)。
|
||||
|
||||
> 如果需要在每轮 rollout 之间额外控制对话历史（例如动态压缩、注入额外提示），推荐直接继承 `MultiTurnScheduler` 并实现 `on_trajectory_start` / `on_turn_end` hook，或重写 `step` / `run` 方法，详见[多轮训练文档](./multi_turn.md#自定义多轮交互逻辑)。
|
||||
|
||||
## 启动训练
|
||||
|
||||
使用内置的 [gym_scheduler](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py) 把 env 串到多轮 rollout 中。
|
||||
|
||||
`GYMScheduler` 基于通用 hook 协议实现：
|
||||
- 继承 `MultiTurnScheduler`，无需自定义 `run` 方法
|
||||
- 实现 `on_trajectory_start`（调用 `env.reset`）和 `on_turn_end`（调用 `env.step`）
|
||||
- 同时适用于 server mode（`run()`）和 colocate mode（`run_multi_turn()`）
|
||||
|
||||
用户自定义的 env 通过 `--external_plugins your_plugin.py` 加载，plugin 里执行 `envs['my_env'] = MyEnv` 完成注册（下文 FrozenLake 示例完整演示）。
|
||||
|
||||
**Colocate 模式**:
|
||||
|
||||
```bash
|
||||
megatron rlhf \
|
||||
--rlhf_type grpo \
|
||||
--vllm_mode colocate \
|
||||
--external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py \
|
||||
--multi_turn_scheduler gym_scheduler \
|
||||
--gym_env frozen_lake \
|
||||
--use_gym_env true \
|
||||
--max_turns 10 \
|
||||
...
|
||||
|
||||
# swift rlhf 同理
|
||||
```
|
||||
|
||||
|
||||
**Server 模式**
|
||||
|
||||
```bash
|
||||
swift rollout \
|
||||
--model xxx \
|
||||
--use_gym_env true \
|
||||
--external_plugins examples/megatron/grpo/multi_turn/frozen_lake_plugin.py \
|
||||
--multi_turn_scheduler gym_scheduler \
|
||||
--gym_env frozen_lake \
|
||||
--max_turns 10
|
||||
|
||||
# trainer 侧需要加 --vllm_server_pass_dataset true，把 env_config 等额外列透传给 rollout 端
|
||||
megatron rlhf --vllm_mode server --vllm_server_pass_dataset true ...
|
||||
# or swift rlhf --vllm_mode server --vllm_server_pass_dataset true ...
|
||||
```
|
||||
|
||||
|
||||
环境选择有两种方式：
|
||||
- 通过 `--gym_env env_name` 全局指定（同一脚本里所有 prompt 共用一个 env）；
|
||||
- 在每行数据的 `env_config.name` 中指定（适用于多环境混合场景，每条数据可指向不同 env，会覆盖 `--gym_env`）。
|
||||
|
||||
## 示例：从零写一个 FrozenLake 环境
|
||||
|
||||
<img src="https://gymnasium.farama.org/_images/frozen_lake.gif" width="220" alt="FrozenLake 环境示意图（来源：Gymnasium 官方文档）" />
|
||||
|
||||
[FrozenLake](https://gymnasium.farama.org/environments/toy_text/frozen_lake/) 是 OpenAI Gym 中的经典任务：智能体从起点出发，需要穿过一片冰湖到达终点，途中要避开冰窟。原始环境如上图所示。下面以纯文本版本（把上图网格直接渲染成 ASCII 字符）为例。
|
||||
|
||||
以下完整代码参考完整代码：[frozen_lake_plugin](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/multi_turn/frozen_lake_plugin.py)。
|
||||
|
||||
**1. 定义 Env**
|
||||
|
||||
每条数据派生一张随机 4x4 地图（随机洞 + 随机 S/G 位置，BFS 校验保证可解）。单元含义：`S` 起点 / `G` 终点 / `H` 冰窟（踩到=失败）/ `F` 安全冰面 / `P` 玩家当前位置。
|
||||
|
||||
```python
|
||||
class FrozenLakeEnv(Env):
|
||||
def __init__(self, env_config):
|
||||
super().__init__(env_config)
|
||||
self.size = int(env_config.get('size', 4))
|
||||
self.p = float(env_config.get('p', 0.8))
|
||||
seed = env_config.get('seed')
|
||||
self.seed = int(seed) if seed is not None else None
|
||||
|
||||
async def reset(self, config: RolloutInferRequest):
|
||||
self.grid = generate_random_map(size=self.size, p=self.p, seed=self.seed)
|
||||
...
|
||||
return observation, {'seed': self.seed}, SYSTEM_PROMPT
|
||||
|
||||
async def step(self, action: Messages):
|
||||
move = _parse_action(action[-1]['content']) # <action>up|down|left|right</action>
|
||||
# 推进一格、判断 G / H；外层 max_turns 由 scheduler 兜底
|
||||
if cell == 'G': return obs, 1.0, True, {'status': 'goal'}
|
||||
if cell == 'H': return obs, 0.0, True, {'status': 'hole'}
|
||||
...
|
||||
```
|
||||
|
||||
**2. GYMScheduler 的 hook 实现**
|
||||
|
||||
框架内置的 `GYMScheduler` 基于多轮 hook 完成了控制逻辑：
|
||||
|
||||
```python
|
||||
class GYMScheduler(MultiTurnScheduler):
|
||||
def on_trajectory_start(self, requests):
|
||||
# 为每个请求创建 env，调用 env.reset，注入初始 observation
|
||||
for req in requests:
|
||||
env = self._create_env(req.data_dict.get('env_config', {}))
|
||||
observation, info, system_message = env.reset(req)
|
||||
req.messages = [system_msg, user_msg(observation)]
|
||||
self._envs[req.uuid] = env
|
||||
|
||||
def on_turn_end(self, req, response_choice, current_turn):
|
||||
# 调用 env.step，累积 reward，返回 done + rollout_infos
|
||||
next_obs, reward, done, info = env.step(deepcopy(req.messages))
|
||||
self._total_rewards[req.uuid] += reward
|
||||
return {
|
||||
'done': done,
|
||||
'rollout_infos': {
|
||||
'total_reward': self._total_rewards[req.uuid],
|
||||
'step_rewards': [...],
|
||||
}
|
||||
}
|
||||
|
||||
def step(self, req, response_choice, current_turn):
|
||||
# 注入下一帧 observation 到 user message
|
||||
if self._pending_obs.get(req.uuid):
|
||||
req.messages.append({'role': 'user', 'content': next_obs})
|
||||
return {'infer_request': req}
|
||||
```
|
||||
|
||||
用户只需实现 Env 接口，无需关心多轮控制细节。
|
||||
|
||||
**3. 注册**
|
||||
|
||||
将 env 类挂到 swift 的 `envs` 注册表里。`--external_plugins` 在训练启动时会 import 该文件，注册随之生效：
|
||||
|
||||
```python
|
||||
# examples/megatron/grpo/multi_turn/frozen_lake_plugin.py
|
||||
from swift.rollout.gym_env import Env, envs
|
||||
|
||||
class FrozenLakeEnv(Env):
|
||||
...
|
||||
|
||||
envs['frozen_lake'] = FrozenLakeEnv
|
||||
```
|
||||
|
||||
**4. 准备数据集**
|
||||
|
||||
数据集在这里仅作占位符处理，数据构造由环境生成，和 `env_config.seed`来控制地图生成的随机性：
|
||||
|
||||
```json
|
||||
{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":0}}
|
||||
{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":1}}
|
||||
...
|
||||
{"messages":[{"role":"user","content":"<placeholder>"}],"env_config":{"seed":127}}
|
||||
```
|
||||
|
||||
**5. （可选）叠加自定义 reward**
|
||||
|
||||
设置 `--use_gym_env true` 后，env 给出的 `total_reward` 会自动作为一路奖励参与训练，无需再写 reward 函数。如果想在此之外再叠加自定义信号（如格式/长度等），通过 `--reward_funcs` 传入即可，gym 奖励会作为额外一列与 reward_funcs 拼在一起，由 `--reward_weights` 统一加权。例如同时启用一个格式校验 reward：
|
||||
|
||||
```bash
|
||||
megatron rlhf ... --use_gym_env true --reward_funcs format --reward_weights 0.2 1.0
|
||||
# reward_weights 末位对应 gym 的 total_reward
|
||||
```
|
||||
|
||||
**6. 训练**
|
||||
|
||||
运行脚本参考：[`examples/megatron/grpo/multi_turn/frozen_lake.sh`](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/multi_turn/frozen_lake.sh)
|
||||
|
||||
|
||||
## OpenEnv 环境训练
|
||||
|
||||
[OpenEnv](https://github.com/huggingface/openenv) 是 HuggingFace 开源的 Agentic RL 环境框架，通过 WebSocket 与环境服务器交互。与上文 FrozenLake 的本地 `Env` 接口不同，OpenEnv 将环境逻辑放在独立的服务进程中，swift 通过 `OpenEnvScheduler` + `OpenEnvWrapper` 与之通信。
|
||||
|
||||
### 架构对比
|
||||
|
||||
| 特性 | 内置 Gym (`GYMScheduler`) | OpenEnv (`OpenEnvScheduler`) |
|
||||
|------|--------------------------|------------------------------|
|
||||
| 环境运行位置 | 训练进程内（Python 对象） | 独立服务器（WebSocket 通信） |
|
||||
| 环境接口 | 继承 `Env`，实现 `reset/step/close` | 服务器提供 HTTP/WebSocket API |
|
||||
| 注册方式 | `--external_plugins` + `envs` 注册表 | `--external_plugins` + `multi_turns` 注册表 |
|
||||
| 适用场景 | 轻量本地环境（FrozenLake 等） | 复杂服务端环境（TextArena、CARLA 等） |
|
||||
| 并发控制 | 无需 | 内置 Semaphore 限制并发连接 |
|
||||
|
||||
### OpenEnvScheduler
|
||||
|
||||
`OpenEnvScheduler` 继承 `GYMScheduler`，将本地 `Env` 替换为 `OpenEnvWrapper`（WebSocket 客户端）。核心设计：
|
||||
|
||||
- **`_create_env`**：创建 `OpenEnvWrapper`，连接 OpenEnv 服务器
|
||||
- **`on_trajectory_start`**：为每个请求创建 wrapper，调用 `reset()`，用 Semaphore 限制并发（默认 4）
|
||||
- **`on_turn_end`**：解析模型输出，调用 `wrapper.step()`，累积奖励
|
||||
- **`parse_action`**（可覆盖）：将模型文本解析为 action dict，默认 `json.loads`
|
||||
- **`format_observation`**（可覆盖）：将服务器返回的 observation 格式化为字符串，默认 `json.dumps`
|
||||
|
||||
用户通过继承 `OpenEnvScheduler` 并覆盖 `parse_action`、`format_observation`、`on_trajectory_start`、`on_turn_end` 来适配具体环境。
|
||||
|
||||
### 示例：Sudoku 环境
|
||||
|
||||
以 TextArena Sudoku 为例，模型需要通过 `[row col number]` 格式下棋，在 9x9 数独棋盘上填入数字。完整代码参考：[sudoku_scheduler.py](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/openenv/sudoku_scheduler.py)。
|
||||
|
||||
**1. 启动 OpenEnv 服务器**
|
||||
|
||||
安装 OpenEnv 和 Sudoku 环境包（textarena 和 nltk 会作为依赖自动安装）：
|
||||
|
||||
```bash
|
||||
pip install openenv
|
||||
pip install git+https://huggingface.co/spaces/openenv/sudoku
|
||||
```
|
||||
|
||||
使用提供的启动脚本启动本地服务器（默认端口 8000）。`MAX_CONCURRENT_ENVS` 需 ≥ 训练时的 `num_generations`：
|
||||
|
||||
```bash
|
||||
TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 python examples/train/grpo/plugin/openenv/start_sudoku_server.py
|
||||
```
|
||||
|
||||
数据集中将 `base_url` 指向本地服务器地址：
|
||||
|
||||
```json
|
||||
{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}
|
||||
```
|
||||
|
||||
**2. 自定义 Scheduler**
|
||||
|
||||
继承 `OpenEnvScheduler`，实现 Sudoku 专用的动作解析、观察格式化和多组件奖励：
|
||||
|
||||
```python
|
||||
from swift.rollout.multi_turn import OpenEnvScheduler
|
||||
|
||||
class SudokuScheduler(OpenEnvScheduler):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._last_content_len = {} # 内容差分跟踪
|
||||
|
||||
async def on_trajectory_start(self, requests):
|
||||
# 创建环境、解析棋盘、生成 hints
|
||||
# hints 包括「保证正确的走法」和候选数字列表
|
||||
...
|
||||
|
||||
async def on_turn_end(self, infer_request, response_choice, current_turn):
|
||||
# 解析 [row col number]，step 环境
|
||||
# 计算 5 路奖励：空格选择 / 合法移动 / 重复惩罚 / 进度 / 正确
|
||||
# 返回更新后的棋盘 + hints 作为下一轮观察
|
||||
...
|
||||
|
||||
def parse_action(self, text):
|
||||
import re
|
||||
match = re.search(r'\[\s*(\d+)\s+(\d+)\s+(\d+)\s*\]', text)
|
||||
if match:
|
||||
row, col, num = match.groups()
|
||||
return {"message": f"[{row} {col} {num}]"}
|
||||
return {"message": "[1 1 1]"}
|
||||
```
|
||||
|
||||
**多组件奖励系统**（参考 [TRL Sudoku 示例](https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_sudoku_grpo.ipynb)）：
|
||||
|
||||
| 奖励组件 | 计算方式 | 作用 |
|
||||
|---------|---------|------|
|
||||
| `empty_cell_reward` | 目标是空格 +1 / 覆盖已有 -1 | 引导模型选择合法位置 |
|
||||
| `valid_move_reward` | 合法新走法 +1 / 警告 -0.5 / 无效 0 | 鼓励合法操作 |
|
||||
| `repetition_reward` | 重复走法指数惩罚（-2^n，上限 -10） | 避免重复 |
|
||||
| `progress_reward` | (已填充 - 初始) / (81 - 初始) | 衡量解题进度 |
|
||||
| `correct_reward` | 环境返回的二值奖励 | 完全解出 |
|
||||
|
||||
组合奖励 = 各组件均值之和，提供比单一二值奖励更密集的学习信号。
|
||||
|
||||
**3. Hints 系统**
|
||||
|
||||
每轮交互中，scheduler 解析当前棋盘状态，为模型提供提示：
|
||||
|
||||
- **GUARANTEED MOVES**：只有一个候选数字的空格（可直接填入）
|
||||
- **Other options**：2-3 个候选数字的空格
|
||||
- **MOVES ALREADY TRIED**：已尝试过的走法（避免重复）
|
||||
|
||||
**4. 准备数据集**
|
||||
|
||||
数据集仅作占位符，实际棋盘由环境服务器生成。`base_url` 指向 OpenEnv 托管地址：
|
||||
|
||||
```json
|
||||
{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}
|
||||
```
|
||||
|
||||
**5. 注册 Scheduler**
|
||||
|
||||
`sudoku_scheduler.py` 末尾已包含注册代码，通过 `--external_plugins` 加载即可：
|
||||
|
||||
```python
|
||||
# sudoku_scheduler.py 末尾
|
||||
from swift.rollout.multi_turn import multi_turns
|
||||
multi_turns['sudoku_scheduler'] = SudokuScheduler
|
||||
```
|
||||
|
||||
**6. 启动训练**
|
||||
|
||||
```bash
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen3.5-4B \
|
||||
--dataset examples/train/grpo/plugin/openenv/sudoku.jsonl \
|
||||
--external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
--enable_thinking false \
|
||||
--max_completion_length 256 \
|
||||
--use_gym_env true \
|
||||
--multi_turn_scheduler sudoku_scheduler \
|
||||
--max_turns 20 \
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate \
|
||||
...
|
||||
```
|
||||
|
||||
运行脚本参考：[`examples/train/grpo/plugin/openenv/run_grpo_sudoku.sh`](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/openenv/run_grpo_sudoku.sh)
|
||||
|
||||
### 注意事项
|
||||
|
||||
- **vLLM 模式**：以上示例使用 `--vllm_mode colocate`，vLLM 与训练共享 GPU。若使用 `--vllm_mode server`，需额外启动 `swift rollout` 作为 vLLM 服务器，且 `--multi_turn_scheduler` 和 `--max_turns` 参数应传给 `swift rlhf` 而非 `swift rollout`。
|
||||
- **并发会话数**：`start_sudoku_server.py` 的 `MAX_CONCURRENT_ENVS` 需 ≥ 训练时的 `num_generations`。默认的 `python -m textarena_env.server.app` 只支持 1 个并发会话。
|
||||
- **enable_thinking**：Sudoku 等环境不需要 CoT 推理，建议设置 `--enable_thinking false` 以减少 token 消耗。
|
||||
- **同步 I/O**：`OpenEnvWrapper` 的 `reset()`/`step()` 是同步 WebSocket 调用。`OpenEnvScheduler` 的子类应使用 `asyncio.to_thread()` 包装这些调用以避免阻塞事件循环。
|
||||
|
||||
|
||||
参考资料:
|
||||
|
||||
- https://gymnasium.farama.org/environments/toy_text/frozen_lake/
|
||||
- https://github.com/alibaba/ROLL/tree/main/roll/pipeline/agentic/env/frozen_lake
|
||||
- https://github.com/huggingface/openenv
|
||||
- https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_sudoku_grpo.ipynb
|
||||
@@ -0,0 +1,11 @@
|
||||
Developer Guide
|
||||
===============
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
loss_types.md
|
||||
multi_turn.md
|
||||
multi_task.md
|
||||
reward_function.md
|
||||
reward_model.md
|
||||
gym_env.md
|
||||
@@ -0,0 +1,131 @@
|
||||
# Loss Types
|
||||
|
||||
GRPO训练支持多种不同的loss类型，主要区别在于归一化的维度和梯度处理方式上有所不同。
|
||||
|
||||
## 损失函数
|
||||
|
||||
token 级别上，GRPO 训练使用以下损失函数
|
||||
|
||||
$$\mathcal{L}_{i,t} = -\min\left(\rho_{i,t} A_{i,t}, \text{clip}(\rho_{i,t}, 1-\epsilon, 1+\epsilon) A_{i,t}\right)$$
|
||||
|
||||
当设置`loss_type cispo`时，使用 cispo 损失
|
||||
|
||||
$$\mathcal{L}_{i,t}^{\text{CISPO}} = -\text{detach}\left(\min(\rho_{i,t}, \epsilon_{\text{high}})\right) \cdot A_{i,t} \cdot \log \pi_\theta(y_{i,t}|y_{i,<t})$$
|
||||
|
||||
当设置`loss_type sapo`时，使用软门控替代硬裁剪，详见 [SAPO](../AdvancedResearch/SAPO.md)
|
||||
|
||||
$$\mathcal{L}_{i,t}^{\text{SAPO}} = -g_{i,t} \cdot A_{i,t}$$
|
||||
|
||||
其中 $g_{i,t} = \sigma(\tau \cdot (\rho_{i,t} - 1))$ 是温度控制的软门控函数。
|
||||
|
||||
其中：
|
||||
- $\rho_{i,t} = \frac{\pi_\theta(y_{i,t}|y_{i,<t})}{\pi_{\theta_{\text{old}}}(y_{i,t}|y_{i,<t})}$ 是重要性采样权重
|
||||
- $A_{i,t}$ 是优势函数
|
||||
- $\epsilon$ 和 $\epsilon_{\text{high}}$ 是clipping参数
|
||||
- $\text{detach}(\cdot)$ 表示该项不参与梯度计算
|
||||
- $\sigma(\cdot)$ 是 sigmoid 函数，$\tau$ 是温度参数
|
||||
|
||||
## GRPO
|
||||
|
||||
`--loss_type grpo`
|
||||
|
||||
GRPO是标准的损失函数实现，对每个样本的token-level损失取平均，然后对所有样本取平均。
|
||||
|
||||
**公式：**
|
||||
|
||||
$$\mathcal{L}_{\text{GRPO}} = \frac{1}{N} \sum_{i=1}^{N} \frac{1}{T_i} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}$$
|
||||
|
||||
其中：
|
||||
- $N$ 是批次中的样本数量
|
||||
- $T_i$ 是第$i$个样本的completion token数量
|
||||
|
||||
**归一化维度：** 样本维度（先对每个样本的所有token取平均，再对所有样本取平均）
|
||||
|
||||
## BNPO (Batch Normalized Policy Optimization)
|
||||
|
||||
`--loss_type bnpo`
|
||||
|
||||
BNPO将所有样本的所有token的损失直接求和，然后除以所有completion token的总数量。
|
||||
|
||||
**公式：**
|
||||
|
||||
$$\mathcal{L}_{\text{BNPO}} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}}{\sum_{i=1}^{N} T_i}$$
|
||||
|
||||
其中：
|
||||
- $N$ 是批次中的样本数量
|
||||
- $T_i$ 是第$i$个样本的completion token数量
|
||||
|
||||
**归一化维度：** Token维度（对所有completion token取平均）
|
||||
|
||||
## DR-GRPO
|
||||
|
||||
`--loss_type dr_grpo`
|
||||
|
||||
DR-GRPO将所有样本的所有token的损失求和，然后除以批次大小乘以最大completion长度。
|
||||
|
||||
**公式：**
|
||||
|
||||
$$\mathcal{L}_{\text{DR-GRPO}} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}}{N \times L_{\text{max}}}$$
|
||||
|
||||
其中：
|
||||
- $N$ 是批次中的样本数量
|
||||
- $T_i$ 是第$i$个样本的completion token数量
|
||||
- $L_{\text{max}}$ 是最大completion长度
|
||||
|
||||
**归一化维度：** 固定维度（批次大小 × 最大completion长度）
|
||||
|
||||
## CISPO
|
||||
|
||||
`--loss_type cispo`
|
||||
|
||||
CISPO损失按所有进程的completion token总数进行归一化。
|
||||
|
||||
**公式：**
|
||||
|
||||
$$\mathcal{L}_{\text{CISPO}} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}^{\text{CISPO}}}{\sum_{\text{all processes}} \sum_{i=1}^{N_p} T_{p,i}}$$
|
||||
|
||||
其中：
|
||||
- $N$ 是当前进程批次中的样本数量
|
||||
- $T_i$ 是第$i$个样本的completion token数量
|
||||
- $N_p$ 是第$p$个进程的样本数量
|
||||
|
||||
**归一化维度：** 全局token维度（跨所有进程的completion token总数）
|
||||
|
||||
## DAPO
|
||||
|
||||
`--loss_type dapo`
|
||||
|
||||
DAPO与BNPO类似，使用token-level归一化，但基于全局数据（多进程）进行归一化。
|
||||
|
||||
**公式：**
|
||||
|
||||
$$\mathcal{L}_{\text{DAPO}} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}}{\sum_{\text{all processes}} \sum_{i=1}^{N_p} T_{p,i}}$$
|
||||
|
||||
其中：
|
||||
- $N$ 是当前进程批次中的样本数量
|
||||
- $T_i$ 是第$i$个样本的completion token数量
|
||||
- $N_p$ 是第$p$个进程的样本数量
|
||||
|
||||
**归一化维度：** 全局token维度（跨所有进程的completion token总数）
|
||||
|
||||
## FIPO
|
||||
|
||||
`--loss_type fipo`
|
||||
|
||||
FIPO 在 DAPO/GRPO 的 clipped policy loss 上引入 Future-KL influence weight。每个 token 的序列级 advantage 会乘以从当前位置到后续 token 的折扣累积 KL 位移得到的权重：
|
||||
|
||||
$$f_{i,t} = \text{clip}\left(\exp\left(\sum_{k=t}^{T_i} \gamma^{k-t} M_{i,k} \Delta \log p_{i,k}\right), 1-\epsilon_f, 1+\epsilon_f\right)$$
|
||||
|
||||
$$\mathcal{L}_{i,t}^{\text{FIPO}} = f_{i,t} \cdot \mathcal{L}_{i,t}$$
|
||||
|
||||
FIPO 的 influence weight 默认不参与梯度计算，并使用与 DAPO 相同的全局 token 归一化。
|
||||
|
||||
**归一化维度:** 全局 token 维度（所有进程的 completion token 总数）
|
||||
|
||||
## SAPO
|
||||
|
||||
`--loss_type sapo`
|
||||
|
||||
SAPO使用温度控制的软门控替代硬裁剪，实现平滑的梯度衰减。归一化方式与GRPO相同。
|
||||
|
||||
详细说明请参考 [SAPO](../AdvancedResearch/SAPO.md)
|
||||
@@ -0,0 +1,51 @@
|
||||
# 多任务训练
|
||||
我们可以在数据集中添加一个用于标识任务类型的列，并在奖励函数/奖励模型插件中根据任务类型进行判断，从而实现多任务训练。假设数据集中包含数学和编程任务，比如：
|
||||
|
||||
```
|
||||
{"query": "Solve the equation x + 2 = 5", "solution": "3", "task": "math"},
|
||||
{"query": "Write a function to calculate the Fibonacci sequence", "solution": "xxx", "task": "code"},
|
||||
{"query": "What is the integral of x^2?", "solution": "xxx", "task": "math"},
|
||||
{"query": "Implement a sorting algorithm in Python", "solution": "xxx", "task": "code"},
|
||||
```
|
||||
|
||||
我们可以设置不同的奖励函数来分别处理数学数据和代码数据，注意数据集中的列会传入奖励函数，所以我们可以通过 `task` 列
|
||||
|
||||
下面是针对不同任务的奖励函数的示例：
|
||||
|
||||
```python
|
||||
from swift.rewards import ORM, orms
|
||||
import random
|
||||
|
||||
# Math-specific reward function
|
||||
class MathRandomReward(ORM):
|
||||
def __call__(self, completions, task, **kwargs):
|
||||
rewards = []
|
||||
for completion, t in zip(completions, task):
|
||||
if t == "math":
|
||||
import random
|
||||
# imple math accuracy logic
|
||||
reward = random.random()
|
||||
rewards.append(reward)
|
||||
else:
|
||||
# Return None for non-math tasks
|
||||
rewards.append(None)
|
||||
return rewards
|
||||
|
||||
# Coding-specific reward function
|
||||
class CodeRandomReward(ORM):
|
||||
def __call__(self, completions, task, **kwargs):
|
||||
rewards = []
|
||||
for prompt, completion, t in zip(prompts, completions, task):
|
||||
if t == "code":
|
||||
# imple coding accuracy logic
|
||||
reward = random.random()
|
||||
rewards.append(reward)
|
||||
else:
|
||||
# Return None for non-coding tasks
|
||||
rewards.append(None)
|
||||
return rewards
|
||||
|
||||
orms['math_reward'] = MathRandomReward
|
||||
orms['code_reward'] = CodeRandomReward
|
||||
```
|
||||
对于非当前任务的数据， 通过返回 None 来处理，从而使得奖励相关仅计算任务内的数据。
|
||||
@@ -0,0 +1,312 @@
|
||||
# 多轮训练
|
||||
|
||||
在强化学习训练场景中，模型采样可能需要与环境进行多轮交互（如工具调用）。这种交互式训练要求模型能够根据环境反馈信息进行连续推理。本文档将详细介绍如何在 GRPO 训练中自定义多轮训练流程。
|
||||
|
||||
> GKD 同样支持多轮训练，与 GRPO 共享同一套 `MultiTurnScheduler` 基础设施。
|
||||
|
||||
以下是多轮训练示例图，模型可能涉及多轮 rollout，包括环境交互、工具调用等步骤：
|
||||
|
||||
|
||||
|
||||
## 多轮规划器 MultiTurnScheduler
|
||||
|
||||
`MultiTurnScheduler` 是一个抽象基类，提供了默认的多轮对话管理逻辑，其工作流程如下图所示：
|
||||
|
||||
<img src="https://raw.githubusercontent.com/modelscope/ms-swift/main/docs/resources/multiturn_pipeline.png " width="300" />
|
||||
|
||||
多轮规划器主要承担两大核心功能：
|
||||
- **终止条件判断**：通过 `check_finished` 方法判断当前轮次推理是否应该结束
|
||||
- **推理请求构造**：通过 `step` 方法构建下一轮推理的请求对象
|
||||
|
||||
抽象基类 `MultiTurnScheduler` 的核心方法如下：
|
||||
|
||||
```python
|
||||
class MultiTurnScheduler(ABC):
|
||||
|
||||
def __init__(self, max_turns: Optional[int] = None, *args, **kwargs):
|
||||
self.max_turns = max_turns
|
||||
|
||||
def on_trajectory_start(self, requests: List['RolloutInferRequest']) -> None:
|
||||
"""在首轮推理前调用，用于初始化轨迹级别状态。
|
||||
|
||||
可在此方法中直接修改 requests（如注入环境初始 observation）。
|
||||
默认实现为空（no-op）。
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_turn_end(self, infer_request: 'RolloutInferRequest',
|
||||
response_choice: 'ChatCompletionResponseChoice',
|
||||
current_turn: int) -> Dict[str, Any]:
|
||||
"""在 assistant 消息追加后、check_finished 前调用。
|
||||
|
||||
用于推进环境状态（如 env.step）并返回每轮元数据。
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: 可选包含以下键：
|
||||
- 'done' (bool): 若存在，将覆盖 check_finished 的结果
|
||||
- 'rollout_infos' (dict): 合并到轨迹累积的额外信息中
|
||||
默认返回空字典（no-op）。
|
||||
"""
|
||||
return {}
|
||||
|
||||
def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice',
|
||||
current_turn: int) -> Dict:
|
||||
"""
|
||||
处理对话轮次之间的转换。
|
||||
|
||||
Args:
|
||||
infer_request: 当前推理请求
|
||||
response_choice: 当前轮次的响应
|
||||
current_turn: 当前轮次数
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: 包含推理结果的字典，结构如下：
|
||||
- infer_request (必需): 下一轮的推理请求对象
|
||||
- response_token_ids (可选): 每个 rollout 轮次的响应 token IDs
|
||||
- response_loss_mask (可选): 每个 rollout 轮次响应的损失掩码
|
||||
- rollout_logprobs (可选): 每个 rollout 轮次的响应对应的 logps
|
||||
- rollout_infos (可选): 额外信息数据
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def check_finished(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice',
|
||||
current_turn: int) -> bool:
|
||||
"""
|
||||
检查多轮 rollout 是否应该结束的默认终止逻辑。
|
||||
|
||||
默认终止条件：
|
||||
1. 当响应达到长度限制时 (finish_reason == 'length')
|
||||
2. 当对话达到最大轮数时 (如果设置了 max_turns)
|
||||
|
||||
Args:
|
||||
infer_request: 推理请求对象
|
||||
response_choice: 包含生成结果的响应选择，包括 finish_reason
|
||||
current_turn: 当前对话轮数
|
||||
|
||||
Returns:
|
||||
bool: True 表示终止对话，False 表示继续
|
||||
"""
|
||||
if response_choice.finish_reason == 'length':
|
||||
return True
|
||||
if self.max_turns and current_turn >= self.max_turns:
|
||||
return True
|
||||
return False
|
||||
```
|
||||
|
||||
`step` 和 `check_finished` 方法接收的参数说明：
|
||||
- **infer_request**: 当前的推理请求
|
||||
- **response_choice**: 当前轮次的推理结果
|
||||
- **current_turn**: 当前推理轮次（从 1 开始）
|
||||
|
||||
<details><summary>入参示例（点击展开）</summary>
|
||||
|
||||
```python
|
||||
infer_request
|
||||
"""
|
||||
RolloutInferRequest(
|
||||
messages=[
|
||||
{'role': 'system', 'content': 'A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>\n'}, {'role': 'user', 'content': 'What is the value of $\\sqrt{36 \\times \\sqrt{16}}$?'},
|
||||
{'role': 'assistant', 'content': 'To find the value of \\(\\sqrt{36 \\times \\sqrt{16}}\\), we will break down the problem step-by-step.\n\nFirst, we need to evaluate the inner square root:\n\\[\n\\sqrt{16}\n\\]\nWe know that:\n\\[\n4^2 = 16 \\implies \\sqrt{16} = 4\n\\]\n\nNext, we substitute this result back into the original expression:\n\\[\n\\sqrt{36 \\times \\sqrt{16}} = \\sqrt{36 \\times 4}\n\\]\n\nNow, we need to evaluate the product inside the square root:\n\\[\n36 \\times 4 = 144\n\\]\n\nSo, the expression simplifies to:\n\\[\n\\sqrt{144}\n\\]\n\nFinally, we determine the square root of 144:\n\\[\n\\sqrt{144} = 12\n\\]\n\nThus, the value of \\(\\sqrt{36 \\times \\sqrt{16}}\\) is:\n\\[\n\\boxed{12}\n\\]'}
|
||||
],
|
||||
images=[],
|
||||
audios=[],
|
||||
videos=[],
|
||||
tools=None,
|
||||
objects={},
|
||||
data_dict={
|
||||
'problem': 'What is the value of $\\sqrt{36 \\times \\sqrt{16}}$?',
|
||||
'solution': "To solve the problem, we need to evaluate the expression \\(\\sqrt{36 \\times \\sqrt{16}}\\).\n\nWe can break down the steps as follows:\n\n1. Evaluate the inner square root: \\(\\sqrt{16}\\).\n2. Multiply the result by 36.\n3. Take the square root of the product obtained in step 2.\n\nLet's compute this step by step using Python code for accuracy.\n```python\nimport math\n\n# Step 1: Evaluate the inner square root\ninner_sqrt = math.sqrt(16)\n\n# Step 2: Multiply the result by 36\nproduct = 36 * inner_sqrt\n\n# Step 3: Take the square root of the product\nfinal_result = math.sqrt(product)\nprint(final_result)\n```\n```output\n12.0\n```\nThe value of \\(\\sqrt{36 \\times \\sqrt{16}}\\) is /\\(\\boxed{12}\\)."
|
||||
}
|
||||
)
|
||||
"""
|
||||
response_choice
|
||||
"""
|
||||
ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=ChatMessage(
|
||||
role='assistant',
|
||||
content='To find the value of \\(\\sqrt{36 \\times \\sqrt{16}}\\), we will break down the problem step-by-step.\n\nFirst, we need to evaluate the inner square root:\n\\[\n\\sqrt{16}\n\\]\nWe know that:\n\\[\n4^2 = 16 \\implies \\sqrt{16} = 4\n\\]\n\nNext, we substitute this result back into the original expression:\n\\[\n\\sqrt{36 \\times \\sqrt{16}} = \\sqrt{36 \\times 4}\n\\]\n\nNow, we need to evaluate the product inside the square root:\n\\[\n36 \\times 4 = 144\n\\]\n\nSo, the expression simplifies to:\n\\[\n\\sqrt{144}\n\\]\n\nFinally, we determine the square root of 144:\n\\[\n\\sqrt{144} = 12\n\\]\n\nThus, the value of \\(\\sqrt{36 \\times \\sqrt{16}}\\) is:\n\\[\n\\boxed{12}\n\\]', tool_calls=None),
|
||||
finish_reason='stop',
|
||||
logprobs=None,
|
||||
messages=None)
|
||||
"""
|
||||
# response_choice.messages will be copied at the end of multi-turn inference.
|
||||
```
|
||||
</details>
|
||||
|
||||
<br>
|
||||
<br>
|
||||
|
||||
默认的 `check_finished` 逻辑会在以下两种情况下停止推理：
|
||||
- 模型回复被截断，即超出了 `max_completion_length`
|
||||
- 模型推理轮数超出了限制的最大轮数
|
||||
|
||||
完整的默认多轮 rollout 逻辑请参考该类的 `run` 方法，我们也可以通过重载`run` 方法来实现自定义多轮逻辑。
|
||||
|
||||
## 设置多轮训练参数
|
||||
|
||||
在 swift rollout 命令中，设置 multi_turn_scheduler 参数指定规划器
|
||||
|
||||
```bash
|
||||
swift rollout \
|
||||
--model Qwen/Qwen3-1.7B \
|
||||
--vllm_use_async_engine true \
|
||||
--multi_turn_scheduler thinking_tips_scheduler \
|
||||
--vllm_max_model_len 32768 \
|
||||
--vllm_gpu_memory_utilization 0.8 \
|
||||
--max_turns 3
|
||||
```
|
||||
|
||||
|
||||
> 通过参数 `external_plugins`，我们可以将本地的多轮规划器注册到 ms-swift 中，具体实现请参考[代码](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py)。
|
||||
|
||||
多轮训练脚本请参考[脚本](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/external/vllm_multi_turn.sh)。
|
||||
|
||||
|
||||
对于多轮 rollout，我们使用 AsyncEngine 来实现高效的批量数据异步多轮采样。AsyncEngine 在多轮推理时能够减少推理过程中的计算气泡：
|
||||
|
||||
<img src="https://raw.githubusercontent.com/modelscope/ms-swift/main/docs/resources/asyncengine.png" width="400" />
|
||||
|
||||
在 `rollout` 命令中使用参数 `vllm_use_async_engine` 来指定 engine 的种类（默认使用 async engine）：
|
||||
|
||||
> 注意: async engine 仅在 server mode 下可用。
|
||||
|
||||
### GYM 环境训练
|
||||
|
||||
如果你的多轮任务可以建模为标准的 gym environment（`reset` / `step` / 环境直接给奖励），推荐直接复用框架内置的 `gym_scheduler`，并通过实现一个 `Env` 子类来描述任务。
|
||||
|
||||
`GYMScheduler` 基于通用 hook 协议实现，无需重载 `run` 方法：
|
||||
- **`on_trajectory_start`**: 调用 `env.reset` 并注入初始 observation 到首轮 user 消息
|
||||
- **`on_turn_end`**: 调用 `env.step` 推进环境，返回 `{'done': bool, 'rollout_infos': dict}`
|
||||
|
||||
这种设计使得 `GYMScheduler` 同时适用于 server mode（`run()`）和 colocate mode（`run_multi_turn()`），用户只需实现 Env 接口即可。
|
||||
|
||||
完整接口、自定义 env 的步骤参考 [GYM 环境训练文档](./gym_env.md)。
|
||||
|
||||
## 高级设置
|
||||
|
||||
### 自定义多轮交互逻辑
|
||||
|
||||
在以上默认逻辑中，我们用一条轨迹来计算多轮 rollout 的损失，这里需要假设多轮交互的过程中，模型的历史信息没有受到改变。
|
||||
|
||||
而在一些多轮场景中，我们可能需要在多轮 rollout 过程中动态地修改模型的历史信息（比如压缩历史信息），此时，我们需要将每轮的 rollout 单独作为一条轨迹进行训练。
|
||||
|
||||
#### 方式一：使用 hook
|
||||
|
||||
```python
|
||||
class CustomScheduler(MultiTurnScheduler):
|
||||
def on_trajectory_start(self, requests):
|
||||
# 首轮推理前初始化（如环境 reset、注入初始状态）
|
||||
for req in requests:
|
||||
req.messages = [system_msg, user_msg(initial_observation)]
|
||||
|
||||
def on_turn_end(self, req, response_choice, current_turn):
|
||||
# 每轮推理后推进状态，返回 done 和 rollout_infos
|
||||
next_obs, reward, done = self.advance_env(req.messages)
|
||||
return {
|
||||
'done': done,
|
||||
'rollout_infos': {'reward': reward, ...}
|
||||
}
|
||||
```
|
||||
|
||||
这种方式同时适用于 server mode 和 colocate mode，无需重载 `run` 方法。
|
||||
|
||||
#### 方式二：重载 run 方法（完全自定义）
|
||||
|
||||
比较常见的一种场景是对于思考类模型，在实际推理过程中，模型通常只会保留最后一轮的思考内容，而忽略历史模型回复中的思考内容。
|
||||
|
||||
对于这类场景，我们需要重写多轮规划器中的交互逻辑，即重载 `run` 方法，从而单独返回每一轮的 Rollout 的结果。
|
||||
|
||||
框架内置的 `ThinkingModelTipsScheduler` 类展示了如何通过重写 `run()` 方法来实现完全自定义的多轮推理逻辑。请参考[内置多轮调度器实现](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py)
|
||||
|
||||
**注意**: 这种情况下，相同轨迹的数据会拆分为多条数据，在奖励相关的处理中，需要对相同轨迹的数据分配同样的reward。
|
||||
|
||||
可以在kwargs中获取 trajectory_inputs 获取完整轨迹的数据，具体实现参考[MultiTurnThinkingTips类](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py)
|
||||
|
||||
### 多模态数据修改
|
||||
在多模态多轮交互场景下，可能需要在对话过程中动态增删或修改多模态数据，并确保这些变更同步至 trainer。
|
||||
|
||||
实现方式：借助 rollout_infos，通过指定键值覆盖原始数据集的多模态内容。
|
||||
|
||||
现已支持覆盖的键：images、audios、videos。
|
||||
|
||||
具体请参考[DeepEyes Schduler](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py#L403-L404)
|
||||
|
||||
### 返回 response token ids
|
||||
在默认的多轮交互流程中，规划器先把模型生成的文本字符串返回给 trainer，trainer 再将其重新 encode 为 token id，用于后续训练。为了避免这一步重复编码的开销，你可以让规划器直接返回 response_token_ids，省去 trainer 侧的再次 encode。
|
||||
|
||||
具体做法如下：
|
||||
|
||||
- 在 response_choice 对象中读取 token_ids 属性，即可获得本次 rollout 生成的 token 序列。
|
||||
- 在 step/run 方法的返回值里加入 response_token_ids，trainer 便能直接使用这些 token id 参与训练，无需重新编码。
|
||||
|
||||
具体实现可以参考[ThinkingModelTipsScheduler](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py)类
|
||||
|
||||
### 损失掩码
|
||||
|
||||
在工具调用或环境交互返回结果时，若需将返回内容作为模型响应的一部分，建议对这些插入内容进行掩码处理，以确保模型在训练过程中不会对这些外部生成的内容计算损失。
|
||||
|
||||
我们可以通过两种方式设置损失掩码
|
||||
|
||||
**第一种：设置 loss_scale**
|
||||
|
||||
ms-swift 提供 loss_scale 参数来对模型回复部分的内容进行损失缩放设置。比如设置`--loss_scale last_round`，可以将非最后一轮的模型回复的损失置零。我们也可以实现自定义 loss_scale，具体请参考[定制化 loss_scale 文档](../../../Customization/Architecture.md#loss-scale)。
|
||||
|
||||
> 注：在GRPO中，loss_scale 只提供掩码功能，不提供缩放功能。
|
||||
|
||||
**第二种：设置loss_mask**
|
||||
|
||||
在`step`或者`run`方法中设置 response_loss_mask, 可以在规划器中自定义损失掩码。前提需要返回response token ids，返回的 response_loss_mask 需要与 response token ids等长。当返回 response_loss_mask 时，loss_scale 参数失效。
|
||||
|
||||
response_loss_mask 返回可以参考[ToolCallScheduler类](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py)
|
||||
|
||||
### 奖励函数相关
|
||||
|
||||
在奖励函数中获取多轮 Rollout 中的信息
|
||||
|
||||
在 `on_turn_end` 方法或 `step/run` 方法中，返回 `rollout_infos` 对象，在奖励函数的 kwargs 中获取 `rollout_infos`：
|
||||
|
||||
```python
|
||||
class Scheduler():
|
||||
def on_turn_end(self, infer_request, response_choice, current_turn):
|
||||
...
|
||||
return {'done': done, 'rollout_infos': extra_dict}
|
||||
|
||||
# 或者在 step 方法中
|
||||
def step(self, infer_request, response_choice, current_turn):
|
||||
...
|
||||
return {'infer_request': infer_request, 'rollout_infos': extra_dict}
|
||||
|
||||
class RewardFunction():
|
||||
def __call__(self, completions, **kwargs):
|
||||
infos = kwargs.get('rollout_infos', {})
|
||||
...
|
||||
```
|
||||
|
||||
### 在 Scheduler 中获取额外的数据集信息
|
||||
|
||||
在训练侧设置参数`--vllm_server_pass_dataset`，可将数据集中的其他列传入多轮规划器。在`infer_request.data_dict`中获取。
|
||||
|
||||
### 训推一致性兼容
|
||||
swift 支持从 vLLM 侧返回 rollout 的 logps 用于纠正训推不一致问题，具体请参考该[文档](../AdvancedResearch/training_inference_mismatch.md)
|
||||
|
||||
在多轮训练中，如果启用了 `rollout_importance_sampling_mode`，框架会自动收集每轮 rollout 的 log probabilities，用于校正训推不一致带来的 off-policy 问题。
|
||||
|
||||
**默认行为**：
|
||||
- 使用默认的 `run` 方法时，框架会自动从 `response_choice.logprobs` 中提取 log probabilities
|
||||
- 这些 logprobs 会与 `response_token_ids` 和 `response_loss_mask` 一起传递给 trainer
|
||||
|
||||
**自定义 Scheduler 的注意事项**：
|
||||
|
||||
如果你在 `step` 方法中修改了 response（如截断、添加内容），需要同步返回对应的 `rollout_logprobs`
|
||||
|
||||
**关键规则**：
|
||||
- `rollout_logprobs` 的长度应该等于 `response_loss_mask` 中值为 1 的数量
|
||||
- 对于 `loss_mask=0` 的 token（如用户添加的提示、工具返回结果），不需要提供 logprobs
|
||||
- 如果 `step` 方法没有返回 `rollout_logprobs`，框架会自动从 `response_choice.logprobs` 中提取
|
||||
|
||||
**重写 `run` 方法的场景**：
|
||||
|
||||
如果你完全重写了 `run` 方法，需要手动收集和传递 `rollout_logprobs`
|
||||
|
||||
具体的实现请参考[内置实现](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py)
|
||||
@@ -0,0 +1,146 @@
|
||||
# 奖励函数
|
||||
## 自定义奖励函数
|
||||
奖励函数接受模型生成的文本 completions 其他数据集中的列以及训练器状态作为参数(kwargs)进行打分, 其中[训练器状态](https://huggingface.co/docs/transformers/main/main_classes/callback#transformers.TrainerState)包含训练的步数等信息。
|
||||
|
||||
注意：模型输入相关的列（比如query，response）会被处理为 messages 键，原数据集中的 assistant response 会被舍弃，请使用额外的列进行保留。
|
||||
相关处理的列名参考[文档](../../../Customization/Custom-dataset.md#query-response格式)
|
||||
|
||||
以下是一个示例，展示了如何实现一个简单的长度奖励函数。该函数会在模型生成的文本长度超过 1024 时，给予 1.0 的奖励信号；否则，奖励信号为 0.0。
|
||||
|
||||
```python
|
||||
from swift.rewards import ORM, orms
|
||||
class DummyLengthRewardFunction(ORM)
|
||||
def __call__(completions, **kwargs):
|
||||
return [1.0 if len(completion) > 1024 else 0.0 for completion in completions]
|
||||
|
||||
orms['dummy']= DummyLengthRewardFunction
|
||||
```
|
||||
|
||||
**获取数据集中的其他列**
|
||||
|
||||
比如奖励函数需要获取数据集`solution`列、当前训练步数和总步数作为辅助计算，以下是两种获取方式
|
||||
|
||||
第一种：在__call__入参中显式定义列名
|
||||
```python
|
||||
def __call__(completions, solution, trainer_state, **kwargs):
|
||||
print(solution)
|
||||
global_step = trainer_state.global_step
|
||||
max_steps = trainer_state.max_steps
|
||||
...
|
||||
```
|
||||
|
||||
第二种：在kwargs中获取
|
||||
```python
|
||||
def __call__(completions, **kwargs):
|
||||
solution = kwargs.get('solution')
|
||||
trainer_state = kwargs.get('trainer_state')
|
||||
global_step = trainer_state.global_step
|
||||
max_steps = trainer_state.max_steps
|
||||
...
|
||||
```
|
||||
|
||||
|
||||
**使用自定义奖励函数**
|
||||
|
||||
可以在`swift/examples/train/grpo/plugin/plugin.py`中加入该奖励函数，使用参数`--external_plugins examples/train/grpo/plugin/plugin.py`进行注册，并通过 reward_funcs 参数进行指定
|
||||
|
||||
执行脚本参考[这里](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/run_external_reward_func.sh)
|
||||
|
||||
## 异步奖励函数
|
||||
|
||||
对于涉及 I/O 操作的奖励函数（如 API 调用、数据库查询等），可以使用异步（async）奖励函数来提高性能。异步奖励函数使用 `asyncio.gather` 并行执行，可以显著加速奖励计算。
|
||||
|
||||
```python
|
||||
from swift.rewards import AsyncORM, orms
|
||||
import asyncio
|
||||
|
||||
class AsyncAPIReward(AsyncORM):
|
||||
async def __call__(self, completions, **kwargs):
|
||||
import aiohttp
|
||||
|
||||
async def score_single(session, text):
|
||||
async with session.post(
|
||||
'https://api.example.com/score',
|
||||
json={'text': text}
|
||||
) as resp:
|
||||
result = await resp.json()
|
||||
return result['score']
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# 使用 asyncio.gather 并行发送所有请求
|
||||
tasks = [score_single(session, c) for c in completions]
|
||||
rewards = await asyncio.gather(*tasks)
|
||||
return list(rewards)
|
||||
|
||||
orms['async_api'] = AsyncAPIReward
|
||||
```
|
||||
|
||||
swift 支持同时使用同步和异步奖励函数。训练器会自动检测奖励函数的类型：
|
||||
- 同步奖励函数按顺序执行
|
||||
- 异步奖励函数使用 `asyncio.gather` 并行执行
|
||||
|
||||
[plugin](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py)文件中提供了一个调用`swift deploy`服务的生成式奖励模型的例子(async_genrm)
|
||||
|
||||
## 内置奖励函数
|
||||
swift内置了五种基于规则的奖励函数(代码见swift/rewards/orm.py)
|
||||
|
||||
| 奖励函数 | 论文 |
|
||||
|----------------|----------------------------------------------------------------------------|
|
||||
| accuracy | [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via RL](https://arxiv.org/abs/2501.12948) |
|
||||
| format | 同上 |
|
||||
| cosine | [Demystifying Long Chain-of-Thought Reasoning in LLMs](https://arxiv.org/abs/2502.03373) |
|
||||
| repetition | 同上 |
|
||||
| soft_overlong | [Decoupled Clip and Dynamic sAmpling Policy Optimization (DAPO)](https://arxiv.org/abs/2503.14476) |
|
||||
|
||||
### 1. **accuracy**
|
||||
|
||||
该函数将模型的生成结果与数据集中的 solution 列进行比较，计算准确率分数。如果生成结果与标准答案一致，则得分为 1.0；否则为 0.0。
|
||||
|
||||
注意：该奖励函数使用`math_verify`库解析生成结果和solution中的答案，可能只适用于特定的数学数据集。
|
||||
|
||||
### 2. **format**
|
||||
|
||||
论文中使用以下system prompt要求模型按照固定格式进行返回
|
||||
```
|
||||
A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>
|
||||
```
|
||||
|
||||
该函数检查模型是否按照 `<think>think content</think><answer>answer content</answer>` 的格式进行生成。如果生成文本符合格式要求，则得分为 1.0；否则为 0.0。
|
||||
|
||||
### 3. **cosine**
|
||||
|
||||
论文发现，仅使用 accuracy 奖励函数进行训练会导致模型的生成长度趋于超长，从而影响训练效果。cosine 奖励函数通过控制模型的生成长度来优化训练过程：
|
||||
|
||||
- 对于生成正确答案的文本，奖励值随长度增加而递减，鼓励模型生成简洁的回答。
|
||||
- 对于生成错误答案的文本，奖励值随长度增加而递增，鼓励模型进行更深入的思考。
|
||||
|
||||
使用余弦函数平滑地调整奖励值，确保奖励变化在合理范围内。余弦函数的参数包括生成文本的长度、最大长度限制以及奖励的最小值和最大值。
|
||||
|
||||
参数
|
||||
- cosine_min_len_value_wrong（默认值：-0.5）：生成错误答案时，最小长度对应的奖励值。
|
||||
- cosine_max_len_value_wrong（默认值：0.0）：生成错误答案时，最大长度对应的奖励值。
|
||||
- cosine_min_len_value_correct（默认值：1.0）：生成正确答案时，最小长度对应的奖励值。
|
||||
- cosine_max_len_value_correct（默认值：0.5）：生成正确答案时，最大长度对应的奖励值。
|
||||
- cosine_max_len（默认值等于模型生成的最大程度）：生成文本的最大长度限制。
|
||||
|
||||
|
||||
### 4. **repetition**
|
||||
|
||||
惩罚模型生成文本中的重复内容，通过检测生成文本中的重复 n-gram 模式来评估重复程度，并给予相应的惩罚。
|
||||
|
||||
函数将生成文本分割为单词，并提取指定大小的 n-gram（默认为 3-gram）。通过统计不同 n-gram 的数量与总 n-gram 数量的比例，计算重复比例。如果生成文本中重复的 n-gram 比例较高，则给予较大的负奖励（惩罚）。惩罚值通过重复比例和最大惩罚值（默认为 -1.0）计算得出。
|
||||
|
||||
参数
|
||||
- repetition_n_grams（默认值：3）：用于检测重复的 n-gram 大小。
|
||||
- repetition_max_penalty（默认值：-1.0）：最大惩罚值，用于控制惩罚的强度。
|
||||
|
||||
### 5. **soft overlong punishment**
|
||||
定义长度惩罚区间。在这个区间内，给予[-1,0]的线性惩罚。
|
||||
|
||||
参数
|
||||
- soft_max_length: 论文中的L_max，模型的最大生成长度，默认等于max_completion_length
|
||||
- soft_cache_length: 论文中的L_cache，控制长度惩罚区间，区间为[soft_max_length-soft_cache_length, soft_max_length]
|
||||
|
||||
## 注意事项
|
||||
|
||||
如果需要在奖励函数中加载模型，默认会使用训练的 deepspeed 插件（transformers逻辑），zero3下可能会导致模型无法正常推理，参考该[issue](https://github.com/modelscope/ms-swift/issues/4580)来跳过 deepspeed 初始化环境
|
||||
@@ -0,0 +1,145 @@
|
||||
# 奖励模型
|
||||
|
||||
默认情况下，奖励模型是指具有分类头数值输出的模型，通常称为输出奖励模型（ORM）。这些模型会对其他模型的输出进行评分，从而生成一个标量值，表示模型响应的质量。
|
||||
|
||||
我们可以通过使用参数 `reward_models` 来加载具有分类头的奖励模型，或者加载经过[奖励建模](../../RLHF.md#rm)训练的奖励模型，进而使用模型的logits作为奖励。
|
||||
|
||||
## 自定义奖励模型
|
||||
对于生成式奖励模型，有两种常见的调用方式：一种是在 Trainer 内部直接使用 reward_model_plugin 定义奖励模型的逻辑，可以使用TransformersEngine对奖励模型进行推理，另一种是通过外部部署的模型服务进行调用。
|
||||
|
||||
- 使用 reward_model_plugin 调用奖励模型时，模型会被内嵌在 Trainer 内部，无需额外占用计算资源。该方式优点是方便集成，但生成速度相对较慢，更适合参数量较小的奖励模型场景。
|
||||
|
||||
- 外部部署奖励模型时，可以通过诸如 swift deploy 或 vllm serve 等命令将模型服务部署于独立设备，大幅提升推理速度，适合参数量较大的模型。但这样需要预留额外的硬件资源。
|
||||
|
||||
### 内部插件
|
||||
|
||||
我们可以在 reward_model_plugin 中灵活地自定义奖励模型的处理逻辑。这使得实现诸如生成式奖励模型等技术成为可能，包括：
|
||||
|
||||
- 自定义模型的系统提示：定义特定的指令和上下文以指导评估过程。
|
||||
- 处理模型交互历史：管理对话上下文，以提供有意义且具有上下文感知的评估。
|
||||
- 定义自定义评估标准：设置独特的标准和度量，用于评估模型的响应，超越默认的准确性和相关性衡量标准。
|
||||
|
||||
通过reward_model_plugin，开发者可以针对其应用的特定需求定制奖励评估过程。这种灵活性允许更细致和有效的基于奖励的训练策略。
|
||||
|
||||
奖励模型通过plugin的`__call__`方法进行调用，该方法接受 `inputs` 作为参数，包含了模型输入输出的 messages 和数据集中的其他列
|
||||
|
||||
```python
|
||||
def __call__(self, inputs):
|
||||
print(inputs)
|
||||
"""
|
||||
[
|
||||
{
|
||||
'messages': [
|
||||
{'role': 'system', 'content': 'system prompt'},
|
||||
{'role': 'query', 'content': 'query'},
|
||||
{'role': 'user', 'content': 'completions1'},
|
||||
],
|
||||
'solution': "abc",
|
||||
},
|
||||
{
|
||||
'messages': [
|
||||
{'role': 'system', 'content': 'system prompt'},
|
||||
{'role': 'query', 'content': 'query'},
|
||||
{'role': 'user', 'content': 'completions2'},
|
||||
],
|
||||
'solution': "abc",
|
||||
}
|
||||
]
|
||||
|
||||
```
|
||||
|
||||
在插件中使用 TransformersEngine 进行奖励模型的推理， 我们只需构造 messages ，并通过 infer 接口调用：
|
||||
```python
|
||||
class RMPlugin(DefaultRMPlugin):
|
||||
|
||||
def __init__(self, model, template):
|
||||
|
||||
super().__init__(model, template)
|
||||
# initilize TransformersEngine to infer
|
||||
self.engine = TransformersEngine(self.model, template=self.template, max_batch_size=0)
|
||||
|
||||
def __call__(self, inputs):
|
||||
system_prompt = ...
|
||||
query = ...
|
||||
messages = [{'role': 'system', 'content': system_prompt}, {'role': 'query', 'content': query}]
|
||||
result = self.engine.infer([messages], self.request_config, use_tqdm=False)
|
||||
rewards = ...
|
||||
return rewards
|
||||
```
|
||||
|
||||
我们在 [rm_plugin.py](https://github.com/modelscope/ms-swift/blob/main/swift/rewards/rm_plugin.py) 中提供了一个简单的生成式奖励模型示例（GenRMPlugin）。
|
||||
|
||||
在 [plugin.py](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py) 中自定义奖励模型插件，并使用 `external_plugins` 参数进行注册。
|
||||
|
||||
|
||||
注意：
|
||||
1. 在 GRPOTrainer 中，reward_model 会依次append到 reward_funcs 中。因此，reward_weights 的顺序对应 [reward_funcs, reward_model]。
|
||||
2. reward_model_plugin 默认为 default，即使用 ORM 处理逻辑。
|
||||
3. 对于参数量较大的模型，TransformersEngine 生成速度较慢，请使用[外部部署](#外部部署)方法
|
||||
|
||||
对于 BERT 这类无法通过 reward_model 加载的模型，我们可以内置在 reward_function 中进行加载，参考[issue](https://github.com/modelscope/ms-swift/issues/4580)
|
||||
|
||||
### 外部部署
|
||||
|
||||
**示例 2：使用 swift deploy 部署奖励模型并进行远程调用**
|
||||
|
||||
这类方法则不需要使用 reward_model_plugin , 而是直接在奖励函数中进行调用即可
|
||||
|
||||
首先用如下命令启动模型服务：
|
||||
|
||||
```bash
|
||||
# 注意部署的设备不要与训练设备重叠
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
swift deploy \
|
||||
--model Qwen/Qwen2.5-72B-Instruct \
|
||||
--vllm_tensor_parallel_size 4
|
||||
|
||||
# [INFO:swift] model_list: ['Qwen2.5-72B-Instruct']
|
||||
# INFO: Started server process [xxxxxx]
|
||||
# INFO: Waiting for application startup.
|
||||
# INFO: Application startup complete.
|
||||
# INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
|
||||
```
|
||||
在奖励函数中通过 OpenAI 库初始化客户端，指定模型服务的地址和端口，示例代码如下：
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
class RMReward(ORM):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
try:
|
||||
self.client = OpenAI(
|
||||
api_key='EMPTY',
|
||||
base_url='http://127.0.0.1:8000/v1', # 若在本地部署则为 127.0.0.1
|
||||
)
|
||||
self.verify_model_name = self.client.models.list().data[0].id
|
||||
except Exception as e:
|
||||
raise RuntimeError('Failed to connect to the model service. Please deploy the model '
|
||||
"using 'swift deploy' or 'vllm serve'.") from e
|
||||
|
||||
|
||||
def __call__(self, completions, messages, **kwargs) -> List[float]:
|
||||
rewards = []
|
||||
for completion, message in zip(completions, messages):
|
||||
rm_prompt = ... # 构建 reward model 的prompt
|
||||
chat_response = self.client.chat.completions.create(
|
||||
model=self.verify_model_name,
|
||||
messages=[
|
||||
{
|
||||
'role': 'system',
|
||||
'content': 'You are a helpful assistant.'
|
||||
},
|
||||
{
|
||||
'role': 'user',
|
||||
'content': rm_prompt
|
||||
},
|
||||
],
|
||||
)
|
||||
response = chat_response.choices[0].message.content.strip()
|
||||
reward = ... # 根据奖励模型生成结果提取奖励值
|
||||
rewards.append(reward)
|
||||
return rewards
|
||||
|
||||
```
|
||||
Reference in New Issue
Block a user