[GRPO] switch grpo liger loss to triton version #4519
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. |
|
@lewtun would you be able to benchmark this branch for memory usage of grpo trainer? |
lewtun
left a comment
lewtun
left a comment
There was a problem hiding this comment.
Thanks a lot for adding this @kashif ! LGTM with a question about adding 1 to logits_to_keep twice. Could you also add your nice plots showing the memory improvement with the Triton loss to the PR description?
Also, do we need to bump the minimum liger version?
| model_inputs[key] = inputs[key] | ||
| model_inputs["use_cache"] = False | ||
| if "logits_to_keep" in self.model_kwarg_keys: | ||
| model_inputs["logits_to_keep"] = logits_to_keep + 1 |
There was a problem hiding this comment.
small suggestion since it wasn't clear to me why we're doing this:
| model_inputs["logits_to_keep"] = logits_to_keep + 1 | |
| # Add 1 because the last logits of the sequence is later excluded | |
| model_inputs["logits_to_keep"] = logits_to_keep + 1 |
| model_inputs["logits_to_keep"] = logits_to_keep + 1 | ||
|
|
||
| logits = model(**model_inputs).logits | ||
| logits = logits[:, -(logits_to_keep + 1) :, :].contiguous() |
There was a problem hiding this comment.
If logits_to_keep is in model_kwarg_keys we add 1 above to model_inputs, but then here we also add 1. Is that a form of double counting or is it correct as currently implemented?
There was a problem hiding this comment.
checking wiht example... let me report back
3 participants
What does this PR do?
Fixes dapo loss type support, but needs to wait for upstream merge and release linkedin/Liger-Kernel#939
Before submitting
Pull Request section?
to it if that's the case.
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.