Validate reward model tokenizer matches policy in XPOTrainer/NashMDTrainer

tests/experimental/test_nash_md_trainer.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import patch
+
 import pytest
 import torch
+from accelerate import Accelerator
 from datasets import DatasetDict, load_dataset
 from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, GenerationConfig
 from transformers.utils import is_peft_available
@@ -195,6 +198,95 @@ 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_vocab_mismatch_raises_when_reward_model_is_wrapped(self):
+        # Regression test: `self.reward_funcs` can already be wrapped (a DeepSpeed engine, or
+        # torch.nn.parallel.DistributedDataParallel under multi-process DDP) by `Accelerator.prepare_model` inside
+        # `super().__init__()`, before the vocab check runs. A wrapper is not itself a `PreTrainedModel`, so the
+        # check must unwrap first or it silently never fires for exactly the training setups (DeepSpeed/DDP) where
+        # catching a mismatch matters most. `Accelerator.prepare_model` doesn't actually wrap on a plain CPU/
+        # single-process test, so reproduce the same wrapper `accelerate.unwrap_model` has to see through in
+        # production by wrapping the reward model in a real (single-process) `DistributedDataParallel`.
+        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")
+        reward_model = self.reward_model
+
+        original_prepare_model = Accelerator.prepare_model
+
+        def fake_prepare_model(self, model, *args, **kwargs):
+            prepared = original_prepare_model(self, model, *args, **kwargs)
+            return torch.nn.parallel.DistributedDataParallel(prepared) if prepared is reward_model else prepared
+
+        store = torch.distributed.HashStore()
+        torch.distributed.init_process_group(backend="gloo", store=store, rank=0, world_size=1)
+        try:
+            with patch.object(Accelerator, "prepare_model", fake_prepare_model):
+                with pytest.raises(ValueError, match="vocabulary"):
+                    NashMDTrainer(
+                        model=self.model,
+                        ref_model=self.ref_model,
+                        reward_funcs=reward_model,
+                        reward_processing_class=mismatched_tokenizer,
+                        args=training_args,
+                        processing_class=self.tokenizer,
+                        train_dataset=dataset,
+                    )
+        finally:
+            torch.distributed.destroy_process_group()
+
+    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_utils.py

@@ -13,10 +13,16 @@
 # limitations under the License.
 
 
+import pytest
 from datasets import Dataset, load_dataset
-from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer
 
-from trl.experimental.utils import DataCollatorForChatML, prepare_peft_model, truncate_dataset
+from trl.experimental.utils import (
+    DataCollatorForChatML,
+    prepare_peft_model,
+    truncate_dataset,
+    validate_reward_processing_class_shares_vocab,
+)
 
 from ..testing_utils import TrlTestCase, require_bitsandbytes, require_peft, require_torch_accelerator
 
@@ -160,6 +166,36 @@ def test_with_extra_column(self):
         assert dataset.to_dict() == expected_output
 
 
+class TestValidateRewardProcessingClassSharesVocab(TrlTestCase):
+    def setup_method(self):
+        self.reward_func = AutoModelForSequenceClassification.from_pretrained(
+            "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", num_labels=1
+        )
+
+    def test_accepts_processor_mixin_reward_processing_class_with_matching_vocab(self):
+        # Regression test for #6951: `NashMDTrainer` defaults `reward_processing_class` to `processing_class`,
+        # which can itself be a `ProcessorMixin` (e.g. for VLM policies), not just a tokenizer. The reward side
+        # must be unwrapped to its `.tokenizer` the same way the policy side already is, or `get_vocab()` raises
+        # `AttributeError` (`ProcessorMixin` doesn't define it) instead of comparing vocabularies -- even when
+        # it's the exact same tokenizer on both sides.
+        processor = AutoProcessor.from_pretrained("trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration")
+
+        validate_reward_processing_class_shares_vocab(
+            self.reward_func, reward_processing_class=processor, policy_processing_class=processor.tokenizer
+        )  # must not raise (and, before the fix, raised AttributeError instead of passing)
+
+    def test_rejects_processor_mixin_reward_processing_class_with_mismatched_vocab(self):
+        # A `ProcessorMixin` on the reward side must still be compared correctly (not skipped, not crash) when its
+        # underlying tokenizer's vocabulary actually differs from the policy's.
+        processor = AutoProcessor.from_pretrained("trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration")
+        mismatched_tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
+
+        with pytest.raises(ValueError, match="vocabulary"):
+            validate_reward_processing_class_shares_vocab(
+                self.reward_func, reward_processing_class=processor, policy_processing_class=mismatched_tokenizer
+            )
+
+
 class TestPreparePeftModel(TrlTestCase):
     @require_peft
     @require_bitsandbytes

