feat: async colocated GRPO with Megatron inference - #2884
Merged
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
tdene
marked this pull request as ready for review
June 22, 2026 16:51
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
Author
|
/ok to test b4265ac |
tdene
force-pushed
the
tde/async_colocated
branch
from
June 23, 2026 15:56
b4265ac to
34c41e9
Compare
Contributor
Author
|
/ok to test 34c41e9 |
tdene
force-pushed
the
tde/async_colocated
branch
from
July 29, 2026 06:53
34c41e9 to
5219117
Compare
cspades
previously approved these changes
Jul 29, 2026
cspades
left a comment
cspades
left a comment
Contributor
There was a problem hiding this comment.
Have some n00b questions but particularly the main RL loop LGTM!
tdene
force-pushed
the
tde/async_colocated
branch
from
August 4, 2026 05:33
5219117 to
0946860
Compare
Contributor
Author
|
/ok to test bd52c73 |
shanmugamr1992
previously approved these changes
Aug 4, 2026
Contributor
Author
|
/ok to test 9ef4202 |
4 tasks
tdene
force-pushed
the
tde/async_colocated
branch
from
August 18, 2026 10:05
607e2c3 to
206cc69
Compare
Contributor
Author
|
/ok to test 206cc69 |
tdene
force-pushed
the
tde/async_colocated
branch
from
August 19, 2026 14:09
206cc69 to
bdf7125
Compare
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
tdene
force-pushed
the
tde/async_colocated
branch
from
August 19, 2026 14:41
bdf7125 to
7b3a56a
Compare
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Contributor
Author
|
/ok to test 3fd09e2 |
Contributor
Author
|
/ok to test b3b51bb |
yuki-97
reviewed
Aug 20, 2026
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Contributor
Author
|
/ok to test 8772e04 |
Contributor
Author
|
/ok to test ab0c137 |
Contributor
Author
|
/ok to test a53fc02 |
4 tasks
yuki-97
approved these changes
Aug 21, 2026
rrs45
added a commit
that referenced
this pull request
Aug 21, 2026
Resolve two conflicts against main's async colocated GRPO work (#2884): - grpo.py: adopt main's refit restructuring (the defer_wake_for_save branch, the dropped NEED_REFIT gate, and the hoisted generation_logger_metrics init) and re-apply the idle/refit_bubble efficiency span over the wake path. The save-bound branch keeps main's 0.0 seed with no span, since a zero-length bubble is only waterfall noise. - vllm_worker.py: keep both additions, native-tracing setup and main's torchcodec video loader registration. Signed-off-by: Raj Singh <rajsin@nvidia.com>
seonjinn
added a commit
to seonjinn/RL
that referenced
this pull request
Aug 26, 2026
Rebuilds the MXFP8 refit-prequantize path on main's CollectiveSenderSpec model (NVIDIA-NeMo#3391): the generation side now owns refit buffer geometry, so the PR's consumer-side buffer_size_bytes plumbing and _resolve_target_packed_tensor_size are dropped in favor of main's interfaces, with the producer-side buffer_size_bytes > 0 guard re-added. Narrows main's grouped-MoE MXFP8 refit guard (NVIDIA-NeMo#3477) to reject only on-the-fly quantization, letting trainer-prequantized E4M3 slabs pass through. Rebuilds the refit trigger on main's unconditional prepare_for_generation structure (NVIDIA-NeMo#2884) while keeping the MegatronGeneration reshard branch. Adds a sender/receiver MXFP8 bit-equivalence test. Signed-off-by: seonjinn <sna@nvidia.com>
4 tasks
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 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.
What does this PR do ?
This PR adds async colocated capability to NeMo RL, which currently only supports async non-colocated and sync colocated.
Example of functionality: https://wandb.ai/nvidia/nemo-rl?nw=i2hxg54gbtq
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information