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
Changes reward computation in training loops for XPO/Nash-MD, including intentional score shifts for same-tokenizer setups; incorrect last-token or chat-template handling would skew optimization but is covered by new tests.

Overview
Fixes reward scoring in XPO and Nash-MD when the policy and reward models use different tokenizers (e.g. Qwen policy + Llama reward), which previously could ignore reward_processing_classes and crash with out-of-vocab IndexError.

Both trainers now score via new get_reward_from_policy_tokens: decode completions with the policy tokenizer, rebuild prompt+completion text (including chat templates), re-encode with reward_processing_classes[0], and take the last non-padding token—aligned with Online DPO and avoiding wrong scores when pad_token_id == eos_token_id at turn boundaries. NashMDTrainer stops overwriting reward_processing_classes with the policy tokenizer and documents the argument.

Behavior note: runs that already used a matching reward tokenizer will see different reward values (decode/re-encode with add_special_tokens=False drops BOS, same as Online DPO). get_reward is refactored to share _compute_reward_logits but unchanged for direct callers.

^{Reviewed by Cursor Bugbot for commit 602bdd0. 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.