NeMo-Gym Integration #4848
Open
NeMo-Gym Integration #4848
Conversation
This file contains hidden or 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
Signed-off-by: Christian Munley <cmunley@nvidia.com>
Signed-off-by: Christian Munley <cmunley@nvidia.com>
Signed-off-by: Christian Munley <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Member
sergiopaniego
left a comment
sergiopaniego
left a comment
There was a problem hiding this comment.
Super cool!
Some initial ideas:
- We could add the training script to the list of examples in here (file).
- Maybe we could rename the training script to be more explanatory
train_multi_agent.pyor something similar. - It would be super cool if we could add a guide to the docs covering this integration (similar to this but specific for NeMo-Gym)
🚀 🚀 🚀 🚀
Signed-off-by: cmunley1 <cmunley@nvidia.com>
… into cmunley1/nemo_gym_on_policy
Signed-off-by: cmunley1 <cmunley@nvidia.com>
1 task
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Signed-off-by: cmunley1 <cmunley@nvidia.com>
Author
|
thanks for reviewing @sergiopaniego what do you think about changes in grpo_trainer and vllm_serve? |
Member
sergiopaniego
left a comment
sergiopaniego
left a comment
There was a problem hiding this comment.
Thanks @cmunley1, I'll review the changes.
There's a conflict on tests/test_vllm_client_server.py, could you take a look at it?
Adding @kashif @qgallouedec here for review.
kashif
reviewed
Jan 23, 2026
Comment on lines
+1246
to
+1247
|
|
||
| uvicorn.run( |
Collaborator
There was a problem hiding this comment.
Suggested change
| uvicorn.run( | |
| # Start the server | |
| uvicorn.run( |
kashif
reviewed
Jan 23, 2026
Member
sergiopaniego
left a comment
sergiopaniego
left a comment
There was a problem hiding this comment.
Could be rename the integration file to nemo_gym.md?
We need to add this new file to _toctree.yml to display it on the documentation.
Co-authored-by: Kashif Rasul <kashif.rasul@gmail.com>
Member
|
Also tagging @lewtun (multi-env) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
This integration supports training language models in NeMo-Gym environments using TRL GRPO. Both single step and multi step tasks are supported, including multi-environment training. NeMo-Gym orchestrates rollouts, returning token ids and logprobs to TRL through the rollout function for training. Currently this integration is only supported through TRL's vllm server mode.
see docs/source/nemo_gym_integration.md for a guide