feat(sc): support colocated MInf

[tdene]

What does this PR do ?

https://wandb.ai/nvidia/nemo-rl?nw=50g85ptkgd4&panelDisplayName=timing%2Ftrain%2Ftotal_step_time&panelSectionName=timing

┌─────────────────────────────────────────────┬───────────┬──────────────┬────────────┬───────────────┐
│                   Metric                    │ H100 colo │ H100 noncolo │ GB200 colo │ GB200 noncolo │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ W&B run                                     │ 8tvgr17u  │ 27nexdg4     │ uaocb97k   │ 8u8vnx6r      │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ GPUs (train / gen)                          │ 8 shared  │ 4 / 4        │ 4 shared   │ 2 / 2         │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Training wall time                          │ 9.9 min   │ 9.5 min      │ 14.7 min   │ 17.5 min      │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Step time, median                           │ 10.32 s   │ 9.87 s       │ 15.30 s    │ 18.15 s       │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Step time, p90                              │ 11.20 s   │ 10.29 s      │ 16.22 s    │ 19.27 s       │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Valid tokens/s/GPU, median                  │ 2220      │ 2310         │ 3026       │ 2536          │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Reward, max                                 │ 0.504     │ 0.492        │ 0.488      │ 0.477         │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Reward, step 50                             │ 0.199     │ 0.217        │ 0.207      │ 0.186         │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ gen_kl_error, median                        │ N/A*      │ N/A*         │ N/A*       │ N/A*          │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ gen_kl_error, healthy-step median†          │ 1.16e-4   │ 1.43e-4      │ 1.23e-4    │ 1.72e-4       │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ token_mult_prob_error, median               │ N/A*      │ N/A*         │ N/A*       │ N/A*          │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ token_mult_prob_error, healthy-step median† │ 1.0035    │ 1.0043       │ 1.0038     │ 1.0045        │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ Poisoned steps*                             │ 31 / 50   │ 26 / 50      │ 29 / 50    │ 28 / 50       │
├─────────────────────────────────────────────┼───────────┼──────────────┼────────────┼───────────────┤
│ CI gate verdict                             │ N/A*      │ N/A*         │ N/A*       │ N/A*          │
└─────────────────────────────────────────────┴───────────┴──────────────┴────────────┴───────────────┘

• N/A: there's a NaN Megatron-Inference bug that will be solved by Avoid emitting NaN / Inf logprobs out of inference NVIDIA/Megatron-LM#6993

† Computed over the non-bugged steps only

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]

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]

/ok to test 53a5b2b

[yuki-97]

Colocated MInf on SC, read as the delta over #3727. The stand-down / deferred-wake shape is a faithful port of async_grpo_train's (#2884, resting on the reshard-on-wake that #3490 put inside prepare_for_generation), and the new unit coverage reads well. Ten comments.

CI is red on this commit. L0_Unit_Tests_Megatron fails with AttributeError: 'NoneType' object has no attribute 'blocks_training' at single_controller.py:911. test_train_pump.py:345 builds the actor with gen_handle=None, which was fine until the train pump started asking the engine before _sync_weights; the two SC fakes this PR does update (test_single_controller.py:759-764, test_sc_checkpointing.py:224-228) cover the other suites, and :955 reaches the same attribute. (The "CI quality check" failure is just the missing CI label.)

Everything on #3727 still applies here, the rebase included: this branch also no longer merges after #3773.

Separately, the PR description is still the untouched template — same ask as #3727: a short summary of the approach, plus convergence evidence (reward, loss, token_mult_prob_error) for the colocated path rather than only the functional test's thresholds.

[tdene]

Team review: 5 specialized reviewers + an adversarial verification pass; every posted finding survived independent re-verification, and the one Major finding was reproduced twice on the repo's own test harness.

Rebase needed: the PR is conflicting with main in tests/unit/test_recipes_and_test_suites.py — this PR's +24 GPU-h bump (4181→4205) vs main's #3678 (4181→4197). Additive resolution: 4197 + 24 = 4221, plus the matching test-function rename.

