Honor reward_processing_classes in XPO and Nash-MD trainers

[22elix3r]

What does this PR do?

Fixes #6951.

XPOTrainer already accepts reward_processing_classes and stores it through OnlineDPOTrainer, but _compute_rewards scored completions with the policy tokenizer's input_ids and pad_token_id. NashMDTrainer used the same scoring path and additionally hardwired reward_processing_classes=processing_class, so a distinct reward tokenizer could not be passed at all.

That silently ignores a user-supplied reward tokenizer. When the reward model has a smaller vocabulary than the policy (Llama reward + Qwen policy in the original report), scoring raises IndexError: index out of range in self far from the cause.

Online DPO already decodes completions and re-tokenizes them with the reward processing class. XPO and Nash-MD now do the same:

1. Decode the completion with the policy tokenizer (skip_special_tokens=True)
2. Rebuild the full prompt+completion text (including the conversational chat-template path)
3. Encode with reward_processing_classes[0] (add_special_tokens=False)
4. Score the last non-padding token using the tokenizer attention_mask

NashMDTrainer now exposes reward_processing_classes instead of overwriting it with the policy tokenizer.

Scoring uses the tokenizer mask rather than get_reward(..., context_length=0). Chat templates often set pad_token_id == eos_token_id and emit that id at turn boundaries; the first-pad rule would otherwise score the prompt.

Same-tokenizer behavior change. Decode/re-encode with add_special_tokens=False drops BOS, matching Online DPO. Existing XPO/Nash-MD runs that already used a matching reward tokenizer will see different scores even though they were not hitting the IndexError. This is intentional.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Discussion: #6951 (comment)

This implements option 2 from the issue (honor the argument the way Online DPO does), rather than dropping the parameter. Docstrings for NashMDTrainer were updated to document the new argument; XPO already documented it.

Tests

• TestGetRewardFromPolicyTokens: policy token ids outside the reward vocab used to crash get_reward and now score after re-tokenization; conversational last-token vs first eos; batched chat padding; same-tokenizer BOS drop
• TestXPOTrainerRewardProcessingClass: XPO scores with a Llama reward tokenizer and a Qwen policy
• TestNashMDTrainerRewardProcessingClass: Nash-MD keeps the passed reward tokenizer instead of replacing it with the policy tokenizer

python -m pytest \
  tests/experimental/test_utils.py::TestGetRewardFromPolicyTokens \
  tests/experimental/test_xpo_trainer.py::TestXPOTrainerRewardProcessingClass \
  tests/experimental/test_nash_md_trainer.py::TestNashMDTrainerRewardProcessingClass

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

@qgallouedec @behroozazarkhalili

---

Note

Medium Risk
Reward scoring path changed for XPO/Nash-MD training, including intentional score shifts when tokenizers match; affects RLHF loss but not auth or data persistence.

Overview
XPO and Nash-MD now score rewards with the reward tokenizer, aligning with Online DPO instead of feeding policy input_ids through get_reward.

Adds get_reward_from_policy_tokens: decode completions with the policy tokenizer, rebuild prompt+completion (including chat templates), re-encode with reward_processing_classes, and take the last non-padding token from the reward model head. Refactors shared backbone scoring into _compute_reward_logits.

NashMDTrainer stops overriding reward_processing_classes with the policy tokenizer and documents the argument. XPOTrainer and NashMDTrainer _compute_rewards both use the new helper.

Behavior note: even when policy and reward tokenizers match, scores can change (decode/re-encode with add_special_tokens=False, last-token vs first-pad scoring for chat where pad equals eos).

Tests cover cross-vocab retokenization, conversational last-token selection, batched chat padding, and trainer integration with Qwen policy + Llama reward.

^{Reviewed by Cursor Bugbot for commit 0d19568. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit f0fa85f. Configure here.}

[behroozazarkhalili]

Thanks for picking this up, and for going with option 2 rather than dropping the parameter. I filed the issue, so I re-checked the implementation against source and measured the parts that a reading alone cannot settle. Everything the PR rests on holds.

Claims I verified

