# 多轮训练

在强化学习训练场景中，模型采样可能需要与环境进行多轮交互（如工具调用）。这种交互式训练要求模型能够根据环境反馈信息进行连续推理。本文档将详细介绍如何在 GRPO 训练中自定义多轮训练流程。

> GKD 同样支持多轮训练，与 GRPO 共享同一套 `MultiTurnScheduler` 基础设施。

以下是多轮训练示例图，模型可能涉及多轮 rollout，包括环境交互、工具调用等步骤：

![多轮示例图](../../../../resources/grpo_multi_turn.png)

## 多轮规划器 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)
