🤖 Properly unwrap torch.compile-ed models in GRPO

[winglian]

What does this PR do?

when using torch compile, there is one more layer to unwrap before we can send the state dict to vlllm

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• 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.
• Did you write any new necessary tests?

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.

[HuggingFaceDocBuilderDev]

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.

[shirinyamani]

That's a very good point! Thanks for pointing it out!, but quick follow-up for my knowledge, I could not find torch.compile() in the main grpo_trainer, could you please help me understand which part of code is specifying the memory hierarchy-aware like torch.compile in your viewpoint ?

bc to the best of my knowledge we could do sth like model = torch.compile(model) but of course we gotta make sure the compatibility with rest of the computation.

@winglian @kashif

[winglian]

The base TrainingArguments from transformers includes a torch_compile option, so you can simply set that on GRPOConfig

That's a very good point! Thanks for pointing it out!, but quick follow-up for my knowledge, I could not find torch.compile() in the main grpo_trainer, could you please help me understand which part of code is specifying the memory hierarchy-aware like torch.compile in your viewpoint ?

bc to the best of my knowledge we could do sth like model = torch.compile(model) but of course we gotta make sure the compatibility with rest of the computation.

[winglian]

@qgallouedec I rebased this so the merge conflict should be resolved. thanks!

[winglian]

I could also move this into the unwrap_model_for_generation function, but I'm not 100% on the deepspeed behavior.

[qgallouedec]

Thanks @winglian, I've added a test, and made sure that it's also compatible with reward models. Can be merged once the CI is green :)