• model_data["raw"] is populated at both call sites, so the new prompts argument is available: xpo_trainer.py:239,250 and nash_md_trainer.py:308,319.
• self.reward_processing_classes[0] is the right index in both trainers. Each rejects more than one reward function and unwraps the list (xpo_trainer.py:176-178, nash_md_trainer.py:230-232), so the list always has length 1.
• Nash-MD really did hardwire the policy tokenizer, and forwarding the argument to the parent is enough: OnlineDPOTrainer.__init__ already auto-loads the reward tokenizer from reward_func.config._name_or_path when it is None, sets pad_token = eos_token if the reward tokenizer has no pad token, and sets reward_func.config.pad_token_id (online_dpo_trainer.py:237-248).

On "matching Online DPO behavior"

I expected a divergence here and did not find one. Online DPO calls the reward model directly with **reward_inputs and reads .logits[:, 0], while get_reward_from_policy_tokens drops the tokenizer's attention_mask and lets get_reward rebuild it as query_responses != pad_token_id. Measured on a tiny Llama sequence-classification head with a right-padded two-row batch, the reconstructed mask equals the tokenizer's mask and the scores are bit-identical, max|A - B| = 0, both when pad_token_id == eos_token_id and when the pad token is distinct.

Routing through get_reward also survives one case the direct call does not. The direct call needs model.config.pad_token_id to be set or it raises Cannot handle batch sizes > 1 if no padding token is defined, and get_reward computes sequence_lengths itself, so it does not.

One thing worth flagging

The PR changes reward scores on the path that already worked, not just the one that crashed, and neither the description nor the tests mention it.

When the policy and reward tokenizers are the same, the old code fed raw policy ids straight to get_reward. The new code decodes with skip_special_tokens=True and re-encodes with add_special_tokens=False, which drops BOS:

OLD ids: [128000, 3923, 374, 220, 17, 10, 17, 30, 1102, 374, 3116, 13]
NEW ids: [        3923, 374, 220, 17, 10, 17, 30, 1102, 374, 3116, 13]

OLD scores: [-0.023193,  0.017456]
NEW scores: [-0.046143, -0.066406]     max|delta| = 0.084

The head is randomly initialized, so treat the magnitude as illustrative rather than as a size estimate for a trained reward model. The point is that the scores move at all for users who were not hitting the bug.

This is not a defect. add_special_tokens=False is what Online DPO does, and matching it is the right call for consistency. But it is a silent change to existing XPO and Nash-MD runs, so it is worth a line in the PR description, and a test that pins the same-tokenizer case would keep it from drifting later. Your three new tests all use a mismatched pair, which is exactly the case that used to raise, so the previously working path is currently uncovered.

Two smaller notes:

• get_reward_from_policy_tokens re-implements the tokenization that OnlineDPOTrainer._calculate_rewards_from_functions does inline, so there are now two copies of that logic that have to stay aligned. That is consistent with how this repo treats trainers, but the two do differ in how they score, so a short comment on the helper pointing at the Online DPO block would help whoever touches one of them next.
• Neither the helper nor the Online DPO original passes truncation, so a re-tokenized sequence longer than the reward model's positions will still fail. Pre-existing and out of scope here, just noting it so it is not mistaken for something this PR introduced.

Nice fix, and thanks for keeping the two siblings aligned with Online DPO instead of removing the argument.

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[22elix3r]

Thanks for the re-check, @behroozazarkhalili — that same-tokenizer BOS drop is a real, silent change and was not called out before.

I pushed a follow-up that:

• Notes the same-tokenizer score change in the PR description (decode/re-encode with add_special_tokens=False drops BOS, matching Online DPO)
• Pins it in test_same_tokenizer_matches_add_special_tokens_false
• Adds a comment on the helper pointing at OnlineDPOTrainer._calculate_rewards_from_functions, including that scoring still uses the XPO/Nash-MD model.score last-token head rather than logits[:, 0]
• Scores chat-formatted sequences from the tokenizer attention_mask (last non-padding token) instead of get_reward(..., context_length=0), which Bugbot flagged for pad == eos turn boundaries
• Wraps the helper docstring to satisfy doc-builder style

Truncation of over-long re-tokenized sequences is still out of scope, as you said.

[22elix3r]

@qgallouedec @albertvillanova @behroozazarkhalili could you take a look when you have a moment?

This is the option-2 fix for #6951: XPO and Nash-MD now honor reward_processing_classes the way Online DPO does, instead of feeding policy token ids into the reward model. CI experimental tests are green; the quality hook failure from the previous push should be gone after the doc-builder wrap.

Happy to iterate on anything that looks off.

[behroozazarkhalili]

