Skip to content

Commit ffb41eb

Browse files
qychen2001hjh0119
andauthored
feat(grpo): add RLSD self-distilled advantage reweighting (#9758)
* feat(grpo): add RLSD self-distilled advantage reweighting Introduce RLSD (Self-Distilled RLVR) token-level advantage reweighting for GRPO, enabled via `--advantage_reweight rlsd`. The per-sequence advantage is redistributed inside each trajectory using the teacher-vs-student logprob gap as a sign-aware, clipped multiplicative weight, without flipping the reward sign. The teacher reuses the OPSD self-distillation forward (current policy conditioned on a privileged `teacher_prompt` column), and a frozen local teacher via `--teacher_model` is also supported. - add `apply_rlsd_reweight` in rl_core/advantage.py - add RLSD args and lambda warmup/decay schedule to the GRPO trainer - add `_check_rlsd` fail-fast validation in RLHFArguments - add unit tests in tests/train/test_rlsd_reweight.py * ci(workflows): add Qoder code review and assistant * docs(grpo): document RLSD advantage reweighting parameters Add the new RLSD (Self-Distilled RLVR) CLI parameters to the English and Chinese command-line references: advantage_reweight, rlsd_lambda, rlsd_reweight_clip_range, rlsd_lambda_warmup_steps, rlsd_lambda_decay_steps and rlsd_negative_only, including defaults, valid ranges, the reweight formula, enable requirements and the two teacher modes. * update core-review.yaml Added separate jobs for trusted and external code reviews in GitHub Actions workflow. --------- Co-authored-by: jinghanhu <hujinghan.hjh@alibaba-inc.com>
1 parent 3470f9c commit ffb41eb

9 files changed

Lines changed: 423 additions & 13 deletions

File tree

.github/workflows/assistant.yaml

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
name: Qoder Assistant
2+
3+
on:
4+
issue_comment:
5+
types: [created]
6+
pull_request_review_comment:
7+
types: [created]
8+
9+
jobs:
10+
qoder-assistant:
11+
if: |
12+
contains(github.event.comment.body, '@qoder') &&
13+
!endsWith(github.event.comment.user.login, '[bot]')
14+
runs-on: ubuntu-latest
15+
permissions:
16+
contents: read
17+
issues: write
18+
pull-requests: write
19+
id-token: write
20+
21+
steps:
22+
- name: Checkout repository
23+
uses: actions/checkout@v4
24+
with:
25+
fetch-depth: 0
26+
27+
- name: Build Arguments
28+
id: build_args
29+
run: |
30+
ARGS="REPO: ${{ github.repository }}
31+
REQUEST_SOURCE: ${{ github.event_name }}
32+
THREAD_ID: ${{ github.event.comment.node_id }}
33+
COMMENT_ID: ${{ github.event.comment.id }}
34+
AUTHOR: ${{ github.event.comment.user.login }}
35+
BODY: ${{ github.event.comment.body }}
36+
URL: ${{ github.event.comment.html_url }}
37+
IS_PR: ${{ github.event.issue.pull_request != null || github.event_name == 'pull_request_review_comment' }}
38+
ISSUE_OR_PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}"
39+
40+
if [ -n "${{ github.event.comment.pull_request_review_id }}" ]; then
41+
ARGS="$ARGS
42+
REVIEW_ID: ${{ github.event.comment.pull_request_review_id }}"
43+
fi
44+
45+
if [ -n "${{ github.event.comment.in_reply_to_id }}" ]; then
46+
ARGS="$ARGS
47+
REPLY_TO_COMMENT_ID: ${{ github.event.comment.in_reply_to_id }}"
48+
fi
49+
50+
echo "args<<EOF" >> $GITHUB_OUTPUT
51+
echo "$ARGS" >> $GITHUB_OUTPUT
52+
echo "EOF" >> $GITHUB_OUTPUT
53+
54+
- name: Run Qoder Assistant
55+
uses: QoderAI/qoder-action@v0
56+
with:
57+
qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }}
58+
prompt: |
59+
/assistant
60+
${{ steps.build_args.outputs.args }}

