Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

🧑‍🍳 Add precompute batch size argument in DPOTrainer for reference model #2426

Merged
merged 6 commits into from
Dec 2, 2024
Merged

🧑‍🍳 Add precompute batch size argument in DPOTrainer for reference model #2426

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
df4104e
added precompute_batch
SwayamInSync Dec 1, 2024
45c2521
review-fixes
SwayamInSync Dec 2, 2024
3ce8394
moving up
SwayamInSync Dec 2, 2024
c0ed5dd
Update trl/trainer/dpo_config.py
SwayamInSync Dec 2, 2024
ae51762
Update trl/trainer/dpo_config.py
qgallouedec Dec 2, 2024
251bdb2
Update trl/trainer/dpo_config.py [ci skip]
qgallouedec Dec 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
added precompute_batch
  • Loading branch information
@SwayamInSync
SwayamInSync committed Dec 1, 2024
commit df4104e7ad2d87e0a7eea9f062531c93f326f6dc
3 changes: 3 additions & 0 deletions trl/trainer/dpo_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,8 @@ class DPOConfig(TrainingArguments):
for saving memory and speeding up training by not computing the logits for all tokens, especially in scenarios
when working with very long prompts where labels are -ignored (-100).
[Read more](https://huggingface.co/docs/transformers/main/model_doc/llama#transformers.LlamaForCausalLM)
precompute_ref_batch_size (`int`, *optional*, defaults to `None`):
SwayamInSync marked this conversation as resolved.
Show resolved Hide resolved
Batch size to use when precomputing reference model log probabilities. Since no gradients need to be stored during precomputation, 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.
"""

learning_rate: float = 1e-6
Expand Down Expand Up @@ -188,6 +190,7 @@ class DPOConfig(TrainingArguments):
rpo_alpha: Optional[float] = None
discopop_tau: float = 0.05
use_num_logits_to_keep: bool = False
precompute_ref_batch_size: Optional[int] = None
SwayamInSync marked this conversation as resolved.
Show resolved Hide resolved

def __post_init__(self):
if self.max_target_length is not None:
Expand Down
6 changes: 4 additions & 2 deletions trl/trainer/dpo_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down