test(online_dpo): pin eval_loss to the batch losses and refuse compute_metrics

The three evaluate tests bounded the magnitude of eval_loss, which a
prediction_step returning any constant satisfies. They now record every
_compute_loss call made during evaluate() and pin eval_loss to the mean
of those losses weighted by batch size, as Trainer aggregates them: the
17-row split leaves a final batch of one example, so an unweighted mean
misses by 0.17 on NashMD. One call per evaluation batch, statistics off.

prediction_step returns no predictions or labels, so a compute_metrics
callback passed to these trainers was accepted and never called.
OnlineDPO, NashMD and XPO now raise a ValueError for it and the
docstrings say so.

Author: behroozazarkhalili

tests/experimental/test_nash_md_trainer.py

@@ -167,11 +167,26 @@ def test_evaluate(self):
             processing_class=self.tokenizer,
         )
 
-        # A hardcoded return value would satisfy a bare membership check, so assert the loss itself. Its sign is
-        # not fixed across these trainers, so bound the magnitude instead: non-zero rules out a stubbed
-        # `prediction_step`, and finite rules out inf and NaN.
+        # A stubbed `prediction_step` returning any constant passes a magnitude bound, so pin `eval_loss` to the losses
+        # `_compute_loss` produced on the evaluation batches, weighted by batch size as `Trainer` does.
+        recorded = []
+        original = trainer._compute_loss
+
+        def spy(model, inputs, log_stats=True):
+            loss = original(model, inputs, log_stats=log_stats)
+            batch_size = len(next(iter(inputs.values())))
+            recorded.append((log_stats, batch_size, loss.item()))
+            return loss
+
+        trainer._compute_loss = spy
         metrics = trainer.evaluate()
-        assert 0 < abs(metrics["eval_loss"]) < float("inf")
+
+        assert len(recorded) == len(trainer.get_eval_dataloader())
+        assert all(log_stats is False for log_stats, _, _ in recorded)
+        expected_loss = sum(batch_size * loss for _, batch_size, loss in recorded) / sum(
+            batch_size for _, batch_size, _ in recorded
+        )
+        assert metrics["eval_loss"] == pytest.approx(expected_loss, abs=1e-5)
 
     def test_evaluate_does_not_pollute_training_stats(self):
         # `self.stats` is averaged and cleared by the training logger, so appending evaluation values to it would

tests/experimental/test_online_dpo_trainer.py

@@ -148,11 +148,41 @@ def test_evaluate(self):
             reward_processing_classes=self.reward_tokenizer,
         )
 
-        # A hardcoded return value would satisfy a bare membership check, so assert the loss itself. Its sign is
-        # not fixed across these trainers, so bound the magnitude instead: non-zero rules out a stubbed
-        # `prediction_step`, and finite rules out inf and NaN.
+        # A stubbed `prediction_step` returning any constant passes a magnitude bound, so pin `eval_loss` to the losses
+        # `_compute_loss` produced on the evaluation batches, weighted by batch size as `Trainer` does.
+        recorded = []
+        original = trainer._compute_loss
+
+        def spy(model, inputs, log_stats=True):
+            loss = original(model, inputs, log_stats=log_stats)
+            batch_size = len(next(iter(inputs.values())))
+            recorded.append((log_stats, batch_size, loss.item()))
+            return loss
+
+        trainer._compute_loss = spy
         metrics = trainer.evaluate()
-        assert 0 < abs(metrics["eval_loss"]) < float("inf")
+
+        assert len(recorded) == len(trainer.get_eval_dataloader())
+        assert all(log_stats is False for log_stats, _, _ in recorded)
+        expected_loss = sum(batch_size * loss for _, batch_size, loss in recorded) / sum(
+            batch_size for _, batch_size, _ in recorded
+        )
+        assert metrics["eval_loss"] == pytest.approx(expected_loss, abs=1e-5)
+
+    def test_compute_metrics_is_rejected(self):
+        # `prediction_step` returns only a loss, so a `compute_metrics` callback would never run; refuse it up front.
+        training_args = OnlineDPOConfig(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="`compute_metrics` is not supported"):
+            OnlineDPOTrainer(
+                model=self.model,
+                reward_funcs=self.reward_model,
+                args=training_args,
+                train_dataset=dataset,
+                processing_class=self.tokenizer,
+                reward_processing_classes=self.reward_tokenizer,
+                compute_metrics=lambda eval_prediction: {},
+            )
 
     def test_evaluate_does_not_pollute_training_stats(self):
         # `self.stats` is averaged and cleared by the training logger, so appending evaluation values to it would

