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

[Kristytiki]

Found at Chai AI Research while training Gemma-4 26B with megatron rlhf --rlhf_type grpo.

Checklist / 检查清单

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

Bug Description / Bug 描述

Checklist

• I have searched existing issues, and this is a new bug report.

Bug Description

In GRPO, response_prefix is appended to the prompt before generation, so vLLM samples after it and the
returned generate_ids do not contain it. Template.decode() then prepends the prefix back onto the decoded
text, and that combined string is encoded as {{RESPONSE}} for the training forward pass — so the prefix
tokens end up with labels set and inside completion_mask.

The policy is therefore scored on tokens it never sampled. vLLM computes log-probs over N tokens, the trainer
over len(prefix) + N. If the model finds the prefix improbable, those few tokens dominate the sequence
log-probability and rollout_correction/ppl_ratio — which should be ~1 for colocated, strictly on-policy
training — blows up.

This affects both swift rlhf and megatron rlhf: the code is in the shared template layer.

Where

1. Generation — prefix goes into the prompt (swift/template/base.py:1449-1451)

elif response_prefix:
    # final round and during inference.
    context_list.append(response_prefix)

2. Decode — prefix is prepended back onto the completion (swift/template/base.py:803-806)

response = self.tokenizer.decode(generate_ids, **kwargs)   # does NOT contain the prefix
response_prefix = self._get_response_prefix(template_inputs)
if first_token and response_prefix:
    response = response_prefix + response                  # now it does

3. Training encode — the whole string becomes the response (swift/template/base.py:1423-1425)

elif response is not None:
    # It is the final round, and the response exists (during training).
    context_list.append('{{RESPONSE}}')     # no prompt-side prefix here

4. Mask is derived from labels (swift/rlhf_trainers/utils.py:1967, 1984)

completion_mask = (rolled_labels != -100)

{{RESPONSE}} is the only context that gets ContextType.RESPONSE in _concat_context_list, so the prefix —
having arrived inside it — is in the loss, the KL term, and the reported perplexity.

Measured impact

google/gemma-4-26B-A4B-it, megatron rlhf --rlhf_type grpo, colocated vLLM, LoRA,
steps_per_generation=4, max_completion_length=80, response_prefix="<|channel>thought\n<channel|>{name}:".
All numbers at step 1, everything else identical:

configuration	ppl_ratio	ess	logged loss
as-is	54.2	0.85	−0.30 … −0.47
response_prefix unset (falls back to non_thinking_prefix)	38.6	0.86	−0.35
base model fine-tuned on that template (prefix in-distribution)	3.1	0.86	−0.35
patched (below)	1.0	0.99	±0.005

Two independent confirmations that the diagnosis is right:

• loss ≈ 0 after the patch. Strictly on-policy with group-centred advantages, the policy-gradient term
  has expectation 0. Every unpatched run sat at −0.3 to −0.47; patched it is ±0.005. The excess was the
  prefix tokens.
• ess 0.85 → 0.99. Sequence-level importance weights were silently discarding ~15% of every batch.

Downstream, unpatched: grad_norm 68–836 with clip_grad=1.0 saturating every step, KL rising monotonically
while reward stayed flat for 130+ steps. Patched, same base model and hyperparameters, reward rose steadily
over 500 steps.

Note on the fallback

Unsetting response_prefix does not remove it. _get_response_prefix (swift/template/base.py:191-203)
falls through to template_meta.non_thinking_prefix when enable_thinking=False, so a prefix is still
prepended — hence 38.6 rather than ~1 in the table above.

Making the prefix cheap is also not a workaround: substituting a short in-distribution prefix got
ppl_ratio to 4.8 but halved the reward, tripled the repetition penalty, and produced malformed text,
because the model was pushed off its chat template.

Suggested fix

Make the training encode treat response_prefix the way generation does — emit it as a prompt-side context
before {{RESPONSE}}, and stop prepending it to the response text. Appended strings get
ContextType.OTHER in _concat_context_list, i.e. the same treatment as the prompt, so they are excluded
from the loss.

The two-line version we are running as a stopgap:

# swift/template/base.py:805  — decode()
if False and response_prefix:          # was: if first_token and response_prefix:
    response = response_prefix + response

# swift/template/base.py:1424 — _swift_encode(), training branch
elif response is not None:
    if response_prefix: context_list.append(response_prefix)
    context_list.append('{{RESPONSE}}')

The invariant to restore: the tokens the policy is scored on during training are exactly the tokens it
sampled during rollout.

Possibly related

#9096 — Gemma-4 GRPO, reward never improved, closed without a working fix. Same template family and the same
response_prefix usage; worth re-checking against this.

How to Reproduce / 如何复现

How to Reproduce

The mechanism only needs three things: response_prefix set, colocated vLLM, and strictly on-policy
updates (steps_per_generation=1, num_iterations=1) so that rollout_correction/ppl_ratio is expected
to be exactly 1. Pick a prefix the model would rarely emit on its own and the ratio scales with how
improbable it is.

swift rlhf \
    --rlhf_type grpo \
    --model <any chat model> \
    --dataset <any prompt dataset> \
    --reward_funcs <any> \
    --response_prefix "<<<ZZZ>>> " \
    --use_vllm true --vllm_mode colocate \
    --steps_per_generation 1 \
    --num_generations 4 \
    --max_completion_length 64 \
    --log_rollout_offpolicy_metrics true \
    --rollout_importance_sampling_mode sequence_mask \
    --logging_steps 1 --max_steps 5

At step 1, rollout_correction/ppl_ratio should be ≈1 (same weights on both sides, nothing off-policy).
Observed instead: a large value that grows with the prefix's improbability. The control is the same run with
a prefix the model emits naturally, or a template with no non_thinking_prefix — that gives ≈1.

Additional Information / 补充信息

No response