Validate reward model tokenizer matches policy in XPOTrainer/NashMDTrainer

tests/experimental/test_nash_md_trainer.py

@@ -195,6 +195,59 @@ def test_train_with_peft_and_ref_model(self):
 
         assert "train_loss" in trainer.state.log_history[-1]
 
+    def test_reward_processing_class_vocab_mismatch_raises(self):
+        # Regression test for #6951: NashMDTrainer feeds token IDs produced by `processing_class` directly to the
+        # reward model, so a `reward_processing_class` with a different vocabulary must be rejected at init time
+        # with a clear error instead of failing later with a cryptic embedding out-of-range error.
+        mismatched_tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
+        training_args = NashMDConfig(output_dir=self.tmp_dir, report_to="none")
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        with pytest.raises(ValueError, match="vocabulary"):
+            NashMDTrainer(
+                model=self.model,
+                ref_model=self.ref_model,
+                reward_funcs=self.reward_model,
+                reward_processing_class=mismatched_tokenizer,
+                args=training_args,
+                processing_class=self.tokenizer,
+                train_dataset=dataset,
+            )
+
+    def test_reward_processing_class_can_be_passed_explicitly(self):
+        # Regression test for #6951: NashMDTrainer previously hardcoded `reward_processing_classes=processing_class`
+        # and didn't expose a `reward_processing_class` argument at all, so callers had no way to point it at the
+        # reward model's own tokenizer.
+        training_args = NashMDConfig(output_dir=self.tmp_dir, report_to="none")
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        trainer = NashMDTrainer(
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_funcs=self.reward_model,
+            reward_processing_class=self.tokenizer,  # shares the policy's vocabulary, so this must be accepted
+            args=training_args,
+            processing_class=self.tokenizer,
+            train_dataset=dataset,
+        )
+        assert trainer.reward_processing_classes == [self.tokenizer]
+
+    def test_reward_processing_class_defaults_to_processing_class(self):
+        # When `reward_processing_class` isn't provided, it should default to `processing_class`, preserving the
+        # trainer's previous behavior.
+        training_args = NashMDConfig(output_dir=self.tmp_dir, report_to="none")
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        trainer = NashMDTrainer(
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_funcs=self.reward_model,
+            args=training_args,
+            processing_class=self.tokenizer,
+            train_dataset=dataset,
+        )
+        assert trainer.reward_processing_classes == [self.tokenizer]
+
     @require_peft
     def test_train_pre_pefted_model_implicit_ref_with_reward_model(self):
         lora_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.1, bias="none", task_type="CAUSAL_LM")

tests/experimental/test_xpo_trainer.py

@@ -145,6 +145,42 @@ def test_train_with_peft_and_ref_model(self):
 
         assert "train_loss" in trainer.state.log_history[-1]
 
+    def test_reward_processing_class_vocab_mismatch_raises(self):
+        # Regression test for #6951: XPOTrainer feeds token IDs produced by `processing_class` directly to the
+        # reward model, so a reward model tokenizer with a different vocabulary must be rejected at init time with
+        # a clear error instead of failing later with a cryptic embedding out-of-range error.
+        mismatched_tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
+        training_args = XPOConfig(output_dir=self.tmp_dir, report_to="none")
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        with pytest.raises(ValueError, match="vocabulary"):
+            XPOTrainer(
+                model=self.model,
+                ref_model=self.ref_model,
+                reward_funcs=self.reward_model,
+                reward_processing_classes=mismatched_tokenizer,
+                args=training_args,
+                processing_class=self.tokenizer,
+                train_dataset=dataset,
+            )
+
+    def test_reward_processing_class_matching_vocab_is_accepted(self):
+        # A reward processing class that does share the policy's vocabulary (e.g. the same tokenizer) must not be
+        # rejected.
+        training_args = XPOConfig(output_dir=self.tmp_dir, report_to="none")
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        trainer = XPOTrainer(
+            model=self.model,
+            ref_model=self.ref_model,
+            reward_funcs=self.reward_model,
+            reward_processing_classes=self.tokenizer,
+            args=training_args,
+            processing_class=self.tokenizer,
+            train_dataset=dataset,
+        )
+        assert trainer.reward_processing_classes == [self.tokenizer]
+
     @require_peft
     def test_train_pre_pefted_model_implicit_ref(self):
         lora_config = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.1, bias="none", task_type="CAUSAL_LM")

trl/experimental/nash_md/nash_md_trainer.py

@@ -37,7 +37,7 @@
 from ...models.utils import unwrap_model_for_generation
 from ...trainer.utils import selective_log_softmax
 from ..online_dpo import OnlineDPOTrainer
-from ..utils import empty_cache, get_reward, truncate_right
+from ..utils import empty_cache, get_reward, truncate_right, validate_reward_processing_class_shares_vocab
 from .nash_md_config import NashMDConfig
 
 
@@ -139,6 +139,15 @@ class NashMDTrainer(OnlineDPOTrainer):
             Processing class used to process the data. If provided, will be used to automatically process the inputs
             for the model, and it will be saved along the model to make it easier to rerun an interrupted training or
             reuse the fine-tuned model.
+        reward_processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*):
+            Processing class for the reward model specified in `reward_funcs`. If set to `None`, it defaults to
+            `processing_class`.
+
+            Note: `NashMDTrainer` scores completions by feeding token IDs produced by `processing_class` directly to
+            the reward model, instead of decoding and re-tokenizing with `reward_processing_class` (unlike
+            [`~trl.experimental.online_dpo.OnlineDPOTrainer`]). Because of this, the reward model must share the exact
+            same vocabulary as the policy model; this is validated at init time and raises a `ValueError` if the
+            vocabularies don't match.
         peft_config ([`~peft.PeftConfig`], *optional*):
             The peft config to use for training.
         compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*):
