Skip to content

feat: async colocated GRPO with Megatron inference - #2884

Open
tdene wants to merge 6 commits into
NVIDIA-NeMo:mainfrom
tdene:tde/async_colocated
Open

feat: async colocated GRPO with Megatron inference#2884
tdene wants to merge 6 commits into
NVIDIA-NeMo:mainfrom
tdene:tde/async_colocated

Conversation

@tdene

@tdene tdene commented Jun 22, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

This PR adds async colocated capability to NeMo RL, which currently only supports async non-colocated and sync colocated.

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

@copy-pr-bot

copy-pr-bot Bot commented Jun 22, 2026

Copy link
Copy Markdown

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
tdene marked this pull request as ready for review June 22, 2026 16:51
@tdene
tdene requested review from a team as code owners June 22, 2026 16:51
@copy-pr-bot

copy-pr-bot Bot commented Jun 22, 2026

Copy link
Copy Markdown

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.

@tdene tdene added the CI:L1 Run doctests, unit tests, and functional tests label Jun 22, 2026
@tdene

tdene commented Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b4265ac

@tdene
tdene force-pushed the tde/async_colocated branch from b4265ac to 34c41e9 Compare June 23, 2026 15:56
@tdene

tdene commented Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 34c41e9

@tdene
tdene force-pushed the tde/async_colocated branch from 34c41e9 to 5219117 Compare July 29, 2026 06:53
@tdene
tdene requested review from a team as code owners July 29, 2026 06:53
cspades
cspades previously approved these changes Jul 29, 2026
View reviewed changes

@cspades cspades left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have some n00b questions but particularly the main RL loop LGTM!

Comment thread nemo_rl/algorithms/grpo.py

assert (not colocated_inference) or (
isinstance(policy_generation, MegatronGeneration)
), "Colocated async GRPO is unsupported for the desired generation backend."

@cspades cspades Jul 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

More helpful user message: "Colocated async GRPO is only supported for the Megatron generation backend." (Add more generation backends to this message as they become supported.)

Comment thread nemo_rl/algorithms/grpo.py
policy_generation.finish_generation()
# Colocated engine stays alive across steps (preserves KV cache).
if not colocated_inference:
policy_generation.finish_generation()

@cspades cspades Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

INFO: Do you happen to know why non-colocated vLLM generation clear its prefix cache while we don't? Is it a performance trade-off to not re-populate the cache?

Comment thread nemo_rl/models/generation/megatron/megatron_worker.py
if not self._inference_engine_initialized:
return
self._sleep()
torch.cuda.synchronize()

@cspades cspades Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

INFO: What's the context behind why this sync existed before? Did we need to make sure all streams on this device have completed their work before we refit in the non-colocated case?

@anlthms anlthms mentioned this pull request Jul 31, 2026
Draft
tdene and others added 4 commits August 4, 2026 00:31
@tdene
feat: async colocated GRPO with Megatron inference
fe9406b 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene @anlthms
Drop unnecessary synchronize
fdbca55 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
@tdene @anlthms
Defer colocated wake past checkpoint saves
0424fa8 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
@tdene @anlthms
Allow inference_optimized during training
0946860 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
@tdene
tdene force-pushed the tde/async_colocated branch from 5219117 to 0946860 Compare August 4, 2026 05:33
@tdene
Account for Megatron-Bridge#5164
bd52c73 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene

tdene commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test bd52c73

shanmugamr1992
shanmugamr1992 previously approved these changes Aug 4, 2026
View reviewed changes
@tdene
Add functional test
9ef4202 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene

tdene commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 9ef4202

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@shanmugamr1992 shanmugamr1992 shanmugamr1992 approved these changes
@cspades cspades cspades left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

CI:L1 Run doctests, unit tests, and functional tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@tdene @cspades @shanmugamr1992