feat: colocated Megatron reshard by tdene · Pull Request #3490 · NVIDIA-NeMo/RL

tdene · 2026-08-04T10:37:35Z

What does this PR do ?

This PR adds the ability to refit during colocated RL.

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 · 2026-08-04T10:37:38Z

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 · 2026-08-04T10:41:14Z

/ok to test fcda426

tdene · 2026-08-04T10:46:50Z

/ok to test d24e7c4

tdene · 2026-08-04T11:54:25Z

/ok to test 4941ebb

tdene · 2026-08-04T12:50:35Z

/ok to test 9fc3e0a

wdykas · 2026-08-04T13:00:52Z

left a few minor comments but looks like it follows the megatron rl pattern and should be good

wdykas

LGTM except a few comments

tdene · 2026-08-04T14:45:51Z

/ok to test d549a67

tdene · 2026-08-04T14:52:04Z

/ok to test 796a168

tdene · 2026-08-04T14:58:04Z

/ok to test 738d84d

tdene · 2026-08-04T17:17:08Z

/ok to test 07f392e

tdene · 2026-08-05T03:46:34Z

/ok to test 7a2653a

tdene · 2026-08-05T07:10:09Z

/ok to test 94f7a94

tdene · 2026-08-05T10:09:35Z

/ok to test 110a8b6

tdene · 2026-08-05T12:36:00Z

/ok to test fe4e7ac

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

The NVIDIA-NeMo#3499 rebase replaced worker_init_timing_metrics with the SetupTimingMetrics dataclass; the eager initial sync is the generation load in that scheme. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

refit_policy_generation no longer hand-rolls the Megatron collective; without a wired synchronizer it fell through to the vLLM broadcast path (AttributeError: model_update_group). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

tdene · 2026-08-17T16:55:08Z

/ok to test 4c3e03c

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

tdene · 2026-08-17T17:40:10Z

/ok to test 3ad74f7

terrykong

Team review of the colocated Megatron reshard. 7 specialised agents plus two independent adversarial passes; 9 candidate findings were dropped as wrong or unfounded before posting.

Overall this reads well and I'd like to see it land. Routing Megatron through WeightSynchronizer removes 7 isinstance(policy_generation, MegatronGeneration) conditions from refit_policy_generation, and I verified the upstream API usage in detail (build_inference_pg_collection, swap_model_weights, the _pg_collection setattr, finalize() idempotency, and the torch_memory_saver pattern all match Megatron-LM / Megatron-Bridge precedent). The moe_router_num_groups reset is a genuinely necessary workaround for InferenceTopKRouter's assert, not a clobber.

Two items worth acting on before merge: the offload_optimizer_for_refit default silently changes behaviour for two shipped nemo_gym configs, and the colocated reshard path lacks the virtual-pipeline guard its nccl_reshard sibling already has. The rest are suggestions and nits.

Reviewed at 3ad74f7 (re-verified after the rebase).

Generated by Claude Code

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

tdene · 2026-08-17T23:33:33Z

/ok to test 5f1f294

tdene · 2026-08-17T23:34:07Z

@terrykong I believe I've addressed all the comments. I also found a couple of extra issues I fixed in 5f1f294

Can you please take another look?

…03-elastic-recovery Part 2 (#3590) landed as a squash, so every file it touched conflicts with itself here. Verified all 17 are byte-identical between the Part 2 branch tip and main, which makes "ours" (that same content plus this branch's delta) the provably correct side. pyrefly.toml was the one real conflict: three commits each appended an allow-list entry, all three kept. The substantive change comes from the colocated Megatron reshard (#3490): mark_stale() is gone from WeightSynchronizer and every implementation, with per-step staleness now owned by the training loop. Nothing on this branch called it, so the removal merged cleanly; the one leftover was a fake in test_reconcile_communicator that still defined it, now dropped so the fake matches the interface it stands in for.

tdene requested review from a team as code owners August 4, 2026 10:37

tdene added the CI:L1 Run doctests, unit tests, and functional tests label Aug 4, 2026

wdykas reviewed Aug 4, 2026

View reviewed changes

Comment thread nemo_rl/models/generation/megatron/megatron_generation.py Outdated

wdykas reviewed Aug 4, 2026

View reviewed changes

Comment thread nemo_rl/models/generation/megatron/megatron_worker.py

wdykas reviewed Aug 4, 2026

View reviewed changes

Comment thread nemo_rl/models/generation/megatron/megatron_worker.py Outdated

wdykas reviewed Aug 4, 2026

View reviewed changes

Comment thread nemo_rl/models/megatron/memory_saver.py Outdated

wdykas approved these changes Aug 4, 2026

View reviewed changes

tdene force-pushed the tde/refit_colocated branch from 738d84d to 07f392e Compare August 4, 2026 17:17

shanmugamr1992 previously approved these changes Aug 4, 2026

View reviewed changes

tdene force-pushed the tde/refit_colocated branch from 07f392e to 7a2653a Compare August 5, 2026 03:46

tdene dismissed shanmugamr1992’s stale review via 94f7a94 August 5, 2026 07:10

tdene and others added 13 commits August 17, 2026 09:32

Add offload_optimizer_for_refit config

75b1c81

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

lint

356d25f

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Fix inference-optimized incompatibility

9ffe449

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Properly fix inference config mismatch

844e964

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Use a QKV-bias-free model in the reshard test

0cd9e5e

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Add refit test that doesn't need TMS

7cf972a

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Minimize the 4n4g reshard recipe config

6fa85ca

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Fix segment-size on nightlies

c7d9592

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Address reviewer comments

7695ba5

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Address review-pr-team comments

5d7f2c4

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Import torch_memory_saver correctly

4c3e03c

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

tdene force-pushed the tde/refit_colocated branch from f81a58d to 4c3e03c Compare August 17, 2026 16:52

lint rebase conflict

3ad74f7

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

terrykong reviewed Aug 17, 2026

View reviewed changes

tdene added 4 commits August 17, 2026 17:38

Address reviewer comments

3ca38b7

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Address reviewer comment

6f63120

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Fix default value of argument

4a21615

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Fix extra issues

5f1f294

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

terrykong enabled auto-merge (squash) August 18, 2026 00:09

terrykong approved these changes Aug 18, 2026

View reviewed changes

terrykong merged commit 3d873eb into NVIDIA-NeMo:main Aug 18, 2026
177 of 180 checks passed

yuki-97 mentioned this pull request Aug 18, 2026

feat(sglang): megatron backend weight refit for sglang rollouts #3612

Merged

yuki-97 mentioned this pull request Aug 26, 2026

feat(sc): support colocated MInf #3730

Open

4 tasks

Uh oh!

Conversation

tdene commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wdykas commented Aug 4, 2026

Uh oh!

wdykas left a comment

Choose a reason for hiding this comment

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 4, 2026

Uh oh!

tdene commented Aug 5, 2026

Uh oh!

tdene commented Aug 5, 2026

Uh oh!

tdene commented Aug 5, 2026

Uh oh!

tdene commented Aug 5, 2026

Uh oh!

tdene commented Aug 17, 2026

Uh oh!

tdene commented Aug 17, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tdene commented Aug 17, 2026

Uh oh!

tdene commented Aug 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tdene commented Aug 4, 2026 •

edited

Loading