feat: Sglang Rollout Refactor

examples/configs/grpo_math_1B_sglang.yaml

@@ -6,23 +6,40 @@ grpo:
 policy:
   generation:
     backend: "sglang"
+    use_async_rollouts: false
     sglang_cfg:
-      # SGLang specific configuration
-      model_path: ${policy.model_name} 
-      gpus_per_server: 1 
+      model_path: ${policy.model_name}
       dtype: ${policy.precision}
-      context_length: 512  # Maximum context length
+      context_length: 512
       allow_auto_truncate: true
-      enable_memory_saver: false
+      tp_size: 1
       dp_size: 1
       pp_size: 1
       ep_size: 1
+      random_seed: 42
       max_running_requests: null
-      mem_fraction_static: 0.7
+      mem_fraction_static: 0.6
       skip_server_warmup: true
       # Piecewise CUDA graph currently crashes with "illegal memory access"
       # (likely torch 2.10 + sglang incompatibility). Keep disabled until upstream fix.
       disable_piecewise_cuda_graph: true
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        # pause generation has no effect here
+        pause_generation_mode: retract
+        num_gpus: 2
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+      sglang_router_config:
+        use_external_router: false
+    colocated:
+      enabled: true
 
 logger:
   wandb_enabled: true
+
+cluster:
+  gpus_per_node: 2
+  num_nodes: 1

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml

@@ -1,11 +1,9 @@
 defaults: ../../grpo_math_1B.yaml
-
 grpo:
   max_num_steps: 450
-
+  val_batch_size: 128
 checkpointing:
   checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang
-
 policy:
   model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
   tokenizer:
@@ -16,36 +14,41 @@ policy:
     enabled: false
   make_sequence_length_divisible_by: 1
   generation:
-    backend: "sglang"
+    backend: sglang
     max_new_tokens: 512
     sglang_cfg:
       model_path: ${policy.model_name}
-      gpus_per_server: 1
       dtype: ${policy.precision}
       context_length: 512
       allow_auto_truncate: true
-      enable_memory_saver: false
+      tp_size: 2
       dp_size: 1
       pp_size: 1
       ep_size: 1
+      random_seed: 42
       max_running_requests: null
-      mem_fraction_static: 0.5
+      mem_fraction_static: 0.6
       skip_server_warmup: true
-      # Piecewise CUDA graphs fail with CUBLAS_STATUS_EXECUTION_FAILED
-      # inside Ray worker forks. See unit test Err 4 for details.
       disable_piecewise_cuda_graph: true
-
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 8
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+      sglang_router_config:
+        use_external_router: false
+        use_distributed_post: true
 data:
   max_input_seq_length: 512
-
 logger:
   log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang
   wandb_enabled: true
   tensorboard_enabled: true
   wandb:
     project: nemo-rl
     name: grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang
-
 cluster:
   gpus_per_node: 8
-

examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang.yaml

@@ -0,0 +1,62 @@
+defaults: ../../grpo_math_1B.yaml
+
+grpo:
+  max_num_steps: 450
+  val_batch_size: 128
+
+checkpointing:
+  checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang
+
+policy:
+  model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  tokenizer:
+    name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  generation:
+    backend: "sglang"
+    max_new_tokens: 512
+    sglang_cfg:
+      model_path: ${policy.model_name}
+      dtype: ${policy.precision}
+      context_length: 512
+      allow_auto_truncate: true
+      tp_size: 4
+      dp_size: 1
+      # pp_size must be 1 (sglang pipeline parallelism not yet supported here).
+      pp_size: 1
+      ep_size: 1
+      random_seed: 42
+      max_running_requests: null
+      mem_fraction_static: 0.6
+      skip_server_warmup: true
+      disable_piecewise_cuda_graph: true
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 16
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+      sglang_router_config:
+        use_external_router: false
+        use_distributed_post: true
+
+data:
+  max_input_seq_length: 512
+
+logger:
+  log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-qwen2.5-math-1.5b-instruct-2n8g-fsdp2tp1-sglang
+
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2

nemo_rl/algorithms/grpo.py

@@ -82,7 +82,8 @@
 )
 from nemo_rl.models.generation.interfaces import GenerationInterface
 from nemo_rl.models.generation.megatron import MegatronGeneration
-from nemo_rl.models.generation.sglang import SGLangConfig, SGLangGeneration
+from nemo_rl.models.generation.sglang.config import SGLangConfig
+from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
 from nemo_rl.models.megatron.router_replay import (
     configure_vllm_for_router_replay,
@@ -719,7 +720,10 @@ def init_vllm():
     def init_sglang():
         """Initialize SGLang generation workers."""
         t0 = time.perf_counter()
-        pg = SGLangGeneration(cluster=inference_cluster, config=generation_config)
+        pg = SGLangGeneration(
+            cluster=inference_cluster,
+            sglang_cfg=generation_config,
+        )
         pg.finish_generation()
         return pg, time.perf_counter() - t0
 
@@ -954,6 +958,9 @@ def init_vllm_then_policy():
             worker_init_timing_metrics=worker_init_timing_metrics,
         )
 
