feat(sc): support non-colocated MInf

docs/guides/async-grpo.md

@@ -207,7 +207,7 @@ If no `replay_buffer.pt` file is found in the latest checkpoint directory, train
 
 4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain; SGLang is in this group today and instead relies on the pause its own weight synchronizer performs around the transfer. Managed Dynamo always drains active trajectories before refit. vLLM requires `async_engine: true`; the Megatron backend is always async-engine.
 
-5. **Recompute KV Cache After Weight Updates**: Set `recompute_kv_cache_after_weight_updates: true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. Other refit paths keep their existing post-update invalidation behavior. When false, in-flight requests retain their pre-update KV cache.
+5. **Recompute KV Cache After Weight Updates**: Set `recompute_kv_cache_after_weight_updates: true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. Other refit paths keep their existing post-update invalidation behavior. When false, in-flight requests retain their pre-update KV cache. On the Megatron generation backend, this must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`; setup errors on a mismatch.
 
 ## Why Importance Sampling Correction Is Required for Async
 

docs/guides/single-controller.md

@@ -25,7 +25,7 @@ uv run examples/run_grpo_single_controller.py --config <your-sc.yaml>
       enabled: true
     ```
 
-2. **Enable vLLM async engine** and **disable colocated inference** (SC drives rollout via `RolloutManager.generate_and_push`, which is only supported on the disaggregated async engine; setup rejects `colocated.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):
 
     ```yaml
     policy:
@@ -40,6 +40,23 @@ 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):
+
+    ```yaml
+    policy:
+      megatron_cfg:
+        enabled: true
+      generation:
+        backend: "megatron"
+        mcore_generation_config:
+          expose_http_server: true  # required for NeMo-Gym rollouts
+        colocated:
+          enabled: false
+          resources:
+            num_nodes: 1
+            gpus_per_node: 1  # inference GPUs; remainder go to training
+    ```
+
 3. **One RL step = one training batch.** The batch a step trains on is the whole step (see `validate_single_controller_config` in [nemo_rl/algorithms/single_controller_utils/config.py](../../nemo_rl/algorithms/single_controller_utils/config.py)). A GRPO step is also one optimizer step; a PPO step is `ppo.ppo_epochs` of them over that same batch.
 
     ```python
@@ -181,8 +198,8 @@ 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: only vLLM is supported and validated; Megatron generation, 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`).
+- Generation backend: vLLM and Megatron generation are supported; 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 and sample filtering — `overlong_filtering`, `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping.
 - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute.

examples/configs/grpo_math_1B.yaml

@@ -62,7 +62,7 @@ grpo:
     # when transient generation errors are expected and acceptable to drop.
     max_generation_failures: 0
     in_flight_weight_updates: false # Set to true to enable in-flight weight updates
-    recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates
+    recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates.
 
 # Reward-zeroing penalties applied to NeMo-Gym rollout results.
 reward_penalties:
@@ -410,7 +410,7 @@ policy:
       enable_chunked_prefill: true  # Split long prefills into chunks for better memory management
       enable_prefix_caching: false  # Reuse KV blocks across requests sharing a prompt prefix.
       max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
-      kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true.
+      kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute".
       materialize_only_last_token_logits: true
       num_speculative_tokens: 0
       logprobs_mode: processed_logprobs  # Return log-probs after sampling processors. Use raw_logprobs for parity with policy recomputation.

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

@@ -0,0 +1,12 @@
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync
+  wandb:
+    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
+policy:
+  generation:
+    backend: megatron
+    mcore_generation_config:
+      kv_cache_management_mode: recompute

examples/nemo_gym/grpo_nanov3.yaml

@@ -242,7 +242,7 @@ policy:
       use_cuda_graphs_for_non_decode_steps: true  # Enable CUDA graphs for prefill/context processing
       enable_chunked_prefill: true
       max_tokens: ${policy.max_total_sequence_length} # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
-      kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true.
+      kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute".
       materialize_only_last_token_logits: true
       num_speculative_tokens: 0
       logprobs_mode: processed_logprobs

examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml

@@ -0,0 +1,78 @@
+# GRPO on the NeMo-Gym workplace-assistant environment via the SingleController path,
+# using non-colocated Megatron Inference.
+# Gym rollouts go through the persistent Megatron engine's OpenAI-compatible server,
+# Smoke-test scale; Qwen3-0.6B on one node with 2 GPUs (1 training + 1 generation).
+# The CI-run variant, tests/functional/grpo_megatron_generation_gym_single_controller.sh,
+# loads this file and overrides only test scale, data paths, and logging; the resolved
+# config stays value-equal to the vLLM SC test (grpo_async_gym_single_controller.sh)
+# everywhere but the generation backend.
+defaults: "grpo_qwen3_30ba3b_instruct.yaml"
+
+grpo:
+  # SC requires one optimizer step per RL step:
+  # num_prompts_per_step * num_generations_per_prompt == policy.train_global_batch_size
+  num_prompts_per_step: 4
+  num_generations_per_prompt: 2
+  max_num_steps: 10  # short demo; raise for a real run
+  # SC does not support validation yet (setup raises when it is enabled).
+  val_period: 0
+  val_at_start: false
+  # The KL term below needs reference logprobs; the base skips them.
+  skip_reference_policy_logprobs_calculation: false
+  # SC replaces the legacy async-GRPO path.
+  async_grpo: null
+
+loss_fn:
+  # A small KL term (the base uses 0) exercises the reference-model path end to end.
+  reference_policy_kl_penalty: 0.01
+  use_importance_sampling_correction: true
+
+policy:
+  model_name: Qwen/Qwen3-0.6B
+  train_global_batch_size: 8
+  # Full workplace-assistant prompts (all tools attached) run past 4k tokens.
+  max_total_sequence_length: 8192
+
+  megatron_cfg:
+    tensor_model_parallel_size: 1
+    expert_model_parallel_size: 1
+    context_parallel_size: 1
+    sequence_parallel: false
+
+  generation:
+    backend: "megatron"
+    mcore_generation_config:
+      # NeMo-Gym drives rollouts through the engine's OpenAI server.
+      expose_http_server: true
+    colocated:
+      enabled: false
+      resources:
+        num_nodes: 1
+        gpus_per_node: 1
+
+# The gym family sets no data_plane block. Only keys without a schema default
+# are set here; see nemo_rl/data_plane/interfaces.py, and the fully documented
+# block in examples/configs/grpo_math_1B.yaml.
+data_plane:
+  enabled: true
+  impl: transfer_queue
+  backend: "simple"
+  claim_meta_poll_interval_s: 0.5
+  simple:
+    num_storage_units: ${mul:2, ${cluster.num_nodes}}  # TQ wants >= 2 per node
+
+async_rl:
+  sampler:
+    name: in_order
+    # 0 = fully synchronous, so the importance sampling correction above is an
+    # inert no-op (all ratios are 1); it is enabled to match the vLLM SC test.
+    max_lookahead_versions: 0
+  min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
+  max_inflight_prompts: ${grpo.num_prompts_per_step}
+  max_buffered_rollouts: ${grpo.num_prompts_per_step}
+
+checkpointing:
+  enabled: false
+
+cluster:
+  gpus_per_node: 2

nemo_rl/algorithms/grpo.py

@@ -1180,18 +1180,26 @@ def _spinup_nemo_gym(base_urls, model_name):
         )
         policy_config["megatron_cfg"]["train_iters"] = total_train_iters
 
-        # When the user opts into recompute-after-refit on the megatron side,
-        # override mcore's kv_cache_management_mode to "recompute" directly.
+    # Megatron generation expresses recompute-after-refit engine-side via
+    # `kv_cache_management_mode="recompute"`; the loop-level flag must agree.
+    if generation_config["backend"] == "megatron":
         async_grpo_config = grpo_config.async_grpo
-        if async_grpo_config.recompute_kv_cache_after_weight_updates:
-            mcore_cfg = policy_config["generation"]["mcore_generation_config"]
-            prior_mode = mcore_cfg.get("kv_cache_management_mode", "persist")
-            if prior_mode != "recompute":
-                print(
-                    f"kv_cache_management_mode overridden '{prior_mode}' -> 'recompute' by "
-                    f"grpo.async_grpo.recompute_kv_cache_after_weight_updates=True."
-                )
-            mcore_cfg["kv_cache_management_mode"] = "recompute"
+        recompute_kv_cache = bool(
+            async_grpo_config is not None
+            and async_grpo_config.recompute_kv_cache_after_weight_updates
+        )
+        kv_cache_mode = generation_config["mcore_generation_config"][
+            "kv_cache_management_mode"
+        ]
+        if recompute_kv_cache != (kv_cache_mode == "recompute"):
+            raise ValueError(
+                "grpo.async_grpo.recompute_kv_cache_after_weight_updates="
+                f"{recompute_kv_cache} conflicts with policy.generation."
+                f"mcore_generation_config.kv_cache_management_mode={kv_cache_mode!r}: "
+                "with policy.generation.backend='megatron' the two must agree. "
+                "Either set the flag to true with kv_cache_management_mode="
+                "'recompute', or leave the flag false with 'persist'/'offload'."
+            )
 
     # Define initialization functions that will be used in all paths
     init_reference_model = loss_config.reference_policy_kl_penalty > 0
@@ -1694,12 +1702,9 @@ def init_dynamo():
         if policy_generation.weight_synchronizer is None:
             init_megatron_weight_synchronizer(policy, policy_generation)
         if enable_nemo_gym:
-            served_urls = policy_generation.dp_openai_server_base_urls
-            if served_urls != [reserved_url]:
-                raise RuntimeError(
-                    "Megatron server came up at a different address than the one "
-                    f"pre-published to NeMo Gym: reserved {reserved_url}, serving {served_urls}."
-                )
+            MegatronGeneration.verify_served_address(
+                policy_generation.dp_openai_server_base_urls, reserved_url
+            )
     # if it is not colocated inference, initialize collective communication for update weights
     elif (
         not colocated_inference

nemo_rl/algorithms/single_controller.py

@@ -89,6 +89,7 @@
 from nemo_rl.experience.failures import RolloutStall
 from nemo_rl.experience.rollout_manager import RolloutOutcome
 from nemo_rl.models.generation.fleet_health import ShardState
+from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmGeneration
 from nemo_rl.models.policy.tq_policy import TQPolicy
@@ -97,7 +98,7 @@
 from nemo_rl.utils.logger import Logger
 from nemo_rl.utils.timer import TimeoutChecker, Timer
 
-Generation = Union[VllmGeneration, SGLangGeneration]
+Generation = Union[VllmGeneration, SGLangGeneration, MegatronGeneration]
 
 # Named `log` rather than `logger` to keep it distinct from the experiment
 # Logger this module also uses as `self._logger`.
@@ -350,8 +351,9 @@ def __init__(
 
     async def run(self) -> dict[str, Any]:
         """Main entry point. Runs until max_train_steps is reached."""
-        # Synchronize weights before starting the pumps
-        await self._sync_weights()
+        # Synchronize weights before starting the pumps, unless setup already delivered them.
+        if self._weight_synchronizer.is_stale:
+            await self._sync_weights()
         self._rollout_manager.set_weight_version(self._trainer_version)
 
         await self._maybe_restore_replay_buffer()