Skip to content

feat(sc): support non-colocated MInf - #3727

Merged
yuki-97 merged 25 commits into
NVIDIA-NeMo:mainfrom
tdene:tde/sc_megatron_generation
Aug 28, 2026
Merged

feat(sc): support non-colocated MInf#3727
yuki-97 merged 25 commits into
NVIDIA-NeMo:mainfrom
tdene:tde/sc_megatron_generation

Conversation

@tdene

@tdene tdene commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

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

  • ...

@tdene
tdene requested review from a team as code owners August 20, 2026 10:30
@copy-pr-bot

copy-pr-bot Bot commented Aug 20, 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 requested a review from a team as a code owner August 20, 2026 12:37
@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Aug 20, 2026
@tdene
tdene force-pushed the tde/sc_megatron_generation branch 2 times, most recently from f37c117 to 4c63c46 Compare August 21, 2026 14:31
@tdene
tdene requested a review from a team as a code owner August 21, 2026 14:31
tdene added 3 commits August 25, 2026 09:30
@tdene
feat(sc): support non-colocated MInf
6b0cdca 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Address PR team review
28c0593 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Cleanup
c9c667a 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
tdene force-pushed the tde/sc_megatron_generation branch from 4c63c46 to c9c667a Compare August 25, 2026 14:47
@tdene tdene added the CI:L1 Run doctests, unit tests, and functional tests label Aug 25, 2026
@tdene

tdene commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c9c667a

yuki-97
yuki-97 reviewed Aug 26, 2026
View reviewed changes

@yuki-97 yuki-97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Non-colocated MInf on SC — the setup wiring and the new unit coverage read well. Fourteen comments.

Rebase first. This no longer merges: #3773 (SC PPO) landed on main and rewrote the same setup.py / single_controller.py / L1_Functional_Tests_SingleController.sh / single-controller.md surface, and renamed test_single_controller_setup.py to test_setup.py — the file this PR adds 242 lines to. Two comments below exist only because of that rebase: the megatron build-task branch bypasses the new critic construction, and the megatron-specific colocated guard is subsumed by a generic one main now applies. Both take two of the new test cases with them.

On the red CI: both lanes hit Gym's 600s endpoint timeout, but differently — on h100 the trainer finishes in ~50s and then nothing is logged for 9.5 minutes with no generation worker ever created; on gb200 the trainer's own init takes 9m45s. Details in the first comment.

Separately: the PR description is still the untouched template. Worth a short summary of the approach, and — since this lands a new generation backend on the SC path — the usual convergence evidence (reward, loss, token_mult_prob_error), rather than only the functional test's median(gen_kl_error) < 1.3 / max(reward) > 0.

Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread tests/functional/grpo_megatron_generation_gym_single_controller.sh Outdated
Comment thread examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml Outdated
Comment thread examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py
yuki-97
yuki-97 reviewed Aug 26, 2026
View reviewed changes
Comment thread tests/functional/L1_Functional_Tests_SingleController.sh
@yuki-97 yuki-97 mentioned this pull request Aug 26, 2026
Open
4 tasks
tdene and others added 8 commits August 26, 2026 08:47
@tdene @yuki-97
Address reviewer comments
6e49356 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
@tdene @yuki-97
Merge branch 'main' into tde/sc_megatron_generation
a3c94e9 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
@tdene
Refactor recompute logic
28218a6 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Address reviewer comments
d3602a3 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Merge branch 'main' into tde/sc_megatron_generation
f2b9a30 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Use example as part of CI
e322324 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Add nightly test
a032e95 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Merge branch 'main' into tde/sc_megatron_generation
3460091 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene

tdene commented Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 3460091

@tdene
tdene force-pushed the tde/sc_megatron_generation branch from 12b504b to 0661efd Compare August 27, 2026 05:14
@tdene

tdene commented Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 0661efd

tdene and others added 4 commits August 27, 2026 08:56
@tdene
fix(megatron): hotfix disable offending tests
a5e2d67 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
(cherry picked from commit c3bf4fb)
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
run the megatron async-gym functional during LFast
ae5c741 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
(cherry picked from commit f8761e0)
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene @claude
test: temporarily run the SC sync math nightlies on GB200
c81afd6 
cw_dfw H100 capacity is unavailable, so the SC sync math pair moves to
the GB200 suite: the vLLM twin already has its 1n4g variant there, and
this adds the matching 1n4g Megatron Inference recipe (chained off the
GB200 SC baseline, same 50-step/180-min sizing as the 1n8g original).
The 1n8g entries are parked in disabled.txt for restoration.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Merge branch 'main' into tde/sc_megatron_generation
0d3bf0c
@tdene

tdene commented Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 0d3bf0c

yuki-97
yuki-97 reviewed Aug 27, 2026
View reviewed changes
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py
Comment thread nemo_rl/models/generation/megatron/megatron_generation.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/setup.py Outdated
Comment thread tests/test_suites/disabled.txt Outdated
tdene added 2 commits August 27, 2026 15:43
@tdene
Address reviewer comments
61f1e2a 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Topology fix
44c13e6 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene

tdene commented Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 44c13e6

tdene added 2 commits August 28, 2026 08:31
@tdene
Merge branch 'main' into tde/sc_megatron_generation
898964c 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Replace _build_trainer_then_megatron_generation
0cc133c 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene

tdene commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 0cc133c

tdene added 2 commits August 28, 2026 11:34
@tdene
Fix tests
83889d5 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Fix NCCL refit TP=1 issue
c3837b7 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
tdene force-pushed the tde/sc_megatron_generation branch from b92f51c to c3837b7 Compare August 28, 2026 16:34
@tdene