@@ -183,6 +192,7 @@ def __init__(
         | FeatureExtractionMixin
         | ProcessorMixin
         | None = None,
+        reward_processing_class: PreTrainedTokenizerBase | None = None,
         peft_config: "PeftConfig | None" = None,
         compute_metrics: Callable[[EvalPrediction], dict] | None = None,
         callbacks: list[TrainerCallback] | None = None,
@@ -198,7 +208,11 @@ def __init__(
             train_dataset=train_dataset,
             eval_dataset=eval_dataset,
             processing_class=processing_class,
-            reward_processing_classes=processing_class,
+            # `NashMDTrainer` feeds token IDs produced by `processing_class` directly to the reward model (see
+            # `_compute_rewards` below), so the reward model's tokenizer must share the same vocabulary as
+            # `processing_class`. Default to `processing_class` itself when the caller doesn't supply one, and
+            # validate the vocabularies match below.
+            reward_processing_classes=reward_processing_class or processing_class,
             peft_config=peft_config,
             compute_metrics=compute_metrics,
             callbacks=callbacks,
@@ -230,6 +244,9 @@ def __init__(
         if len(self.reward_funcs) != 1:
             raise ValueError("NashMDTrainer only supports one reward function/model.")
         self.reward_funcs = self.reward_funcs[0]
+        validate_reward_processing_class_shares_vocab(
+            self.reward_funcs, self.reward_processing_classes[0], self.processing_class
+        )
 
     @property
     def mixture_coef(self):

trl/experimental/utils.py

@@ -793,6 +793,38 @@ def get_reward(
     )
 
 
+def validate_reward_processing_class_shares_vocab(
+    reward_func: Any,
+    reward_processing_class: PreTrainedTokenizerBase | None,
+    policy_processing_class: PreTrainedTokenizerBase | ProcessorMixin,
+) -> None:
+    """
+    Validate that a model-based reward function's tokenizer shares its vocabulary with the policy's tokenizer.
+
+    Some trainers (e.g. `XPOTrainer`, `NashMDTrainer`) score completions with [`get_reward`], which feeds token IDs
+    produced by the *policy*'s tokenizer directly to the reward model, instead of decoding and re-tokenizing with the
+    reward model's own tokenizer (as [`~trl.experimental.online_dpo.OnlineDPOTrainer`] does). This is only correct if
+    the reward model uses the exact same vocabulary (token -> ID mapping) as the policy model; otherwise the reward
+    model scores the wrong tokens, which silently produces meaningless rewards or raises a cryptic out-of-range error
+    deep inside the model's embedding lookup. This raises an informative error at init time instead.
+    """
+    if not isinstance(reward_func, PreTrainedModel) or reward_processing_class is None:
+        return  # only model-based reward functions read token IDs produced by the policy's tokenizer this way
+    policy_tokenizer = (
+        policy_processing_class.tokenizer
+        if isinstance(policy_processing_class, ProcessorMixin)
+        else policy_processing_class
+    )
+    if reward_processing_class.get_vocab() != policy_tokenizer.get_vocab():
+        raise ValueError(
+            "The reward model's tokenizer vocabulary does not match the policy model's tokenizer vocabulary. This "
+            "trainer feeds token IDs produced by the policy's tokenizer directly to the reward model (it does not "
+            "decode and re-tokenize per reward model), so the reward model must share the exact same vocabulary as "
+            "the policy model. Use a reward model that shares a tokenizer with the policy model, or pass a "
+            "`reward_processing_classes` whose vocabulary matches `processing_class`."
+        )
+
+
 def prepare_model_for_kbit_training(model, use_gradient_checkpointing=True, gradient_checkpointing_kwargs=None):
     r"""
     Prepare a k-bit quantized transformers model for training (PEFT/QLoRA).

trl/experimental/xpo/xpo_trainer.py

@@ -36,7 +36,7 @@
 from ...models.utils import unwrap_model_for_generation
 from ...trainer.utils import selective_log_softmax
 from ..online_dpo import OnlineDPOTrainer
-from ..utils import empty_cache, get_reward, truncate_right
+from ..utils import empty_cache, get_reward, truncate_right, validate_reward_processing_class_shares_vocab
 from .xpo_config import XPOConfig
 
 
@@ -82,6 +82,12 @@ class XPOTrainer(OnlineDPOTrainer):
 
             If set to `None`, the tokenizer for each model-based reward function is automatically loaded using
             [`~transformers.AutoTokenizer.from_pretrained`].
+
+            Note: `XPOTrainer` scores completions by feeding token IDs produced by `processing_class` directly to the
+            reward model, instead of decoding and re-tokenizing with `reward_processing_classes` (unlike
+            [`~trl.experimental.online_dpo.OnlineDPOTrainer`]). Because of this, the reward model must share the exact
+            same vocabulary as the policy model; this is validated at init time and raises a `ValueError` if the
+            vocabularies don't match.
         peft_config ([`~peft.PeftConfig`], *optional*):
             The peft config to use for training.
         compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*):
@@ -176,6 +182,9 @@ def __init__(
         if len(self.reward_funcs) != 1:
             raise ValueError("XPOTrainer only supports one reward function/model.")
         self.reward_funcs = self.reward_funcs[0]
+        validate_reward_processing_class_shares_vocab(
+            self.reward_funcs, self.reward_processing_classes[0], self.processing_class
+        )
 
     @property
     def alpha(self):