Verified during review (beyond the inline comments): all 9 replied review threads deliver what they claim at this head; full pre-commit (ruff, ruff-format, pyrefly, recipe minimize-check) passes on the PR files; the nightly budget bump is exact (+24 = 1n8g × 180 min; DRYRUN total 4200 ≤ 4205); the new functional script's log-grep target is emitted by production (megatron_policy_worker.py:3303); the new unit tests' event-order assertions match the production call order, and the four changed SC test files pass locally (179 passed).

Design points worth keeping as-is (deliberately not comments): the serial trainer→generation build seam (merging it with the vLLM path would be conditional soup); blocks_training/wake_carries_weight_updates declared on GenerationInterface with fail-safe defaults instead of hasattr probing; the save choreography is a faithful port of async_grpo_train's, with the once-bound check_save and the hoisted metrics tail genuine improvements over the original; and the whole choreography is unit-testable without Ray/GPU, which the new tests exploit well.

Generated by Claude Code

[tdene]

/ok to test a5a7375

[tdene]

/ok to test 5d12742

[tdene]

/ok to test 331383a

[tdene]

/ok to test 0118405

[tdene]

/ok to test e68a99a

[tdene]

/ok to test 3d9d2c3

[tdene]

/ok to test fd0c068

[tdene]

/ok to test 1e833e1

[tdene]

/ok to test 77e9e62

[yuki-97]

Scope: Re-review at head 77e9e62 — the colocated MInf path end to end (stand-down/wake, deadline crediting, interaction with checkpoint saves), plus an audit of all eleven self-review threads against the current tree rather than against their replies.

Checked:

• Stand-down / wake pairing — every arm that clears the gate has a matching reopen: _sync_weights' tail on ordinary steps, the deferred post-save wake on save-bound ones. Standing down implies the refit block is entered, so there is no path that freezes without reaching a reopen. The single unpaired arm is loop_will_exit, which is deliberate and now pinned by the timeout_save_exits leg.

• Save vs refit ordering — confirmed deliberate rather than an inconsistency. Colocated runs save then wake: waking first would hold the inference model resident across a save that onloads model and optimizer, for an engine that cannot serve during it anyway. Every other backend runs refit then save: deferring would serve stale weights for the whole save, and offload_before_refit's finalize_async_save would then block on the async write. The predicate choosing between them is capability-based, so a future backend with the same property lands on the right arm without a backend check.

• Deadline crediting — the three _Deadline sites map one-to-one onto the three RolloutTimeouts fields, with env_s deliberately uncredited since an environment step is not blocked by the stand-down, and suspend/resume wired to the stand-down and to both wake paths. One deadline is not credited and cannot be from here: the router's backend_timeout_s is an aiohttp total enforced in the router actor's own process. Not raised as a comment, since generation_router.enabled is false in every config this PR ships and both rollout deadlines default to null.

• Colocated weight path — there is no weight transfer to review: the synchronizer's colocated branch skips the collective entirely and the wake carries the weights, by shared tensors or by reshard into the inference model. That is what wake_carries_weight_updates() licenses, and what makes deferring the refit past a save sound in the first place.

Remaining comments are on a correctness guard in the deadline suspend, when the clocks freeze relative to the stand-down, naming and comment density around the stand-down block, a test gap on the narrowed colocated reject, and parking the new nightlies until the NaN fix lands.

[yuki-97]