tests/experimental/test_xpo_trainer.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import patch
+
 import pytest
+import torch
+from accelerate import Accelerator
 from datasets import DatasetDict, load_dataset
 from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
 from transformers.utils import is_peft_available
@@ -145,6 +149,79 @@ 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]
+
+    def test_reward_processing_class_vocab_mismatch_raises_when_reward_model_is_wrapped(self):
+        # Regression test: `self.reward_funcs` can already be wrapped (a DeepSpeed engine, or
+        # torch.nn.parallel.DistributedDataParallel under multi-process DDP) by `Accelerator.prepare_model` inside
+        # `super().__init__()`, before the vocab check runs. A wrapper is not itself a `PreTrainedModel`, so the
+        # check must unwrap first or it silently never fires for exactly the training setups (DeepSpeed/DDP) where
+        # catching a mismatch matters most. `Accelerator.prepare_model` doesn't actually wrap on a plain CPU/
+        # single-process test, so reproduce the same wrapper `accelerate.unwrap_model` has to see through in
+        # production by wrapping the reward model in a real (single-process) `DistributedDataParallel`, instead of
+        # a hand-rolled stand-in that `extract_model_from_parallel` wouldn't recognize either way.
+        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")
+        reward_model = self.reward_model
+
+        original_prepare_model = Accelerator.prepare_model
+
+        def fake_prepare_model(self, model, *args, **kwargs):
+            prepared = original_prepare_model(self, model, *args, **kwargs)
+            return torch.nn.parallel.DistributedDataParallel(prepared) if prepared is reward_model else prepared
+
+        store = torch.distributed.HashStore()
+        torch.distributed.init_process_group(backend="gloo", store=store, rank=0, world_size=1)
+        try:
+            with patch.object(Accelerator, "prepare_model", fake_prepare_model):
+                with pytest.raises(ValueError, match="vocabulary"):
+                    XPOTrainer(
+                        model=self.model,
+                        ref_model=self.ref_model,
+                        reward_funcs=reward_model,
+                        reward_processing_classes=mismatched_tokenizer,
+                        args=training_args,
+                        processing_class=self.tokenizer,
+                        train_dataset=dataset,
+                    )
+        finally:
+            torch.distributed.destroy_process_group()
+
     @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,11 @@ def __init__(
         if len(self.reward_funcs) != 1:
             raise ValueError("NashMDTrainer only supports one reward function/model.")
         self.reward_funcs = self.reward_funcs[0]
+        # self.reward_funcs may already be wrapped (DeepSpeed engine / DDP) by the super().__init__() call above,
+        # which would make it fail the PreTrainedModel check below; unwrap first so the check still fires.
+        validate_reward_processing_class_shares_vocab(
+            self.accelerator.unwrap_model(self.reward_funcs), self.reward_processing_classes[0], self.processing_class
+        )
 
     @property
     def mixture_coef(self):

trl/experimental/utils.py

@@ -793,6 +793,43 @@ def get_reward(
     )
 
 
+def validate_reward_processing_class_shares_vocab(
+    reward_func: Any,
+    reward_processing_class: PreTrainedTokenizerBase | ProcessorMixin | 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
+    )
+    reward_tokenizer = (
+        reward_processing_class.tokenizer
+        if isinstance(reward_processing_class, ProcessorMixin)
+        else reward_processing_class
+    )
+    if reward_tokenizer.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).