tdene commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c3837b7

@tdene
Revert temporary test changes
df3df79 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene

tdene commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test df3df79

tdene added 2 commits August 28, 2026 15:03
@tdene
Expand comment
f2fc2e8 
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
Merge remote-tracking branch 'origin/main' into tde/sc_megatron_gener…
399a162 
…ation

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

tdene commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 399a162

@yuki-97
yuki-97 enabled auto-merge (squash) August 28, 2026 20:13
@yuki-97
yuki-97 merged commit dd4b988 into NVIDIA-NeMo:main Aug 28, 2026
101 checks passed
@tianyi-zhang-02 tianyi-zhang-02 mentioned this pull request Aug 29, 2026
Open
asolergi-nv added a commit that referenced this pull request Aug 29, 2026
@asolergi-nv
Merge upstream/main into feat/sc-resiliency-04-restart-readmission
a53e2eb 
PR3 (#3591) was SQUASH-merged into main as b3b6713, so none of its commits are
ancestors of main while PR4 still carries all of them. Git therefore sees PR3's whole
diff as independently added on both sides, which is why all 16 conflicts name b3b6713
and why the PR showed CONFLICTING despite the content being identical.

That made the classification, not the content, the work. For each conflicted file: is
main's version byte-identical to PR3's head (3d9ce21), and does PR4 add anything beyond
it? Three groups fell out.

GROUP A -- pure squash artefacts, resolved by taking OURS (10 files)
  fleet_health.py, collective_weight_synchronizer.py, membership.py,
  nccl_reshard_weight_synchronizer.py, grpo_sc_generation_shard_recovery.sh,
  test_watchdog_pump.py, test_membership.py, test_reconcile_communicator.py,
  test_reshard_rebuild.py, test_weight_synchronizer.py
  main == PR3 exactly and no other PR touched them, so PR4's side is main's content plus
  PR4's delta. Taking ours loses nothing.

GROUP B -- PR4 contributes nothing, resolved by taking THEIRS (2 files)
  single_controller_utils/setup.py  (#3480, #3727, #3821 on top of PR3)
  tests/unit/single_controller/test_refit_recovery.py  (#3480 on top of PR3)

GROUP C -- genuine merges (4 files), one per upstream PR below.

The six upstream PRs that contributed real content, and what each needed:

  #3480 recover replay buffer from native TQ checkpoints
        single_controller.py: rollout_recovery imports. Kept alongside ours.
        setup.py, test_refit_recovery.py, L1 harness: group B / additive.
  #3765 log toolcall and thinktag violation rate
        single_controller.py: VIOLATION_TAG_KEYS. Auto-merged, verified present.
  #3727 support non-colocated MInf
        single_controller.py: MegatronGeneration import, kept alongside ours.
        L1 harness: grpo_megatron_generation_gym_single_controller.sh entry.
  #3821 warm-start the value model from a critic-pretrain checkpoint
        config.py: the max_num_epochs validator. Ours only adds restart_dead_shards to
        FleetHealthConfig, so both survive; verified the field landed in the right class
        and the validator is intact.
  #3655 nemo-lens telemetry
        vllm_generation.py: the @trace_fn decorator on generate. Ours adds restart_shard
        in a different region; both kept.
  #3839 pause generation during in-flight refit
        vllm_generation.py: pause_generation_for_refit / resume_generation_after_refit.
        Auto-merged, verified present -- worth knowing it exists, since it pauses engines
        around a refit and this PR restarts them.

Verified after resolving: no conflict markers; all four lint hooks clean (the single
pyrefly error is the pre-existing unrelated transfer_queue import); 1122 unit tests pass;
both submodule pointers and uv.lock/pyproject byte-identical to main.

Both sides' work was checked individually rather than assumed: EngineSupervisor wiring,
restart_dead_shards, restart_shard, recreate_worker, desired_membership and the report_refit
call on our side; the six items above on main's.

Note for anyone reproducing locally: #3655 adds a nemo-lens dependency that the pre-merge
container image does not carry, so tests fail at import with ModuleNotFoundError: nemo
until the venv is refreshed. Plain upstream/main fails the same way in that image; it is
not a merge defect.

Signed-off-by: asolergibert <asolergibert@nvidia.com>
@tdene tdene mentioned this pull request Aug 31, 2026
Merged
4 tasks
ZhiyuLi-Nvidia added a commit that referenced this pull request Sep 4, 2026
@ZhiyuLi-Nvidia @claude
test: raise the nightly compute cap to 4266 GPU hours
5827223 
The three nightly entries this branch adds bring the suite from 4138 to
4266 GPU hours:

  +32  vlm_grpo-nemotron-omni-30ba3b-clevr-2n8g-megatron-tp8ep8.v1
  +32  vlm_grpo-nemotron-omni-30ba3b-clevr-2n8g-...v1-tq_mooncake
  +64  vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple

Two of those are the VLM data-plane coverage itself; the third is the
2-node base the mooncake wrapper delegates to, which exists because the
recipe does not fit on one node.

Follows the convention for adding nightly coverage -- the cap has moved
4048 -> 4139 -> 4149 -> 4157 -> 4181 -> 4197 as recipes landed (#2744,
#3773, #3768, #3727, #3678).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yuki-97 yuki-97 yuki-97 approved these changes

Assignees

No one assigned

Labels

CI:L1 Run doctests, unit tests, and functional tests Documentation Improvements or additions to documentation

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@tdene @yuki-97