How to log loss terms in the DPO calculation?

[mynameismon]

I want to log both the reference chosen rewards and the model chosen rewards, and similarly for the reference rejected rewards and model rejected rewards. Is there any way to do that?

[Mahesh-Paul]

In TRL's DPOTrainer, you can log individual loss and margin components by either customizing the evaluation step or overriding get_batch_loss_metrics.

1. Default Metrics via Callbacks / WandB

DPOTrainer computes implicit metrics internally during training and evaluation:

• rewards/chosen
• rewards/rejected
• rewards/accuracies
• rewards/margins
• logps/chosen
• logps/rejected

Make sure you pass your logger in TrainingArguments:

training_args = DPOConfig(
    ...,
    logging_steps=10,
    report_to=["wandb"],  # or "tensorboard"
)


2. Custom Logging via Method Override

from trl import DPOTrainer

class CustomDPOTrainer(DPOTrainer):
    def get_batch_loss_metrics(self, model, batch, train_eval="train"):
        loss, metrics = super().get_batch_loss_metrics(model, batch, train_eval=train_eval)
        
        # Log additional customized metrics
        prefix = "eval_" if train_eval == "eval" else ""
        metrics[f"{prefix}custom_loss_term"] = loss.detach().item()
        
        return loss, metrics