Skip to content

fix(rlhf): align rollout response prefixes during training - #10107

Merged
hjh0119 merged 2 commits into
modelscope:mainfrom
taking-lying-flat:fix/grpo-response-prefix-from-rollout
Sep 17, 2026
Merged

hjh0119 merged 2 commits into
modelscope:mainfrom
taking-lying-flat:fix/grpo-response-prefix-from-rollout

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

PR type

  • Bug Fix
  • New Feature
  • Document Updates
  • More Models or Datasets Support

PR information

GRPO/GKD can reconstruct a sampled response with a different response prefix from generation. Training previously selected thinking/non-thinking prefixes directly from template metadata, bypassing the template's actual prefix resolver.

Trigger Previous training behavior
Global custom response_prefix Injects the default prefix instead of the configured prefix.
Per-sample chat_template_kwargs.response_prefix Ignores the sample override.
Explicit response_prefix="" May inject the default even though generation uses no prefix.

Resolve the training prefix through Template._get_response_prefix with the sample's chat-template settings, including teacher views and requests. This also respects model-specific resolver overrides. Default thinking/non-thinking configurations continue to use the template's existing rules.

For example, tokens sampled after prompt + CUSTOM: must be trained with that same context. The prefix belongs to the prompt and is excluded from completion loss. The original sampled response IDs are reused directly; only the prompt prefix is tokenized.

When reinserting the prefix, recognize an already inserted prefix by both its IDs and zero loss mask, so an actual sampled repetition is preserved. Copy the outer token/mask lists to avoid modifying stored rollout IDs during repeated encoding.

Text, image, and video samples share this response encoding path. Visual preprocessing and teacher image selection retain their existing behavior. The diff is limited to swift/rlhf_trainers/utils.py and swift/rlhf_trainers/gkd_helpers.py (25 additions, 27 deletions). No test code is committed.

Fixes #10106.

Experiment results

Validation on the final diff:

  • 19 regression tests passed for prefix selection, masking, sampled repetitions, repeated encoding, and teacher response encoding.
  • 60 real-tokenizer cases passed with local Qwen3-0.6B and Qwen3.5-0.8B: default/thinking/custom/empty/sample overrides, flat/nested IDs, sampled repetitions, and conversation history. Training inputs match the generation prompt plus sampled IDs; completion labels and stored IDs remain correct.
  • 42 Qwen3.5-0.8B multimodal cases passed: one/two images, video frames, and zero/one/two teacher images. Checks cover visual tensors and masking, response alignment, teacher requests, completion counts, and position-ID lengths.
  • Formatting, import ordering, static lint, and git diff --check passed.

Validation scripts and tests were run locally and are not included in this PR. These checks do not constitute an end-to-end vLLM or Megatron training run.

Signed-off-by: taking-lying-flat <1615405@qq.com>
Signed-off-by: taking-lying-flat <1615405@qq.com>
@taking-lying-flat
taking-lying-flat force-pushed the fix/grpo-response-prefix-from-rollout branch from 21d897d to 6ba1186 Compare September 16, 2026 14:16
@hjh0119
hjh0119 merged commit b2732cb into modelscope:main Sep 17, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

GRPO: response_prefix is prompt at rollout time but completion at training time

2 participants