DPO trainer supports num_logits_to_keep to save memory#2129
DPO trainer supports num_logits_to_keep to save memory#2129kashif merged 55 commits intohuggingface:mainfrom
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@xyangk thank you! can you run the |
|
Sure! |
…ep_necessary_logits
…ep_necessary_logits
…ep_necessary_logits
|
Implemented the calculation of num_logits_to_keep on the new master code, fixed a bug in the previous calculation, and used more rigorous test code. |
|
Hi @kashif , I wanted to check in on this PR. Please let me know if there's anything further I can clarify or update. Thank you for your time and review! |
|
@xyangk can you use the transformers helper |
I've implemented the suggested changes. Thank you for the feedback! |
|
|
||
| # num_logits_to_keep is supported since transformers v4.45.0 | ||
| if self.use_num_logits_to_keep: | ||
| check_min_version("4.45.0") |
There was a problem hiding this comment.
You don't need it
https://github.com/huggingface/trl/blob/main/setup.py#L82
There was a problem hiding this comment.
Sorry I missed that the transformers version was already updated. Since the current PR is already merged, I'll create a new PR to remove the version check. Thank you for cathing this!
) * Support num_logits_to_keep, which computes necessary logits in the forward pass. * update doc * bug fix * update * check is model supports num_logits_to_keep * ruff format * update test file * peft model support * test passed * update * apply use_num_logits_to_keep * fix num_logits_to_keep compute bug * compare all outputs * pytest * pass test * use check_min_version * format * test_dpo_trainer_use_num_logits_to_keep passed * add some comments --------- Co-authored-by: Kashif Rasul <kashif.rasul@gmail.com>
4 participants
What does this PR do?
Since transformers v4.45.0 supports
num_logits_to_keepin the forward pass(check this PR), it can be useful for saving memory and speeding up training by computing logits only for the last completion tokens, especially in scenarios with very long prompts where labels are ignored (-100).I implemented the logic to calculate
num_logits_to_keepbased on the labels and tested this on DPO, resulting in substantial memory savings.Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.