+        # Capture rollout TP size on the policy once; refit calls no longer need it.
+        policy.set_rollout_num_gpus_per_engine(policy_generation.num_gpus_per_engine)
+
         print(
             f"  ✓ Using SGLang backend for generation with {policy_config['model_name']}",
             flush=True,
@@ -1425,21 +1432,26 @@ def _apply_configured_message_level_advantage_penalties(
 def _should_use_async_rollouts(master_config: MasterConfig) -> bool:
     """Determine if async rollouts should be used based on the configuration.
 
-    Returns True if async_engine is enabled.
+    SGLang only uses async rollouts when explicitly configured with
+    ``policy.generation.use_async_rollouts``. vLLM and Megatron use async
+    rollouts when their respective ``async_engine`` config is enabled.
     """
     generation_config = master_config.policy["generation"]
     if generation_config is None:
         return False
     backend = generation_config.get("backend", "")
 
+    if backend == "sglang":
+        return bool(generation_config.get("use_async_rollouts", False))
+
     if backend == "vllm":
-        vllm_cfg = generation_config.get("vllm_cfg", {})
-        return vllm_cfg.get("async_engine", False)
-    elif backend == "megatron":
+        return bool(generation_config.get("vllm_cfg", {}).get("async_engine", False))
+
+    if backend == "megatron":
         mcore_cfg = generation_config.get("mcore_generation_config", {})
         return mcore_cfg.get("async_engine", False)
-    else:
-        return False
+
+    return False
 
 
 def _should_use_nemo_gym(master_config: MasterConfig) -> bool:
@@ -1616,15 +1628,10 @@ def refit_policy_generation(
                 )
 
             if isinstance(policy_generation, SGLangGeneration):
-                sglang_url_to_gpu_uuids = (
-                    policy_generation.get_sglang_url_to_gpu_uuids()
-                )
-                # Stream weights via HTTP
-                flush_success = policy_generation.invalidate_kv_cache()
-                if not flush_success:
-                    print("SGLang KV cache invalidation failed before weight update. ")
+                # Stream weights to colocated SGLang engines via CUDA IPC over HTTP.
                 futures_train = policy.stream_weights_via_http(
-                    sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids,
+                    rollout_engine_urls=policy_generation.get_rollout_engine_urls(),
+                    buffer_size_bytes=buffer_size_bytes,
                 )
                 # Wait for all workers to complete
                 ray.get(futures_train)
@@ -2066,7 +2073,7 @@ def grpo_train(
                                 if "full_result" in key:
                                     rollout_metrics.pop(key)
 
-                    # Use async rollouts if vLLM async engine is enabled
+                    # Use async rollouts when enabled by config/backend defaults.
                     elif _should_use_async_rollouts(master_config):
                         (
                             repeated_batch,
@@ -2870,7 +2877,7 @@ def validate(
 
             additional_metrics_to_report = dict()
             # Generate responses (updates the LLMMessageLogType in batch_with_msg_logs)
-            # Use async rollouts if vLLM async engine is enabled
+            # Use async rollouts when enabled by config/backend defaults.
             # We cascade NeMo-Gym first since NeMo-Gym also uses async rollouts.
             if _should_use_nemo_gym(master_config):
                 generation_config = master_config.policy["generation"]
@@ -3053,9 +3060,15 @@ def async_grpo_train(
         master_config: Master configuration
         max_trajectory_age_steps: Maximum age (in training steps) for trajectories to be used in training
     """
-    # Ensure we are running with a compatible async generation backend
-    assert _should_use_async_rollouts(master_config), (
-        "Async GRPO requires an async generation engine. "
+    # Ensure we are running with a compatible async generation backend.
+    # Async GRPO (with in-flight weight updates) supports vLLM and Megatron;
+    # SGLang async rollouts do not support the async GRPO replay path.
+    generation_config = master_config.policy["generation"]
+    backend = generation_config.get("backend", "") if generation_config else ""
+    assert backend in ("vllm", "megatron") and _should_use_async_rollouts(
+        master_config
+    ), (
+        "Async GRPO requires an async vLLM or Megatron generation engine. "
         "Set either policy.generation.vllm_cfg.async_engine=true (vLLM) or "
         "policy.generation.mcore_generation_config.async_engine=true (Megatron)."
     )

nemo_rl/algorithms/ppo.py

@@ -68,7 +68,8 @@
     run_multi_turn_rollout,
 )
 from nemo_rl.models.generation.interfaces import GenerationInterface
-from nemo_rl.models.generation.sglang import SGLangConfig, SGLangGeneration
+from nemo_rl.models.generation.sglang.config import SGLangConfig
+from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface
@@ -511,7 +512,10 @@ def init_vllm():
     def init_sglang():
         """Initialize SGLang generation workers."""
         t0 = time.perf_counter()
-        pg = SGLangGeneration(cluster=inference_cluster, config=generation_config)
+        pg = SGLangGeneration(
+            cluster=inference_cluster,
+            sglang_cfg=generation_config,
+        )
         pg.finish_generation()
         return pg, time.perf_counter() - t0