Validate reward model tokenizer matches policy in XPOTrainer/NashMDTrainer - #6955
Validate reward model tokenizer matches policy in XPOTrainer/NashMDTrainer#6955amanyagami wants to merge 3 commits into
Conversation
…ainer XPOTrainer and NashMDTrainer score completions via get_reward(), which runs the reward model directly on the *policy*-tokenized input_ids, unlike OnlineDPOTrainer, which decodes and re-tokenizes completions per reward function. This is only correct if the reward model shares the exact same vocabulary as the policy model; otherwise the reward model scores the wrong tokens, silently producing meaningless rewards or crashing with a cryptic embedding out-of-range error. XPOTrainer already accepted and documented reward_processing_classes, but never used it anywhere, with no indication that it did nothing. NashMDTrainer didn't accept it at all and hardcoded reward_processing_classes=processing_class in its call to OnlineDPOTrainer.__init__, forcing the policy tokenizer onto every reward model with no way to override it and no validation. Implementing full decode/re-tokenize like OnlineDPOTrainer would require reworking the padding/masking/EOS-detection logic that get_reward() shares between both trainers -- a much larger, harder to validate change for two experimental trainers. Instead, add validate_reward_processing_class_shares_vocab() (trl/experimental/utils.py) and call it from both trainers' __init__ after the reward model is resolved, raising an informative ValueError on a vocabulary mismatch instead of letting it fail later. Also add the missing reward_processing_class argument to NashMDTrainer (defaulting to processing_class, preserving its previous behavior when omitted), and document the shared-vocab constraint in both docstrings. Fixes huggingface#6951
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
❌ 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 84ddbd8. Configure here.
… unwrap ProcessorMixin symmetrically Two real bugs in the initial fix, both caught by automated review on PR huggingface#6955: 1. self.reward_funcs can already be wrapped (DeepSpeed engine / DDP) by Accelerator.prepare_model inside super().__init__() by the time the vocab check runs in XPOTrainer/NashMDTrainer's own __init__. A wrapper is not a PreTrainedModel, so validate_reward_processing_class_shares_vocab's guard silently skipped the check under exactly the training setups (DeepSpeed, multi-process DDP) where catching a mismatch matters most. Fixed by unwrapping with self.accelerator.unwrap_model() before the check, at both call sites. Regression tests reproduce this with a real single-process DistributedDataParallel wrapping (not a hand-rolled stand-in, since accelerate's extract_model_from_parallel only recognizes real wrapper types) -- confirmed to fail without the unwrap and pass with it. 2. The vocab check unwrapped a ProcessorMixin on the policy side but not the reward side, so NashMDTrainer defaulting reward_processing_class to a ProcessorMixin processing_class (VLM policies) raised AttributeError from get_vocab() instead of comparing vocabularies -- even for the same processor object on both sides. Fixed by unwrapping reward_processing_class the same way. Regression tests cover both the matching-vocab (must not raise) and mismatched-vocab (must still raise ValueError, not AttributeError) cases directly against the helper. Fixes the two issues raised by Cursor Bugbot on 84ddbd8.
|
Both issues flagged by Cursor Bugbot are fixed in fbfba76:
Both new regression tests were confirmed to fail without their respective fix and pass with it. Full |
1 participant
What does this PR do?
XPOTrainerandNashMDTrainer(both subclasses ofexperimental.online_dpo.OnlineDPOTrainer) score completionswith the
get_reward()helper, which runs the reward model directly on the policy-tokenizedinput_ids:Unlike
OnlineDPOTrainer._calculate_rewards_from_functions, which decodes completions with the policy's tokenizerand re-tokenizes them per reward function with
reward_processing_classbefore scoring,XPOTrainerandNashMDTrainernever decode/re-tokenize: they feed the policy's token IDs straight into the reward model's forwardpass. That's only correct if the reward model shares the exact same vocabulary (token -> ID mapping) as the policy
model — otherwise the reward model scores the wrong tokens, which either silently produces meaningless rewards or
crashes deep inside an embedding lookup with an unhelpful
IndexError.On top of that:
XPOTraineraccepts and documentsreward_processing_classes, but never actually uses it anywhere in_compute_rewards— passing a reward tokenizer that differs from the policy's just wasn't a way to avoid thebug, and there was no signal to the user that it silently does nothing.
NashMDTrainerdoesn't acceptreward_processing_classesat all, and hardcodesreward_processing_classes=processing_classin its call tosuper().__init__(), so it always feeds thepolicy's tokenizer to the (potentially different) reward model, with no way to override it and no validation that
this is actually safe.
Why this fix (see issue for the 3 options considered)
The issue lists three options: (1) implement full decode/re-tokenize per reward model like
OnlineDPOTrainerdoes,(2) add validation + document the shared-vocab constraint, or (3) something else. I went with (2):
_compute_rewardsin both trainers usesget_reward(), a single-model helper built around rawquery_responsestoken IDs (used for computing sequence lengths, building the attention mask from
pad_token_id, etc.) — notOnlineDPOTrainer's per-reward-function decode/apply-chat-template/re-tokenize path. Rewriting_compute_rewardsto decode and re-tokenize would touch the padding/masking/EOS-detection logic in both trainers and is a much larger,
higher-risk diff for two experimental trainers I can't validate on GPU. Since the issue explicitly allows the
safer, minimal option when a full rewrite is too large, I added a validation check + documentation instead of
reimplementing the reward pipeline.
Changes
validate_reward_processing_class_shares_vocab()totrl/experimental/utils.py: given a reward function,its processing class, and the policy's processing class, it's a no-op for non-model reward functions and raises a
clear
ValueErrorif a model-based reward function's tokenizer vocabulary doesn't match the policy's.XPOTrainer.__init__: call the new validation after resolvingself.reward_funcsto a single model. Documentedthe shared-vocab constraint in the
reward_processing_classesdocstring entry.NashMDTrainer.__init__: added the missingreward_processing_classargument (singular, since NashMD onlysupports one reward model) so callers can now point it at the reward model's own tokenizer. When not provided, it
still defaults to
processing_class, preserving the existing default behavior exactly. Also call the newvalidation, and documented the constraint in the docstring.
Why this doesn't break existing behavior
For both trainers, when
reward_processing_class(es)isn't explicitly overridden with something incompatible, thevalidation either compares a tokenizer to itself (
NashMDTrainer's unchanged default) or compares theauto-loaded reward-model tokenizer to the policy tokenizer, which only raises when they were already silently
mismatched (previously broken, just with a worse failure mode later). A genuinely matching setup is unaffected.
Tests
Added regression tests in
tests/experimental/test_xpo_trainer.pyandtests/experimental/test_nash_md_trainer.py:test_reward_processing_class_vocab_mismatch_raises(both trainers): passing a reward tokenizer with a differentvocabulary (a tiny GPT-2 tokenizer against a Qwen2 policy) raises
ValueErrormentioning "vocabulary".test_reward_processing_class_matching_vocab_is_accepted(XPO) /test_reward_processing_class_can_be_passed_explicitly(NashMD): explicitly passing a matching processing class is accepted and stored.
test_reward_processing_class_defaults_to_processing_class(NashMD): omittingreward_processing_classstilldefaults to
processing_class, confirming the previous default behavior is preserved.I confirmed all 5 new tests fail on
main(2 withTypeError: unexpected keyword argumentfor NashMD, since theparameter didn't exist; the 2 mismatch tests fail because nothing validates the mismatch) and pass with the fix.
Test output (after fix)
Full suite for both trainers (includes the existing training/PEFT tests, which run actual tiny-model training steps
on CPU, in addition to the 5 new reward-tokenizer tests):
Just the new reward-tokenizer tests:
I confirmed these 5 fail on
main(git stashthe source changes, keep the tests): 2 withTypeError: unexpected keyword argument 'reward_processing_class'forNashMDTrainer(the parameter didn't exist), and the 2 mismatchtests with no exception raised at all (nothing validated the mismatch).
pre-commit run --files trl/experimental/utils.py trl/experimental/xpo/xpo_trainer.py trl/experimental/nash_md/nash_md_trainer.py tests/experimental/test_xpo_trainer.py tests/experimental/test_nash_md_trainer.pypasses (ruff check, ruff format, doc-builder style check all green).
Fixes #6951
Before submitting
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.
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.
Note
Low Risk
Changes are limited to experimental trainers: init-time validation and docs, with no change to reward scoring logic beyond rejecting already-invalid tokenizer pairings.
Overview
Fixes #6951 by failing fast when
XPOTrainerandNashMDTrainerwould score completions with policy token IDs on a reward model whose vocabulary does not match the policy tokenizer (unlikeOnlineDPOTrainer, these trainers useget_reward()on rawinput_idswithout decode/re-tokenize).Adds
validate_reward_processing_class_shares_vocab()intrl/experimental/utils.py: for model-based reward functions it comparesget_vocab()on policy and reward sides, unwrappingProcessorMixinto.tokenizeron both sides.XPOTrainerandNashMDTrainercall it aftersuper().__init__(), usingaccelerator.unwrap_model()on the reward model so the check still runs when the reward model is wrapped by DDP/DeepSpeed.NashMDTrainergains optionalreward_processing_class(default remainsprocessing_classviareward_processing_class or processing_class). Docstrings document the shared-vocabulary requirement forreward_processing_classes/reward_processing_class.Regression tests cover vocab mismatch
ValueError, matching vocab acceptance, NashMD defaults, DDP-wrapped reward models (patchedAccelerator.prepare_model), andProcessorMixinpaths intest_utils.py.Reviewed by Cursor Bugbot for commit 786e110. Bugbot is set up for automated code reviews on this repo. Configure here.