diff --git a/tests/experimental/test_nash_md_trainer.py b/tests/experimental/test_nash_md_trainer.py index e2ce0ca8234..8178d4936e1 100644 --- a/tests/experimental/test_nash_md_trainer.py +++ b/tests/experimental/test_nash_md_trainer.py @@ -225,3 +225,50 @@ def test_train_pre_pefted_model_implicit_ref_with_reward_model(self): trainer.train() assert "train_loss" in trainer.state.log_history[-1] + + +class TestNashMDTrainerRewardProcessingClass(TrlTestCase): + def setup_method(self): + self.policy_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" + self.reward_id = "trl-internal-testing/tiny-LlamaForCausalLM-3.2" + self.model = AutoModelForCausalLM.from_pretrained(self.policy_id, dtype="float32") + self.ref_model = AutoModelForCausalLM.from_pretrained(self.policy_id) + self.reward_model = AutoModelForSequenceClassification.from_pretrained(self.reward_id, num_labels=1) + self.tokenizer = AutoTokenizer.from_pretrained(self.policy_id) + self.reward_tokenizer = AutoTokenizer.from_pretrained(self.reward_id) + self.tokenizer.pad_token = self.tokenizer.eos_token + if self.reward_tokenizer.pad_token is None: + self.reward_tokenizer.pad_token = self.reward_tokenizer.eos_token + + def test_accepts_reward_processing_classes(self): + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + training_args = NashMDConfig(output_dir=self.tmp_dir, report_to="none", per_device_train_batch_size=1) + trainer = NashMDTrainer( + model=self.model, + ref_model=self.ref_model, + reward_funcs=self.reward_model, + args=training_args, + processing_class=self.tokenizer, + reward_processing_classes=self.reward_tokenizer, + train_dataset=dataset, + ) + + assert trainer.reward_processing_classes[0].pad_token_id == self.reward_tokenizer.pad_token_id + + prompt = "hello" + completion = " world" + device = trainer.accelerator.device + prompt_ids = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")["input_ids"].to(device) + completion_ids = self.tokenizer(completion, add_special_tokens=False, return_tensors="pt")["input_ids"].to( + device + ) + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + data = { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "raw": [prompt], + } + model_scores, mixture_scores = trainer._compute_rewards(data, data, prompt_ids.shape[1]) + assert model_scores.shape == (1,) + assert mixture_scores.shape == (1,) + assert torch.isfinite(model_scores).all() diff --git a/tests/experimental/test_utils.py b/tests/experimental/test_utils.py index 58ee79cc6ad..dfbebb0d7ee 100644 --- a/tests/experimental/test_utils.py +++ b/tests/experimental/test_utils.py @@ -13,10 +13,20 @@ # limitations under the License. +import pytest +import torch from datasets import Dataset, load_dataset +from torch import nn from transformers import AutoTokenizer -from trl.experimental.utils import DataCollatorForChatML, prepare_peft_model, truncate_dataset +from trl.data_utils import apply_chat_template +from trl.experimental.utils import ( + DataCollatorForChatML, + get_reward, + get_reward_from_policy_tokens, + prepare_peft_model, + truncate_dataset, +) from ..testing_utils import TrlTestCase, require_bitsandbytes, require_peft, require_torch_accelerator @@ -182,3 +192,233 @@ def test_qlora_bf16_yields_uniform_dtype(self): fp32 = [name for name, param in model.named_parameters() if param.dtype == torch.float32] assert fp32 == [], f"expected no float32 params after prepare_peft_model, got e.g. {fp32[:5]}" + + +class _TinyBackbone(nn.Module): + def __init__(self, vocab_size, hidden_size): + super().__init__() + self.embed = nn.Embedding(vocab_size, hidden_size) + + def forward( + self, + input_ids, + attention_mask=None, + position_ids=None, + return_dict=True, + output_hidden_states=True, + use_cache=False, + ): + hidden = self.embed(input_ids) + + class Output: + pass + + output = Output() + output.hidden_states = (hidden, hidden) + return output + + +class _TinyRewardModel(nn.Module): + def __init__(self, vocab_size, hidden_size=8): + super().__init__() + self.base_model_prefix = "model" + self.config = type("Config", (), {"vocab_size": vocab_size, "hidden_size": hidden_size})() + self.model = _TinyBackbone(vocab_size, hidden_size) + self.score = nn.Linear(hidden_size, 1, bias=False) + + +class _IndexReportingBackbone(nn.Module): + def __init__(self, vocab_size, hidden_size): + super().__init__() + self.embed = nn.Embedding(vocab_size, hidden_size) + + def forward( + self, + input_ids, + attention_mask=None, + position_ids=None, + return_dict=True, + output_hidden_states=True, + use_cache=False, + ): + hidden = self.embed(input_ids).clone() + token_index = torch.arange(input_ids.size(1), device=input_ids.device, dtype=hidden.dtype) + hidden[..., 0] = token_index + output = type("Output", (), {})() + output.hidden_states = (hidden, hidden) + return output + + +class _IndexReportingRewardModel(nn.Module): + """`score` returns the token index so tests can see which position was selected.""" + + def __init__(self, vocab_size, hidden_size=8): + super().__init__() + self.base_model_prefix = "model" + self.config = type("Config", (), {"vocab_size": vocab_size, "hidden_size": hidden_size})() + self.model = _IndexReportingBackbone(vocab_size, hidden_size) + self.score = nn.Linear(hidden_size, 1, bias=False) + with torch.no_grad(): + self.score.weight.zero_() + self.score.weight[0, 0] = 1.0 + + +class TestGetRewardFromPolicyTokens(TrlTestCase): + def setup_method(self): + self.policy_tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5") + self.reward_tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-LlamaForCausalLM-3.2") + if self.policy_tokenizer.pad_token is None: + self.policy_tokenizer.pad_token = self.policy_tokenizer.eos_token + if self.reward_tokenizer.pad_token is None: + self.reward_tokenizer.pad_token = self.reward_tokenizer.eos_token + reward_vocab = max(self.reward_tokenizer.vocab_size, len(self.reward_tokenizer)) + self.reward_model = _TinyRewardModel(reward_vocab) + + def test_retokenizes_policy_ids_outside_reward_vocab(self): + prompt = "hello" + completion = " world" + prompt_ids = self.policy_tokenizer(prompt, add_special_tokens=False, return_tensors="pt")["input_ids"] + completion_ids = self.policy_tokenizer(completion, add_special_tokens=False, return_tensors="pt")["input_ids"] + query_responses = torch.cat([prompt_ids, completion_ids], dim=1) + context_length = prompt_ids.shape[1] + + oob_id = self.reward_model.config.vocab_size + if oob_id < self.policy_tokenizer.vocab_size: + query_responses = query_responses.clone() + query_responses[0, -1] = oob_id + with pytest.raises((IndexError, RuntimeError)): + get_reward(self.reward_model, query_responses, self.reward_tokenizer.pad_token_id, context_length) + + scores = get_reward_from_policy_tokens( + self.reward_model, + query_responses, + context_length, + [prompt], + self.policy_tokenizer, + self.reward_tokenizer, + ) + assert scores.shape == (1,) + assert torch.isfinite(scores).all() + + def test_conversational_prompts(self): + prompt = [{"role": "user", "content": "Hi"}] + completion = "Hello" + completion_ids = self.policy_tokenizer(completion, add_special_tokens=False, return_tensors="pt")["input_ids"] + context_length = 1 + query_responses = torch.cat([torch.zeros((1, context_length), dtype=torch.long), completion_ids], dim=1) + + scores = get_reward_from_policy_tokens( + self.reward_model, + query_responses, + context_length, + [prompt], + self.policy_tokenizer, + self.reward_tokenizer, + ) + assert scores.shape == (1,) + assert torch.isfinite(scores).all() + + def test_conversational_scores_last_token_not_prompt_eos(self): + tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5") + tokenizer.pad_token = tokenizer.eos_token + assert tokenizer.pad_token_id == tokenizer.eos_token_id + + prompt = [{"role": "user", "content": "Hi"}] + completion = "Hello there" + example = {"messages": prompt + [{"role": "assistant", "content": completion}]} + text = apply_chat_template(example, tokenizer)["text"] + encoded = tokenizer(text, add_special_tokens=False, return_tensors="pt") + token_ids = encoded["input_ids"][0] + pad_positions = (token_ids == tokenizer.pad_token_id).nonzero(as_tuple=False) + last_index = token_ids.size(0) - 1 + assert pad_positions.numel() > 0, "chat template must emit eos/pad at turn boundaries for this test" + first_pad_index = int(pad_positions[0]) + assert first_pad_index < last_index, "first pad/eos must not be the completion end" + + vocab = max(tokenizer.vocab_size, len(tokenizer), int(token_ids.max()) + 1) + reward_model = _IndexReportingRewardModel(vocab) + completion_ids = tokenizer(completion, add_special_tokens=False, return_tensors="pt")["input_ids"] + context_length = 2 + query_responses = torch.cat([torch.zeros((1, context_length), dtype=torch.long), completion_ids], dim=1) + + scores = get_reward_from_policy_tokens( + reward_model, + query_responses, + context_length, + [prompt], + tokenizer, + tokenizer, + ) + scored_index = int(scores.item()) + assert scored_index == last_index + assert scored_index != first_pad_index - 1 + + def test_batched_chat_padding_scores_per_row_last_token(self): + tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5") + tokenizer.pad_token = tokenizer.eos_token + prompts = [ + [{"role": "user", "content": "Hi"}], + [{"role": "user", "content": "A much longer user question for padding"}], + ] + completions = ["Yes", "This is a longer assistant reply"] + expected = [] + for prompt, completion in zip(prompts, completions, strict=True): + text = apply_chat_template( + {"messages": prompt + [{"role": "assistant", "content": completion}]}, tokenizer + )["text"] + token_ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][0] + expected.append(token_ids.size(0) - 1) + + vocab = max(tokenizer.vocab_size, len(tokenizer)) + reward_model = _IndexReportingRewardModel(vocab) + completion_ids = tokenizer(completions, add_special_tokens=False, padding=True, return_tensors="pt")[ + "input_ids" + ] + context_length = 1 + query_responses = torch.cat( + [torch.zeros((len(prompts), context_length), dtype=torch.long), completion_ids], dim=1 + ) + + scores = get_reward_from_policy_tokens( + reward_model, + query_responses, + context_length, + prompts, + tokenizer, + tokenizer, + ) + assert [int(x) for x in scores.tolist()] == expected + + def test_same_tokenizer_matches_add_special_tokens_false(self): + tokenizer = self.reward_tokenizer + if tokenizer.bos_token_id is None: + pytest.skip("tokenizer needs a BOS token to pin the Online DPO encoding") + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + prompt = "The answer is 2 + 2?" + completion = " that is four." + with_bos = tokenizer(prompt + completion, add_special_tokens=True, return_tensors="pt")["input_ids"][0] + without_bos = tokenizer(prompt + completion, add_special_tokens=False, return_tensors="pt")["input_ids"][0] + assert with_bos[0] == tokenizer.bos_token_id + assert without_bos[0] != tokenizer.bos_token_id + assert with_bos.tolist() != without_bos.tolist() + + vocab = max(tokenizer.vocab_size, len(tokenizer), int(with_bos.max()) + 1) + reward_model = _IndexReportingRewardModel(vocab) + + prompt_ids = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")["input_ids"] + prompt_ids = torch.cat([torch.tensor([[tokenizer.bos_token_id]], dtype=prompt_ids.dtype), prompt_ids], dim=1) + completion_ids = tokenizer(completion, add_special_tokens=False, return_tensors="pt")["input_ids"] + query_responses = torch.cat([prompt_ids, completion_ids], dim=1) + + scores = get_reward_from_policy_tokens( + reward_model, + query_responses, + prompt_ids.shape[1], + [prompt], + tokenizer, + tokenizer, + ) + assert int(scores.item()) == without_bos.size(0) - 1 + assert int(scores.item()) != query_responses.size(1) - 1 diff --git a/tests/experimental/test_xpo_trainer.py b/tests/experimental/test_xpo_trainer.py index ac64d6633aa..2be685c9e1d 100644 --- a/tests/experimental/test_xpo_trainer.py +++ b/tests/experimental/test_xpo_trainer.py @@ -13,6 +13,7 @@ # limitations under the License. import pytest +import torch from datasets import DatasetDict, load_dataset from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer from transformers.utils import is_peft_available @@ -173,3 +174,50 @@ def test_train_pre_pefted_model_implicit_ref(self): trainer.train() assert "train_loss" in trainer.state.log_history[-1] + + +class TestXPOTrainerRewardProcessingClass(TrlTestCase): + def setup_method(self): + self.policy_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" + self.reward_id = "trl-internal-testing/tiny-LlamaForCausalLM-3.2" + self.model = AutoModelForCausalLM.from_pretrained(self.policy_id, dtype="float32") + self.ref_model = AutoModelForCausalLM.from_pretrained(self.policy_id) + self.reward_model = AutoModelForSequenceClassification.from_pretrained(self.reward_id, num_labels=1) + self.tokenizer = AutoTokenizer.from_pretrained(self.policy_id) + self.reward_tokenizer = AutoTokenizer.from_pretrained(self.reward_id) + self.tokenizer.pad_token = self.tokenizer.eos_token + if self.reward_tokenizer.pad_token is None: + self.reward_tokenizer.pad_token = self.reward_tokenizer.eos_token + + def test_compute_rewards_with_distinct_reward_tokenizer(self): + dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train") + training_args = XPOConfig(output_dir=self.tmp_dir, report_to="none", per_device_train_batch_size=1) + trainer = XPOTrainer( + model=self.model, + ref_model=self.ref_model, + reward_funcs=self.reward_model, + args=training_args, + processing_class=self.tokenizer, + reward_processing_classes=self.reward_tokenizer, + train_dataset=dataset, + ) + + assert trainer.reward_processing_classes[0].pad_token_id == self.reward_tokenizer.pad_token_id + + prompt = "hello" + completion = " world" + device = trainer.accelerator.device + prompt_ids = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")["input_ids"].to(device) + completion_ids = self.tokenizer(completion, add_special_tokens=False, return_tensors="pt")["input_ids"].to( + device + ) + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + data = { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "raw": [prompt], + } + model_scores, ref_scores = trainer._compute_rewards(data, data, prompt_ids.shape[1]) + assert model_scores.shape == (1,) + assert ref_scores.shape == (1,) + assert torch.isfinite(model_scores).all() diff --git a/trl/experimental/nash_md/nash_md_trainer.py b/trl/experimental/nash_md/nash_md_trainer.py index e824775f14f..f22bc00db6a 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_from_policy_tokens, truncate_right from .nash_md_config import NashMDConfig @@ -139,6 +139,14 @@ 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_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + + If set to `None`, the tokenizer for each model-based reward function is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. peft_config ([`~peft.PeftConfig`], *optional*): The peft config to use for training. compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): @@ -183,6 +191,7 @@ def __init__( | FeatureExtractionMixin | ProcessorMixin | None = None, + reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, peft_config: "PeftConfig | None" = None, compute_metrics: Callable[[EvalPrediction], dict] | None = None, callbacks: list[TrainerCallback] | None = None, @@ -198,7 +207,7 @@ def __init__( train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=processing_class, - reward_processing_classes=processing_class, + reward_processing_classes=reward_processing_classes, peft_config=peft_config, compute_metrics=compute_metrics, callbacks=callbacks, @@ -322,12 +331,23 @@ def _process_completions(self, model_output, mixture_output, prompts): return model_data, mixture_data def _compute_rewards(self, model_data, mixture_data, context_length): + reward_processing_class = self.reward_processing_classes[0] with torch.no_grad(): - _, model_scores, _ = get_reward( - self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length + model_scores = get_reward_from_policy_tokens( + self.reward_funcs, + model_data["input_ids"], + context_length, + model_data["raw"], + self.processing_class, + reward_processing_class, ) - _, mixture_scores, _ = get_reward( - self.reward_funcs, mixture_data["input_ids"], self.processing_class.pad_token_id, context_length + mixture_scores = get_reward_from_policy_tokens( + self.reward_funcs, + mixture_data["input_ids"], + context_length, + mixture_data["raw"], + self.processing_class, + reward_processing_class, ) # Apply EOS penalty if needed diff --git a/trl/experimental/utils.py b/trl/experimental/utils.py index 26edd89abd1..1b34e3f42f9 100644 --- a/trl/experimental/utils.py +++ b/trl/experimental/utils.py @@ -743,6 +743,24 @@ def first_true_indices(bools: torch.Tensor, dtype=torch.long) -> torch.Tensor: return torch.min(zero_or_index, dim=-1).values +def _compute_reward_logits( + model: torch.nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor +) -> torch.Tensor: + attention_mask = attention_mask.bool() + position_ids = attention_mask.cumsum(1) - attention_mask.long() # exclusive cumsum + lm_backbone = getattr(model, model.base_model_prefix) + masked_input_ids = torch.masked_fill(input_ids, ~attention_mask, 0) + output = lm_backbone( + input_ids=masked_input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + output_hidden_states=True, + use_cache=False, # otherwise mistral-based RM would error out + ) + return model.score(output.hidden_states[-1]) + + def get_reward( model: torch.nn.Module, query_responses: torch.Tensor, pad_token_id: int, context_length: int ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -769,18 +787,7 @@ def get_reward( The lengths of the sequences in the query responses. """ attention_mask = query_responses != pad_token_id - position_ids = attention_mask.cumsum(1) - attention_mask.long() # exclusive cumsum - lm_backbone = getattr(model, model.base_model_prefix) - input_ids = torch.masked_fill(query_responses, ~attention_mask, 0) - output = lm_backbone( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - return_dict=True, - output_hidden_states=True, - use_cache=False, # otherwise mistral-based RM would error out - ) - reward_logits = model.score(output.hidden_states[-1]) + reward_logits = _compute_reward_logits(model, query_responses, attention_mask) sequence_lengths = first_true_indices(query_responses[:, context_length:] == pad_token_id) - 1 + context_length # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454 return ( @@ -793,6 +800,65 @@ def get_reward( ) +def get_reward_from_policy_tokens( + model: torch.nn.Module, + query_responses: torch.Tensor, + context_length: int, + prompts: list, + policy_processing_class: PreTrainedTokenizerBase, + reward_processing_class: PreTrainedTokenizerBase, +) -> torch.Tensor: + """ + Score policy-token sequences with a reward model that may use a different tokenizer. + + Completions are decoded with `policy_processing_class` and re-tokenized with `reward_processing_class`, matching + [`experimental.online_dpo.OnlineDPOTrainer`]. Scoring uses the tokenizer `attention_mask` and the last non-padding + token. This is required for chat templates: `pad_token_id` is often `eos_token_id` and appears at turn boundaries, + so [`get_reward`]'s "first pad id" rule would score the prompt. + + Args: + model (`torch.nn.Module`): + The reward model used to compute the scores. + query_responses (`torch.Tensor`): + Policy-token sequences of shape `(batch_size, sequence_length)` containing the prompt followed by the + completion. + context_length (`int`): + Number of prompt tokens in `query_responses`. + prompts (`list`): + Raw prompts corresponding to each row, either strings or conversational message lists. + policy_processing_class ([`~transformers.PreTrainedTokenizerBase`]): + Tokenizer that produced `query_responses`. + reward_processing_class ([`~transformers.PreTrainedTokenizerBase`]): + Tokenizer of the reward model. + + Returns: + `torch.Tensor`: + Reward scores of shape `(batch_size,)`. + """ + # Tokenization matches `OnlineDPOTrainer._calculate_rewards_from_functions` (`skip_special_tokens=True`, + # `add_special_tokens=False`). That drops BOS on the already-working same-tokenizer path. Scoring differs: Online + # DPO reads SequenceClassification `logits[:, 0]`; we keep the XPO/Nash-MD `model.score` last-token head. + completions = policy_processing_class.batch_decode(query_responses[:, context_length:], skip_special_tokens=True) + if is_conversational({"prompt": prompts[0]}): + completion_messages = [[{"role": "assistant", "content": completion}] for completion in completions] + examples = [ + {"messages": prompt + completion} for prompt, completion in zip(prompts, completion_messages, strict=True) + ] + texts = [apply_chat_template(example, reward_processing_class)["text"] for example in examples] + else: + texts = [prompt + completion for prompt, completion in zip(prompts, completions, strict=True)] + + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + device = query_responses.device + reward_ids = reward_inputs["input_ids"].to(device=device) + attention_mask = reward_inputs["attention_mask"].to(device=device) + reward_logits = _compute_reward_logits(model, reward_ids, attention_mask) + sequence_lengths = attention_mask.sum(dim=1).clamp(min=1) - 1 + return reward_logits[torch.arange(reward_logits.size(0), device=device), sequence_lengths].squeeze(-1) + + 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..5f6115e1219 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_from_policy_tokens, truncate_right from .xpo_config import XPOConfig @@ -253,12 +253,23 @@ def _process_completions(self, model_output, ref_output, prompts): return model_data, ref_data def _compute_rewards(self, model_data, ref_data, context_length): + reward_processing_class = self.reward_processing_classes[0] with torch.no_grad(): - _, model_scores, _ = get_reward( - self.reward_funcs, model_data["input_ids"], self.processing_class.pad_token_id, context_length + model_scores = get_reward_from_policy_tokens( + self.reward_funcs, + model_data["input_ids"], + context_length, + model_data["raw"], + self.processing_class, + reward_processing_class, ) - _, ref_scores, _ = get_reward( - self.reward_funcs, ref_data["input_ids"], self.processing_class.pad_token_id, context_length + ref_scores = get_reward_from_policy_tokens( + self.reward_funcs, + ref_data["input_ids"], + context_length, + ref_data["raw"], + self.processing_class, + reward_processing_class, ) # Apply EOS penalty if needed