feat(sc): setup + entrypoint (#3266)

Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: ruit <ruit@nvidia.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: ruit <ruit@nvidia.com>

Author: 3 people

.github/workflows/cicd-main.yml

@@ -731,6 +731,8 @@ jobs:
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_PPO
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
+          - script: L1_Functional_Tests_SingleController
+            runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_Eval
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_Other_1
@@ -800,6 +802,8 @@ jobs:
             runner: ${{ vars.GB200_RUNNER }}
           - script: L1_Functional_Tests_PPO
             runner: ${{ vars.GB200_RUNNER }}
+          - script: L1_Functional_Tests_SingleController
+            runner: ${{ vars.GB200_RUNNER }}
           - script: L1_Functional_Tests_Eval
             runner: ${{ vars.GB200_RUNNER }}
           - script: L1_Functional_Tests_Other_1
@@ -873,6 +877,8 @@ jobs:
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_PPO
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
+          - script: L1_Functional_Tests_SingleController
+            runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_Eval
             runner: ${{ needs.org-member-pre-flight.outputs.runner_prefix }}
           - script: L1_Functional_Tests_Other_1

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -0,0 +1,62 @@
+# SingleController variant of grpo_math_1B.yaml.
+defaults: grpo_math_1B.yaml
+
+# SC does not support validation yet.
+grpo:
+  val_period: 0
+
+async_rl:
+  sampler:
+    name: in_order
+    # How far generation may run ahead of the trainer.
+    max_lookahead_versions: 1
+  # Recompute generation KV caches after each weight update.
+  recompute_kv_cache_after_weight_updates: false
+  # Min ready groups the streaming trainer waits for before dispatching a batch.
+  min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
+  # Cap on in-flight generate_and_push calls in the rollout pump.
+  max_inflight_prompts: ${grpo.num_prompts_per_step}
+  # Cap on unconsumed rollout groups buffered in the DataPlane (backpressure).
+  max_buffered_rollouts: 64
+  # Enable per-rollout diagnostic prints (prompt content / completion previews).
+  diagnostics: false
+
+# SC does not support checkpointing yet.
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-single-controller
+
+policy:
+  dtensor_cfg:
+    enabled: false
+  megatron_cfg:
+    enabled: true
+
+  generation:
+    port_range_low: 11001
+    port_range_high: 15000
+    mcore_generation_config:
+      unified_memory_level: 0
+    vllm_cfg:
+      async_engine: true
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 1
+        num_nodes: 1
+
+logger:
+  wandb:
+    name: grpo-single-controller-dev
+  swanlab:
+    name: grpo-single-controller-dev
+  mlflow:
+    run_name: grpo-single-controller-dev
+
+data_plane:
+  enabled: true
+
+cluster:
+  gpus_per_node: 2
+  master_port_range_low: 25000
+  master_port_range_high: 28000

examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml

@@ -0,0 +1,30 @@
+# SingleController variant of grpo-llama3.1-8b-instruct-2n8g-async-1off.yaml.
+defaults: ./performance/grpo-llama3.1-8b-instruct-2n8g-async-1off.yaml
+
+logger:
+  log_dir: logs/grpo-llama3.1-8b-instruct-2n8g-async-1off-sc
+  wandb:
+    name: grpo-llama3.1-8b-instruct-2n8g-async-1off-sc
+
+# SC does not support validation and checkpointing yet.
+grpo:
+  val_period: 0
+
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-llama3.1-8b-instruct-2n8g-async-1off-sc
+
+# TransferQueue data plane is mandatory for the SingleController path.
+data_plane:
+  enabled: true
+
+# SC async-RL runtime knobs.
+async_rl:
+  sampler:
+    name: windowed
+    # Matches grpo.async_grpo.max_trajectory_age_steps=1.
+    max_staleness_versions: 1
+  # Min ready groups the streaming trainer waits for before dispatching a batch.
+  min_groups_for_streaming_train: 16
+  max_inflight_prompts: ${grpo.num_prompts_per_step}  # match grpo-llama3.1-8b-instruct-2n8g-async-1off
+  max_buffered_rollouts: 128

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

@@ -0,0 +1,49 @@
+# SingleController + Megatron variant of grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3.yaml.
+defaults: ./grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.v3.yaml
+
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-sc
+  wandb:
+    name: grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-sc
+
+# SC does not support validation and checkpointing yet.
+grpo:
+  val_period: 0
+
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-sc
+
+# TransferQueue data plane is mandatory for the SingleController path.
+data_plane:
+  enabled: true
+
+# SC async-RL runtime knobs.
+async_rl:
+  sampler:
+    name: in_order
+    max_lookahead_versions: 0
+  # Preserve sync-mode cache invalidation after every weight update.
+  recompute_kv_cache_after_weight_updates: true
+  # Min ready groups the streaming trainer waits for before dispatching a batch.
+  min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
+  max_inflight_prompts: ${grpo.num_prompts_per_step}
+  # In-order with zero lookahead requires one prompt batch of capacity.
+  max_buffered_rollouts: 32
+
+policy:
+  dtensor_cfg:
+    enabled: false
+  megatron_cfg:
+    enabled: true
+    scheduler:
+      lr_warmup_iters: 50
+  generation:
+    vllm_cfg:
+      async_engine: true
+    colocated:
+      enabled: false
+      resources:
+        # 4 GPUs for inference; remaining 4 GPUs on the node go to training.
+        gpus_per_node: 4
+        num_nodes: 1

examples/run_grpo_single_controller.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Async GRPO launcher driven by the SingleController actor.
+
+Builds the full SC actor args driver-side via setup_single_controller and hands them
+to SingleControllerActor. Mirrors run_grpo.py for config loading so the same YAML
+files apply. data_plane.enabled=true is mandatory.
+"""
+
+import argparse
+import os
+import pprint
+import sys
+
+import ray
+from omegaconf import OmegaConf
+
+from nemo_rl.algorithms.single_controller import SingleControllerActor
+from nemo_rl.algorithms.single_controller_utils import (
+    MasterConfig,
+    setup_single_controller,
+)
+from nemo_rl.algorithms.utils import get_tokenizer
+from nemo_rl.distributed.virtual_cluster import init_ray
+from nemo_rl.models.generation import configure_generation_config
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
+from nemo_rl.utils.logger import get_next_experiment_dir
+
+# Drop examples/ from sys.path so examples/nemo_gym/ (no __init__.py) doesn't
+# shadow the real nemo_gym package as a namespace package.
+current_dir = os.path.dirname(os.path.abspath(__file__))
+while current_dir in sys.path:
+    sys.path.remove(current_dir)
+
+
+def parse_args() -> tuple[argparse.Namespace, list[str]]:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Run async GRPO training via SingleController"
+    )
+    parser.add_argument(
+        "--config", type=str, default=None, help="Path to YAML config file"
+    )
+    args, overrides = parser.parse_known_args()
+    return args, overrides
+
+
+def main() -> None:
+    """Main entry point."""
+    register_omegaconf_resolvers()
+    args, overrides = parse_args()
+
+    if not args.config:
+        args.config = os.path.join(
+            os.path.dirname(__file__),
+            "configs",
+            "grpo_math_1B_megatron_single_controller.yaml",
+        )
+
+    config = load_config(args.config)
+    print(f"Loaded configuration from: {args.config}")
+
+    if overrides:
+        print(f"Overrides: {overrides}")
+        config = parse_hydra_overrides(config, overrides)
+
+    config = OmegaConf.to_container(config, resolve=True)
+    config = MasterConfig(**config)
+    print("Applied CLI overrides")
+
+    dp_cfg = config.data_plane
+    if not dp_cfg.get("enabled", False):
+        raise ValueError(
+            "run_grpo_single_controller requires data_plane.enabled=true. "
+            "Use examples/run_grpo.py for the legacy / sync paths."
+        )
+
+    print("Final config:")
+    pprint.pprint(config)
+
+    config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"])
+    print(f"📊 Using log directory: {config.logger['log_dir']}")
+    if config.checkpointing["enabled"]:
+        print(
+            f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}"
+        )
+
+    init_ray()
+
+    tokenizer = get_tokenizer(config.policy["tokenizer"])
+    assert config.policy["generation"] is not None, (
+        "A generation config is required for SC-driven async GRPO"
+    )
+    has_refit_draft_weights = bool(config.policy["draft"]["enabled"])
+    megatron_cfg = config.policy.get("megatron_cfg") or {}
+    trains_mtp = bool(megatron_cfg.get("mtp_num_layers"))
+    config.policy["generation"] = configure_generation_config(
+        config.policy["generation"],
+        tokenizer,
+        has_refit_draft_weights=has_refit_draft_weights,
+        trains_mtp=trains_mtp,
+    )
+
+    actor_args = setup_single_controller(config, tokenizer)
+
+    print("🚀 Launching SingleControllerActor")
+    sc = SingleControllerActor.remote(master_config=config, actor_args=actor_args)
+    try:
+        result = ray.get(sc.run.remote())
+        print(f"SC run complete: {result}")
+    finally:
+        # Drain env actors before generation to avoid in-flight requests during shutdown.
+        for env_name, handle in actor_args.env_handles.items():
+            try:
+                ray.get(handle.shutdown.remote())
+            except Exception as e:
+                print(f"Env {env_name!r} shutdown failed: {e}")
+
+        for resource_name, resource in (
+            ("Generation", actor_args.gen_handle),
+            ("Trainer", actor_args.trainer_handle),
+        ):
+            try:
+                resource.shutdown()
+            except Exception as e:
+                print(f"{resource_name} shutdown failed: {e}")
+
+
+if __name__ == "__main__":
+    main()