given that they'll have nan which will be fixed in later PR, let's move the new tests (also gb200's) to disable.txt for now, and move them back when new PR lands.

[tdene]

Done

[tdene]

Resolved in 8e9f622

[yuki-97]

Both colocated recipes are parked now, thanks. One recipe this PR adds is still live: nightly_gb200.txt:53 registers grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh.

It runs the same Megatron-inference path as the two parked scripts and gates on exactly the same three metrics -- gen_kl_error median, step-50, reward > 0.2. Only the colocation mode differs.

Suggest fix: park it with the other two, or say why noncolocated is not exposed to the inf/nan logprobs.

[yuki-97]

nit: self._gen is always not None in SC, so feels no need to check.

It is annotated non-Optional at single_controller.py:249, and this same function calls it unguarded further down — :2136 would AttributeError on every non-colocated step if it could be None. Two other guards have the same shape and could go with it: :1820 and :2049.

Suggested change

	if self._gen is not None and self._gen.blocks_training():
	if self._gen.blocks_training():

Separately, while you are here: this block runs 9 lines of comment against 3 of code, and some of it restates the code. Both halves could halve —

:1823-1826 — :1824 is what the three lines below it already say. The other two facts (why standing down mid-loop is safe, and where the gate reopens) are the ones a reader cannot get locally, and they fit in two lines:

                    # Safe mid-loop: colocated steps are assembled whole, so the loop closes after this.
                    # The gate reopens at the post-step _sync_weights wake, or after the save on save-bound steps.

:1830-1834 — five lines restating each other for the one call they annotate:

                        # Deadline clocks measure inference service time, not wall clock:
                        # the switch to training must not tick them down.

Entirely your call, keeping them as is is fine too.

[tdene]

Accepted your suggestions

[tdene]

Resolved in 8e9f622

[yuki-97]

Both comment halves landed, thanks -- the guard half didn't. :1831 is still if self._gen is not None and self._gen.blocks_training(): at 0d2ada9, and the two siblings at :1820 and :2050 are unchanged.

self._gen is still non-Optional at :249 and still dereferenced unguarded at :2134 and :2142, so all three checks stay tautological.

[tdene]

My bad. I thought I did this. I must have lost the patch somehow.

I'll do it again properly when I'm not tired & rushed.

[tdene]

/ok to test 8e9f622

[tdene]

/ok to test 0d2ada9

[yuki-97]

Re-review of the fix round (8e9f622 -> 0d2ada9).

Three of the five open threads are fixed and replied to inline. Two landed only partially:

• single_controller.py self._gen guards -- both comment halves applied, the guard half not.
• Nightly parking -- both colocated recipes are parked, but nightly_gb200.txt:53 still registers a recipe this PR adds, gated identically to the two parked ones.

Three new comments below, on the un-park rationale, the -sync naming under lag-4, and the deadline-registry wiring.

#3570 landed in 4213e08 and this branch conflicts with it -- both add tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh and the same nightly_gb200.txt line, with different gates. Needs a rebase.

[yuki-97]

The un-park condition points at a gate neither parked script has: 77e9e62 replaced the token_mult_prob_error gates with gen_kl_error in all four scripts, and this comment landed afterwards in 8e9f622.

Both parked recipes now gate only on median(data["train/gen_kl_error"]) < 1.3, which the PR table shows is N/A for the same NaN reason — so the rename keeps the rationale true.

Suggested change

	# which poisons the token_mult_prob_error median gate on roughly half the steps.
	# which poisons the gen_kl_error median gate on roughly half the steps.

[yuki-97]

This flips a pre-existing -sync nightly to lag-4, and the two colocated children inherit it — so four recipes carry "sync" in their filename, wandb.name and log_dir while running off-policy.

The family root pins max_lookahead_versions: 0 deliberately ("In-order with zero lookahead requires one prompt batch of capacity"), and on main this file had no async_rl block at all.

Suggest fix: rename the four to -lag4, or say why -sync should mean lookahead 4 here.

[yuki-97]

Nothing pins that a real RolloutManager hands its registry to the impl — delete this kwarg and the whole unit suite stays green, so the feature could ship silently disabled.

The registry test builds its own registry by hand, all four impl tests set _deadline_registry = None, and the SC tests mock the manager. The functional suite will not catch it either: all three timeouts default to null, and a None deadline is a no-op with or without the registry.

Suggest fix: build a real RolloutManager in tests/unit/experience/test_rollout_manager.py and assert the impl received the manager's own registry.

[yuki-97]

Both comment halves landed, thanks -- the guard half didn't. :1831 is still if self._gen is not None and self._gen.blocks_training(): at 0d2ada9, and the two siblings at :1820 and :2050 are unchanged.

self._gen is still non-Optional at :249 and still dereferenced unguarded at :2134 and :2142, so all three checks stay tautological.

[yuki-97]

Both colocated recipes are parked now, thanks. One recipe this PR adds is still live: nightly_gb200.txt:53 registers grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh.

It runs the same Megatron-inference path as the two parked scripts and gates on exactly the same three metrics -- gen_kl_error median, step-50, reward > 0.2. Only the colocation mode differs.

Suggest fix: park it with the other two, or say why noncolocated is not exposed to the inf/nan logprobs.