Repository: huggingface/trl
Pull request: 🧑‍🍳 Add precompute batch size argument in `DPOTrainer` for reference model
# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly. They may suggest changes to make the code even better.
-->

<!-- Remove if not applicable -->

Fixes #2421 


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [ ] Did you write any new necessary tests?


## Who can review?
@qgallouedec 

Problem statement:
Adding precompute batch size argument in DPOTrainer for reference model
### Feature request

Proposing adding a new configuration parameter `precompute_ref_batch_size` to allow users to specify a different (likely larger) batch size specifically for the reference model precomputation phase. This would:

1. Speed up the precomputation phase by processing more examples per batch
2. Make better use of available GPU memory since no gradients need to be stored
3. Maintain backward compatibility by defaulting to current behavior if not specified

The change would affect:
- DPOConfig: Add new optional parameter `precompute_ref_batch_size`
- `get_train_dataloader()` and `get_eval_dataloader()`: Use the new batch size when precomputing reference probabilities

This is particularly useful for large-scale DPO training where the precomputation phase can be a significant bottleneck.

### Motivation

Currently when `precompute_ref_log_probs=True` the `DPOTrainer` class uses the `per_device_train_batch_size` or `per_device_eval_batch_size` respectively during training or evaluation for generating the logprobs.
But for efficiency this batch size can be more than the used for training one since this step does not require gradient computation and storage in memory. Adding a configurable `precompute_ref_batch_size` parameter would allow users to optimize this preprocessing step by using larger batch sizes while maintaining memory efficiency.

### Your contribution

Yes, I can help by submitting a PR following the contribution guidelines.

Hints from issue discussion:
Thanks for this suggestion @SwayamInSync!
Do you have any idea of the gain in speed?
If you've a working implementation, feel free to submit a PR so that we can test and discuss the code
> Thanks for this suggestion @SwayamInSync! Do you have any idea of the gain in speed? If you've a working implementation, feel free to submit a PR so that we can test and discuss the code

Made a PR at #2426 
From a quick test on my settings I can fit only a batch size upto 8 (get OOM elsewise) but with this new parameter, inference batch size can go upto 32 (instead of same as train, so pretty better than before I guess)

Commit message: Update trl/trainer/dpo_config.py

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

Patch to review:
diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py
index 1e6e8e67ad..ea9c916d76 100644
--- a/tests/test_dpo_trainer.py
+++ b/tests/test_dpo_trainer.py
@@ -350,6 +350,40 @@ def test_dpo_trainer_with_ref_model_is_model(self):
                     train_dataset=dummy_dataset["train"],
                 )
 
+    def test_precompute_ref_batch_size(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = DPOConfig(
+                output_dir=tmp_dir,
+                per_device_train_batch_size=2,
+                precompute_ref_log_probs=True,
+                precompute_ref_batch_size=4,
+                report_to="none",
+            )
+
+            dummy_dataset = load_dataset("trl-internal-testing/zen", "standard_preference")
+
+            trainer = DPOTrainer(
+                model=self.model,
+                ref_model=self.ref_model,
+                args=training_args,
+                processing_class=self.tokenizer,
+                train_dataset=dummy_dataset["train"],
+                eval_dataset=dummy_dataset["test"],
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # check the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                # check the params have changed - ignore 0 biases
+                if param.sum() != 0:
+                    self.assertFalse(torch.allclose(param, new_param, rtol=1e-12, atol=1e-12))
+
     @require_peft
     def test_dpo_trainer_without_providing_ref_model_with_lora(self):
         from peft import LoraConfig
diff --git a/trl/trainer/dpo_config.py b/trl/trainer/dpo_config.py
index ea4a176aa1..17c48e1555 100644
--- a/trl/trainer/dpo_config.py
+++ b/trl/trainer/dpo_config.py
@@ -94,6 +94,10 @@ class DPOConfig(TrainingArguments):
         precompute_ref_log_probs (`bool`, *optional*, defaults to `False`):
             Whether to precompute reference model log probabilities for training and evaluation datasets. This is
             useful when training without the reference model to reduce the total GPU memory needed.
+        precompute_ref_batch_size (`int`, *optional*, defaults to `None`):
+            Batch size to use when precomputing reference model log probabilities. This can be set higher than the
+            training batch size to speed up preprocessing. If `None``, defaults to `per_device_train_batch_size` for
+            training and `per_device_eval_batch_size` for evaluation.
         dataset_num_proc (`Optional[int]`, *optional*, defaults to `None`):
             Number of processes to use for processing the dataset.
         model_init_kwargs (`Optional[dict[str, Any]]`, *optional*, defaults to `None`):
@@ -173,6 +177,7 @@ class DPOConfig(TrainingArguments):
     disable_dropout: bool = True
     generate_during_eval: bool = False
     precompute_ref_log_probs: bool = False
+    precompute_ref_batch_size: Optional[int] = None
     dataset_num_proc: Optional[int] = None
     model_init_kwargs: Optional[dict[str, Any]] = None
     ref_model_init_kwargs: Optional[dict[str, Any]] = None
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
index c1f2776511..ea3f24a39d 100644
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -684,8 +684,9 @@ def get_train_dataloader(self) -> DataLoader:
         """
 
         if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs:
+            batch_size = self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size
             dataloader_params = {
-                "batch_size": self.args.per_device_train_batch_size,
+                "batch_size": batch_size,
                 "collate_fn": self.data_collator,
                 "num_workers": self.args.dataloader_num_workers,
                 "pin_memory": self.args.dataloader_pin_memory,
@@ -737,8 +738,9 @@ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoa
         eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
 
         if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs:
+            batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size
             dataloader_params = {
-                "batch_size": self.args.per_device_eval_batch_size,
+                "batch_size": batch_size,
                 "collate_fn": self.data_collator,
                 "num_workers": self.args.dataloader_num_workers,
                 "pin_memory": self.args.dataloader_pin_memory,


Code review comments:
diff_hunk: @@ -94,6 +94,10 @@ class DPOConfig(TrainingArguments):
         precompute_ref_log_probs (`bool`, *optional*, defaults to `False`):
             Whether to precompute reference model log probabilities for training and evaluation datasets. This is
             useful when training without the reference model to reduce the total GPU memory needed.
+        precompute_ref_batch_size (`int`, *optional*, defaults to `None`):
+            Batch size to use when precomputing reference model log probabilities. This can be set higher than the
+            training batch size to speed up preprocessing. If `None``, defaults to `per_device_train_batch_size` for
path: trl/trainer/dpo_config.py
text: @user1:
```suggestion
            training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for
```