.github/workflows/code-review.yaml

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
name: Qoder Auto Code Review
2+
3+
on:
4+
pull_request_target:
5+
types: [opened, synchronize, reopened]
6+
7+
jobs:
8+
qoder-review-trusted:
9+
if: contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.pull_request.author_association)
10+
runs-on: ubuntu-latest
11+
permissions:
12+
contents: read
13+
pull-requests: write
14+
id-token: write
15+
16+
steps:
17+
- name: Checkout PR head
18+
uses: actions/checkout@v4
19+
with:
20+
fetch-depth: 0
21+
ref: ${{ github.event.pull_request.head.sha }}
22+
23+
- name: Run Qoder Code Review
24+
uses: QoderAI/qoder-action@v0
25+
with:
26+
qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }}
27+
prompt: |
28+
/review-pr
29+
REPO:${{ github.repository }} PR_NUMBER:${{ github.event.pull_request.number }}
30+
31+
qoder-review-external:
32+
if: ${{ !contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.pull_request.author_association) }}
33+
runs-on: ubuntu-latest
34+
environment: qoder-review
35+
permissions:
36+
contents: read
37+
pull-requests: write
38+
id-token: write
39+
40+
steps:
41+
- name: Checkout PR head
42+
uses: actions/checkout@v4
43+
with:
44+
fetch-depth: 0
45+
ref: ${{ github.event.pull_request.head.sha }}
46+
47+
- name: Run Qoder Code Review
48+
uses: QoderAI/qoder-action@v0
49+
with:
50+
qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }}
51+
prompt: |
52+
/review-pr
53+
REPO:${{ github.repository }} PR_NUMBER:${{ github.event.pull_request.number }}

docs/source/Instruction/Command-line-parameters.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -679,6 +679,12 @@ reward模型参数将在PPO、GRPO中使用;teacher模型参数在GKD与GRPO
679679
- 注意:`gdpo` 模式不支持 `kl_in_reward=True`,若同时设置会自动将 `kl_in_reward` 设为 `False`
680680
- GDPO 适用于多奖励优化场景:当使用多个奖励函数时,GDPO 会对每个奖励函数分别在组内进行标准化(减均值、除标准差),然后使用 `reward_weights` 进行加权求和,最后再进行批次级别的标准化。这种方式可以更好地保留各个奖励的相对差异,避免不同奖励组合坍塌成相同的 advantage 值。
681681
- teacher_kl_coef: OPD-RL中teacher_kl的系数,即 `adv_t = base_adv + teacher_kl_coef * teacher_kl`。默认为 1.0。
682+
- advantage_reweight: 优势重加权算法,默认为None(不启用)。可选项为 `rlsd`([RLSD, Self-Distilled RLVR](https://arxiv.org/abs/2604.03128)):利用 teacher 与 student 在相同采样 token 上的对数概率差异,将 GRPO 的序列级优势在 token 维度重新分配,即 `delta_t = logP_T(y_t) - logP_S(y_t)`、`w_t = exp(sign(A) * delta_t)`、`reweight = (1 - λ) + λ * clip(w_t)`、`A_hat_t = A * reweight`;reweight 恒为正,不会翻转环境奖励的符号。启用后需满足:已设置 `reward_funcs`(reweight 方向依赖优势符号 `sign(A)`),且数据集包含 `teacher_prompt` 列(特权 prompt,如 问题 + 参考解答/标准答案);与 `use_liger_kernel`、`loss_type in [real, fipo]`、`off_policy_sequence_mask_delta`、`teacher_model_server` 不兼容。teacher 模式:不设置 `--teacher_model` 时为动态自蒸馏(teacher = 当前策略在特权 prompt 上打分);设置 `--teacher_model <ckpt>` 时为冻结本地 teacher(在 no_grad 下打分,训练中不更新)。
683+
- rlsd_lambda: RLSD 混合权重,`reweight = (1 - λ) + λ * clip(w_t)``0` 退化为普通 GRPO,`1` 为完整 RLSD 重加权。取值需在 [0, 1],默认为 0.5。
684+
- rlsd_reweight_clip_range: 证据权重裁剪范围 eps_w,将 `w_t = exp(sign(A) * delta_t)` 裁剪到 `[1 - eps_w, 1 + eps_w]`,避免 teacher/student 对数概率差异造成过大的倍率。需 >= 0,默认为 0.2。
685+
- rlsd_lambda_warmup_steps: 将 λ 从 0 线性 warmup 到 `rlsd_lambda` 的步数。默认为 0(不进行 warmup)。
686+
- rlsd_lambda_decay_steps: 在 warmup 之后,将 λ 从 `rlsd_lambda` 线性衰减到 0 的步数。默认为 0(不衰减,λ 保持恒定)。
687+
- rlsd_negative_only: 是否仅对优势为负(`A < 0`,即错误回答)的序列进行重加权;优势非负的序列保留普通 GRPO 优势。默认为 False。
682688
- sync_ref_model: 是否定期同步ref_model,默认为False。
683689
- ref_model_mixup_alpha: 控制在更新过程中model和先前ref_model之间的混合。更新公式为 $π_{ref} = α * π_θ + (1 - α) * π_{ref_{prev}}$。默认为0.6。
684690
- ref_model_sync_steps: 同步频率,默认为512。

docs/source_en/Instruction/Command-line-parameters.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -696,6 +696,12 @@ The hyperparameters for the reward function can be found in the [Built-in Reward
696696
- Note: `gdpo` mode does not support `kl_in_reward=True`. If both are set, `kl_in_reward` will be automatically set to `False`.
697697
- GDPO is designed for multi-reward optimization: When using multiple reward functions, GDPO normalizes each reward function separately within groups (subtract mean, divide by std), then performs weighted aggregation using `reward_weights`, and finally applies batch-level normalization. This approach better preserves the relative differences between rewards and prevents different reward combinations from collapsing into identical advantage values.
698698
- teacher_kl_coef: Coefficient for teacher KL in OPD-RL, i.e. `adv_t = base_adv + teacher_kl_coef * teacher_kl`. Default is 1.0.
699+
- advantage_reweight: Advantage reweighting algorithm. Default is None (disabled). Option: `rlsd` ([RLSD, Self-Distilled RLVR](https://arxiv.org/abs/2604.03128)), which redistributes the per-sequence GRPO advantage across tokens using the teacher-vs-student log-probability gap on the same sampled tokens, i.e. `delta_t = logP_T(y_t) - logP_S(y_t)`, `w_t = exp(sign(A) * delta_t)`, `reweight = (1 - λ) + λ * clip(w_t)`, `A_hat_t = A * reweight`. The reweight is strictly positive and never flips the sign of the environment reward. When enabled it requires `reward_funcs` (the reweight direction uses `sign(A)`) and a `teacher_prompt` column in the dataset (the privileged prompt, e.g. question + reference solution / ground-truth answer); it is incompatible with `use_liger_kernel`, `loss_type in [real, fipo]`, `off_policy_sequence_mask_delta`, and `teacher_model_server`. Teacher modes: without `--teacher_model` it runs dynamic self-distillation (teacher = current policy scored on the privileged prompt); with `--teacher_model <ckpt>` it uses a frozen local teacher (scored under no_grad, never updated during training).
700+
- rlsd_lambda: RLSD mixing weight, `reweight = (1 - λ) + λ * clip(w_t)`. `0` reduces to plain GRPO, `1` is full RLSD reweighting. Must be in [0, 1]. Default is 0.5.
701+
- rlsd_reweight_clip_range: Evidence-weight clip range eps_w. Clips `w_t = exp(sign(A) * delta_t)` to `[1 - eps_w, 1 + eps_w]` to prevent overly large multipliers from the teacher/student log-prob gap. Must be >= 0. Default is 0.2.
702+
- rlsd_lambda_warmup_steps: Number of steps to linearly warm up λ from 0 to `rlsd_lambda`. Default is 0 (no warmup).
703+
- rlsd_lambda_decay_steps: Number of steps to linearly decay λ from `rlsd_lambda` to 0 after warmup. Default is 0 (no decay; λ stays constant).
704+
- rlsd_negative_only: Whether to reweight only sequences with negative advantage (`A < 0`, i.e. incorrect responses); sequences with non-negative advantage keep plain GRPO advantages. Default is False.
699705
- sync_ref_model: Whether to synchronize the reference model. Default is False.
700706
- ref_model_mixup_alpha: The Parameter controls the mix between the current policy and the previous reference policy during updates. The reference policy is updated according to the equation: $π_{ref} = α * π_θ + (1 - α) * π_{ref_{prev}}$. Default is 0.6.
701707
- ref_model_sync_steps:The parameter determines how frequently the current policy is synchronized with the reference policy. Default is 512.

swift/arguments/rlhf_args.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -568,6 +568,62 @@ def _check_grpo(self):
568568
raise NotImplementedError('Currently, async_generate is not supported with multi-turn functionality.')
569569

570570
self._check_opd_rl()
571+
self._check_rlsd()
572+
573+
def _check_rlsd(self):
574+
"""Validate RLSD (Self-Distilled RLVR) advantage reweighting parameters.
575+
576+
RLSD reweights the per-sequence GRPO advantage per token using the teacher-vs-student
577+
logprob gap. It relies on an OPSD teacher forward driven by a per-sample ``teacher_prompt``
578+
column (privileged context, e.g. question + reference solution / GT answer). It also needs
579+
an environment reward (``reward_funcs``) because the reweight direction uses ``sign(A)``.
580+
581+
Two teacher modes are supported (mirroring the reference's ``freeze_teacher_model`` flag):
582+
583+
* **Dynamic self-distillation** (no ``--teacher_model``, default): teacher = current policy
584+
evaluated on the privileged prompt (paper's same-model, two-contexts Δ_t; equivalent to
585+
the reference's ``freeze_teacher_model=false``).
586+
* **Frozen local teacher** (``--teacher_model <ckpt>``): a separate model scores the
587+
privileged prompt under ``no_grad`` and is never in the optimizer, so it stays frozen
588+
throughout training. Point ``--teacher_model`` at the initial policy checkpoint to
589+
reproduce the reference's ``freeze_teacher_model=true`` (frozen θ₀ anchor).
590+
"""
591+
if self.advantage_reweight != 'rlsd':
592+
return
593+
if not (0.0 <= self.rlsd_lambda <= 1.0):
594+
raise ValueError(f'rlsd_lambda must be in [0, 1], got {self.rlsd_lambda}.')
595+
if self.rlsd_reweight_clip_range < 0:
596+
raise ValueError(f'rlsd_reweight_clip_range (eps_w) must be >= 0, got {self.rlsd_reweight_clip_range}.')
597+
if not self.reward_funcs:
598+
raise ValueError('advantage_reweight=rlsd requires reward_funcs: the reweight direction uses '
599+
'sign(A) of the GRPO advantage, which needs an environment reward.')
600+
if self.use_liger_kernel:
601+
raise ValueError('advantage_reweight=rlsd is not compatible with use_liger_kernel '
602+
'(the fused loss path bypasses the per-token RLSD reweight).')
603+
if self.loss_type in ['real', 'fipo']:
604+
raise ValueError(f'advantage_reweight=rlsd does not support loss_type={self.loss_type!r} '
605+
'(it reduces the advantage to a per-sequence scalar).')
606+
if self.off_policy_sequence_mask_delta is not None:
607+
raise ValueError('advantage_reweight=rlsd does not support off_policy_sequence_mask_delta.')
608+
# Local frozen teacher (--teacher_model) is allowed: it drives the OPSD teacher forward on
609+
# the privileged prompt via the same code path as OPD-RL, and use_rlsd bypasses OPD-RL's
610+
# additive term (grpo_trainer.py:266) so there is no double-count. The API-teacher path
611+
# (--teacher_model_server) scores tokens over a remote inference call whose tokenizer
612+
# alignment with the student is not guaranteed, so keep it gated for RLSD's per-token reweight.
613+
if self.teacher_model_server is not None:
614+
raise ValueError('advantage_reweight=rlsd does not support --teacher_model_server (remote API '
615+
'teacher). Use --teacher_model <local ckpt> for a frozen local teacher, or omit '
616+
'both to use dynamic self-distillation (teacher = current policy).')
617+
if self.beta not in (0, 0.0):
618+
logger.warning(f'advantage_reweight=rlsd: the reference RLSD config disables KL (beta=0), '
619+
f'but beta={self.beta}. Consider setting --beta 0.')
620+
if self.teacher_model is not None:
621+
logger.info(f'advantage_reweight=rlsd: frozen-teacher mode — --teacher_model={self.teacher_model!r} scores '
622+
f'the `teacher_prompt` column under no_grad and stays frozen throughout training (matches the '
623+
f'RLSD reference freeze_teacher_model=true when this checkpoint is the initial policy).')
624+
else:
625+
logger.info('advantage_reweight=rlsd: dynamic self-distillation mode — expecting a `teacher_prompt` '
626+
'column in the dataset; teacher = current policy on the privileged prompt.')
571627

572628
def _check_opd_rl(self):
573629
"""Fail-fast OPD-RL (teacher distillation on GRPO) parameter compatibility.

swift/rl_core/advantage.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,58 @@ def expand_advantage_to_per_token(
279279
return per_token_adv
280280

281281

282+
def apply_rlsd_reweight(
283+
base_advantages: torch.Tensor,
284+
completion_mask: torch.Tensor,
285+
teacher_per_token_logps: torch.Tensor,
286+
policy_per_token_logps: torch.Tensor,
287+
lam: float,
288+
clip_range: float,
289+
negative_only: bool = False,
290+
) -> torch.Tensor:
291+
"""RLSD (Self-Distilled RLVR) token-level advantage reweighting.
292+
293+
Redistributes the per-sequence GRPO advantage *inside* each trajectory using the
294+
teacher-vs-student log-prob gap, without ever flipping the sign of the environment
295+
reward (the reweight is strictly positive). Mirrors ``_build_stgca_advantages`` in
296+
the reference implementation (RLSD/verl/workers/actor/dp_opsd_actor.py)::
297+
298+
delta_t = stop_grad(logP_T(y_t) - logP_S(y_t))
299+
w_t = exp(sign(A) * delta_t)
300+
reweight = (1 - lam) + lam * clip(w_t, 1 - clip_range, 1 + clip_range)
301+
A_hat_t = A * stop_grad(reweight)
302+
303+
``lam`` mixes between pure GRPO (``lam=0`` -> plain broadcast) and full RLSD reweighting
304+
(``lam=1``). The teacher is the "informed self": the same policy conditioned on the
305+
ground-truth answer, scoring the *same* sampled tokens (``teacher_per_token_logps``).
306+
``policy_per_token_logps`` is the student side (the batch-time old logps).
307+
308+
Args:
309+
base_advantages: ``[B]`` per-sequence GRPO advantage.
310+
completion_mask: ``[B, T]`` response-token mask.
311+
teacher_per_token_logps: ``[B, T]`` teacher logp on sampled tokens.
312+
policy_per_token_logps: ``[B, T]`` student (old) logp on the same tokens.
313+
lam: Effective mixing weight (already schedule-adjusted).
314+
clip_range: Evidence-weight clip epsilon ``eps_w``.
315+
negative_only: When True, only reweight sequences with ``A < 0`` (incorrect
316+
responses); ``A >= 0`` sequences keep pure GRPO advantages.
317+
318+
Returns:
319+
``[B, T]`` per-token reweighted advantage (masked outside the response).
320+
"""
321+
mask = completion_mask
322+
base_bt = base_advantages.unsqueeze(1).expand_as(mask)
323+
delta = (teacher_per_token_logps.detach() - policy_per_token_logps.detach()) * mask
324+
sign_a = torch.sign(base_bt)
325+
weights = torch.exp(sign_a * delta) * mask
326+
clipped = torch.clamp(weights, min=1.0 - clip_range, max=1.0 + clip_range)
327+
reweight = (1.0 - lam) + lam * clipped
328+
if negative_only:
329+
seq_neg = (base_advantages < 0).float().unsqueeze(1)
330+
reweight = seq_neg * reweight + (1.0 - seq_neg)
331+
return base_bt * reweight.detach() * mask
332+
333+
282334
@dataclass
283335
class RewardMetrics:
284336
"""Reward statistics for logging."""

swift/rlhf_trainers/args_mixin.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,17 @@ class GRPOArgumentsMixin(RolloutTrainerArgumentsMixin):
436436
# enabled when a teacher (teacher_model / teacher_model_server) is set on a GRPO run.
437437
teacher_kl_coef: float = 1.0
438438

439+
# RLSD (Self-Distilled RLVR), https://arxiv.org/abs/2604.03128
440+
# Token-level advantage reweighting using the teacher-vs-student logprob gap. Reuses the OPSD
441+
# self-distillation teacher forward (teacher = current policy conditioned on the ground-truth
442+
# answer via a per-sample ``teacher_prompt`` column). Set ``advantage_reweight='rlsd'`` to enable.
443+
advantage_reweight: Optional[Literal['rlsd']] = None
444+
rlsd_lambda: float = 0.5 # mixing weight: 0 -> pure GRPO, 1 -> full RLSD reweighting
445+
rlsd_reweight_clip_range: float = 0.2 # eps_w: clip evidence weight to [1-eps_w, 1+eps_w]
446+
rlsd_lambda_warmup_steps: int = 0 # linear warmup of lambda from 0 to rlsd_lambda
447+
rlsd_lambda_decay_steps: int = 0 # linear decay of lambda to 0 over this many steps
448+
rlsd_negative_only: bool = False # only reweight sequences with advantage < 0
449+
439450
# REAL https://arxiv.org/abs/2602.05630
440451
real_tau: float = 0.5
441452

0 commit comments

Comments
 (0)