Checked the follow-up. The turn-boundary fix is right, and it closes a gap my earlier measurement could not have caught.

I had reported that routing through get_reward was numerically identical to Online DPO's direct call, max|A - B| = 0. That result holds, but only for the sequences I tested, which were plain strings with no EOS inside them. Bugbot's case is different and it reproduces:

ids (pad == eos, eos at turn boundary index 3): [10, 11, 12, 128009, 20, 21, 22, 23]

get_reward "first pad id" rule  -> scores index 2  (value 12)   mid-prompt
attention-mask last real token  -> scores index 7  (value 23)   correct

first_true_indices(ids == pad_token_id) - 1 stops at the first EOS, so with a chat template the score comes from the prompt and the whole completion is discarded. A second row with no internal EOS agrees between the two rules, which is exactly why my earlier probe showed no difference. Scoring from the tokenizer attention_mask is the correct rule here.

Two notes on the new code:

• Factoring the backbone call into _compute_reward_logits and having get_reward call it keeps the two paths on one implementation, so the position ids, the masked fill and the use_cache=False workaround cannot drift apart. Worth keeping.
• The docstring now states that scoring uses the model.score last-token head rather than logits[:, 0]. That is the accurate description. The two agree numerically on a standard sequence-classification reward model, but they are different code paths and the difference matters if someone passes a model where they diverge.

test_same_tokenizer_matches_add_special_tokens_false covers the case I raised, which was the one gap in the original tests since all three used a mismatched tokenizer pair.

Nothing further from me. Agreed that truncation of over-long re-tokenized sequences is a separate concern and not this PR's job.

[22elix3r]

Rebased this branch onto current main and resolved the merge conflicts.

The only conflicted files were trl/experimental/utils.py and tests/experimental/test_utils.py (trainers auto-merged). I kept get_reward_from_policy_tokens and the _compute_reward_logits refactor, and folded in the overlapping main changes:

• XPO/NashMD perf: avoid output_hidden_states when only last_hidden_state is used #4755: _compute_reward_logits now uses output_hidden_states=False and output.last_hidden_state instead of hidden_states[-1]. Dummy reward backbones in the tests set last_hidden_state so they match that path.
• Require peft>=0.13.0 and drop the 0.12 version guards #7102: dropped the peft 0.12 version guards around get_peft_model.
• Remove PPOTrainer #7020: kept TestReferenceModel in tests/experimental/test_utils.py alongside the new TestGetRewardFromPolicyTokens tests.

XPO and Nash-MD still import get_reward_from_policy_tokens; Nash-MD still forwards reward_processing_classes rather than hardwiring processing_class. Targeted tests for the new helper and both trainer classes passed locally.

[22elix3r]

Rebased onto current main — conflicts are resolved and the PR is mergeable.

@qgallouedec @albertvillanova @behroozazarkhalili @kashif could you take another look when you have a chance?

[behroozazarkhalili]

Re-checked the rebased head 0d19568. Verdict: the rebase is clean and the conflict resolution kept everything the earlier rounds settled. Nothing new from me; this is ready from my side.

What I checked

• All three touched test files on a CPU compute node against the PR head: 35 passed, 1 skipped, 0 failed (tests/experimental/test_utils.py, test_xpo_trainer.py, test_nash_md_trainer.py).
• ruff check and ruff format --check at the CI-pinned 0.13.3 pass on all six changed files; the pinned doc-builder style check passes.
• git merge-tree onto current main reports no conflicts (4 commits behind, both auto-mergeable).

Conflict resolution

The #4755 fold is correct: _compute_reward_logits calls the backbone with output_hidden_states=False and reads output.last_hidden_state, and get_reward now routes through the same helper, so the two scoring paths share one backbone call as before. get_reward_from_policy_tokens still scores from the tokenizer attention_mask (last non-padding token), which is the turn-boundary fix from the previous round, and both trainers still import it. Nash-MD forwards reward_processing_classes to the parent rather than hardwiring processing_class.

Unchanged from earlier rounds

The two notes I left before still hold and need no action: the model.score last-token head is the accurate description of what is scored, and truncation of over-long re-tokenized sequences remains a separate concern.

[albertvillanova]

Thanks for taking this on, and sorry it ends this way. The two trainers this targets have since been removed, XPO in #7256 and Nash-MD in #7255, so the files it edits no longer exist and it can no longer be merged. Closing on that basis rather than on anything about the change itself.