Document how to change the training objective by subclassing a trainer (#6918)

Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>

Author: sergiopaniego

docs/source/customization.md

@@ -111,3 +111,67 @@ training_args = DPOConfig(
     gradient_accumulation_steps=8,
 )
 ```
+
+## Change the training objective
+
+Subclassing a trainer is the simplest way to customize training. The loss is defined in the `compute_loss` method, so to train with a different objective, subclass the trainer and override it. Data preparation, generation, logging, and checkpointing are inherited, so the subclass holds only what actually changes.
+
+```python
+import torch
+
+from trl import DPOTrainer
+from trl.trainer.utils import selective_log_softmax
+
+
+class MyDPOTrainer(DPOTrainer):
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # In this example: the hinge loss from https://huggingface.co/papers/2309.06657
+        input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
+        shift_labels, shift_completion_mask = input_ids[:, 1:], inputs["completion_mask"][:, 1:]
+        logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
+        with torch.no_grad():
+            ref_logits = self.ref_model(input_ids=input_ids, attention_mask=attention_mask).logits
+
+        # Sum the log-probs over the completion tokens, for the policy and for the reference model
+        logps = (selective_log_softmax(logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
+        ref_logps = (selective_log_softmax(ref_logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
+
+        chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
+        ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)
+        delta = (chosen_logps - ref_chosen_logps) - (rejected_logps - ref_rejected_logps)
+        return torch.relu(1 - self.beta * delta).mean()
+```
+
+### Add parameters to the config
+
+Subclass the config to declare the parameters your loss needs:
+
+```python
+from dataclasses import dataclass, field
+
+from trl import DPOConfig
+
+
+@dataclass
+class MyDPOConfig(DPOConfig):
+    my_coef: float = field(default=0.1, metadata={"help": "Coefficient of the custom term."})
+```
+
+### Change the batch format
+
+When the loss needs inputs that the default collator doesn't produce, pass your own collator, and override `_prepare_dataset` to leave the dataset untouched when it is already in the expected format:
+
+```python
+class MyDPOTrainer(DPOTrainer):
+    def _prepare_dataset(self, dataset, *args, **kwargs):
+        return dataset
+
+
+trainer = MyDPOTrainer(..., data_collator=my_collator)
+```
+
+### Complete example
+
+The block-diffusion SFT example [`examples/sft_diffusion_gemma/sft_diffusion_gemma.py`](https://github.com/huggingface/trl/blob/main/examples/sft_diffusion_gemma/sft_diffusion_gemma.py) combines the three: `DiffusionGemmaSFTConfig` extends [`SFTConfig`] with the canvas and corruption parameters, and `DiffusionGemmaSFTTrainer` extends [`SFTTrainer`] and replaces the autoregressive cross-entropy with a block-diffusion denoising objective. Neither requires a change to the library.
+
+[Antidoom](https://github.com/Liquid4All/antidoom), an open-source tool from [Liquid AI](https://huggingface.co/LiquidAI), extends [`DPOTrainer`] with Final Token Preference Optimization ([Antislop](https://huggingface.co/papers/2510.15061)), a preference loss over a single token position that reduces repetition loops in reasoning models.