feat(sc): support colocated MInf

docs/guides/single-controller.md

@@ -25,7 +25,7 @@ uv run examples/run_grpo_single_controller.py --config <your-sc.yaml>
       enabled: true
     ```
 
-2. **Pick a generation backend** and **disable colocated inference** (setup rejects `colocated.enabled: true`). With vLLM, enable the async engine (SC drives rollout via `RolloutManager.generate_and_push`, which is only supported on the disaggregated async engine):
+2. **Pick a generation backend**. With vLLM, **disable colocated inference** (setup rejects `colocated.enabled: true` for every backend except Megatron) and enable the async engine (SC drives rollout via `RolloutManager.generate_and_push`, which is only supported on the disaggregated async engine):
 
     ```yaml
     policy:
@@ -40,7 +40,8 @@ uv run examples/run_grpo_single_controller.py --config <your-sc.yaml>
             gpus_per_node: 4  # inference GPUs; remainder go to training
     ```
 
-    Megatron generation is also supported, non-colocated only. It requires the Megatron trainer (`policy.megatron_cfg.enabled: true`) and NeMo-Gym rollouts additionally require `policy.generation.mcore_generation_config.expose_http_server: true`. The exemplar — a NeMo-Gym run with the OpenAI server exposed — lives at [examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml](../../examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml):
+    Megatron generation is also supported, colocated or non-colocated. It requires the Megatron trainer (`policy.megatron_cfg.enabled: true`) and NeMo-Gym rollouts additionally require `policy.generation.mcore_generation_config.expose_http_server: true`. Colocated (`colocated.enabled: true`) additionally requires `async_rl.min_groups_for_streaming_train == num_prompts_per_step`: the engine stands down for the whole train step, so each step must be assembled as one full batch before training takes the GPUs. Per-request deadlines (`generation_timeout_s`, `rollout_timeout_s`) exclude the time the engine is stood down: in-flight requests freeze with their clocks suspended, and the clocks resume on the post-step wake.
+    The non-colocated exemplar — a NeMo-Gym run with the OpenAI server exposed — lives at [examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml](../../examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml); the colocated exemplar at [examples/configs/grpo_math_1B_megatron_generation_colocated_single_controller.yaml](../../examples/configs/grpo_math_1B_megatron_generation_colocated_single_controller.yaml):
 
     ```yaml
     policy:
@@ -51,7 +52,7 @@ uv run examples/run_grpo_single_controller.py --config <your-sc.yaml>
         mcore_generation_config:
           expose_http_server: true  # required for NeMo-Gym rollouts
         colocated:
-          enabled: false
+          enabled: false  # set true for colocated (no resources split needed)
           resources:
             num_nodes: 1
             gpus_per_node: 1  # inference GPUs; remainder go to training