tests/experimental/test_xpo_trainer.py

@@ -117,11 +117,26 @@ def test_evaluate(self):
             processing_class=self.tokenizer,
         )
 
-        # A hardcoded return value would satisfy a bare membership check, so assert the loss itself. Its sign is
-        # not fixed across these trainers, so bound the magnitude instead: non-zero rules out a stubbed
-        # `prediction_step`, and finite rules out inf and NaN.
+        # A stubbed `prediction_step` returning any constant passes a magnitude bound, so pin `eval_loss` to the losses
+        # `_compute_loss` produced on the evaluation batches, weighted by batch size as `Trainer` does.
+        recorded = []
+        original = trainer._compute_loss
+
+        def spy(model, inputs, log_stats=True):
+            loss = original(model, inputs, log_stats=log_stats)
+            batch_size = len(next(iter(inputs.values())))
+            recorded.append((log_stats, batch_size, loss.item()))
+            return loss
+
+        trainer._compute_loss = spy
         metrics = trainer.evaluate()
-        assert 0 < abs(metrics["eval_loss"]) < float("inf")
+
+        assert len(recorded) == len(trainer.get_eval_dataloader())
+        assert all(log_stats is False for log_stats, _, _ in recorded)
+        expected_loss = sum(batch_size * loss for _, batch_size, loss in recorded) / sum(
+            batch_size for _, batch_size, _ in recorded
+        )
+        assert metrics["eval_loss"] == pytest.approx(expected_loss, abs=1e-5)
 
     def test_evaluate_does_not_pollute_training_stats(self):
         # `self.stats` is averaged and cleared by the training logger, so appending evaluation values to it would

trl/experimental/nash_md/nash_md_trainer.py

@@ -142,8 +142,8 @@ class NashMDTrainer(OnlineDPOTrainer):
         peft_config ([`~peft.PeftConfig`], *optional*):
             The peft config to use for training.
         compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*):
-            The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to
-            metric values.
+            Not supported: evaluation computes only a loss, since `prediction_step` returns no predictions or labels
+            for the callback. Passing a callable raises a `ValueError`.
         callbacks (`list[transformers.TrainerCallback]`):
             The callbacks to use for training.
         optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):
@@ -189,6 +189,11 @@ def __init__(
         optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
         preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
     ) -> None:
+        if compute_metrics is not None:
+            raise ValueError(
+                "`compute_metrics` is not supported: `prediction_step` returns only a loss, so the callback would never "
+                "be called. Read `eval_loss` from `evaluate()` instead."
+            )
         super().__init__(
             model=model,
             ref_model=ref_model,

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -149,8 +149,8 @@ class OnlineDPOTrainer(_BaseTrainer):
         peft_config ([`~peft.PeftConfig`], *optional*):
             PEFT configuration used to wrap the model. If `None`, the model is not wrapped.
         compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*):
-            The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to
-            metric values.
+            Not supported: evaluation computes only a loss, since `prediction_step` returns no predictions or labels
+            for the callback. Passing a callable raises a `ValueError`.
         callbacks (`list[transformers.TrainerCallback]`):
             The callbacks to use for training.
         optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):
@@ -403,6 +403,11 @@ def __init__(
             args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {}
             args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False)
 
+        if compute_metrics is not None:
+            raise ValueError(
+                "`compute_metrics` is not supported: `prediction_step` returns only a loss, so the callback would never "
+                "be called. Read `eval_loss` from `evaluate()` instead."
+            )
         super().__init__(
             model=model,
             args=args,

trl/experimental/xpo/xpo_trainer.py

@@ -85,8 +85,8 @@ class XPOTrainer(OnlineDPOTrainer):
         peft_config ([`~peft.PeftConfig`], *optional*):
             The peft config to use for training.
         compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*):
-            The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to
-            metric values.
+            Not supported: evaluation computes only a loss, since `prediction_step` returns no predictions or labels
+            for the callback. Passing a callable raises a `ValueError`.
         callbacks (`list[transformers.TrainerCallback]`):
             The callbacks to use for training.
         optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`):
@@ -131,6 +131,11 @@ def __init__(
         optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
         preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
     ) -> None:
+        if compute_metrics is not None:
+            raise ValueError(
+                "`compute_metrics` is not supported: `prediction_step` returns only a loss, so the callback would never "
+                "be called. Read `eval_loss` from `evaluate()` instead."
+            )
         super().__init__(
             model=model,
             ref_model=ref_model,