diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index bd9b8b3f12c..19addcf444a 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -524,6 +524,53 @@ def test_train_with_f_divergence(self, f_divergence_type): new_param = trainer.model.get_parameter(n) assert not torch.equal(param, new_param), f"Parameter {n} has not changed." + @pytest.mark.parametrize( + "f_divergence_type", + ["forward_kl", "js_divergence", "alpha_divergence"], + ) + def test_apo_down_loss_is_f_divergence_invariant(self, f_divergence_type): + # apo_down implements Eq. 8 of the APO paper (https://huggingface.co/papers/2408.06266), which is + # defined purely on the raw reverse-KL log-ratios r_theta, exactly like apo_zero. Its loss must + # therefore be identical regardless of `f_divergence_type`. Previously the rejected term used the + # f-divergence-transformed `delta_score` while the chosen term used the raw log-ratio, so a + # non-default `f_divergence_type` silently changed the loss (it only matched by accident under the + # default `reverse_kl`). See #6441. + dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") + + training_args = DPOConfig( + output_dir=self.tmp_dir, + loss_type="apo_down", + per_device_train_batch_size=4, + bf16=False, # precision-independent loss-math check; keep fp32 so it runs on CPU too + report_to="none", + ) + trainer = DPOTrainer( + model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + args=training_args, + train_dataset=dataset, + ) + trainer.model.eval() + + # The reference model is a frozen snapshot of the initial policy, so an untrained policy has zero + # log-ratios and `delta_score` collapses to zero under every `f_divergence_type`, which would hide + # the bug. Perturb the policy weights (the reference stays fixed) so the log-ratios are non-zero. + torch.manual_seed(0) + with torch.no_grad(): + for param in trainer.model.parameters(): + param.add_(torch.randn_like(param) * 0.05) + + # Reuse a single fixed batch so only `f_divergence_type` varies between the two loss computations. + batch = next(iter(trainer.get_train_dataloader())) + batch = trainer._prepare_inputs(batch) + + with torch.no_grad(): + trainer.f_divergence_type = "reverse_kl" + reference_loss = trainer.compute_loss(trainer.model, batch) + trainer.f_divergence_type = f_divergence_type + other_loss = trainer.compute_loss(trainer.model, batch) + + torch.testing.assert_close(other_loss, reference_loss) + def test_train_with_explicit_ref_model(self): dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train") diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py index fac06a9dcde..4b6cf993241 100644 --- a/trl/trainer/dpo_trainer.py +++ b/trl/trainer/dpo_trainer.py @@ -1507,8 +1507,11 @@ def _compute_loss(self, model, inputs, return_outputs): # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266) # Use this loss when you believe the chosen outputs are worse than your model's default output. # Decrease chosen likelihood and decrease rejected likelihood more + # Both terms use the raw reverse-KL log-ratios (the paper's r_theta), matching apo_zero + # and remaining invariant to f_divergence_type. Using delta_score mixed a raw chosen + # term with an f-divergence-transformed rejected term (#6441). losses_chosen = torch.sigmoid(self.beta * chosen_logratios) - losses_rejected = 1 - torch.sigmoid(self.beta * delta_score) + losses_rejected = 1 - torch.sigmoid(self.beta * (chosen_logratios - rejected_logratios)) per_sequence_loss = losses_chosen + losses_rejected elif loss_type == "discopop":