diff --git a/docs/source/customization.md b/docs/source/customization.md index e3051b7c61e..652bef03f34 100644 --- a/docs/source/customization.md +++ b/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.