This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,63 @@
|
||||
# On-Policy RL Meets Off-Policy Experts: Harmonizing SFT and RL via Dynamic Weighting (CHORD)
|
||||
|
||||
This document describes the CHORD algorithm proposed in the paper [On-Policy RL Meets Off-Policy Experts: Harmonizing SFT and RL via Dynamic Weighting](https://arxiv.org/abs/2508.11408). The core idea of CHORD is to dynamically integrate expert data (SFT) into reinforcement learning by a dual control mechanism: a global weight μ plus a token-level weight φ, thereby balancing imitation and exploration.
|
||||
|
||||
## Algorithm Overview
|
||||
CHORD mixes training by introducing the SFT loss into the GRPO loss. The overall objective is:
|
||||
|
||||
$$
|
||||
\mathcal{L}_{\text{CHORD}} = (1 - \mu) \cdot \mathcal{L}_{\text{GRPO}} + \mu \cdot \mathcal{L}_{\text{SFT}}
|
||||
$$
|
||||
|
||||
where:
|
||||
- $\mathcal{L}_{\text{GRPO}}$: on-policy RL loss based on on-policy samples (similar to PPO).
|
||||
- $\mathcal{L}_{\text{SFT}}$: supervised fine-tuning (SFT) loss.
|
||||
- $\mu \in [0, 1]$: global balancing coefficient that controls the contribution of the SFT signal to the overall gradient.
|
||||
|
||||
### Configuration (data and batch sizes)
|
||||
We can implement CHORD training based on GRPO training.
|
||||
|
||||
CHORD requires specifying an additional SFT dataset and batch size at training time:
|
||||
- `chord_sft_dataset`: the SFT dataset that provides expert data.
|
||||
- `chord_sft_per_device_train_batch_size`: SFT mini-batch size per device.
|
||||
|
||||
---
|
||||
|
||||
## Two CHORD Variants
|
||||
|
||||
The paper proposes two variants: CHORD-μ and CHORD-φ.
|
||||
|
||||
### CHORD-μ
|
||||
CHORD-μ gradually decays μ during training to transition from imitating experts toward autonomous exploration.
|
||||
|
||||
Parameters:
|
||||
- `chord_mu_peak`: the peak value of μ.
|
||||
- `chord_mu_valley`: the final decayed value of μ.
|
||||
- `chord_mu_warmup_steps`: number of training steps to ramp μ up to the peak.
|
||||
- `chord_mu_decay_steps`: number of training steps to decay μ from peak to valley.
|
||||
|
||||
### CHORD-φ (Token-level weighting)
|
||||
CHORD-φ uses a token-wise weighting function φ to dynamically control each expert token's gradient contribution.
|
||||
|
||||
Definition of φ:
|
||||
$$
|
||||
\phi(y_t^\star, \pi_\theta) = p_t \cdot (1 - p_t)
|
||||
$$
|
||||
|
||||
where:
|
||||
- $p_t = \pi_\theta(y_t^\star \mid x, y_{<t}^\star)$ is the model's current predicted probability of the expert token.
|
||||
- When $p_t \approx 0.5$ (model uncertainty), φ is maximal → emphasize tokens the model is uncertain about.
|
||||
- When $p_t \approx 0$ or $p_t \approx 1$, φ → 0 → avoid overemphasizing tokens that are already certain or impossible.
|
||||
|
||||
Parameter to enable φ weighting:
|
||||
- `chord_enable_phi_function: bool = False`
|
||||
- Set to `True` to enable token-wise weight φ.
|
||||
|
||||
Note: If using a constant μ, set `chord_mu_peak` and `chord_mu_valley` to the same value.
|
||||
|
||||
<details>
|
||||
<summary>Code implementation of μ scheduling and loss computation</summary>
|
||||
See the `GRPOTrainer` method `_compute_chord_loss`.
|
||||
</details>
|
||||
|
||||
Training reference script: https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/internal/chord.sh
|
||||
@@ -0,0 +1,62 @@
|
||||
# Clipped Importance Sampling Policy Optimization (CISPO)
|
||||
|
||||
Clipped Importance Sampling Policy Optimization (CISPO) is a reinforcement learning algorithm proposed in the [MiniMax-M1](https://arxiv.org/abs/2506.13585) paper. Compared to GRPO (Group Relative Policy Optimization), CISPO clips the importance sampling weights themselves.
|
||||
|
||||
## Algorithm Overview
|
||||
|
||||
For clarity, we explain CISPO by contrasting it with GRPO.
|
||||
|
||||
GRPO limits the magnitude of policy updates by clipping the policy ratio. Its loss function is:
|
||||
|
||||
$$
|
||||
\mathcal{L}_{\text{GRPO}}(\theta) = -\mathbb{E}\left[\min\left(r_t(\theta) \cdot \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \cdot \hat{A}_t\right)\right]
|
||||
$$
|
||||
|
||||
where $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}$ is the importance sampling ratio.
|
||||
|
||||
When handling long reasoning chains, this clipping approach can lead to the following issues:
|
||||
|
||||
**Gradient Suppression of Critical Tokens**: In complex reasoning tasks, certain critical low-probability tokens (such as *However, Recheck, Wait, Aha*) are crucial for triggering deep thinking and reasoning error correction. These tokens have low probability in the old policy $\pi_{\theta_{\text{old}}}$. When the new policy attempts to increase their probability, it results in a large policy ratio $r_t(\theta)$, and GRPO's clipping mechanism will discard these tokens.
|
||||
|
||||
|
||||
### CISPO's Solution
|
||||
|
||||
The core idea of CISPO is to clip the importance sampling weights while preserving gradient updates. Specifically, CISPO's loss function is:
|
||||
|
||||
$$
|
||||
\mathcal{L}_{\text{CISPO}}(\theta) = -\mathbb{E}\left[\text{detach}\left(\min(r_t(\theta), \epsilon_{\text{high}})\right) \cdot \hat{A}_t \cdot \log \pi_\theta(a_t|s_t)\right]
|
||||
$$
|
||||
|
||||
where $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}$ is the importance sampling ratio.
|
||||
|
||||
**Key Mechanisms**:
|
||||
- Clip the importance sampling weights: $\min(r_t(\theta), \epsilon_{\text{high}})$
|
||||
- **Detach operation**: The clipped weights do not participate in gradient computation and serve as constant coefficients
|
||||
- Gradients come from the $\log \pi_\theta(a_t|s_t)$ term, ensuring all tokens contribute gradients
|
||||
|
||||
|
||||
## Implementation Details
|
||||
|
||||
The pseudo-code implementation of CISPO is as follows:
|
||||
|
||||
```python
|
||||
log_ratio = per_token_logps - old_per_token_logps
|
||||
importance_weights = torch.exp(log_ratio) # r_t(θ) = π_θ / π_θ_old
|
||||
|
||||
clamped_ratios = torch.clamp(importance_weights, max=epsilon_high).detach()
|
||||
|
||||
per_token_loss = -clamped_ratios * advantages.unsqueeze(1) * per_token_logps
|
||||
```
|
||||
|
||||
## Parameter Configuration
|
||||
|
||||
CISPO training can be enabled based on `GRPOTrainer` by setting the following parameters:
|
||||
|
||||
```bash
|
||||
--loss_type cispo
|
||||
--epsilon_high 5.0
|
||||
```
|
||||
|
||||
> Compared to other algorithms, cispo generally uses a larger value for epsilon_high. The minimax paper does not provide specific parameter settings; the value used here refers to the experimental setup in the paper [ScaleRL](https://arxiv.org/pdf/2510.13786).
|
||||
|
||||
For other training parameters, refer to the [GRPO parameter documentation](../../Command-line-parameters.md#grpo-arguments).
|
||||
@@ -0,0 +1,89 @@
|
||||
# DAPO: An Open-Source LLM Reinforcement Learning System at Scale
|
||||
|
||||
[Decoupled Clip and Dynamic sAmpling Policy Optimization (DAPO)](https://arxiv.org/abs/2503.14476) introduces several tricks based on GRPO, including:
|
||||
- [Clip Higher](#clip-higher)
|
||||
- [Dynamic Sampling](#dynamic-sampling)
|
||||
- [Token level Loss](#token-level-loss)
|
||||
- [Overlong Filtering](#overlong-filtering)
|
||||
- [Soft Overlong Punishment](#soft-overlong-punishment)
|
||||
|
||||
## Clip Higher
|
||||
PPO and GRPO use symmetric clipping ranges (e.g., ±0.2) to limit the magnitude of policy updates. While this ensures stability, it also restricts the model's exploratory capabilities. Specifically, when certain tokens have extremely low probabilities under the old policy, even if the current gradient indicates they should be reinforced (A > 0), the maximum increase is strictly limited.
|
||||
|
||||
DAPO employs an asymmetric clipping range, raising the upper clipping limit to encourage exploration:
|
||||
- The upper bound (encouragement side) is relaxed to 0.28.
|
||||
- The lower bound (suppression side) remains unchanged at 0.2.
|
||||
|
||||
In GRPO, the default symmetric clipping range is set using `epsilon`.
|
||||
|
||||
Parameters:
|
||||
- `epsilon_high` sets the upper clipping range, while `epsilon` serves as the lower clipping range.
|
||||
|
||||
## Dynamic Sampling
|
||||
GRPO samples multiple responses per question to compute inter-group advantages:
|
||||
|
||||
$$
|
||||
\hat{A}_{i,t} = \frac{R_i - \text{mean}(\{R_j\}_{j=1}^G)}{\text{std}(\{R_j\}_{j=1}^G)}
|
||||
$$
|
||||
|
||||
However, when all generated outputs {o_i} receive the same reward, the inter-group advantage becomes zero, leading to vanishing gradients and reduced training efficiency.
|
||||
|
||||
DAPO addresses this issue with a dynamic sampling strategy:
|
||||
- Skips data with zero inter-group reward standard deviation during sampling.
|
||||
- Continues generating samples until the batch is filled.
|
||||
|
||||
Parameters:
|
||||
- `dynamic_sample true` enables dynamic sampling.
|
||||
- `max_resample_times` sets the maximum number of resampling attempts.
|
||||
|
||||
## Token level Loss
|
||||
GRPO normalizes losses at the sentence level, which introduces bias based on response length.
|
||||
|
||||
DAPO uses token-level normalization to avoid this bias in loss calculation.
|
||||
|
||||
Parameters:
|
||||
- `loss_type bnpo/dapo` enables token-level normalization.
|
||||
|
||||
> For the loss_type formula, please refer to the [documentation](../DeveloperGuide/loss_types.md).
|
||||
|
||||
## Overlong Filtering
|
||||
DAPO argues that forcibly truncated responses contain high reward noise, making it difficult for the model to distinguish between quality issues and length issues. To address this, DAPO filters out truncated data during training, excluding it from loss computation.
|
||||
|
||||
Parameters:
|
||||
- `overlong_filter` enables filtering of overly long samples.
|
||||
|
||||
## Soft Overlong Punishment
|
||||
Language models often struggle with controlling output length:
|
||||
- Overly long outputs may be truncated, leading to incorrect judgments of valid content.
|
||||
- Unconstrained length generation affects practicality and computational efficiency.
|
||||
|
||||
DAPO designs a three-stage length penalty function:
|
||||
|
||||
$$
|
||||
R_{\text{length}}(L) =
|
||||
\begin{cases}
|
||||
0, & L \leq L_{\text{max}} - L_{\text{cache}} \\[10pt]
|
||||
\dfrac{(L_{\text{max}} - L_{\text{cache}}) - L}{L_{\text{cache}}}, & L_{\text{max}} - L_{\text{cache}} < L \leq L_{\text{max}} \\[10pt]
|
||||
-1, & L > L_{\text{max}}
|
||||
\end{cases}
|
||||
$$
|
||||
|
||||
When the length falls within the interval $(L_{\text{max}} - L_{\text{cache}} < L \leq L_{\text{max}})$, a linearly increasing penalty is applied. For lengths $(L > L_{\text{max}})$, the maximum penalty (-1) is imposed.
|
||||
|
||||
Parameters:
|
||||
- `reward_funcs soft_overlong` enables this reward function.
|
||||
- `soft_max_length` sets L_max, which defaults to the model's maximum output length (`max_completion_length`).
|
||||
- `soft_cache_length` sets L_cache.
|
||||
|
||||
## Parameter Settings
|
||||
In summary, the following parameters can be set based on GRPOTrainer to implement DAPO training.
|
||||
|
||||
| Parameter | Type | Value |
|
||||
|-----------------------|-----------|-------------|
|
||||
| `--loss_type` | `str` | `bnpo`/`dapo`|
|
||||
| `--epsilon_high` | `float` | `0.28` |
|
||||
| `--dynamic_sample` | `bool` | `true` |
|
||||
| `--max_resample_times`| `int` | `3` |
|
||||
| `--overlong_filter` | `bool` | `true` |
|
||||
| `--reward_funcs` | `str` | `soft_overlong`|
|
||||
| `--soft_cache_length` | `int` | `4096` |
|
||||
@@ -0,0 +1,55 @@
|
||||
# FIPO: Future-KL Influenced Policy Optimization
|
||||
|
||||
Author: [li2zhi](https://github.com/li2zhi)
|
||||
|
||||
[FIPO](https://arxiv.org/abs/2603.19835) is a value-free RL method for eliciting longer and deeper reasoning. It keeps the GRPO/DAPO training scaffold, but changes how token-level policy updates are weighted: instead of applying one sequence-level advantage uniformly to every token, FIPO uses a discounted Future-KL signal to estimate whether the future trajectory after each token is being reinforced or suppressed.
|
||||
|
||||
## Core Idea
|
||||
|
||||
In GRPO/DAPO, tokens in the same response usually share the same sequence-level advantage:
|
||||
|
||||
$$
|
||||
\hat{A}_{i,t} = \hat{A}_{i}
|
||||
$$
|
||||
|
||||
This is simple and stable, but the credit assignment is coarse. FIPO starts from the signed log-probability shift between the current policy and the old policy:
|
||||
|
||||
$$
|
||||
\Delta \log p_t = \log \pi_\theta(y_t \mid x, y_{<t}) -
|
||||
\log \pi_{\mathrm{old}}(y_t \mid x, y_{<t})
|
||||
$$
|
||||
|
||||
A positive value means the token probability is being increased by the current update, while a negative value means it is being suppressed. FIPO then accumulates this signal from the current token to the end of the response:
|
||||
|
||||
$$
|
||||
\mathrm{FutureKL}_t =
|
||||
\sum_{k=t}^{T}\gamma^{k-t} M_k \Delta \log p_k
|
||||
$$
|
||||
|
||||
where $M_k$ is the completion mask and $\gamma = 2^{-1 / \text{decay\_rate}}$. A larger `decay_rate` gives farther future tokens more influence; a smaller value makes the signal more local. FIPO maps the Future-KL value into a bounded influence weight:
|
||||
|
||||
$$
|
||||
f_t = \mathrm{clip}(\exp(\mathrm{FutureKL}_t), 1-\epsilon_f, 1+\epsilon_f)
|
||||
$$
|
||||
|
||||
The original advantage is then replaced by a future-aware advantage:
|
||||
|
||||
$$
|
||||
\tilde{A}_{i,t} = \hat{A}_{i} \cdot f_{i,t}
|
||||
$$
|
||||
|
||||
## Parameters
|
||||
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ------------------------- | ------- | ------- |---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `--loss_type` | `str` | `grpo` | Set to`fipo` to enable FIPO loss |
|
||||
| `--delta` | `float` | `None` | When enabled, it is used for both Future-KL high-IS-ratio token filtering and the main-loss dual-clip upper bound, and should be greater than `1 + epsilon_high`. Set it to `10.0` to match the official 32B script |
|
||||
| `--fipo_decay_rate` | `float` | `32.0` | Half-life parameter for Future-KL; the actual discount is`2 ** (-1 / fipo_decay_rate)` |
|
||||
| `--fipo_clip_range` | `float` | `0.2` | Influence weight clipping range;`0.2` clips to `[0.8, 1.2]` |
|
||||
| `--fipo_clip_high_only` | `bool` | `true` | If`true`, clips the weight to `[1.0, 1.0 + fipo_clip_range]` |
|
||||
| `--fipo_safety_threshold` | `float` | `4.0` | Caps the FIPO weight to `[0.8, 1.0]` for negative-advantage tokens whose IS ratio exceeds this threshold |
|
||||
|
||||
## Training Example
|
||||
|
||||
[swift](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/internal/fipo.sh)
|
||||
@@ -0,0 +1,66 @@
|
||||
# Group Sequence Policy Optimization (GSPO)
|
||||
|
||||
In [Group Sequence Policy Optimization](https://arxiv.org/abs/2507.18071), it is pointed out that GRPO computes importance sampling weights at the token level. However, this approach is problematic: since each token is only sampled once, it cannot realize effective distribution correction, and instead introduces high-variance noise during training, which can easily lead to unstable gradient estimates and even training collapse. Therefore, the paper argues that the unit of the objective function should be consistent with that of the reward. Since the reward is typically given at the sequence level (i.e., for the entire generated response), it is more reasonable to perform off-policy correction and optimization at the sequence level rather than the token level.
|
||||
|
||||
Below are the three main strategies for computing importance sampling weights:
|
||||
|
||||
1. GRPO
|
||||
GRPO computes the importance sampling ratio independently for each token, as follows:
|
||||
|
||||
$$
|
||||
w^{\mathrm{GRPO}}_{i,t} = \frac{\pi_\theta (y_{i, t} \mid x, y_{i, <t})}{\pi_{\theta_{\mathrm{old}}} (y_{i, t} \mid x, y_{i, <t})}
|
||||
$$
|
||||
|
||||
2. GSPO (Sequence-Level)
|
||||
GSPO calculates the importance sampling ratio at the sequence level, given by:
|
||||
|
||||
$$
|
||||
w^{\mathrm{GSPO}}_{i} = \left[ \frac{\pi_\theta (y_i \mid x)}{\pi_{\theta_{\mathrm{old}}} (y_i \mid x)} \right]^{\frac{1}{|y_i|}}
|
||||
= \exp\left( \frac{1}{|y_i|} \sum_{t=1}^{|y_i|} \log \frac{\pi_\theta (y_{i, t} \mid x, y_{i, <t})}{\pi_{\theta_{\mathrm{old}}} (y_{i, t} \mid x, y_{i, <t})} \right)
|
||||
$$
|
||||
|
||||
3. GSPO-token
|
||||
GSPO-token combines both sequence-level and token-level importance sampling:
|
||||
|
||||
$$
|
||||
w_{i, t}^{\mathrm{GSPO-token}} = \mathrm{sg}\left[w_i^{\mathrm{GSPO}}\right] \cdot \frac{\pi_{\theta}(y_{i, t} \mid x, y_{i, < t})}{\mathrm{sg}\left[\pi_{\theta}(y_{i, t} \mid x, y_{i, < t})\right]}
|
||||
$$
|
||||
|
||||
where $\mathrm{sg}[\cdot]$ denotes stop-gradient (detach()).
|
||||
|
||||
> **NOTE:** According to gradient analysis (i.e., Eqs. (11) and (18) in the paper), when the advantage for each token is identical, GSPO-token is equivalent to GSPO. In the current implementation of GRPO, all token advantages are normalized based on the sentence-level reward within each group. Therefore, in this setting, GSPO-token and GSPO are theoretically equivalent. However, GSPO-token provides support for future fine-grained (token-level) advantages.
|
||||
|
||||
Pseudo-code implementation:
|
||||
```python
|
||||
log_ratio = per_token_logps - old_per_token_logps
|
||||
# GRPO
|
||||
log_importance_weights = log_ratio
|
||||
|
||||
# GSPO (Sequence-Level)
|
||||
seq_weight = (log_ratio * mask).sum(-1) / mask.sum(-1)
|
||||
log_importance_weights = seq_weight.unsqueeze(-1) # (B,1)
|
||||
|
||||
# GSPO-token
|
||||
seq_weight = (log_ratio * mask).sum(-1) / mask.sum(-1)
|
||||
log_importance_weights = seq_weight.detach().unsqueeze(-1) + (per_token_logps - per_token_logps.detach())
|
||||
|
||||
importance_weights = torch.exp(log_importance_weights)
|
||||
```
|
||||
|
||||
Based on GRPO training, you can select different algorithms via the `--importance_sampling_level` argument:
|
||||
|
||||
- `importance_sampling_level token` (default, GRPO implementation)
|
||||
- `importance_sampling_level sequence` (GSPO)
|
||||
- `importance_sampling_level sequence_token` (GSPO-token)
|
||||
|
||||
|
||||
Other hyperparameters in the paper
|
||||
```bash
|
||||
--epsilon 3e-4 # from paper section 5.1
|
||||
--epsilon_high 4e-4 # from paper section 5.1
|
||||
--gradient_accumulation_steps 8
|
||||
--steps_per_generation 32 # from paper section 5.1 (each batch of rollout data is partitioned into four minibatches for gradient updates)
|
||||
--beta 0 # zero kl regularization https://github.com/volcengine/verl/pull/2775#issuecomment-3131807306
|
||||
```
|
||||
|
||||
For training, you can refer to [this script](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/internal/gspo.sh).
|
||||
@@ -0,0 +1,93 @@
|
||||
# Rewards as Labels: Revisiting RLVR from a Classification Perspective
|
||||
|
||||
Author: [li2zhi](https://github.com/li2zhi)
|
||||
|
||||
[Rewards as Labels: Revisiting RLVR from a Classification Perspective](https://arxiv.org/abs/2602.05630) proposes a reformulation of GRPO by treating rewards as labels and performing **in-group classification** instead of advantage estimation. This converts the policy optimization problem into a classification problem, thereby addressing two key issues in the GRPO loss:
|
||||
- **Gradient Misassignment** for positive samples
|
||||
- **Gradient Domination** for negative samples
|
||||
|
||||
## Background and Motivation
|
||||
|
||||
GRPO Objective
|
||||
|
||||
$$
|
||||
J_{\mathrm{GRPO}}(\theta)=\mathbb{E}_{q,o\sim\pi_{\mathrm{od}}(\cdot|q)}\left[\frac{1}{|o|}\sum_{t=1}^{|o|}\left(\min\left(\rho_tA_t,\mathrm{clip}(\rho_t,1-\epsilon,1+\epsilon)A_t\right)\right)\right]
|
||||
$$
|
||||
|
||||
where:
|
||||
- $\rho_t = \frac{\pi_\theta(o_t|q)}{\pi_{\mathrm{old}}(o_t|q)}$ is the probability ratio
|
||||
- $A_t$ is the advantage function
|
||||
|
||||
The corresponding gradient is:
|
||||
|
||||
$$
|
||||
\nabla_{\theta} J_{\mathrm{GRPO}} = \mathbb { E } \left[ \frac { 1 } { | o | } \sum _ { t = 1 } ^ { | o | } \mathbb { I } _ { \mathrm { clip } } \cdot A _ { t } e ^ { s _ { t } } \nabla _ { \theta } \log \pi _ { \theta } \left( o _ { t } | q \right) \right]
|
||||
$$
|
||||
|
||||
where:
|
||||
- $s_t = \log \frac{\pi_\theta(o_t|q)}{\pi_{\mathrm{old}}(o_t|q)}$ is the relative log-probability
|
||||
- $\mathbb{I}_{\mathrm{clip}}$ is the clipping indicator
|
||||
|
||||
Thus, the per-token gradient weight in GRPO is:
|
||||
|
||||
$$
|
||||
|\mathcal{W}_{\mathrm{GRPO}}|=\left\{ \begin{array} {ll}\left|A\cdot e^s\right|, & \mathrm{if~}\mathbb{I}_{\mathrm{clip}}=1, \\ 0, & \text{otherwise.} \end{array}\right.
|
||||
$$
|
||||
|
||||
|
||||
|
||||
1. **Gradient Misassignment (Positive Samples)**：
|
||||
For positive samples, as the relative log-probability $s$ decreases, the gradient magnitude also decreases.
|
||||
This is counterintuitive: tokens that the model is less confident about but correct should receive larger updates. However, GRPO assigns more weight to already confident tokens, causing under-trained tokens to receive insufficient learning signal.
|
||||
|
||||
2. **Gradient Domination (Negative Samples)**：
|
||||
For negative samples, as $s$ decreases, the gradient magnitude increases exponentially.
|
||||
This leads to a situation where a few overconfident incorrect tokens dominate the gradient, overwhelming other negative signals within the same group. Due to the absence of an upper bound, this may result in unstable and excessively large parameter updates.
|
||||
|
||||
To address the above issues, REAL treats rewards directly as labels and performs **group-wise classification training**.
|
||||
|
||||
|
||||
|
||||
The classification logit for each sample is defined as:
|
||||
|
||||
$$
|
||||
\bar{s}^k=\frac{1}{|o^k|}\sum_{t=1}^{|o^k|}\left(\log\frac{\pi_\theta(o_t^k\mid q)}{\pi_{\mathrm{old}}(o_t^k\mid q)}\right)
|
||||
$$
|
||||
|
||||
- $\bar{s}^k > 0$: The sample is more likely under the current policy than the old policy → the model tends to **promote** this sample
|
||||
- $\bar{s}^k < 0$: The sample is less likely under the current policy → the model tends to **suppress** this sample
|
||||
|
||||
Loss Function
|
||||
|
||||
$$
|
||||
\mathcal{L}_{REAL}=\log\left(1+\sum_{\mathcal{O}_+}e^{-\bar{s}^i/\tau}\right)+\log\left(1+\sum_{\mathcal{O}_-}e^{\bar{s}^j/\tau}\right)
|
||||
$$
|
||||
|
||||
Gradient Properties
|
||||
|
||||
$$
|
||||
|\mathcal{W}_{\mathrm{REAL}}|=
|
||||
\begin{cases}
|
||||
\frac{1}{\tau}\frac{1}{1+C_{+}e^{\bar{s}^{k}/\tau}}, & r=1 \\
|
||||
\\
|
||||
\frac{1}{\tau}\frac{1}{1+C_{-}e^{-\bar{s}^{k}/\tau}}, & r=0 & & &
|
||||
\end{cases}
|
||||
$$
|
||||
|
||||
## Parameter Settings
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|--------------------------------------------------------------------|
|
||||
| `--loss_type` | `str` | - | Set to `real` |
|
||||
| `--real_tau` | `float` | `0.5` | Temperature parameter controlling decision boundary sharpness |
|
||||
|
||||
Training Script Reference
|
||||
|
||||
[swift](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/internal/real.sh)
|
||||
|
||||
## Important Notes
|
||||
|
||||
When configuring training parameters, ensure that:
|
||||
- `per_device_train_batch_size` is divisible by `num_generations`
|
||||
|
||||
This guarantees that each training batch contains complete groups, which is required for correct in-group classification.
|
||||
@@ -0,0 +1,86 @@
|
||||
# REINFORCE++: An Efficient RLHF Algorithm with Robustness to Both Prompt and Reward Models
|
||||
|
||||
[REINFORCE++ Baseline](https://arxiv.org/abs/2501.03262) is a simplified version of the REINFORCE++ algorithm, designed for outcome rewards (response-level scalar rewards). Similar to GRPO, it samples multiple model outputs for each prompt and uses an intra-group baseline to estimate advantages. The key difference lies in the statistics used for normalization.
|
||||
|
||||
## Algorithm Overview
|
||||
For clarity, we explain REINFORCE++ Baseline by contrasting it with GRPO (Group Relative Policy Optimization).
|
||||
|
||||
Both GRPO and REINFORCE++ Baseline estimate advantages via intra-group comparisons. Their main differences are:
|
||||
|
||||
### Difference 1: Statistics Used for Normalization
|
||||
|
||||
**GRPO (Group Relative Policy Optimization)**
|
||||
|
||||
For each prompt, GRPO generates $G$ response samples and normalizes using the **mean and standard deviation of all samples within the group**:
|
||||
|
||||
$$
|
||||
\hat{A}_{i} = \frac{R_i - \text{mean}(\{R_j\}_{j=1}^G)}{\text{std}(\{R_j\}_{j=1}^G)}
|
||||
$$
|
||||
|
||||
When `scale_rewards='batch'` is set, it uses the **batch-level std of original rewards**:
|
||||
|
||||
$$
|
||||
\hat{A}_{i} = \frac{R_i - \text{mean}(\{R_j\}_{j=1}^G)}{\text{std}(\{R_j\}_{j=1}^{N})}
|
||||
$$
|
||||
|
||||
where $N$ is the total number of samples in the batch.
|
||||
|
||||
**REINFORCE++ Baseline**
|
||||
|
||||
For each prompt, REINFORCE++ generates $G$ response samples, subtracts the group mean, and then normalizes using the **standard deviation of the group-mean-subtracted rewards**:
|
||||
|
||||
$$
|
||||
\begin{align}
|
||||
\tilde{A}_{i} &= R_i - \text{mean}(\{R_j\}_{j=1}^G) \\
|
||||
\hat{A}_{i} &= \frac{\tilde{A}_{i}}{\text{std}(\{\tilde{A}_k\}_{k=1}^{N})}
|
||||
\end{align}
|
||||
$$
|
||||
|
||||
where $N$ is the total number of samples in the batch.
|
||||
|
||||
**Key Difference**:
|
||||
- **GRPO**: Uses the std of **original rewards $R$** for normalization
|
||||
- **REINFORCE++**: Uses the std of **group-mean-subtracted rewards $\tilde{A}$** for normalization
|
||||
|
||||
### Difference 2: KL Divergence Regularization
|
||||
|
||||
Similar to RLOO, REINFORCE++ Baseline integrates KL divergence directly into the reward:
|
||||
|
||||
$$
|
||||
R'_i = R_i - \beta \cdot \text{KL}(\pi_\theta || \pi_{\text{ref}})
|
||||
$$
|
||||
|
||||
where $\beta$ is the KL divergence weight coefficient (corresponding to the parameter `beta`), and $\pi_{\text{ref}}$ is the reference policy.
|
||||
|
||||
## Parameter Configuration
|
||||
|
||||
We can implement REINFORCE++ Baseline training by configuring the following parameters with `GRPOTrainer`:
|
||||
|
||||
```bash
|
||||
--advantage_estimator reinforce_plus_plus
|
||||
--scale_rewards batch
|
||||
--kl_in_reward true
|
||||
```
|
||||
|
||||
For training examples, please refer to this [script](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/internal/reinforce_plus_plus.sh)
|
||||
|
||||
### Key Parameter Descriptions
|
||||
|
||||
- **`--advantage_estimator`**: Selects the advantage estimation method
|
||||
- `grpo` (default): Uses the std of original rewards for normalization
|
||||
- `reinforce_plus_plus`: Uses the std of group-mean-subtracted rewards for normalization
|
||||
|
||||
- **`--kl_in_reward`**: Controls where the KL divergence regularization term is applied
|
||||
- `false`: KL divergence is an independent regularization term in the loss function (GRPO default)
|
||||
- `true`: KL divergence is subtracted directly from the reward (REINFORCE++ original implementation)
|
||||
|
||||
- **`--scale_rewards`**: Controls the normalization method
|
||||
- `group` (default): Intra-group normalization
|
||||
- `batch`: Global batch-level normalization (REINFORCE++ original implementation)
|
||||
- `none`: No normalization
|
||||
|
||||
- **`--num_generations`**: Number of samples generated per prompt ($G$)
|
||||
|
||||
- **`--beta`**: KL divergence regularization coefficient ($\beta$)
|
||||
|
||||
For other parameters, please refer to [GRPO Parameters](../../Command-line-parameters.md#grpo-arguments)
|
||||
@@ -0,0 +1,93 @@
|
||||
# REINFORCE Leave-One-Out (RLOO)
|
||||
|
||||
[REINFORCE Leave-One-Out (RLOO)](https://arxiv.org/abs/2402.14740) is a reinforcement learning algorithm based on the classic REINFORCE policy-gradient method. It constructs an unbiased advantage baseline via the Leave-One-Out (LOO) technique.
|
||||
|
||||
## Algorithm Overview
|
||||
|
||||
For clarity, we explain RLOO by contrasting it with GRPO (Group Relative Policy Optimization).
|
||||
|
||||
Both GRPO and RLOO estimate advantages via intra-group comparisons to avoid the high variance of a global baseline. Their core differences are mainly in the following aspects:
|
||||
|
||||
### Difference 1: How the Advantage Baseline Is Constructed
|
||||
|
||||
**1. GRPO (Group Relative Policy Optimization)**
|
||||
|
||||
For each prompt, GRPO generates $G$ response samples and normalizes rewards using the group mean and standard deviation:
|
||||
|
||||
$$
|
||||
\hat{A}_{i} = \frac{R_i - \text{mean}(\{R_j\}_{j=1}^G)}{\text{std}(\{R_j\}_{j=1}^G)}
|
||||
$$
|
||||
|
||||
Where:
|
||||
- $R_i$ is the reward of the $i$-th sample
|
||||
- $\text{mean}(\{R_j\}_{j=1}^G) = \frac{1}{G}\sum_{j=1}^G R_j$ is the group mean
|
||||
- $\text{std}(\{R_j\}_{j=1}^G)$ is the group standard deviation
|
||||
|
||||
**2. RLOO (REINFORCE Leave-One-Out)**
|
||||
|
||||
For each prompt, RLOO generates $K$ response samples and constructs the baseline via Leave-One-Out, i.e., for the $i$-th sample, the baseline is the mean of the other $K-1$ samples:
|
||||
|
||||
$$
|
||||
\hat{A}_{i} = R_i - \frac{1}{K-1}\sum_{j \neq i} R_j
|
||||
$$
|
||||
|
||||
This can be equivalently rewritten as:
|
||||
|
||||
$$
|
||||
\hat{A}_{i} = \frac{K}{K-1} \left(R_i - \bar{R}\right)
|
||||
$$
|
||||
|
||||
where $\bar{R} = \frac{1}{K}\sum_{j=1}^K R_j$ is the group mean reward.
|
||||
|
||||
> Note: We use $K$ here to match the notation in the paper. It has the same meaning as $G$ in GRPO and corresponds to the configuration parameter `num_generations`.
|
||||
|
||||
**Why Leave-One-Out?**
|
||||
|
||||
The key advantage is unbiasedness. For the $i$-th sample, its reward $R_i$ is independent of the baseline $\frac{1}{K-1}\sum_{j \neq i} R_j$, hence the advantage estimate is unbiased. In contrast, using the mean including itself as the baseline introduces bias.
|
||||
|
||||
### Difference 2: How KL Regularization Is Applied
|
||||
|
||||
To prevent the policy from drifting too far from the reference policy, both algorithms introduce KL divergence regularization, but in different ways:
|
||||
|
||||
**GRPO**: Adds KL divergence as an independent regularization term to the [loss](../GetStarted/GRPO.md#algorithm-overview):
|
||||
|
||||
$$
|
||||
\mathcal{L}(\theta) = -\mathbb{E}\left[\hat{A}_i \log \pi_\theta(a_i|s_i)\right] + \beta \cdot \text{KL}(\pi_\theta \Vert \pi_{\text{ref}})
|
||||
$$
|
||||
|
||||
**RLOO**: Integrates KL divergence directly into the reward, constructing a modified reward:
|
||||
|
||||
$$
|
||||
R'_i = R_i - \beta \cdot \text{KL}(\pi_\theta \Vert \pi_{\text{ref}})
|
||||
$$
|
||||
|
||||
where $\beta$ is the KL coefficient (parameter `beta`), and $\pi_{\text{ref}}$ is the reference policy (typically an SFT model or the initial policy).
|
||||
|
||||
## Parameter Configuration
|
||||
|
||||
RLOO training can be enabled based on `GRPOTrainer` by setting the following parameters:
|
||||
|
||||
```bash
|
||||
# Basic RLOO configuration
|
||||
--advantage_estimator rloo # Use RLOO's leave-one-out advantage estimator
|
||||
--kl_in_reward true # Integrate KL divergence into the reward (default for RLOO)
|
||||
```
|
||||
|
||||
You can refer to this [script](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/internal/rloo.sh) for training.
|
||||
|
||||
### Important Parameters
|
||||
|
||||
- **`--advantage_estimator`**: Choose the advantage estimator
|
||||
- `grpo` (default): standardize using group mean and standard deviation
|
||||
- `rloo`: construct the baseline via Leave-One-Out
|
||||
|
||||
- **`--kl_in_reward`**: Controls where the KL term is applied
|
||||
- `false`: KL as a separate regularization term in the loss (GRPO style)
|
||||
- `true`: subtract KL directly from the reward to form a modified reward (RLOO style)
|
||||
|
||||
- **`--num_generations`**: Number of samples per prompt, i.e., $K$
|
||||
|
||||
- **`--beta`**: KL regularization coefficient $\beta$
|
||||
- Controls how conservatively the policy updates
|
||||
|
||||
Other parameters are consistent with the [GRPO arguments](../../Command-line-parameters.md#grpo-arguments).
|
||||
@@ -0,0 +1,91 @@
|
||||
# Soft Adaptive Policy Optimization (SAPO)
|
||||
|
||||
[Soft Adaptive Policy Optimization (SAPO)](https://arxiv.org/abs/2511.20347) addresses the issues caused by hard clipping in GRPO by proposing a temperature-controlled soft gate mechanism that smoothly attenuates off-policy updates while preserving useful learning signals.
|
||||
|
||||
## Background and Motivation
|
||||
|
||||
When training LLMs with reinforcement learning, GRPO handles off-policy training by computing token-level importance sampling ratios:
|
||||
|
||||
$$
|
||||
r_t = \frac{\pi_\theta(y_t|x, y_{<t})}{\pi_{\theta_{\mathrm{old}}}(y_t|x, y_{<t})}
|
||||
$$
|
||||
|
||||
However, token-level importance sampling ratios often exhibit high variance, which can be exacerbated in the following cases:
|
||||
- **Long text generation**
|
||||
- **MoE model routing heterogeneity**: The old-policy model during sampling and the training model may use different expert routing, significantly amplifying logps differences
|
||||
|
||||
To address this, GRPO uses hard clipping to limit the magnitude of policy updates:
|
||||
|
||||
$$
|
||||
L^{\mathrm{GRPO}} = -\min\left( r_t \cdot A, \mathrm{clip}(r_t, 1-\epsilon, 1+\epsilon) \cdot A \right)
|
||||
$$
|
||||
|
||||
**The Dilemma of Hard Clipping**: Hard clipping struggles to balance stability and learning efficiency—too strict clipping limits the number of effective samples, while too loose clipping introduces noisy gradients from off-policy samples, leading to training instability.
|
||||
|
||||
## SAPO Method
|
||||
|
||||
SAPO uses a temperature-controlled sigmoid soft gate function to replace hard clipping, achieving smooth gradient attenuation.
|
||||
|
||||
### Soft Gate Function
|
||||
|
||||
The core of SAPO is using the sigmoid function to apply soft gating on the importance sampling ratio:
|
||||
|
||||
For positive advantages ($A > 0$), use positive gating:
|
||||
|
||||
$$
|
||||
g^{+}_t = \sigma\left( \tau_{\mathrm{pos}} \cdot (r_t - 1) \right) \cdot \frac{4}{\tau_{\mathrm{pos}}}
|
||||
$$
|
||||
|
||||
For negative advantages ($A < 0$), use negative gating:
|
||||
|
||||
$$
|
||||
g^{-}_t = \sigma\left( \tau_{\mathrm{neg}} \cdot (r_t - 1) \right) \cdot \frac{4}{\tau_{\mathrm{neg}}}
|
||||
$$
|
||||
|
||||
where:
|
||||
- $\sigma(\cdot)$ is the sigmoid function
|
||||
- $\tau_{\mathrm{pos}}$ and $\tau_{\mathrm{neg}}$ are temperature parameters that control the gate function slope
|
||||
- $r_t$ is the importance sampling ratio
|
||||
|
||||
### SAPO Loss Function
|
||||
|
||||
$$
|
||||
L^{\mathrm{SAPO}} = -g_t \cdot A
|
||||
$$
|
||||
|
||||
where $g_t = g^{+}_t$ when $A > 0$, $g_t = g^{-}_t$ when $A < 0$.
|
||||
|
||||
### Temperature Parameters
|
||||
|
||||
The temperature parameter $\tau$ controls the decay rate of the soft gate function—larger values result in faster decay.
|
||||
|
||||
|
||||
|
||||
The paper points out that positive advantages increase the logit of sampled tokens while decreasing the logits of all unsampled tokens; negative advantages do the opposite, increasing the logits of many unsampled tokens, which may spread to a large number of irrelevant tokens and introduce instability. Therefore, the paper recommends setting $\tau_\text{neg} > \tau_\text{pos}$ to make the gradient decay faster for tokens with negative rewards, improving training stability and performance.
|
||||
|
||||
The paper recommends default values of $\tau_{\mathrm{pos}} = 1.0$ and $\tau_{\mathrm{neg}} = 1.05$.
|
||||
|
||||
## Parameter Settings
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `--loss_type` | `str` | - | Set to `sapo` |
|
||||
| `--tau_pos` | `float` | `1.0` | Temperature parameter for positive advantages, controls gate slope |
|
||||
| `--tau_neg` | `float` | `1.05` | Temperature parameter for negative advantages, controls gate slope |
|
||||
|
||||
```bash
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--loss_type sapo \
|
||||
--tau_pos 1.0 \
|
||||
--tau_neg 1.05 \
|
||||
# ... other parameters
|
||||
```
|
||||
|
||||
Example training scripts:
|
||||
|
||||
- [swift](https://github.com/modelscope/ms-swift/blob/main/examples/train/grpo/internal/sapo.sh)
|
||||
- [megatron swift](https://github.com/modelscope/ms-swift/blob/main/examples/megatron/grpo/sapo.sh)
|
||||
|
||||
> The soft gate mechanism of SAPO only takes effect during off-policy training.
|
||||
> The importance sampling granularity in SAPO is at the token level (i.e., importance_sampling_level defaults to token), which conflicts with GSPO.
|
||||
@@ -0,0 +1,93 @@
|
||||
# DeepEyes: Incentivizing "Thinking with Images" via Reinforcement Learning
|
||||
|
||||
## Principle Introduction
|
||||
|
||||
The [DeepEyes paper](https://arxiv.org/abs/2505.14362) proposes a method that enables models to "think with images" (image-assisted reasoning) by leveraging reinforcement learning. This approach achieves emergent model capabilities through end-to-end reinforcement learning, without requiring additional SFT (Supervised Fine-Tuning) steps. The model is equipped with built-in image localization capabilities and can actively invoke an "image zoom-in tool": during inference, the model automatically selects specific regions within an image for zooming and cropping, and then chains the processed region information into downstream reasoning, thus achieving multi-step visual-text reasoning.
|
||||
|
||||
|
||||
|
||||
## Best Practices
|
||||
|
||||
**Dataset Download and Registration**
|
||||
|
||||
Download the official DeepEyes training datasets locally:
|
||||
```bash
|
||||
# modelscope
|
||||
modelscope download --dataset Lixiang/ChenShawn-DeepEyes-Datasets-47k
|
||||
|
||||
# huggingface
|
||||
huggingface-cli download ChenShawn/DeepEyes-Datasets-47k --repo-type=dataset
|
||||
```
|
||||
|
||||
There are three parquet files in the dataset. Register them in the `swift/dataset/data/dataset_info.json` file. For each, rename the `prompt` column to `messages`:
|
||||
|
||||
```json
|
||||
{
|
||||
"ms_dataset_id": "path/to/data_0.1.2_visual_toolbox_v2.parquet",
|
||||
"columns": {
|
||||
"prompt": "messages"
|
||||
}
|
||||
},
|
||||
{
|
||||
"ms_dataset_id": "path/to/data/data_thinklite_reasoning_acc.parquet",
|
||||
"columns": {
|
||||
"prompt": "messages"
|
||||
}
|
||||
},
|
||||
{
|
||||
"ms_dataset_id": "path/to/data/data_v0.8_visual_toolbox_v2.parquet",
|
||||
"columns": {
|
||||
"prompt": "messages"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
For registering the reward function and tool call logic locally as used in the paper, you can refer to the [DeepEyes implementation example](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py).
|
||||
|
||||
**Deploying the Evaluation Model**
|
||||
|
||||
The reward function in DeepEyes relies on a generative reward model to compare model outputs with ground-truth answers. For efficiency, it is recommended to deploy the reward model as a service.
|
||||
|
||||
Assuming you use the Qwen2.5-VL-72B-Instruct model for evaluation, refer to the following deployment command:
|
||||
|
||||
```bash
|
||||
# 4*80G
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
swift deploy \
|
||||
--model Qwen/Qwen2.5-VL-72B-Instruct \
|
||||
--infer_backend vllm \
|
||||
--vllm_tensor_parallel_size 4 \
|
||||
```
|
||||
|
||||
In the plugin file, you can use the OpenAI interface to call the model; see the [Reward Model documentation](../DeveloperGuide/reward_model.md#external-deployment).
|
||||
|
||||
The complete training script can be found at https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/deepeyes/deepeyes.sh
|
||||
|
||||
## Implementation Details
|
||||
|
||||
The [DeepEyes implementation example](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py), which references the [official implementation](https://github.com/Visual-Agent/DeepEyes/blob/main/verl/utils/reward_score/vl_agent.py), provides sample code for a DeepEyes training plugin, covering the logic of the reward function and multi-turn interaction calls.
|
||||
|
||||
**Dataset Info** is shown below:
|
||||
|
||||
| Dataset File Name | data_source | Scoring Function | Tool Call |
|
||||
|------------------------------------------|-----------------------|-----------------------------------|------------------------|
|
||||
| data_v0.8_visual_toolbox_v2.parquet | chart | vl_agent.compute_score | True (image_zoom_in_tool) |
|
||||
| data_0.1.2_visual_toolbox_v2.parquet | vstar | vl_agent.compute_score | True (image_zoom_in_tool) |
|
||||
| data_thinklite_reasoning_acc.parquet | thinklite_eureka | vl_agent.compute_score_math | False |
|
||||
|
||||
**Note**: When processing image inputs, multimodal large models may perform preprocessing (such as cropping or resizing limited by the `max_pixels` parameter). When using the image zoom-in tool (`image_zoom_in_tool`), the model will output the cropped bbox based on the input image. Therefore, it is necessary to ensure the image fed into the tool has already been preprocessed. The following example shows how this is implemented in the Qwen2.5-VL series:
|
||||
|
||||
```python
|
||||
from qwen_vl_utils import fetch_image
|
||||
# At this point, images have not yet been preprocessed
|
||||
infer_request.images
|
||||
# Load as PIL.Image and apply cropping (same as using the MAX_PIXELS environment variable)
|
||||
img = fetch_image({'image': load_pil_image(infer_request.images[0])})
|
||||
```
|
||||
|
||||
**Tool Reward**
|
||||
|
||||
According to the paper, if the final answer is correct and the trajectory uses at least one tool, a tool reward is given. To prevent the model from generating invalid tool calls, we determine this based on the number of images, not just on the presence of `<tool_call>` tokens.
|
||||
```python
|
||||
tool_reward = 1.0 if num_image > 1 and acc_reward > 0.5 else 0.0
|
||||
```
|
||||
@@ -0,0 +1,28 @@
|
||||
# Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning
|
||||
|
||||
The [paper](https://arxiv.org/abs/2506.01939) finds that when training large language models for reasoning abilities with methods such as RLVR, the key to learning progress lies in a small fraction of high-entropy "minority tokens," rather than the majority of low-entropy tokens.
|
||||
|
||||
The paper demonstrates that within the token distribution during model reasoning, only a few high-entropy tokens play a dominant role. These tokens typically appear at critical junctures where the reasoning or decision path diverges the most (e.g., tokens like "wait," "since," etc.), determining whether the model can master complex reasoning tasks. In contrast, most low-entropy tokens contribute little to the model's reasoning ability. The paper proposes computing policy gradients exclusively on high-entropy tokens, discarding gradients for low-entropy tokens.
|
||||
|
||||
The formula for token entropy is as follows:
|
||||
|
||||
$
|
||||
H_t := -\sum_{j=1}^{V} p_{t,j} \log p_{t,j}, \qquad \text{where } (p_{t,1}, \cdots, p_{t,V}) = \mathbf{p}_t = \pi_\theta(\cdot | \mathbf{q}, \mathbf{o}_{<t}) = \text{Softmax}\left(\frac{\mathbf{z}_t}{T}\right)
|
||||
$
|
||||
|
||||
Where:
|
||||
- $\pi_\theta$: The model parameterized by $\theta$;
|
||||
- $\mathbf{q}$: The input query;
|
||||
- $\mathbf{o}_{<t} = (o_1, o_2, \cdots, o_{t-1})$: The sequence of tokens generated prior to timestep $t$;
|
||||
- $V$: Vocabulary size;
|
||||
- $\mathbf{z}_t \in \mathbb{R}^V$: The pre-softmax logits at timestep $t$;
|
||||
- $\mathbf{p}_t \in \mathbb{R}^V$: The model's output probability distribution over the vocabulary;
|
||||
- $T \in \mathbb{R}$: The decoding temperature, controlling the smoothness of the distribution.
|
||||
|
||||
Object of entropy computation: $H_t$ is the entropy of the token generation distribution $\mathbf{p}_t$, which measures the uncertainty in the policy $\pi_\theta$ under the given context $(\mathbf{q}, \mathbf{o}_{<t})$.
|
||||
|
||||
> "Token entropy" $H_t$ always refers to the uncertainty of the generation distribution $\mathbf{p}_t$ at position $t$, rather than a property of the token $o_t$ itself. In other words, $H_t$ is the entropy of the distribution $\mathbf{p}_t$ at position $t$, and is independent of the sampled token $o_t$.
|
||||
|
||||
In practice, during GRPO training, the top_entropy_quantile parameter can be used to control the percentile threshold for entropy filtering. In the experiments from the paper, this parameter is set to 0.2, meaning that only the top 20% of tokens (with the highest entropy) at each sequence position are used for optimization in each batch.
|
||||
|
||||
By setting the parameter `log_entropy`, you can record the changes in entropy during training; see the [documentation](../GetStarted/GRPO.md#logged-metrics) for reference.
|
||||
@@ -0,0 +1,19 @@
|
||||
Advanced Research
|
||||
===============
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
entropy_mask.md
|
||||
CISPO.md
|
||||
DAPO.md
|
||||
deepeyes.md
|
||||
FIPO.md
|
||||
GSPO.md
|
||||
CHORD.md
|
||||
RLOO.md
|
||||
REINFORCEPP.md
|
||||
REAL.md
|
||||
router_replay.md
|
||||
SAPO.md
|
||||
training_inference_mismatch.md
|
||||
treepo.md
|
||||
@@ -0,0 +1,104 @@
|
||||
# Router Replay (R2/R3)
|
||||
|
||||
**TL;DR**: In RL training of MoE models, routing inconsistency between the training engine and the inference engine can significantly amplify training-inference mismatch, and even cause training collapse. Router Replay eliminates this inconsistency by replaying fixed routing masks during the training forward pass. Depending on the replay source, there are two strategies: R2 (Vanilla Routing Replay) and R3 (Rollout Routing Replay).
|
||||
|
||||
## Background
|
||||
|
||||
### Three Policies in MoE RL
|
||||
|
||||
In GRPO training of MoE models, there are three distinct policy stages that share the same model weights but may differ in routing behavior:
|
||||
|
||||
| Policy | Notation | Routing Result | Description |
|
||||
|--------|----------|---------------|-------------|
|
||||
| **Training Policy** | $\pi_\theta$ | $e^{\pi}_t$ | The model during gradient updates |
|
||||
| **Old Policy** | $\pi_{\theta_{\text{old}}}$ | $e^{\pi}_{\text{old},t}$ | The model state before batch updates |
|
||||
| **Rollout Policy** | $\mu_{\theta_{\text{old}}}$ | $e^{\mu}_{\text{old},t}$ | The sampling policy in the inference engine (e.g., vLLM), with the same weights as old policy, but different routing due to kernel implementation differences, precision, etc. |
|
||||
|
||||
Here, $\pi_{\theta_{\text{old}}}$ and $\mu_{\theta_{\text{old}}}$ have identical weights at sampling time, but due to implementation differences between the inference and training engines (e.g., operator implementations), routing results may differ even for the same input.
|
||||
|
||||
### Decomposition of Training-Inference Mismatch
|
||||
|
||||
According to the [paper](https://arxiv.org/abs/2507.18071), the token-level importance sampling ratio can be decomposed into two factors:
|
||||
|
||||
$$
|
||||
\frac{\pi_\theta(y_t|x, y_{<t})}{\mu_{\theta_{\text{old}}}(y_t|x, y_{<t})} = \underbrace{\frac{\pi_{\theta_{\text{old}}}(y_t|x, y_{<t})}{\mu_{\theta_{\text{old}}}(y_t|x, y_{<t})}}_{\text{training-inference discrepancy}} \times \underbrace{\frac{\pi_\theta(y_t|x, y_{<t})}{\pi_{\theta_{\text{old}}}(y_t|x, y_{<t})}}_{\text{policy staleness}}
|
||||
$$
|
||||
|
||||
For MoE models, expert routing is deeply coupled with both factors:
|
||||
|
||||
- **Training-inference discrepancy**: Routing inconsistency between the training and inference engines ($e^{\pi}_{\text{old},t} \neq e^{\mu}_{\text{old},t}$) amplifies output divergence
|
||||
- **Policy staleness**: As mini-batch updates proceed, routing also drifts ($e^{\pi}_t \neq e^{\pi}_{\text{old},t}$), further deviating from the sampling policy
|
||||
|
||||
## R2: Vanilla Routing Replay
|
||||
|
||||
The core idea of R2 is to **replay the routing determined by the old policy in the training engine** ($e^{\pi}_{\text{old},t}$) during gradient updates.
|
||||
|
||||
### Principle
|
||||
|
||||
During the training forward pass, first run a forward pass with the old policy weights to record the expert indices selected by each MoE Router layer, then force the training model $\pi_\theta$ to use these indices in its forward pass:
|
||||
|
||||
$$
|
||||
g_{\text{replay},i} = \frac{I^{\pi}_{\text{old},i} \cdot \exp(s_{\text{train},i})}{\sum_j I^{\pi}_{\text{old},j} \cdot \exp(s_{\text{train},j})}
|
||||
$$
|
||||
|
||||
where $I^{\pi}_{\text{old}}$ is the old policy's routing mask and $s_{\text{train}}$ is the router logits computed by the training model. The softmax still operates on the training logits, so gradients can flow back to the router weights normally.
|
||||
|
||||
### Properties
|
||||
|
||||
| Scenario | Behavior |
|
||||
|----------|----------|
|
||||
| **First mini-batch** (on-policy) | $\theta = \theta_{\text{old}}$, so $e^{\pi}_t = e^{\pi}_{\text{old},t}$, **target policy unchanged** (no bias) |
|
||||
| **Subsequent mini-batches** (off-policy) | $\theta \neq \theta_{\text{old}}$, so $e^{\pi}_t \neq e^{\pi}_{\text{old},t}$, **target policy changed** (biased), but policy staleness is controlled |
|
||||
|
||||
## R3: Rollout Routing Replay
|
||||
|
||||
The core idea of R3 is to **replay the routing determined by the rollout policy in the inference engine** ($e^{\mu}_{\text{old},t}$) during the training forward pass.
|
||||
|
||||
### Principle
|
||||
|
||||
During sampling in the inference engine (e.g., vLLM), additionally record the expert indices (routing mask) for each token at every MoE Router layer, then pass these masks to the training engine and force $\pi_\theta$ to use them in its forward pass:
|
||||
|
||||
$$
|
||||
g_{\text{replay},i} = \frac{I^{\mu}_{\text{old},i} \cdot \exp(s_{\text{train},i})}{\sum_j I^{\mu}_{\text{old},j} \cdot \exp(s_{\text{train},j})}
|
||||
$$
|
||||
|
||||
### Compatibility with Other Methods
|
||||
|
||||
- R3 is **orthogonal** to **GSPO** and can be combined for further improvement
|
||||
- R3 combined with **TIS** may not provide additional gains (R3 already eliminates inconsistency at the source; TIS's additional correction may be redundant)
|
||||
- Router Replay and **Clipping** are both essential in off-policy training
|
||||
|
||||
## Router Mask Caching
|
||||
|
||||
The R3 paper also proposes that routing masks can be cached alongside the KV Cache: for the same prefix tokens, the MoE Router output is deterministic, so routing masks can be stored and reused together with the prefix KVCache. This is particularly important in multi-turn Agent scenarios (tool calling), avoiding the need to re-prefill the prefix to obtain routing masks, with an overall rollout latency overhead of less than 3%.
|
||||
|
||||
## Swift Implementation
|
||||
|
||||
### Parameters
|
||||
|
||||
Select the routing replay strategy via the `--router_replay_mode` parameter:
|
||||
|
||||
| Value | Description |
|
||||
|-------|-------------|
|
||||
| `disabled` (default) | No routing replay |
|
||||
| `R2` | Vanilla Routing Replay: record old policy routing in the training engine and replay |
|
||||
| `R3` | Rollout Routing Replay: export routing masks from the inference engine and replay in training |
|
||||
|
||||
Environment requirements:
|
||||
|
||||
- R3 requires vLLM ≥ 0.14.0 to support returning `routed_experts` information.
|
||||
- Router Replay is currently only available with the Megatron backend, requiring megatron-core ≥ 0.16.0.
|
||||
|
||||
### Relationship with Training-Inference Correction
|
||||
|
||||
Router Replay and the importance sampling (IS) correction described in [Training-Inference Mismatch](training_inference_mismatch.md) are complementary:
|
||||
|
||||
- **IS correction**: Corrects probability divergence at the loss level via weighting
|
||||
- **Router Replay**: Eliminates the source of divergence at the model architecture level by fixing routing
|
||||
|
||||
## References
|
||||
|
||||
1. [Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers](https://arxiv.org/abs/2510.11370)
|
||||
2. [Group Sequence Policy Optimization](https://arxiv.org/abs/2507.18071)
|
||||
3. [Stabilizing Reinforcement Learning with LLMs: Formulation and Practices](https://arxiv.org/abs/2512.01374)
|
||||
4. [Megatron Core Router Replay Design Document](https://docs.nvidia.com/megatron-core/developer-guide/nightly/api-guide/router_replay.html)
|
||||
@@ -0,0 +1,236 @@
|
||||
# Training-Inference-Mismatch
|
||||
|
||||
**TL;DR**: While GRPO introduces vLLM to accelerate the sampling process, it also introduces Training-Inference Mismatch issues that may affect training stability. This document explains the background, causes, and solutions to this problem.
|
||||
|
||||
## Background
|
||||
|
||||
### Basic Assumptions of GRPO
|
||||
|
||||
The training objective of GRPO (Group Relative Policy Optimization) can be expressed as:
|
||||
|
||||
$$
|
||||
\mathcal{L}_{\text{GRPO}} = - \mathbb{E}_{y \sim \pi_\theta} \left[ \min \left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t \right) \right]
|
||||
$$
|
||||
|
||||
Where:
|
||||
- $r_t(\theta) = \frac{\pi_\theta(y_t|x, y_{<t})}{\pi_{\theta_{\text{old}}}(y_t|x, y_{<t})}$ is the importance sampling ratio
|
||||
- $\hat{A}_t$ is the advantage function, calculated based on reward and group baseline
|
||||
- $\epsilon$ is the clipping parameter
|
||||
|
||||
**Core Assumption**: Samples $y$ are drawn from the policy $\pi_\theta$. In practice, this means:
|
||||
1. The rollout model and the training model (policy model) should be **the same model** $\pi_\theta$
|
||||
2. The probability distributions of both models should be **exactly identical**, i.e., $\pi_{\text{rollout}} = \pi_\theta$
|
||||
|
||||
### Deviation After Introducing vLLM
|
||||
|
||||
GRPO's training speed is largely constrained by the sampling process (rollout). To accelerate this, training frameworks introduce high-performance inference engines (such as vLLM) for sampling. The **ideal assumption** is that through weight synchronization, vLLM maintains consistency with the training model, i.e., $\pi_{\text{vLLM}} \equiv \pi_\theta$.
|
||||
|
||||
However, in practice, even with fully synchronized weights, due to differences in operator implementations, the probability distributions still deviate:
|
||||
|
||||
$$
|
||||
\pi_{\text{vLLM}}(y|x) \neq \pi_\theta(y|x)
|
||||
$$
|
||||
|
||||
At this point, the actual training objective becomes:
|
||||
|
||||
$$
|
||||
\mathcal{L} = - \mathbb{E}_{y \sim \pi_{\text{vLLM}}} \left[ \min \left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t \right) \right]
|
||||
$$
|
||||
|
||||
Where samples come from $\pi_{\text{vLLM}}$, but gradients are computed based on $\pi_\theta$. This **violates the algorithm's on-policy assumption**, introducing training-inference mismatch issues and potentially causing performance degradation.
|
||||
|
||||
## Solution
|
||||
|
||||
To address training-inference mismatch, we can introduce **Importance Sampling (IS)** correction mechanisms.
|
||||
|
||||
### Importance Sampling Correction
|
||||
|
||||
The basic idea of importance sampling is: when samples come from distribution $q$ rather than target distribution $p$, we can correct the expectation calculation by introducing weights:
|
||||
|
||||
$$
|
||||
\mathbb{E}_{x \sim p} [f(x)] = \mathbb{E}_{x \sim q} \left[ \frac{p(x)}{q(x)} \cdot f(x) \right]
|
||||
$$
|
||||
|
||||
Applied to the GRPO scenario, the corrected loss function is:
|
||||
|
||||
$$
|
||||
\mathcal{L}_{\text{corrected}} = - \mathbb{E}_{y \sim \pi_{\text{vLLM}}} \left[ w(x, y) \cdot \min \left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t \right) \right]
|
||||
$$
|
||||
|
||||
Where $w(x, y)$ is the importance sampling weight used to correct the distribution bias between vLLM and the training model.
|
||||
|
||||
Importance sampling weights can be computed and applied at different granularities:
|
||||
|
||||
1. **Token-Level**
|
||||
|
||||
Compute the importance sampling ratio at each token:
|
||||
|
||||
$$
|
||||
w_{i,t}^{\text{token}} = \frac{\pi_\theta(y_{i,t}|x, y_{i,<t})}{\pi_{\text{vLLM}}(y_{i,t}|x, y_{i,<t})}
|
||||
$$
|
||||
|
||||
2. **Sequence-Level**
|
||||
|
||||
Compute the sequence-level importance sampling ratio, then broadcast to each token:
|
||||
|
||||
$$
|
||||
w_i^{\text{seq}} = \left[ \frac{\pi_\theta(y_i|x)}{\pi_{\text{vLLM}}(y_i|x)} \right]^{\frac{1}{|y_i|}} = \exp\left( \frac{1}{|y_i|} \sum_{t=1}^{|y_i|} \log \frac{\pi_\theta(y_{i,t}|x, y_{i,<t})}{\pi_{\text{vLLM}}(y_{i,t}|x, y_{i,<t})} \right)
|
||||
$$
|
||||
|
||||
### Stability Control: Truncate vs. Mask
|
||||
|
||||
Excessively large importance sampling weights can cause gradient explosion and destabilize training. Therefore, weight control is necessary:
|
||||
|
||||
#### 1. Truncate
|
||||
|
||||
Truncate the importance sampling weight to the $[0, \tau]$ interval:
|
||||
|
||||
$$
|
||||
w_{\text{truncate}} = \min(w, \tau)
|
||||
$$
|
||||
|
||||
This method retains all samples but limits their influence.
|
||||
|
||||
#### 2. Mask
|
||||
|
||||
Discard token/sequence data where weights exceed the threshold:
|
||||
|
||||
$$
|
||||
w_{\text{mask}} = \begin{cases}
|
||||
w & \text{if } w \leq \tau \\
|
||||
0 & \text{otherwise}
|
||||
\end{cases}
|
||||
$$
|
||||
|
||||
|
||||
### Four Correction Modes
|
||||
|
||||
Combining granularity and control strategies, there are four correction modes (selected via `--rollout_importance_sampling_mode` parameter):
|
||||
|
||||
| Mode | Description |
|
||||
|------|-------------|
|
||||
| `token_truncate` | Token-level truncation |
|
||||
| `token_mask` | Token-level masking |
|
||||
| `sequence_truncate` | Sequence-level truncation |
|
||||
| `sequence_mask` | Sequence-level masking |
|
||||
|
||||
The threshold is set via the `--rollout_importance_sampling_threshold` parameter.
|
||||
|
||||
## Metrics
|
||||
|
||||
To monitor the degree of training-inference mismatch during training, we add the following metrics to the logs (prefixed with `rollout_correction/`):
|
||||
|
||||
### 1. KL Divergence
|
||||
|
||||
KL divergence measures the deviation between the rollout policy and the training policy. Both metrics estimate $\text{KL}(\pi_{\text{vLLM}} \| \pi_\theta)$
|
||||
|
||||
**Direct estimator `kl`**:
|
||||
|
||||
$$
|
||||
\text{KL}(\pi_{\text{vLLM}} \| \pi_\theta) = \mathbb{E}_{\pi_{\text{vLLM}}}\left[ \log \frac{\pi_{\text{vLLM}}}{\pi_\theta} \right]
|
||||
$$
|
||||
|
||||
**K3 estimator `k3_kl`**:
|
||||
|
||||
$$
|
||||
\text{KL}(\pi_{\text{vLLM}} \| \pi_\theta) \approx \mathbb{E}_{\pi_{\text{vLLM}}}\left[ \rho - \log \rho - 1 \right], \quad \rho = \frac{\pi_\theta}{\pi_{\text{vLLM}}}
|
||||
$$
|
||||
|
||||
The K3 estimator is more numerically stable when KL values are small and is always non-negative.
|
||||
|
||||
### 2. Perplexity (PPL)
|
||||
|
||||
Perplexity measures the model's prediction uncertainty for a sequence:
|
||||
|
||||
$$
|
||||
\text{PPL} = \exp\left( -\frac{1}{|y|} \sum_{t=1}^{|y|} \log p(y_t) \right)
|
||||
$$
|
||||
|
||||
Related metrics:
|
||||
- `training_ppl` / `training_log_ppl`: Training policy PPL and its logarithm
|
||||
- `rollout_ppl` / `rollout_log_ppl`: Rollout policy PPL and its logarithm
|
||||
- `log_ppl_diff`: Log PPL difference, positive value means training policy assigns lower probability
|
||||
- `log_ppl_abs_diff`: Absolute log PPL difference
|
||||
- `log_ppl_diff_max` / `log_ppl_diff_min`: Max/min of log PPL difference
|
||||
- `ppl_ratio`: PPL ratio $\frac{\text{PPL}_{\text{training}}}{\text{PPL}_{\text{rollout}}}$
|
||||
|
||||
### 3. χ² Divergence (Chi-squared Divergence)
|
||||
|
||||
χ² divergence measures the variance of importance sampling weights:
|
||||
|
||||
$$
|
||||
\chi^2(\pi_\theta \| \pi_{\text{vLLM}}) = \mathbb{E}_{\pi_{\text{vLLM}}}\left[ \rho^2 \right] - 1, \quad \rho = \frac{\pi_\theta}{\pi_{\text{vLLM}}}
|
||||
$$
|
||||
|
||||
- `chi2_token`: Token-level χ² divergence, $\mathbb{E}[\rho_t^2] - 1$
|
||||
- `chi2_seq`: Sequence-level χ² divergence (geometric mean based), $\mathbb{E}[\rho_{\text{geo}}^2] - 1$, where $\rho_{\text{geo}} = \exp(\frac{1}{T}\sum_t \log \rho_t)$
|
||||
|
||||
Higher χ² divergence indicates larger IS weight variance and less stable training. `chi2_seq` uses geometric mean instead of product, making it comparable in scale to `chi2_token`.
|
||||
|
||||
### 4. Effective Sample Size (ESS)
|
||||
|
||||
Effective sample size measures the number of samples that actually contribute after importance sampling:
|
||||
|
||||
$$
|
||||
\text{ESS} = \frac{1}{\mathbb{E}\left[\left(\frac{w}{\mathbb{E}[w]}\right)^2\right]}
|
||||
$$
|
||||
|
||||
A larger ESS value (closer to 1) indicates more uniform importance sampling weight distribution and higher sample utilization efficiency. When all weights are equal (on-policy), ESS = 1; when weights differ significantly (severely off-policy), ESS becomes small.
|
||||
|
||||
### 5. IS Weight Statistics
|
||||
|
||||
- `is_weight_mean`: Average importance sampling weight, ideal value is 1.0
|
||||
- `clipped_frac`: Fraction of samples that were truncated or masked
|
||||
|
||||
|
||||
## Usage
|
||||
|
||||
### Logging Diagnostic Metrics Only
|
||||
|
||||
If you only want to monitor the degree of training-inference mismatch without enabling importance sampling correction, you can set:
|
||||
|
||||
```
|
||||
--log_rollout_offpolicy_metrics true
|
||||
```
|
||||
|
||||
This will log all diagnostic metrics (KL, PPL, χ², etc.) without modifying the loss function.
|
||||
|
||||
### Enabling Importance Sampling Correction
|
||||
|
||||
Enable the correction mechanism with the following parameters:
|
||||
|
||||
```
|
||||
--rollout_importance_sampling_mode (default None)
|
||||
--rollout_importance_sampling_threshold (default 2)
|
||||
```
|
||||
|
||||
When `rollout_importance_sampling_mode` is set, diagnostic metrics are automatically logged without needing to set `log_rollout_offpolicy_metrics`.
|
||||
|
||||
## Off-Policy Sequence Masking
|
||||
|
||||
In addition to importance sampling correction, you can use **Off-Policy Sequence Masking** to address training-inference mismatch. This technique comes from the [DeepSeek-V3.2 paper](https://arxiv.org/abs/2512.02556).
|
||||
|
||||
### Principle
|
||||
|
||||
The core idea of Off-Policy Sequence Masking is: when the current policy deviates significantly from the old policy (rollout or old policy), directly discard (mask) that sequence from loss computation. This approach specifically targets **sequences with negative advantage**, as these are more likely to cause training instability when policy shift is large.
|
||||
|
||||
Specifically, for each sequence, compute:
|
||||
|
||||
$$
|
||||
\delta_i = \frac{1}{|y_i|} \sum_{t=1}^{|y_i|} \bigl( \log \pi_{\text{old}}(y_{i,t}|x, y_{i,<t}) - \log \pi_\theta(y_{i,t}|x, y_{i,<t}) \bigr)
|
||||
$$
|
||||
|
||||
Sequence $i$ will be masked when the following conditions are met (mean taken over completion tokens, i.e., where `completion_mask=1`):
|
||||
1. $\delta_i > \tau$
|
||||
2. **AND** $\hat{A}_i < 0$
|
||||
|
||||
Where:
|
||||
- $\pi_{\text{old}}$ preferentially uses `rollout_per_token_logps` (logprobs from rollout/behavior policy); if unavailable, falls back to `old_per_token_logps`
|
||||
- $\tau$ is the user-set threshold (`--off_policy_sequence_mask_delta`, default None = disabled)
|
||||
|
||||
## References
|
||||
|
||||
1. https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda
|
||||
2. https://fengyao.notion.site/off-policy-rl
|
||||
3. https://github.com/volcengine/verl/blob/main/verl/trainer/ppo/rollout_corr_helper.py
|
||||
4. [DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models](https://arxiv.org/abs/2512.02556)
|
||||
@@ -0,0 +1,35 @@
|
||||
# TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling
|
||||
|
||||
Author: [li2zhi](https://github.com/li2zhi)
|
||||
|
||||
## Principle Introduction
|
||||
[TreePO paper](https://arxiv.org/abs/2508.17445) proposes a tree-structured modeling method. This method organizes sequence generation into a segmented tree structure search. Through dynamic branching, backtracking, and early termination mechanisms, it significantly improves the reuse rate of the key-value cache, thereby reducing computational overhead, while maintaining or even enhancing the diversity of exploration.
|
||||
|
||||
|
||||
|
||||
## Implementation Details
|
||||
[TreePO implementation example](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/treepo/tree_rollout_plugin.py), which references the [official implementation](https://github.com/multimodal-art-projection/TreePO/blob/main/recipe/treepo/vllm_rollout_tree.py) provides sample code for a TreePO training plugin，covering logic related to multi-round interactions, termination judgment, and branch rollback.
|
||||
|
||||
**Note:** In actual use, you need to rewrite the logic of methods such as step and check_finished according to your own scenario requirements to ensure that they can execute as expected in the custom scenario.
|
||||
For information on the design and use of custom rewards, you can refer to the implementation of [DeepEyes](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py).
|
||||
|
||||
The complete training script can be found at [script](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo/plugin/treepo/tree_rollout.sh).
|
||||
|
||||
## Test Data
|
||||
> model: Qwen/Qwen2.5-0.5B
|
||||
> dataset: AI-MO/NuminaMath-TIR
|
||||
> subset size: 1,000 samples
|
||||
> 1 GPU for training, 1 GPU for inference
|
||||
|
||||
| \ | batch_size | num_generation | max_tree_depth | global_step | total inference calls | saving ratio | train_speed(iter/s) | improvement rate |
|
||||
| ----------------------- | ---------- | -------------- | -------------- | ----------- | --------------------- | ------------ | ------------------- | ---------------- |
|
||||
| original implementation | 8 | 8 | 4 | 200 | 5965 | 0.00% | 0.292436 | 0.00% |
|
||||
| tree(max_divergence=3) | 8 | 8 | 4 | 200 | 3678 | 38.34% | 0.31819 | 8.81% |
|
||||
| | | | | | | | | |
|
||||
| original implementation | 8 | 8 | 5 | 105 | 4312 | 0.00% | 0.261324 | 0.00% |
|
||||
| tree(max_divergence=2) | 8 | 8 | 5 | 105 | 2513 | 52.69% | 0.336639 | 28.82% |
|
||||
| tree(max_divergence=3) | 8 | 8 | 5 | 105 | 2990 | 30.66% | 0.308791 | 18.16% |
|
||||
| | | | | | | | | |
|
||||
| original implementation | 8 | 8 | 6 | 105 | 5202 | 0.00% | 0.24832 | 0.00% |
|
||||
| tree(max_divergence=2) | 8 | 8 | 6 | 105 | 3348 | 35.64% | 0.27755 | 11.77% |
|
||||
| tree(max_divergence=3) | 8 | 8 | 6 | 105 | 3888 | 25.26% | 0.272339 | 9.67% |
|
||||
@@ -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:
|
||||
|
||||
|
||||
|
||||
## 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
|
||||
```
|
||||
@@ -0,0 +1,401 @@
|
||||
# GRPO
|
||||
|
||||
[GRPO (Group Relative Policy Optimization)](https://arxiv.org/abs/2402.03300) leverages intra-group relative advantage calculations to replace the independent value model in the PPO algorithm and directly incorporates KL divergence penalties into the loss function to improve training stability.
|
||||
|
||||
## Algorithm Overview
|
||||
|
||||
GRPO Objective Function is defined as
|
||||
$
|
||||
{\scriptstyle
|
||||
\begin{aligned}
|
||||
\mathcal{J}_{G R P O}(\theta) & =\mathbb{E}_{\left[q \sim P(Q),\left\{o_i\right\}_{i=1}^G \sim \pi_{\theta_{o l d}}(O \mid q)\right]} \\
|
||||
& \frac{1}{G} \sum_{i=1}^G \frac{1}{\left|o_i\right|} \sum_{t=1}^{\left|o_i\right|}\left\{\min \left[\frac{\pi_\theta\left(o_{i, t} \mid q, o_{i,<t}\right)}{\pi_{\theta_{o l d}}\left(o_{i, t} \mid q, o_{i,<t}\right)} \hat{A}_{i, t}, \operatorname{clip}\left(\frac{\pi_\theta\left(o_{i, t} \mid q, o_{i,<t}\right)}{\pi_{\theta_{o l d}}\left(o_{i, t} \mid q, o_{i,<t}\right)}, 1-\varepsilon, 1+\varepsilon\right) \hat{A}_{i, t}\right]-\beta \mathbb{D}_{K L}\left[\pi_\theta| | \pi_{r e f}\right]\right\}
|
||||
\end{aligned}
|
||||
}
|
||||
$
|
||||
|
||||
The advantage function is defined as
|
||||
|
||||
$
|
||||
\hat{A}_{i,t} = \frac{R_i - \text{mean}(\{R_j\}_{j=1}^G)}{\text{std}(\{R_j\}_{j=1}^G)}
|
||||
$
|
||||
|
||||
|
||||
<details> <summary>GRPO Algorithm Pseudocode</summary>
|
||||
|
||||
```python
|
||||
# ========== 1. Rollout Generation Phase ==========
|
||||
prompt = "Question: Which is bigger? 9.11 or 9.9?"
|
||||
|
||||
# Generate multiple completions through parallel sampling
|
||||
completions = rollout_function(
|
||||
model=current_policy_model,
|
||||
prompt=prompt,
|
||||
num_generations=8, # Hyperparameter: number of samples per prompt
|
||||
temperature=1.0 # Hyperparameter: sampling diversity
|
||||
)
|
||||
"""
|
||||
completions = [
|
||||
(completion 1) "The larger number is 9.9...",
|
||||
(completion 2) "9.11 is bigger than...",
|
||||
...
|
||||
(completion 8) "After calculation, 9.9..."
|
||||
]
|
||||
"""
|
||||
|
||||
# ========== 2. Reward Calculation Phase ==========
|
||||
# Evaluate generated completions using reward model
|
||||
rewards = reward_function(
|
||||
completions=completions,
|
||||
ground_truth="9.9" # Expected correct answer
|
||||
)
|
||||
"""
|
||||
rewards = [
|
||||
(reward 1) 1.0, # Correct answer
|
||||
(reward 2) 0.0, # Incorrect
|
||||
...
|
||||
(reward 8) 1.0 # Correct
|
||||
]
|
||||
"""
|
||||
|
||||
# Normalize rewards to advantages
|
||||
rewards_mean = mean(rewards) # μ = 0.5
|
||||
rewards_std = std(rewards) # σ = 0.25
|
||||
advantages = (rewards - rewards_mean) / (rewards_std + 1e-8) # Standardization
|
||||
"""
|
||||
advantages = [
|
||||
(advantage 1) 2.0, # (1.0 - 0.5)/0.25
|
||||
(advantage 2) -2.0,
|
||||
...
|
||||
(advantage 8) 2.0
|
||||
]
|
||||
"""
|
||||
|
||||
# ========== 3. Policy Optimization Phase ==========
|
||||
# Get token-level log probabilities from different models
|
||||
current_logps = get_per_token_logps(current_policy_model, prompt, completions) # π_θ
|
||||
old_logps = get_per_token_logps(old_policy_model, prompt, completions) # π_θ_old
|
||||
ref_logps = get_per_token_logps(reference_model, prompt, completions) # π_ref
|
||||
|
||||
# PPO Clipped Objective
|
||||
is_ratio = exp(current_logps - old_logps) # Importance sampling ratio: e^(π_θ - π_θ_old)
|
||||
clipped_ratio = clip(is_ratio, 1-ε, 1+ε) # ε=0.2 typically
|
||||
|
||||
# Policy gradient term (dual form)
|
||||
policy_loss = -mean(
|
||||
minimum(is_ratio * advantages, # Unclipped objective
|
||||
clipped_ratio * advantages) # Clipped objective
|
||||
)
|
||||
|
||||
# KL Divergence Penalty (K3 estimator)
|
||||
# KL(π_θ||π_ref) ≈ e^(logπ_ref - logπ_θ) - (logπ_ref - logπ_θ) - 1
|
||||
kl_penalty = beta * mean(
|
||||
exp(ref_logps - current_logps) -
|
||||
(ref_logps - current_logps) - 1
|
||||
)
|
||||
|
||||
# Total Loss = Policy Loss + KL Penalty
|
||||
total_loss = policy_loss + kl_penalty
|
||||
|
||||
# ========== 4. Update Rule ==========
|
||||
# Apply gradient descent to minimize total_loss
|
||||
optimizer.zero_grad()
|
||||
total_loss.backward()
|
||||
optimizer.step()
|
||||
```
|
||||
</details>
|
||||
|
||||
|
||||
For training script examples, refer to [examples](https://github.com/modelscope/ms-swift/tree/main/examples/train/grpo).
|
||||
|
||||
For GRPO parameters, refer to the [documentation](../../../Instruction/Command-line-parameters.md#grpo-arguments)
|
||||
|
||||
## Cluster Support
|
||||
|
||||
|
||||
|
||||
The GRPO training framework supports integration with high-performance inference engines (e.g., vLLM) to accelerate the sampling process, offering the following two deployment modes:
|
||||
|
||||
### 1. Colocate (Internal) Mode
|
||||
Training and inference share GPU resources, with the inference service launched internally within the Trainer.
|
||||
|
||||
Startup parameters
|
||||
```bash
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate
|
||||
```
|
||||
|
||||
#### Memory Optimization Solutions in Colocate Mode
|
||||
|
||||
When running in Colocate mode, out-of-memory (OOM) issues may frequently occur. Below are several effective memory optimization methods and parameter configurations:
|
||||
|
||||
1. Reduce the vllm_gpu_memory_utilization parameter.
|
||||
|
||||
2. During the training phase, release the GPU memory occupied by vLLM:
|
||||
|
||||
|
||||
```bash
|
||||
--sleep_level 1
|
||||
```
|
||||
|
||||
3. During the vLLM inference phase, release the GPU memory occupied by the model and optimizer:
|
||||
|
||||
```bash
|
||||
--offload_optimizer true \
|
||||
--offload_model true \
|
||||
```
|
||||
|
||||
4. Use Tensor Parallelism in vLLM:
|
||||
|
||||
```bash
|
||||
--vllm_tensor_parallel_size [tp_size]
|
||||
```
|
||||
|
||||
5. Gather model weights in batches (when synchronizing vLLM weights under zero3):
|
||||
|
||||
```bash
|
||||
--move_model_batches [批次数量]
|
||||
```
|
||||
|
||||
6. Store Megatron exported HF format weights for vLLM updates in CPU main memory to reduce GPU memory usage:
|
||||
|
||||
```bash
|
||||
--offload_bridge true
|
||||
```
|
||||
|
||||
### 2. Async(External) Mode
|
||||
|
||||
Training and inference resources are separated, with a dedicated inference server deployed.
|
||||
|
||||
Use the `swift rollout` command to deploy the vLLM server (currently only supports vLLM backend):
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
swift rollout \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--vllm_tensor_parallel_size 2 \
|
||||
--vllm_data_parallel_size 1
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
swift rollout \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--vllm_tensor_parallel_size 2 \
|
||||
--vllm_data_parallel_size 1
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
swift rollout \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--vllm_tensor_parallel_size 2 \
|
||||
--vllm_data_parallel_size 2
|
||||
```
|
||||
For more rollout parameters, refer to the [vllm arguments](../../../Instruction/Command-line-parameters.md#vllm-arguments) and [rollout arguments](../../../Instruction/Command-line-parameters.md#rollout-arguments)
|
||||
|
||||
Note: When set `vllm_use_async_engine`, enabling only DP (Data Parallelism) may cause errors. [Related issue](https://github.com/vllm-project/vllm/issues/18567). If errors occur, try enabling both TP (Tensor Parallelism) and DP or upgrading vLLM.
|
||||
|
||||
To configure the external vLLM server during training, use the following parameters:
|
||||
|
||||
```bash
|
||||
--use_vllm true \
|
||||
--vllm_mode server \
|
||||
--vllm_server_host <server_IP> \
|
||||
--vllm_server_port <service_port> \
|
||||
--vllm_server_timeout <timeout> \
|
||||
```
|
||||
|
||||
### Weight-Sync Acceleration
|
||||
|
||||
Setting the following parameters optimizes weight synchronization speed for LoRA training by syncing only the LoRA adapter weights instead of the full model weights.
|
||||
|
||||
> Note: This synchronization method may slightly impact vLLM inference speed.
|
||||
|
||||
```bash
|
||||
# rollout(server mode)
|
||||
swift rollout \
|
||||
--vllm_enable_lora true \
|
||||
--vllm_max_lora_rank xxx # match the lora_rank in the training script
|
||||
...
|
||||
|
||||
# grpo(colocate mode)
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--vllm_mode colocate \
|
||||
--vllm_enable_lora true \
|
||||
...
|
||||
|
||||
# megatron grpo(colocate mode)
|
||||
swift megatron rlhf \
|
||||
--rlhf_type grpo \
|
||||
--vllm_mode colocate \
|
||||
--vllm_enable_lora true \
|
||||
...
|
||||
```
|
||||
|
||||
**Multimodal ViT LoRA Sync:** If ViT LoRA is enabled during training (`freeze_vit false`),
|
||||
tower/connector LoRA support must also be enabled on the vLLM side.
|
||||
|
||||
pass via `vllm_engine_kwargs`:
|
||||
|
||||
```bash
|
||||
--vllm_engine_kwargs '{"enable_tower_connector_lora": true}'
|
||||
```
|
||||
|
||||
This is an experimental vLLM feature, currently supporting models such as Qwen2.5-VL and Qwen3-VL.
|
||||
For model-specific support details, see the [vLLM documentation](https://docs.vllm.ai/en/latest/features/lora/)
|
||||
and the [vLLM issue](https://github.com/vllm-project/vllm/issues/31479).
|
||||
|
||||
## logged metrics
|
||||
- completions/mean_length: The average length of generated completions.
|
||||
- completions/min_length: The minimum length among generated completions.
|
||||
- completions/max_length: The maximum length among generated completions.
|
||||
- completions/clipped_ratio: The proportion of completions that were truncated due to length limits.
|
||||
- reward/{reward_func_name}/mean: The average reward value for a specific reward function.
|
||||
- reward/{reward_func_name}/std: The standard deviation of the reward for a specific reward function.
|
||||
> Note: These two metrics are calculated across all completions.
|
||||
- reward: The overall average reward after applying reward_weights.
|
||||
- reward_std: The standard deviation of the overall reward within each batch after applying reward_weights.
|
||||
> Note: These two metrics are first computed within each group and then averaged (for mean/std) across groups.
|
||||
- frac_reward_zero_std: The proportion of samples in a generation batch where the reward standard deviation is zero, meaning there is almost no diversity in answers for that prompt (i.e., the rewards of all completions are same).
|
||||
- kl: The average KL divergence between the model and the reference model on completions. This is logged only if beta is nonzero.
|
||||
- clip_ratio/region_mean: The average proportion of tokens clipped by the CLIP operator across different sentences.
|
||||
- clip_ratio/low_mean: The average proportion of tokens clipped by the lower CLIP bound across different sentences.
|
||||
- clip_ratio/low_min: The minimum proportion of tokens clipped by the lower CLIP bound across different sentences.
|
||||
- clip_ratio/high_mean: The average proportion of tokens clipped by the upper CLIP bound across different sentences.
|
||||
- clip_ratio/high_max: The maximum proportion of tokens clipped by the upper CLIP bound across different sentences.
|
||||
> Note: If `overlong_filter` is enabled, the kl and clip_ratio metrics will exclude overlength samples.
|
||||
|
||||
If the `log_entropy` parameter is set, additional entropy-related metrics will be logged, including:
|
||||
- entropy/mean: the average entropy across different sentences
|
||||
- entropy/max: the maximum entropy among different sentences
|
||||
- entropy/min: the minimum entropy among different sentences
|
||||
> Note: Here, sentence entropy refers to the mean entropy of tokens in each completion.
|
||||
|
||||
If `top_entropy_quantile` is set to a value smaller than 1.0, the entropy threshold value will also be recorded:
|
||||
- entropy/threshold: Tokens with entropy below this value will be excluded from the loss calculation.
|
||||
|
||||
Training-inference consistency metrics, prefixed with rollout_correction, requires setting `log_rollout_offpolicy_metrics=true` or `rollout_importance_sampling_mode`:
|
||||
- `kl` / `k3_kl`: KL divergence between training policy and rollout policy (direct estimator / K3 estimator)
|
||||
- `training_ppl` / `rollout_ppl`: Perplexity of training policy and rollout policy
|
||||
- `log_ppl_diff`: Log PPL difference, reflects the degree of distribution shift
|
||||
- `ppl_ratio`: PPL ratio
|
||||
- `chi2_token` / `chi2_seq`: Token/Sequence-level χ² divergence
|
||||
|
||||
IS correction metrics (requires setting `rollout_importance_sampling_mode`):
|
||||
- `is_weight_mean`: Average importance sampling weight
|
||||
- `ess`: Effective Sample Size
|
||||
- `clipped_frac`: Fraction of samples that were truncated or masked
|
||||
|
||||
> For detailed explanation of training-inference consistency metrics, please refer to [Training-Inference-Mismatch](../AdvancedResearch/training_inference_mismatch.md)
|
||||
|
||||
If `log_completions` is set, the training dynamics will be saved in the output directory, including:
|
||||
- step: The training step at the time of logging.
|
||||
- prompt: The model input.
|
||||
- completion: The model's sampled answer.
|
||||
- {reward_func_name}: The specific reward(s).
|
||||
- entropy: The average token entropy (recorded if `log_entropy` is set).
|
||||
|
||||
Setting `report_to wandb/swanlab` will send training dynamics table to the respective platform.
|
||||
|
||||
If you want to log extra columns in the Table, populate the `metrics_to_gather` dictionary inside `GRPOTrainer._generate_and_score_completions`.
|
||||
|
||||
The trainer automatically detects and logs the following keys:
|
||||
|
||||
- image: image inputs for vision models(wandb only).
|
||||
- solution: the solution column from the dataset.
|
||||
|
||||
## FAQ
|
||||
|
||||
**1. Loss Equals Zero / Approaches Zero / Is Negative During Training**
|
||||
|
||||
This is normal behavior. For reference, see [issue](https://github.com/huggingface/open-r1/issues/239#issuecomment-2646297851).
|
||||
|
||||
---
|
||||
|
||||
**2. num_generations / Batch Size Related**
|
||||
|
||||
In GRPO, the batch size is measured in terms of completions (i.e., model-generated outputs). For example, setting `per_device_train_batch_size=8` means that each GPU processes 8 completions for loss calculation during training.
|
||||
|
||||
During the training phase, the total effective batch size in a full gradient accumulation step equals:
|
||||
|
||||
```python
|
||||
effective_batch_size = num_processes * per_device_train_batch_size * gradient_accumulation_steps
|
||||
```
|
||||
|
||||
During the sampling phase, the total batch size (completion-level) depends on the following:
|
||||
|
||||
- If generation_batch_size is set, the total equals generation_batch_size.
|
||||
- If steps_per_generation is set, the total equals per_device_train_batch_size * steps_per_generation * num_processes.
|
||||
- By default, steps_per_generation is set to gradient_accumulation_steps, and generation_batch_size equals the per_device_train_batch_size * steps_per_generation * num_processes = per_device_train_batch_size * gradient_accumulation_steps * num_processes = effective_batch_size.
|
||||
|
||||
During evaluation, the number of completions equals:
|
||||
|
||||
```
|
||||
num_processes * per_device_eval_batch_size
|
||||
```
|
||||
|
||||
The parameter `num_generations` must be divisible by the total batch size used in sampling and evaluation to ensure even distribution across devices.
|
||||
|
||||
**Example**
|
||||
|
||||
- num_processes = 8
|
||||
- per_device_train_batch_size = 4
|
||||
- gradient_accumulation_steps = 8
|
||||
- generation_batch_size = 512
|
||||
- num_generations = 64
|
||||
|
||||
|
||||
1. Total prompts needed for sampling: 512 / 64 = 8
|
||||
2. Generate 512 responses from the model per sampling step
|
||||
3. Model update batch size: 8 * 4 * 8 = 256
|
||||
|
||||
**3. Why did KL result in NaN?**
|
||||
|
||||
With `overlong_filter` enabled, all completions on a certain GPU were truncated.
|
||||
|
||||
**4. How is the training steps calculated?**
|
||||
|
||||
Refer to [issue](https://github.com/modelscope/ms-swift/issues/3912).
|
||||
|
||||
**5. Why is the clip ratio always 0?**
|
||||
|
||||
The core purpose of the clip mechanism is to limit the magnitude of policy updates to prevent policy performance collapse due to excessively large updates (i.e., a drastic decline in performance after policy updates). The specific formula for the clip operation is as follows:
|
||||
|
||||
$$
|
||||
L_{\text{CLIP}}(\theta) = \mathbb{E}_{t} \left[ \min\left(r_{t}(\theta) \hat{A}_{t}, \text{clip}(r_{t}(\theta), 1 - \epsilon, 1 + \epsilon) \hat{A}_{t} \right) \right]
|
||||
$$
|
||||
|
||||
Where: $r_{t}(\theta) = \frac{\pi_{\theta}(a_{t} \mid s_{t})}{\pi_{\text{old}}(a_{t} \mid s_{t})}$ is the importance sampling ratio, which measures the difference between the new and old policy. $\hat{A}_{t}$ is the advantage function, representing the relative return of the action. $\epsilon$ is used to limit the deviation range of $r_{t}(\theta)$.
|
||||
|
||||
In the on-policy training process, since each update uses data generated by the latest policy, the new and old policies are the same, i.e., $\pi_{\theta} = \pi_{\text{old}}$.
|
||||
|
||||
Thus, the importance sampling ratio is always 1, and the clip operation does not take effect.
|
||||
|
||||
The algorithm becomes off-policy (near-on-policy) under the following parameter settings:
|
||||
1. num_iterations > 1, or
|
||||
2. gradient_accumulation_steps % steps_per_generation != 0
|
||||
|
||||
Refer to [issue](https://github.com/huggingface/open-r1/issues/239#issuecomment-2646297851).
|
||||
|
||||
**6. How to set the training `mini-batch size`**
|
||||
|
||||
In GRPO training, we can configure mini-batch updates in the following two ways:
|
||||
|
||||
- Set `generation_batch_size` to be an integer multiple of the training global batch size (effective_batch_size).
|
||||
- Or set `steps_per_generation` to be an integer multiple of `gradient_accumulation_steps`.
|
||||
|
||||
Typical configuration example:
|
||||
- When configured with:
|
||||
steps_per_generation = 16, gradient_accumulation_steps = 8, mini_batch_size = steps_per_generation / gradient_accumulation_steps = 2. The results from 1 rollout will be split into 2 mini-batch updates.
|
||||
|
||||
**7. Difference between swift deploy and swift rollout**
|
||||
|
||||
- swift deploy is primarily used for model deployment and inference. It supports various engines such as PT, vLLM, and SGLang, and is compatible with streaming inference as well as the OpenAI API format.
|
||||
|
||||
- swift rollout, on the other hand, is dedicated to GRPO rollout acceleration. Currently, it only supports the vLLM engine and comes with built-in automatic weight synchronization.
|
||||
|
||||
**8. How to disable the KL loss term**
|
||||
|
||||
Set the parameter `--beta 0` to disable KL loss calculation. The reference model (ref model) will not be loaded in this case.
|
||||
|
||||
|
||||
## RL WeChat Group
|
||||
|
||||
<img src="https://raw.githubusercontent.com/modelscope/ms-swift/main/docs/resources/wechat/grpo.png" width="250">
|
||||
@@ -0,0 +1,6 @@
|
||||
Get Started
|
||||
===============
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
GRPO.md
|
||||
@@ -0,0 +1,19 @@
|
||||
GRPO
|
||||
===============
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
:caption: Get Started
|
||||
|
||||
GetStarted/index.rst
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
:caption: Developer Guide
|
||||
|
||||
DeveloperGuide/index.rst
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
:caption: Advanced Research
|
||||
|
||||
AdvancedResearch/index.rst
|
||||
Reference in New Issue
Block a user