@@ -238,7 +239,7 @@ The SC path is still under active development. Feature gaps are tracked in [issu
   Gym rollouts; multimodal/VLM MOPD is not yet supported. See
   [Multi-Teacher On-Policy Distillation](../about/algorithms/mopd.md#running-mopd).
 - Train backend: only Megatron is supported and validated; the AutoModel training path has not been tested on SC.
-- Generation backend: vLLM and Megatron generation are supported; SGLang and TRT-LLM have not been tested on SC.
+- Generation backend: vLLM and Megatron generation are supported (Megatron in both non-colocated and colocated modes); SGLang and TRT-LLM have not been tested on SC.
 - Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`); checkpointing is.
 - (PPO) Rollout drop budgets — `async_rl.rollout_failure.max_skipped_prompts` and `max_consecutive_dropped_prompts` must both be `0`. A drop shortens the step, and the critic shards it against the configured `value.train_global_batch_size` rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO.
 - Reward shaping — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and `overlong_filtering` are supported.

examples/configs/grpo_math_1B_megatron_generation_colocated_single_controller.yaml

@@ -0,0 +1,9 @@
+defaults: grpo_math_1B_megatron_single_controller.yaml
+policy:
+  generation:
+    backend: megatron
+    colocated:
+      enabled: true
+      resources:
+        gpus_per_node: null
+        num_nodes: null

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.yaml

@@ -0,0 +1,16 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron-single-controller-sync.yaml
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync
+async_rl:
+  sampler:
+    max_lookahead_versions: 4
+  max_buffered_rollouts: 160
+policy:
+  generation:
+    backend: megatron
+    mcore_generation_config:
+      kv_cache_management_mode: recompute

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-single-controller-sync.yaml

@@ -0,0 +1,14 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.yaml
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-single-controller-sync
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-single-controller-sync
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-single-controller-sync
+policy:
+  generation:
+    colocated:
+      enabled: true
+      resources:
+        gpus_per_node: null
+        num_nodes: null

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync.yaml

@@ -5,6 +5,10 @@ logger:
     name: grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync
 checkpointing:
   checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync
+async_rl:
+  sampler:
+    max_lookahead_versions: 4
+  max_buffered_rollouts: 160
 policy:
   generation:
     backend: megatron

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-single-controller-sync.yaml

@@ -0,0 +1,14 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync.yaml
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-single-controller-sync
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-single-controller-sync
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-single-controller-sync
+policy:
+  generation:
+    colocated:
+      enabled: true
+      resources:
+        gpus_per_node: null
+        num_nodes: null

nemo_rl/algorithms/async_utils/staleness_sampler.py

@@ -677,8 +677,10 @@ class CustomSamplerConfig(BaseModel, extra="allow"):
 def required_buffer_capacity_for_config(
     cfg: SamplerConfig,
     groups_per_step: int,
-) -> Optional[int]:
-    """Return a built-in sampler's required capacity without constructing it."""
+    *,
+    min_groups_for_streaming_train: int,
+) -> int:
+    """Return a sampler config's required capacity without constructing it."""
     if isinstance(cfg, ReadyFirstSamplerConfig):
         return _gated_required_buffer_capacity(
             groups_per_step,
@@ -694,7 +696,9 @@ def required_buffer_capacity_for_config(
             groups_per_step,
             gate_window=cfg.peak_lookahead_versions,
         )
-    return None
+    if isinstance(cfg, WindowedSamplerConfig):
+        return min_groups_for_streaming_train
+    return groups_per_step
 
 
 def _custom_sampler_class(cfg: CustomSamplerConfig) -> type:

nemo_rl/algorithms/single_controller.py

@@ -1773,6 +1773,7 @@ async def _train_pump(self) -> None:
                         max_prompt_groups = target_groups - groups_dispatched
                         if max_prompt_groups <= 0:
                             break
+                        # For a colocated engine this is max_prompt_groups, pinned inside setup.
                         min_prompt_groups = min(
                             self._async_cfg.min_groups_for_streaming_train,
                             max_prompt_groups,
@@ -1819,6 +1820,20 @@ async def _train_pump(self) -> None:
                     if groups_dispatched == 0 and self._gen is not None:
                         await asyncio.to_thread(self._gen.snapshot_step_metrics)
 
+                    # Whole steps are assembled for colocated engines so the loop closes after this.
+                    # In-flight requests freeze, then continue on fresh weights.
+                    # The post-step `_sync_weights` wake reopens the gate,
+                    # except on save-bound steps, where the wake is deferred until after the save.
+                    if self._gen is not None and self._gen.blocks_training():
+                        self._rollout_permitted.clear()
+                        await asyncio.to_thread(self._gen.finish_generation)
+                        # Deadline clocks are meant to be in terms of the inference clock-time.
+                        # However, naively, we measure them against wall clock-time.
+                        # This is incorrect in colocated. Every time we switch to training,
+                        # we unfairly penalize the deadline clocks.
+                        # Instead, freeze the deadline clocks while we are training.
+                        self._rollout_manager.suspend_request_deadlines()
+
                     # ---- 2. Prepare the batch ----
                     # Compute prev_logprobs / ref_logprobs
                     if (
@@ -2078,46 +2093,6 @@ async def _train_pump(self) -> None:
                     for step, promoted in self._batch_promotions.items()
                     if step > version_during_step
                 }
-
-                # ---- 6. Refit the model ----
-                with self._timer.time("weight_sync"):
-                    calibration_data = (
-                        BatchedDataDict.from_batches(calibration_batches)
-                        if calibration_batches
-                        else None
-                    )
-                    # Critic warmup doesn't need refit, and the version still advances.
-                    aborted_stale_inflight_groups = 0
-                    if is_policy_training_step:
-                        aborted_stale_inflight_groups = await self._sync_weights(
-                            calibration_data=calibration_data
-                        )
-                    self._retune_lookahead_versions()
-                    self._rollout_manager.set_weight_version(self._trainer_version)
-                    step_metrics.update(
-                        {
-                            "evicted_stale_prompt_groups": evicted_stale_prompt_groups,
-                            "aborted_stale_inflight_groups": aborted_stale_inflight_groups,
-                            # Non-zero means this step trained on a smaller batch than
-                            # num_prompts_per_step, which any comparison of step metrics
-                            # across steps has to account for.
-                            "dropped_prompt_groups": dropped_prompt_groups,
-                            # Groups filled by a spare prompt this step waited on --
-                            # either one it lost itself, or one it lent to an earlier
-                            # step and was repaid for. Non-zero here with zero above is
-                            # the healthy shape of on_dropped_prompt="replace": the
-                            # batch stayed whole, and the cost was the wall-clock spent
-                            # waiting on the spare.
-                            "replaced_prompt_groups": replaced_prompt_groups,
-                            # Groups this step lost and filled by borrowing finished work
-                            # from a later step. The better shape of the same thing: the
-                            # batch stayed whole and nothing waited for it, with the
-                            # repayment showing up as a replacement on the lender.
-                            "promoted_prompt_groups": promoted_prompt_groups,
-                        }
-                    )
-
-                # Checkpointing (mirrors async_grpo_train's save block).
                 # What the step actually trained on, which is num_prompts_per_step only
                 # when nothing was dropped. Counted from the dispatch tally rather than
                 # derived from the shortfall so the figure does not depend on the
@@ -2142,16 +2117,87 @@ async def _train_pump(self) -> None:
                         and self._train_steps % ft_save_period == 0
                     )
                 )
+                # Call once per step and reuse the bool.
                 should_save_by_timeout = self._timeout.check_save()
+                will_save_checkpoint = self._master_config.checkpointing[
+                    "enabled"
+                ] and (should_save_by_step or should_save_by_timeout)
+                # A save-bound colocated step defers the wake until after the save.
+                defer_wake_for_save = (
+                    will_save_checkpoint
+                    and self._gen.blocks_training()
+                    and self._gen.wake_carries_weight_updates()
+                )
 
-                if self._master_config.checkpointing["enabled"] and (
-                    should_save_by_step or should_save_by_timeout
-                ):
+                # ---- 6. Refit the model ----
+                # Critic warmup doesn't need a refit and the version still advances.
+                # But a colocated engine that was stood down still needs to be woken up via reshard.
+                aborted_stale_inflight_groups = 0
+                if is_policy_training_step or self._gen.blocks_training():
+                    if defer_wake_for_save:
+                        # Wake-deferral (colocated): the engine is about to be saved; let it sleep.
+                        # Record `weight_sync` for consistency in reports.
+                        with self._timer.time("weight_sync"):
+                            pass
+                        with self._timer.time("offload_before_refit"):
+                            await asyncio.to_thread(self._trainer.offload_before_refit)
+                    else:
+                        with self._timer.time("weight_sync"):
+                            calibration_data = (
+                                BatchedDataDict.from_batches(calibration_batches)
+                                if calibration_batches
+                                else None
+                            )
+                            aborted_stale_inflight_groups = await self._sync_weights(
+                                calibration_data=calibration_data,
+                            )
+                self._retune_lookahead_versions()
+                self._rollout_manager.set_weight_version(self._trainer_version)
+                step_metrics.update(
+                    {
+                        "evicted_stale_prompt_groups": evicted_stale_prompt_groups,
+                        "aborted_stale_inflight_groups": aborted_stale_inflight_groups,
+                        # Non-zero means this step trained on a smaller batch than
+                        # num_prompts_per_step, which any comparison of step metrics
+                        # across steps has to account for.
+                        "dropped_prompt_groups": dropped_prompt_groups,
+                        # Groups filled by a spare prompt this step waited on --
+                        # either one it lost itself, or one it lent to an earlier
+                        # step and was repaid for. Non-zero here with zero above is
+                        # the healthy shape of on_dropped_prompt="replace": the
+                        # batch stayed whole, and the cost was the wall-clock spent
+                        # waiting on the spare.
+                        "replaced_prompt_groups": replaced_prompt_groups,
+                        # Groups this step lost and filled by borrowing finished work
+                        # from a later step. The better shape of the same thing: the
+                        # batch stayed whole and nothing waited for it, with the
+                        # repayment showing up as a replacement on the lender.
+                        "promoted_prompt_groups": promoted_prompt_groups,
+                    }
+                )
+
+                # Checkpointing (mirrors async_grpo_train's save block).
+                if will_save_checkpoint:
                     with self._timer.time("checkpointing"):
                         await self._save_checkpoint(
                             step_metrics,
                             is_policy_training_step=is_policy_training_step,
                         )
+                    if defer_wake_for_save:
+                        # The save is done; wake the engine unless the loop is about to exit.
+                        # Deliberately a bare wake, not `_sync_weights`.
+                        loop_will_exit = is_last_step or should_save_by_timeout
+                        if not loop_will_exit:
+                            with self._timer.time("weight_sync"):
+                                # The save onloaded model+optimizer; generation needs offload.
+                                await asyncio.to_thread(
+                                    self._trainer.offload_after_refit
+                                )
+                                await asyncio.to_thread(
+                                    self._gen.prepare_for_generation
+                                )
+                                self._rollout_permitted.set()
+                                self._rollout_manager.resume_request_deadlines()
 
             timing_metrics: dict[str, float] = self._timer.get_timing_metrics(
                 reduction_op="sum"
@@ -2776,10 +2822,10 @@ async def _save_checkpoint(
     ) -> None:
         """Write a full checkpoint for the just-finished train step.
 
-        Everything except the (possibly async) policy weight write must be
-        on disk before begin_finalization; rollouts keep running throughout.
-        The policy optimizer is skipped during critic warmup -- it has never
-        stepped.
+        Everything except the (possibly async) policy weight write must be on disk
+        before `begin_finalization`. Non-colocated engines keep serving rollouts throughout.
+        Colocated engines are stood down for the save.
+        The policy optimizer is skipped during critic warmup: it has never stepped.
         """
         save_state = self._save_state
         # SC has no validation loop yet; drop the default sentinel instead of
@@ -3047,9 +3093,9 @@ async def _sync_weights(
     ) -> int:
         """Pause new rollout dispatches, synchronize weights, resume.
 
-        SC owns the pause gate; in-flight generations continue through the
-        refit — vLLM V1 async engine supports weight updates during pending
-        requests.
+        SC owns the pause gate. vLLM serves through the refit; it supports live weight updates.
+        A colocated engine is instead already stood down: the synchronizer's sync is the wake,
+        and step 4 resumes dispatch.
 
         Flow:
           1. _rollout_permitted.clear()  — no new dispatches
@@ -3160,6 +3206,7 @@ async def _sync_weights(
 
         print(f"  _sync_weights: sync done in {elapsed:.3f}s", flush=True)
         self._rollout_permitted.set()
+        self._rollout_manager.resume_request_deadlines()
         return aborted_stale_inflight_groups
 
     async def _value_stage(self, meta: KVBatchMeta) -> KVBatchMeta: