Skip to content

perf: GRPO runs a redundant pre-update policy forward under the default setup #10358

Description

@rrrxxx0510

Checklist / 检查清单

  • I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues,确认这是一个新的 bug report。

Bug Description / Bug 描述

grpo_trainer.py always runs an extra forward to fill old_per_token_logps. With the default setup (num_iterations=1, rollout batch consumed within one optimizer step) the pre-update policy and the current policy are the same weights, so the result is bit-for-bit identical to per_token_logps.detach() — the PPO ratio is 1 and the clip term is a no-op. The forward is pure overhead.

A probe confirmed this: before the change old_per_token_logps is non-None and torch.equal against per_token_logps.detach() is True with max_abs_diff=0.0 on every step. After skipping it, all logged metrics (loss, grad_norm, completions/{mean,min,max}_length, clipped_ratio, reward, reward_std, kl, every clip_ratio/*) and the peak memory are unchanged.

The loss site already falls back to per_token_logps.detach() when the field is None, so the field is only needed where something reads it without that fallback: kl_in_reward, teacher KL (RLSD/OPD), the Liger path and rollout importance correction, and async_generate.

How to Reproduce / 如何复现

Any default GRPO run (--rlhf_type grpo, no --num_iterations). Cost of the removed forward, measured on an H200 with Qwen2.5-0.5B (one forward over the completion batch plus its logps):

batch x seqlen tokens time peak
8 x 1024 8192 65.2 ms 5.80 GiB
16 x 2048 32768 109.2 ms 20.31 GiB

Additional Information / 补充信息

ms-swift main (4.6.0.dev0).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions