From 84ddbd80075f93ebedf8116875d79f02a96318c2 Mon Sep 17 00:00:00 2001 From: amanyagami <2amansingh2@gmail.com> Date: Thu, 27 Aug 2026 20:46:52 -0700 Subject: [PATCH 1/2] Validate reward model tokenizer matches policy in XPOTrainer/NashMDTrainer 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 #6951 --- tests/experimental/test_nash_md_trainer.py | 53 +++++++++++++++++++++ tests/experimental/test_xpo_trainer.py | 36 ++++++++++++++ trl/experimental/nash_md/nash_md_trainer.py | 21 +++++++- trl/experimental/utils.py | 32 +++++++++++++ trl/experimental/xpo/xpo_trainer.py | 11 ++++- 5 files changed, 150 insertions(+), 3 deletions(-) diff --git a/tests/experimental/test_nash_md_trainer.py b/tests/experimental/test_nash_md_trainer.py index e2ce0ca8234..db180d49b00 100644 --- a/tests/experimental/test_nash_md_trainer.py +++ b/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") diff --git a/tests/experimental/test_xpo_trainer.py b/tests/experimental/test_xpo_trainer.py index ac64d6633aa..459b7caee84 100644 --- a/tests/experimental/test_xpo_trainer.py +++ b/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") diff --git a/trl/experimental/nash_md/nash_md_trainer.py b/trl/experimental/nash_md/nash_md_trainer.py index e824775f14f..a9def2db6eb 100644 --- a/trl/experimental/nash_md/nash_md_trainer.py +++ b/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): diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index 26edd89abd1..fe7b5635e3c 100644 --- a/trl/experimental/utils.py +++ b/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). diff --git a/trl/experimental/xpo/xpo_trainer.py b/trl/experimental/xpo/xpo_trainer.py index 5edbbc3dc20..5ee20c00b40 100644 --- a/trl/experimental/xpo/xpo_trainer.py +++ b/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): From fbfba76859646bcf64be59fc9e68c266295fbaab Mon Sep 17 00:00:00 2001 From: amanyagami <2amansingh2@gmail.com> Date: Thu, 27 Aug 2026 21:59:04 -0700 Subject: [PATCH 2/2] Address Cursor Bugbot review: unwrap reward model before vocab check, unwrap ProcessorMixin symmetrically Two real bugs in the initial fix, both caught by automated review on PR #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. --- tests/experimental/test_nash_md_trainer.py | 39 ++++++++++++++++++++ tests/experimental/test_utils.py | 40 +++++++++++++++++++- tests/experimental/test_xpo_trainer.py | 41 +++++++++++++++++++++ trl/experimental/nash_md/nash_md_trainer.py | 4 +- trl/experimental/utils.py | 9 ++++- trl/experimental/xpo/xpo_trainer.py | 4 +- 6 files changed, 131 insertions(+), 6 deletions(-) diff --git a/tests/experimental/test_nash_md_trainer.py b/tests/experimental/test_nash_md_trainer.py index db180d49b00..77e42db8b42 100644 --- a/tests/experimental/test_nash_md_trainer.py +++ b/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 @@ -214,6 +217,42 @@ def test_reward_processing_class_vocab_mismatch_raises(self): 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 diff --git a/tests/experimental/test_utils.py b/tests/experimental/test_utils.py index 58ee79cc6ad..6de5cfe5e8e 100644 --- a/tests/experimental/test_utils.py +++ b/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 diff --git a/tests/experimental/test_xpo_trainer.py b/tests/experimental/test_xpo_trainer.py index 459b7caee84..d25e0bbe07e 100644 --- a/tests/experimental/test_xpo_trainer.py +++ b/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 @@ -181,6 +185,43 @@ def test_reward_processing_class_matching_vocab_is_accepted(self): ) 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") diff --git a/trl/experimental/nash_md/nash_md_trainer.py b/trl/experimental/nash_md/nash_md_trainer.py index a9def2db6eb..6d8b4b069c2 100644 --- a/trl/experimental/nash_md/nash_md_trainer.py +++ b/trl/experimental/nash_md/nash_md_trainer.py @@ -244,8 +244,10 @@ 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.reward_funcs, self.reward_processing_classes[0], self.processing_class + self.accelerator.unwrap_model(self.reward_funcs), self.reward_processing_classes[0], self.processing_class ) @property diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index fe7b5635e3c..58a9a85fc60 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -795,7 +795,7 @@ def get_reward( def validate_reward_processing_class_shares_vocab( reward_func: Any, - reward_processing_class: PreTrainedTokenizerBase | None, + reward_processing_class: PreTrainedTokenizerBase | ProcessorMixin | None, policy_processing_class: PreTrainedTokenizerBase | ProcessorMixin, ) -> None: """ @@ -815,7 +815,12 @@ def validate_reward_processing_class_shares_vocab( if isinstance(policy_processing_class, ProcessorMixin) else policy_processing_class ) - if reward_processing_class.get_vocab() != policy_tokenizer.get_vocab(): + 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 " diff --git a/trl/experimental/xpo/xpo_trainer.py b/trl/experimental/xpo/xpo_trainer.py index 5ee20c00b40..2874a8662b3 100644 --- a/trl/experimental/xpo/xpo_trainer.py +++ b/trl/experimental/xpo/xpo_trainer.py @@ -182,8 +182,10 @@ def __init__( if len(self.reward_funcs) != 1: raise ValueError("XPOTrainer 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.reward_funcs, self.reward_processing_classes[0], self.processing_class + self.accelerator.unwrap_model(self.reward_funcs), self.reward_processing_classes[0], self.processing_class ) @property