chore: import upstream snapshot with attribution
Lint test / lint (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,361 @@
# GYM Environment Training
GYM-style environment training wraps the "model → environment → reward" chain behind an abstract interface, letting the LLM interact with the environment as an Agent over multiple turns. The reward of each step is produced directly by the environment, so you don't need a separate reward function to infer it from the trajectory. This document first introduces the interface, then walks through a complete custom example (FrozenLake) showing how to plug it into training.
## Gym interface
GYM originates from the [Gymnasium library](https://github.com/Farama-Foundation/Gymnasium). In ms-swift we define the following interface:
```python
class Env(ABC):
def __init__(self, env_config):
"""env_config comes from the env_config column of each dataset row and carries initialization arguments."""
self.env_config = env_config
@abstractmethod
async def reset(self, config: RolloutInferRequest) -> Tuple[str, Dict[str, Any], str]:
"""
Returns:
- observation: sent to the model as the first user message
- info: debug/log information, recorded in completions.jsonl
- system_message: system prompt for this trajectory
"""
pass
@abstractmethod
async def step(self, action: Messages) -> Tuple[str, float, bool, Dict[str, Any]]:
"""
Args:
action: the complete conversation messages so far; the last one is the model's latest reply
Returns:
- next_observation: next user message
- reward: reward for the current step
- done: whether the trajectory is finished
- info: debug/log information
"""
pass
@abstractmethod
async def close(self):
"""Release resources."""
pass
```
The `RolloutInferRequest` received by `reset` contains the dataset row's `messages`, `data_dict` (extra columns including `env_config`), etc. See the [input example](./multi_turn.md#multiturnscheduler) for the full structure.
> If you need extra control over the conversation history between turns (e.g. dynamic compression, injecting hints), subclass `MultiTurnScheduler` and implement `on_trajectory_start` / `on_turn_end` hooks, or override `step` / `run` — see the [multi-turn doc](./multi_turn.md#customising-the-interaction-logic).
## Launching training
Use the built-in [gym_scheduler](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py) to wire the env into multi-turn rollout.
`GYMScheduler` is based on the generic hook protocol:
- Inherits `MultiTurnScheduler` — no need to override the `run` method
- Implements `on_trajectory_start` (calls `env.reset`) and `on_turn_end` (calls `env.step`)
- Works with both server mode (`run()`) and colocate mode (`run_multi_turn()`)
User-defined envs are loaded via `--external_plugins your_plugin.py`; the plugin runs `envs['my_env'] = MyEnv` to register them (the FrozenLake example below demonstrates the full pattern).
The built-in `GYMScheduler` completes the control logic via hooks:
```python
class GYMScheduler(MultiTurnScheduler):
def on_trajectory_start(self, requests):
# Create an env for each request, call env.reset, inject initial 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):
# Call env.step, accumulate reward, return 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):
# Inject the next observation into a user message
if self._pending_obs.get(req.uuid):
req.messages.append({'role': 'user', 'content': next_obs})
return {'infer_request': req}
```
**Colocate mode**:
```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 works the same way
```
**Server mode**
```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
# On the trainer side, add --vllm_server_pass_dataset true so the env_config column reaches the rollout server.
megatron rlhf --vllm_mode server --vllm_server_pass_dataset true ...
# or swift rlhf --vllm_mode server --vllm_server_pass_dataset true ...
```
Two ways to select the environment:
- Set it globally via `--gym_env env_name` (recommended — one env for the whole script);
- Or specify it per dataset row via `env_config.name` (for mixed-env workloads; overrides `--gym_env`).
## Example: writing a FrozenLake environment from scratch
<img src="https://gymnasium.farama.org/_images/frozen_lake.gif" width="220" alt="FrozenLake environment (image from Gymnasium docs)" />
[FrozenLake](https://gymnasium.farama.org/environments/toy_text/frozen_lake/) is a classic task from OpenAI Gym: the agent starts at the start cell, must cross a frozen lake to reach the goal, and avoid holes along the way. The original environment is illustrated above. The walkthrough below uses a text-only version of it (the same grid rendered as ASCII).
Full source: [frozen_lake_plugin](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/multi_turn/frozen_lake_plugin.py).
**1. Define the Env**
Each dataset row produces a freshly generated random 4x4 map (random holes + random S/G positions, BFS-validated to be solvable). Cell meanings: `S` start / `G` goal / `H` hole (stepping in = fail) / `F` safe ice / `P` player's current position.
```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>
# Advance one cell, check G / H; the outer max_turns is enforced by the scheduler.
if cell == 'G': return obs, 1.0, True, {'status': 'goal'}
if cell == 'H': return obs, 0.0, True, {'status': 'hole'}
...
```
**2. Register**
Hook the env class into swift's `envs` registry. `--external_plugins` imports the file at startup, so the registration takes effect automatically:
```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
```
**3. Prepare the dataset**
The dataset is just a placeholder here — the actual data is constructed by the env, with `env_config.seed` controlling map-generation randomness:
```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}}
```
**4. (Optional) Blend in extra rewards**
With `--use_gym_env true`, the env-provided `total_reward` is automatically added as one reward column — no reward function is required. To mix in additional signals (e.g. format/length checks), just pass them via `--reward_funcs`; the gym reward is appended as an extra column and blended with the reward_funcs through `--reward_weights`. For example, also enabling a format reward:
```bash
megatron rlhf ... --use_gym_env true --reward_funcs format --reward_weights 0.2 1.0
# the last entry of reward_weights corresponds to the gym total_reward
```
**5. Train**
Runnable script: [`examples/megatron/grpo/multi_turn/frozen_lake.sh`](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/multi_turn/frozen_lake.sh)
During training, observe `rollout_infos.num_turns` (steps per trajectory) and the reward mean in the logs. `--log_completions true` writes full conversations to `completions.jsonl`, so you can verify the model outputs in the `<action>...</action>` format turn by turn.
References:
- https://gymnasium.farama.org/environments/toy_text/frozen_lake/
- https://github.com/alibaba/ROLL/tree/main/roll/pipeline/agentic/env/frozen_lake
- [OpenEnv](https://github.com/huggingface/openenv)
- [TRL Sudoku GRPO Example](https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_sudoku_grpo.ipynb)
## OpenEnv Environment Training
[OpenEnv](https://github.com/huggingface/openenv) is an open-source Agentic RL environment framework by HuggingFace that communicates with environment servers via WebSocket. Unlike the local `Env` interface used by FrozenLake above, OpenEnv places environment logic in a separate server process, and swift communicates with it through `OpenEnvScheduler` + `OpenEnvWrapper`.
### Architecture Comparison
| Feature | Built-in Gym (`GYMScheduler`) | OpenEnv (`OpenEnvScheduler`) |
|---------|------------------------------|------------------------------|
| Environment location | In-process (Python object) | Standalone server (WebSocket) |
| Environment interface | Subclass `Env`, implement `reset/step/close` | Server provides HTTP/WebSocket API |
| Registration | `--external_plugins` + `envs` registry | `--external_plugins` + `multi_turns` registry |
| Use case | Lightweight local envs (FrozenLake, etc.) | Complex server envs (TextArena, CARLA, etc.) |
| Concurrency control | Not needed | Built-in Semaphore for connection limiting |
### OpenEnvScheduler
`OpenEnvScheduler` extends `GYMScheduler`, replacing the local `Env` with `OpenEnvWrapper` (a WebSocket client). Key design:
- **`_create_env`**: Creates an `OpenEnvWrapper` connected to the OpenEnv server
- **`on_trajectory_start`**: Creates a wrapper per request, calls `reset()`, uses Semaphore to limit concurrency (default 4)
- **`on_turn_end`**: Parses model output, calls `wrapper.step()`, accumulates reward
- **`parse_action`** (overridable): Converts model text to action dict, default `json.loads`
- **`format_observation`** (overridable): Converts server observation to string, default `json.dumps`
Users subclass `OpenEnvScheduler` and override `parse_action`, `format_observation`, `on_trajectory_start`, and `on_turn_end` to adapt to specific environments.
### Example: Sudoku Environment
Using TextArena Sudoku as an example, the model places numbers on a 9x9 Sudoku grid via `[row col number]` format. Full code: [sudoku_scheduler.py](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/openenv/sudoku_scheduler.py).
**1. Start OpenEnv Server**
Install OpenEnv and the Sudoku environment package (textarena and nltk are installed automatically as dependencies):
```bash
pip install openenv
pip install git+https://huggingface.co/spaces/openenv/sudoku
```
Use the provided startup script to start the local server (default port 8000). `MAX_CONCURRENT_ENVS` must be ≥ `num_generations` used in training:
```bash
TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 python examples/train/grpo/plugin/openenv/start_sudoku_server.py
```
> The default `python -m textarena_env.server.app` only supports 1 concurrent session, which is insufficient for GRPO's parallel multi-generation sampling. `start_sudoku_server.py` lifts this restriction by setting `SUPPORTS_CONCURRENT_SESSIONS`.
Point `base_url` to the local server in your dataset:
```json
{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}
```
**2. Custom Scheduler**
Subclass `OpenEnvScheduler` to implement Sudoku-specific action parsing, observation formatting, and multi-component rewards:
```python
from swift.rollout.multi_turn import OpenEnvScheduler
class SudokuScheduler(OpenEnvScheduler):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._last_content_len = {} # Content diff tracking
async def on_trajectory_start(self, requests):
# Create env, parse board, generate hints
# hints include 'guaranteed moves' and candidate numbers
...
async def on_turn_end(self, infer_request, response_choice, current_turn):
# Parse [row col number], step env
# Compute 5-component reward: empty_cell / valid_move / repetition / progress / correct
# Return updated board + hints as next observation
...
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]"}
```
**Multi-component reward system** (adapted from [TRL Sudoku example](https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_sudoku_grpo.ipynb)):
| Reward component | Calculation | Purpose |
|-----------------|-------------|---------|
| `empty_cell_reward` | Targets empty cell +1 / overwrites -1 | Guide model to valid positions |
| `valid_move_reward` | Valid new move +1 / warning -0.5 / invalid 0 | Encourage legal moves |
| `repetition_reward` | Exponential penalty for repeats (-2^n, cap -10) | Avoid repetition |
| `progress_reward` | (filled - initial) / (81 - initial) | Measure solving progress |
| `correct_reward` | Binary reward from environment | Puzzle fully solved |
Combined reward = sum of component averages, providing denser learning signal than a single binary reward.
**3. Hints System**
At each turn, the scheduler parses the current board state and provides hints to the model:
- **GUARANTEED MOVES**: Cells with only one candidate (can be filled directly)
- **Other options**: Cells with 2-3 candidates
- **MOVES ALREADY TRIED**: Previously attempted moves (to avoid repetition)
This significantly reduces exploration difficulty and enables the model to make more valid moves.
**4. Prepare Dataset**
The dataset serves as a placeholder; actual boards are generated by the environment server. Point `base_url` to the OpenEnv hosted address:
```json
{"messages":[{"role":"user","content":"Play"}],"env_config":{"name":"openenv","base_url":"http://127.0.0.1:8000"}}
```
**5. Register Scheduler**
`sudoku_scheduler.py` includes registration code at the end, loaded via `--external_plugins`:
```python
# End of sudoku_scheduler.py
from swift.rollout.multi_turn import multi_turns
multi_turns['sudoku_scheduler'] = SudokuScheduler
```
**6. Start Training**
```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 \
...
```
Runnable script: [`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)
### Notes
1. **vLLM mode**: The example above uses `--vllm_mode colocate`, where vLLM and training share the same GPUs. If using `--vllm_mode server`, you need to start `swift rollout` separately as the vLLM server, and `--multi_turn_scheduler` / `--max_turns` should be passed to `swift rlhf`, not `swift rollout`.
2. **Server concurrency**: `start_sudoku_server.py`'s `MAX_CONCURRENT_ENVS` must be ≥ `num_generations` used in training. The default `python -m textarena_env.server.app` only supports 1 concurrent session.
3. **Content diff**: Environments like TextArena return cumulative messages (full history each turn). The scheduler tracks `_last_content_len` to return only the new portion, preventing context length explosion.
4. **First-turn timing**: `on_trajectory_start` is called BEFORE the first rollout, ensuring the model sees the actual environment observation (e.g., Sudoku board) rather than the placeholder text from the dataset.
5. **enable_thinking**: When using Qwen3.5 series models, set `--enable_thinking false` to skip `<think>` block generation.
6. **Sync I/O**: `OpenEnvWrapper`'s `reset()`/`step()` are synchronous WebSocket calls. `OpenEnvScheduler` subclasses should wrap these calls with `asyncio.to_thread()` to avoid blocking the event loop.
@@ -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 training supports multiple loss types, with the main differences being the normalization dimension and gradient handling.
## Loss Function
At the token level, GRPO training uses the following loss function:
$$\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)$$
When setting `loss_type cispo`, the CISPO loss is used:
$$\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})$$
When setting `loss_type sapo`, soft gating replaces hard clipping, see [SAPO](../AdvancedResearch/SAPO.md)
$$\mathcal{L}_{i,t}^{\text{SAPO}} = -g_{i,t} \cdot A_{i,t}$$
where $g_{i,t} = \sigma(\tau \cdot (\rho_{i,t} - 1))$ is the temperature-controlled soft gate function.
where:
- $\rho_{i,t} = \frac{\pi_\theta(y_{i,t}|y_{i,<t})}{\pi_{\theta_{\text{old}}}(y_{i,t}|y_{i,<t})}$ is the importance sampling weight
- $A_{i,t}$ is the advantage function
- $\epsilon$ and $\epsilon_{\text{high}}$ are the clipping parameters
- $\text{detach}(\cdot)$ indicates that this term does not participate in gradient computation
- $\sigma(\cdot)$ is the sigmoid function, $\tau$ is the temperature parameter
## GRPO
`--loss_type grpo`
GRPO is the standard loss function implementation that averages the token-level losses for each sample, then averages across all samples.
**Formula:**
$$\mathcal{L}_{\text{GRPO}} = \frac{1}{N} \sum_{i=1}^{N} \frac{1}{T_i} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}$$
where:
- $N$ is the number of samples in the batch
- $T_i$ is the number of completion tokens for the $i$-th sample
**Normalization Dimension:** Sample dimension (first average over tokens for each sample, then average over all samples)
## BNPO (Batch Normalized Policy Optimization)
`--loss_type bnpo`
BNPO sums all token losses from all samples and then divides by the total number of completion tokens.
**Formula:**
$$\mathcal{L}_{\text{BNPO}} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}}{\sum_{i=1}^{N} T_i}$$
where:
- $N$ is the number of samples in the batch
- $T_i$ is the number of completion tokens for the $i$-th sample
**Normalization Dimension:** Token dimension (average over all completion tokens)
## DR-GRPO
`--loss_type dr_grpo`
DR-GRPO sums all token losses from all samples and then divides by the batch size multiplied by the maximum completion length.
**Formula:**
$$\mathcal{L}_{\text{DR-GRPO}} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{T_i} \mathcal{L}_{i,t}}{N \times L_{\text{max}}}$$
where:
- $N$ is the number of samples in the batch
- $T_i$ is the number of completion tokens for the $i$-th sample
- $L_{\text{max}}$ is the maximum completion length
**Normalization Dimension:** Fixed dimension (batch size × maximum completion length)
## CISPO
`--loss_type cispo`
CISPO loss is normalized by the total number of completion tokens across all processes.
**Formula:**
$$\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}}$$
where:
- $N$ is the number of samples in the current process batch
- $T_i$ is the number of completion tokens for the $i$-th sample
- $N_p$ is the number of samples for the $p$-th process
**Normalization Dimension:** Global token dimension (total completion tokens across all processes)
## DAPO
`--loss_type dapo`
DAPO is similar to BNPO, using token-level normalization, but based on global data (multi-process) normalization.
**Formula:**
$$\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}}$$
where:
- $N$ is the number of samples in the current process batch
- $T_i$ is the number of completion tokens for the $i$-th sample
- $N_p$ is the number of samples for the $p$-th process
**Normalization Dimension:** Global token dimension (total completion tokens across all processes)
## FIPO
`--loss_type fipo`
FIPO adds a Future-KL influence weight on top of the DAPO/GRPO clipped policy loss. The sequence-level advantage for each token is weighted by the discounted accumulated KL shift from the current token to future tokens:
$$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}$$
The FIPO influence weight is detached by default and uses the same global token normalization as DAPO.
**Normalization Dimension:** Global token dimension (total completion tokens across all processes)
## SAPO
`--loss_type sapo`
SAPO uses temperature-controlled soft gating instead of hard clipping to achieve smooth gradient attenuation. The normalization method is the same as GRPO.
For details, please refer to [SAPO](../AdvancedResearch/SAPO.md)
@@ -0,0 +1,55 @@
# Multi-Task Training
We can add a column to the dataset that indicates the task type, and then use this information in the reward function or reward model plugin to determine which task is being processed. This allows us to implement multi-task training. For example, suppose our dataset contains both math and programming tasks like the following:
```json
[
{"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"}
]
```
We can set up different reward functions to handle math and code data separately. Note that the columns in the dataset will be passed to the reward function, so we can use the `task` column to distinguish between tasks.
Below are examples of reward functions tailored for different tasks:
```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
# Implement 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":
# Implement 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
```
For data that does not belong to the current task, we handle it by returning None, ensuring that the reward calculation only applies to data within the designated task.
@@ -0,0 +1,323 @@
# Multi-turn Training
In reinforcement-learning scenarios, the model may need to interact with the environment over multiple turns (e.g., tool calls).
This interactive training requires the model to carry out continuous reasoning based on the feedback from the environment.
This document explains in detail how to customise the multi-turn training workflow in GRPO training.
> GKD also supports multi-turn training, sharing the same `MultiTurnScheduler` infrastructure as GRPO.
The figure below shows a typical multi-turn training process, where the model may perform several rollout rounds that include environment interaction, tool calls, and so on:
![Multi-turn example](../../../../resources/grpo_multi_turn.png)
## MultiTurnScheduler
`MultiTurnScheduler` is an abstract base class that provides the default multi-turn dialogue-management logic.
Its workflow is illustrated below:
<img src="https://raw.githubusercontent.com/modelscope/ms-swift/main/docs/resources/multiturn_pipeline.png" width="300" />
The scheduler is responsible for two core functions:
- **Termination check** — decide whether the current turn of inference should stop via `check_finished`.
- **Inference request construction** — build the request object for the next turn via `step`.
Key methods of the abstract base class `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:
"""Called before the first inference turn to initialize trajectory-level state.
This method can directly modify requests (e.g., inject initial environment observation).
Default is no-op.
"""
pass
def on_turn_end(self, infer_request: 'RolloutInferRequest',
response_choice: 'ChatCompletionResponseChoice',
current_turn: int) -> Dict[str, Any]:
"""Called after the assistant message is appended and before check_finished.
Used to advance the environment state (e.g., env.step) and return per-turn metadata.
Returns:
Dict[str, Any]: may optionally contain:
- 'done' (bool): if present, overrides the result of check_finished
- 'rollout_infos' (dict): merged into the accumulated trajectory info
Default returns an empty dict (no-op).
"""
return {}
def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice',
current_turn: int) -> Dict:
"""
Handle the transition between dialogue turns.
Args:
infer_request: current inference request
response_choice: response of the current turn
current_turn: current turn index (starting from 1)
Returns:
Dict[str, Any]: a dict containing the result of this turn
- infer_request (required): the inference request for the next turn
- response_token_ids (optional): token IDs of each rollout response
- response_loss_mask (optional): loss mask of each rollout response
- rollout_logprobs (optional): token logps of each rollout response
- rollout_infos (optional): extra information
"""
raise NotImplementedError
def check_finished(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice',
current_turn: int) -> bool:
"""
Default termination logic for multi-turn rollout.
Termination conditions:
1. The response is truncated (finish_reason == 'length').
2. The dialogue reaches the maximum number of turns (if max_turns is set).
Args:
infer_request: the inference request
response_choice: response choice containing finish_reason
current_turn: current turn index
Returns:
bool: True to stop, False to continue
"""
if response_choice.finish_reason == 'length':
return True
if self.max_turns and current_turn >= self.max_turns:
return True
return False
```
Arguments passed to `step` and `check_finished`:
- **infer_request**: current inference request
- **response_choice**: inference result of the current turn
- **current_turn**: current turn index (starting from 1)
<details><summary>Input example (click to expand)</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>
The default `check_finished` logic stops inference in the following cases:
- The model reply is truncated, i.e. exceeds `max_completion_length`.
- The number of inference turns exceeds the specified maximum.
For the full default multi-turn rollout logic, see the `run` method of the class.
You can override `run` to implement a completely custom workflow.
## Setting multi-turn parameters
Specify the scheduler via `multi_turn_scheduler` in the `swift rollout` command:
```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
```
> With the `external_plugins` argument you can register your own local scheduler with ms-swift.
> Refer to the [plugin code](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py).
A full multi-turn training script can be found [here](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/external/vllm_multi_turn.sh).
For multi-turn rollout we use `AsyncEngine` to perform efficient batched asynchronous sampling.
AsyncEngine reduces compute bubbles in multi-turn inference:
<img src="https://raw.githubusercontent.com/modelscope/ms-swift/main/docs/resources/asyncengine.png" width="400" />
Use the `vllm_use_async_engine` argument in the `rollout` command to specify the engine type (async is the default).
> Note: The async engine is only available in server mode.
### GYM environment training
If your multi-turn task can be modeled as a standard gym environment (`reset` / `step` / reward produced by the env directly), use the built-in `gym_scheduler` and implement an `Env` subclass to describe the task.
`GYMScheduler` is based on the generic hook protocol and does not require overriding the `run` method:
- **`on_trajectory_start`**: calls `env.reset` and injects the initial observation into the first user message
- **`on_turn_end`**: calls `env.step` to advance the environment, returns `{'done': bool, 'rollout_infos': dict}`
This design makes `GYMScheduler` compatible with both server mode (`run()`) and colocate mode (`run_multi_turn()`) — users only need to implement the `Env` interface.
See the [GYM environment training doc](./gym_env.md) for the full interface, the steps to define a custom env, and a minimal end-to-end example with no external dependencies (FrozenLake — runs out of the box on Megatron in colocate mode).
## Advanced topics
### Customising the interaction logic
In the default logic we treat the whole multi-turn rollout as one trajectory when computing the loss.
This assumes the model's history is not modified during interaction.
In some scenarios you may need to dynamically change the history during rollout (e.g., compressing context).
In that case each turn should be treated as a separate trajectory.
#### Approach 1: Using hooks
```python
class CustomScheduler(MultiTurnScheduler):
def on_trajectory_start(self, requests):
# Initialise before the first turn (e.g., env.reset, inject initial state)
for req in requests:
req.messages = [system_msg, user_msg(initial_observation)]
def on_turn_end(self, req, response_choice, current_turn):
# Advance state after each turn, return done and rollout_infos
next_obs, reward, done = self.advance_env(req.messages)
return {
'done': done,
'rollout_infos': {'reward': reward, ...}
}
```
This approach works with both server mode and colocate mode, and does not require overriding the `run` method.
#### Approach 2: Overriding the `run` method (fully custom)
A common scenario is for "thinking" models: during real inference the model keeps only the last reasoning step and discards previous ones.
For such cases override the `run` method in your scheduler to return the result for each rollout turn individually.
The built-in `ThinkingModelTipsScheduler` shows how to fully customise multi-turn inference by overriding `run()`.
See the implementation in [multi_turn.py](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py).
**NOTE**: In this scenario, the data for a single trajectory is split into multiple records. When computing rewards, you must assign the same reward to every record that belongs to the same trajectory.
The complete trajectory can be accessed via `trajectory_inputs` in `kwargs`.
For a concrete implementation, see the [MultiTurnThinkingTips class](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py)
### Multimodal Data Override
In multimodal, multi-turn interactions, you may need to dynamically add, delete, or modify multimodal data during the conversation and ensure these changes are synchronized to the trainer.
Implementation: Use `rollout_infos` to override the original multimodal content in the dataset by specifying the corresponding keys.
Supported override keys: images, audios, videos.
For details, see [DeepEyes Scheduler](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py#L403-L404).
### Returning response token IDs
In the default workflow the scheduler returns text, the trainer re-encodes it to token IDs for training.
To avoid this extra encoding, have the scheduler return `response_token_ids` directly.
Steps:
- Read the `token_ids` attribute from `response_choice` to obtain the sequence.
- Include `response_token_ids` in the dict returned by `step` / `run`; the trainer can then use them directly.
For a concrete implementation, refer to the [ThinkingModelTipsScheduler class](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py)
### Loss mask
When the environment or a tool call returns content that becomes part of the model response, you may want to mask it so the model is not penalised on externally generated tokens.
You can set the loss mask in two ways.
**1. Using `loss_scale`**
ms-swift provides the `loss_scale` parameter to scale or mask parts of the response.
For example, `--loss_scale last_round` zeroes out the loss for all but the last round.
Custom `loss_scale` can also be implemented; see the [customisation guide](../../../Customization/Architecture.md#loss-scale).
> Note: In GRPO, `loss_scale` serves only as a mask; it does not scale the loss.
**2. Using `loss_mask`**
In `step` or `run`, set `response_loss_mask` to define a custom mask.
This requires returning `response_token_ids`; the mask must be the same length.
When `response_loss_mask` is provided, `loss_scale` is ignored.
For how to return response_loss_mask, see the [ToolCallScheduler class](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py)
### Reward-function related tips
**Accessing multi-turn rollout information in a reward function**
Return a `rollout_infos` object from `step` / `run`, then read it from `kwargs` in the reward function:
```python
class Scheduler():
def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice',
current_turn: int) -> Dict:
...
return {'infer_request': infer_request, 'rollout_infos': extra_dict}
class RewardFunction():
def __call__(self, completions, **kwargs):
infos = kwargs.get('rollout_infos', {})
...
```
### Accessing additional dataset information in scheduler
Set `--vllm_server_pass_dataset` on the training side to pass other dataset columns to the scheduler.
They can be read from `infer_request.data_dict`.
### Training-Inference-Mismatch
Swift supports returning rollout logprobs from the vLLM side to address training-inference mismatch issues. For details, please refer to this [document](../AdvancedResearch/training_inference_mismatch.md).
In multi-turn training, if `rollout_importance_sampling_mode` is enabled, the framework automatically collects log probabilities from each rollout turn to correct off-policy issues.
**Default Behavior**:
- When using the default `run` method, the framework automatically extracts log probabilities from `response_choice.logprobs`
- These logprobs are passed to the trainer along with `response_token_ids` and `response_loss_mask`
**Notes for Custom Schedulers**:
If you modify the response in your `step` method (e.g., truncation, adding content), you need to return the corresponding `rollout_logprobs`:
**Key Rules**:
- The length of `rollout_logprobs` should equal the count of 1s in `response_loss_mask`
- For tokens with `loss_mask=0` (e.g., user-added prompts, tool return results), no logprobs are needed
- If `step` does not return `rollout_logprobs`, the framework will automatically extract them from `response_choice.logprobs`
**When Overriding the `run` Method**:
If you completely override the `run` method, you need to manually collect and pass `rollout_logprobs`
For implementation, please refer to [here](https://github.com/modelscope/ms-swift/blob/main/swift/rollout/multi_turn.py)
@@ -0,0 +1,144 @@
# Reward Function
## Custom Reward Function
The reward function takes as arguments (via kwargs) the model-generated completions, other columns from the dataset, and the training state, and calculates a reward score. The [trainer state](https://huggingface.co/docs/transformers/main/main_classes/callback#transformers.TrainerState) includes information such as the current training step.
Note: The columns related to model input (such as query and response) are converted to the messages key. The original assistant response in the dataset will be discarded, so please use extra columns if you wish to retain it.
The relevant column names for processing can be found in the [document](../../../Customization/Custom-dataset.md#Query-Response)
Below is an example illustrating how to implement a simple length-based reward function. This function assigns a reward of 1.0 if the length of the generated completion exceeds 1024, and 0.0 otherwise.
```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
```
**Accessing Other Columns in the Dataset**
For example, if the reward function needs to access the solution column from the dataset, as well as the current training step and the total number of steps for calculation, there are two ways to retrieve these values:
Explicitly define the column name in the __call__ parameters:
```python
def __call__(completions, solution, trainer_state, **kwargs):
print(solution)
global_step = trainer_state.global_step
max_steps = trainer_state.max_steps
...
```
Retrieve it from 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
...
```
**Using Custom Reward Functions**
You can add the reward function in [plugin program](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py), register it using the parameter `--external_plugins examples/train/grpo/plugin/plugin.py`, and specify it via the `reward_funcs` parameter.
For execution scripts, refer to [here](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/run_external_reward_func.sh).
## Async Reward Functions
For reward functions involving I/O operations (such as API calls, database queries, etc.), you can use asynchronous (async) reward functions to improve performance. Async reward functions are executed in parallel using `asyncio.gather`, which can significantly speed up reward computation.
```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:
# Use asyncio.gather to send all requests in parallel
tasks = [score_single(session, c) for c in completions]
rewards = await asyncio.gather(*tasks)
return list(rewards)
orms['async_api'] = AsyncAPIReward
```
Swift supports using both synchronous and asynchronous reward functions simultaneously. The trainer automatically detects the type of reward function:
- Synchronous reward functions are executed sequentially
- Asynchronous reward functions are executed in parallel using `asyncio.gather`
The [plugin](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py) file provides an example of a generative reward model (async_genrm) that calls the `swift deploy` service.
## Built-in Reward Functions
Swift includes five rule-based reward functions (code can be found in swift/rewards/orm.py).
| Reward Function | Paper |
|----------------|----------------------------------------------------------------------------|
| accuracy | [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via RL](https://arxiv.org/abs/2501.12948) |
| format | Same as above |
| cosine | [Demystifying Long Chain-of-Thought Reasoning in LLMs](https://arxiv.org/abs/2502.03373) |
| repetition | Same as above |
| soft_overlong | [Decoupled Clip and Dynamic sAmpling Policy Optimization (DAPO)](https://arxiv.org/abs/2503.14476) |
### 1. **accuracy**
This function compares the model's generated output with the solution column in the dataset to calculate an accuracy score. If the generated output matches the reference answer, the score is 1.0; otherwise, it is 0.0.
Note: This reward function uses the `math_verify` library to parse the generated output and the solution, which may only be applicable to specific mathematical datasets.
### 2. **format**
The paper uses the following system prompt to require the model to return responses in a fixed format:
```
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>
```
This function checks whether the model generates text in the format `<think>think content</think><answer>answer content</answer>`. If the generated text meets the format requirements, the score is 1.0; otherwise, it is 0.0.
### 3. **cosine**
The paper found that using only the accuracy reward function for training could lead to excessively long generated outputs, thereby affecting training effectiveness. The cosine reward function optimizes the training process by controlling the length of the model's outputs:
- For texts with correct answers, the reward decreases as the length increases, encouraging the model to generate concise responses.
- For texts with incorrect answers, the reward increases as the length increases, encouraging the model to think more deeply.
A cosine function is used to smoothly adjust the reward value, ensuring the changes remain within a reasonable range. The parameters of the cosine function include the length of the generated text, the maximum length limit, and the minimum and maximum reward values.
Parameters:
- cosine_min_len_value_wrong (default: -0.5): The reward value for the minimum length when the answer is incorrect.
- cosine_max_len_value_wrong (default: 0.0): The reward value for the maximum length when the answer is incorrect.
- cosine_min_len_value_correct (default: 1.0): The reward value for the minimum length when the answer is correct.
- cosine_max_len_value_correct (default: 0.5): The reward value for the maximum length when the answer is correct.
- cosine_max_len (default equals the model's maximum generation length): The maximum length limit for the generated text.
### 4. **repetition**
Penalizes repetitive content in the model's generated text by detecting repeated n-gram patterns and applying corresponding penalties.
The function splits the generated text into words and extracts n-grams of a specified size (default: 3-grams). By calculating the ratio of unique n-grams to the total number of n-grams, it determines the repetition rate. If the repetition rate is high, a larger negative reward (penalty) is applied. The penalty value is calculated based on the repetition rate and the maximum penalty value (default: -1.0).
Parameters:
- repetition_n_grams (default: 3): The size of n-grams used to detect repetition.
- repetition_max_penalty (default: -1.0): The maximum penalty value, controlling the penalty strength.
### 5. **soft overlong punishment**
Defines a length penalty interval. Within this interval, a linear penalty in the range [-1, 0] is applied.
Parameters:
- soft_max_length: L_max in the paper, the model's maximum generation length, defaulting to max_completion_length.
- soft_cache_length: L_cache in the paper, controlling the length penalty interval, which is [soft_max_length - soft_cache_length, soft_max_length].
## Notes
If a model needs to be loaded in the reward function, the training DeepSpeed plugin (transformers logic) will be used by default. Under Zero3, this may cause the model to fail to perform inference properly. Refer to this [issue](https://github.com/modelscope/ms-swift/issues/4580) to skip the DeepSpeed initialization environment.
@@ -0,0 +1,142 @@
# Reward Model
By default, a reward model refers to a model with a classification head that outputs numeric values, usually called an Output Reward Model (ORM). These models score the outputs from other models and produce a scalar value representing the quality of the model response.
You can load reward models with a classification head using the `reward_models` parameter, or load reward models trained by [reward modeling](../../RLHF.md#rm), and then use the model's logits as rewards.
## Custom Reward Models
For generative reward models, there are two common ways to use them: one is by directly defining the reward model logic inside the Trainer via the `reward_model_plugin`, and then using TransformersEngine for inference; the other is to call an externally deployed model service.
- Using `reward_model_plugin`, the reward model will be embedded within the Trainer and does not require additional computational resources. The advantage of this approach is ease of integration, but generation speed is relatively slow, making it more suitable for small-parameter reward models.
- When deploying reward models externally, you can use commands like `swift deploy` or `vllm serve` to deploy the model service on an independent device to greatly improve inference speed, which is more suitable for large models. However, this approach requires reserving extra hardware resources.
### Internal Plugin
You can flexibly customize the reward model processing logic inside `reward_model_plugin`. This enables implementations such as generative reward models, including:
- Custom model system prompts: define specific instructions and context to guide the evaluation process.
- Handling model interaction history: manage dialog context to allow meaningful and context-aware evaluation.
- Defining custom evaluation metrics: set unique criteria and measures for response evaluation, beyond the default accuracy and relevance checks.
With `reward_model_plugin`, developers can tailor the reward evaluation process for specific application needs. This flexibility allows for more fine-grained and effective reward-based training strategies.
The reward model is called via the plugin's `__call__` method, which takes `inputs` as a parameter. `inputs` contains the messages of model input/output and other columns from the dataset.
```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",
}
]
"""
```
When using TransformersEngine in the plugin for reward model inference, you only need to construct messages and call the infer interface:
```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
```
We provide a simple example of a generative reward model (`GenRMPlugin`) in [rm_plugin.py](https://github.com/modelscope/ms-swift/blob/main/swift/rewards/rm_plugin.py).
You can customize your reward model plugin in [plugin.py](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/plugin/plugin.py) and register it using the `external_plugins` parameter.
Note:
1. In `GRPOTrainer`, the reward_model will be appended to reward_funcs one by one. Therefore, the order of `reward_weights` corresponds to `[reward_funcs, reward_model]`.
2. The default for `reward_model_plugin` is `default`, which uses ORM logic.
3. For models with a large number of parameters, TransformersEngine generation is slow. Please use [external deployment](#external-deployment).
For models like BERT that cannot be loaded by `reward_model`, you can load them inside `reward_function`, see [issue](https://github.com/modelscope/ms-swift/issues/4580).
### External Deployment
This approach does not require the `reward_model_plugin` and can be called directly in the reward function.
First, use the following command to start the model service:
```bash
# Note: Do not overlap deployment devices with training devices
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)
```
In the reward function, initialize the client using the OpenAI library and specify the address and port of the model service. Example:
```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 if deployed locally
)
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 = ... # Construct the prompt for the reward model
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 = ... # Extract the reward value from the result
rewards.append(reward)
return rewards
```