GRPO: Scalable training with one LLM/node

[jglaser]

What does this PR do?

Truly scalable training with GRPO + 1 local vllm process per node. Also works with FSDP. Tested on 256 nodes of Frontier (2048 GPUs).

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. Scaling bottleneck in GRPO Training. #3258 Support FSDP #3259
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?
  • developing FSDP or multinode tests is challenging and might be addressed best via a separate PR. For now we would like to make sure that existing single node tests are not breaking. Suggestions welcome!

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.
@qgallouedec @binary-husky

Caveat:

• currently disables an optimization for multiple completions as that one does not seem to be robust w.r.t. duplicate user inputs

[qgallouedec]

Thanks @jglaser

I'm not sure to understand why FSDP requires to have one vLLM instance per node?

[jglaser]

Thanks @jglaser

I'm not sure to understand why FSDP requires to have one vLLM instance per node?

It does not... FSDP changes and vllm scaling in this PR are not strictly related - however they arose in the same stream of work, as I was trying to train a 14B model which also required sharding (in addition to data parallelism).

If the FSDP feature complicates review unnecessarily, this can be factored out into a separate PR. Suggestions?

[qgallouedec]

Ok, it makes more sense. To make the review easier can you split into two separate PRs 🙏

[LeonEricsson]

hey @jglaser, just checking in. really appreciate the work on this! there's been some interest in getting the FSDP part merged. would you be cool with someone helping to split into separate PRs as suggested?

[jglaser]

hey @jglaser, just checking in. really appreciate the work on this! there's been some interest in getting the FSDP part merged. would you be cool with someone helping to split into separate PRs as suggested?

working on #3354

[qgallouedec]

closed by #3260