GRPO: Scalable training with one LLM/node #3186
Open
+134
−53
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Temporarily disable optimations for multiple generations from the same prompt - Compute number of local proceses using collective - Strip FSDP checkpointing artifact from parameter names
relevant for multiple processes posting requests to the same server
Use an intra-node communicator to avoid sending large global messages
This was referenced Apr 4, 2025
This was referenced Apr 8, 2025
Open
|
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? |
|
Ok, it makes more sense. To make the review easier can you split into two separate PRs 🙏 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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
Pull Request section?
to it if that's the case. Scaling bottleneck in GRPO Training. #3258 Support FSDP #3259
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.
@qgallouedec @binary-husky
Caveat: