feat(sglang): megatron backend weight refit for sglang rollouts

examples/configs/grpo_math_1B_sglang.yaml

@@ -22,7 +22,7 @@ policy:
       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
+      cuda_graph_backend_prefill: breakable
       disable_cuda_graph: false
       sglang_server_config:
         needs_offload: true
@@ -32,6 +32,8 @@ policy:
         pause_generation_mode: retract
         num_gpus: 2
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        # "ipc" for colocated inference, "broadcast" for disaggregate GPUs.
+        weight_transfer_mode: ipc
       sglang_router_config:
         use_external_router: false
     colocated:

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

@@ -0,0 +1,63 @@
+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-1n4g-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: 1
+      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
+      cuda_graph_backend_prefill: breakable
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 4
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
+      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-1n4g-fsdp2tp1-sglang
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl-sglang-1n4g
+    name: grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang
+
+cluster:
+  gpus_per_node: 4
+  num_nodes: 1

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

@@ -0,0 +1,72 @@
+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-1n4g-megatrontp1-sglang
+
+policy:
+  model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  tokenizer:
+    name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  optimizer: null
+  megatron_cfg:
+    enabled: true
+    tensor_model_parallel_size: 1
+    pipeline_model_parallel_size: 1
+    scheduler:
+      lr_warmup_iters: 50
+  dtensor_cfg:
+    enabled: false
+  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: 1
+      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
+      cuda_graph_backend_prefill: breakable
+      disable_cuda_graph: false
+      sglang_server_config:
+        needs_offload: true
+        cpu_weight_backup: true
+        sglang_server_concurrency: 1024
+        pause_generation_mode: retract
+        num_gpus: 4
+        num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
+      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-1n4g-megatrontp1-sglang
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl-sglang-1n4g
+    name: grpo-qwen2.5-math-1.5b-instruct-1n4g-megatrontp1-sglang
+
+cluster:
+  gpus_per_node: 4
+  num_nodes: 1

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

@@ -29,7 +29,7 @@ policy:
       max_running_requests: null
       mem_fraction_static: 0.6
       skip_server_warmup: true
-      disable_piecewise_cuda_graph: true
+      cuda_graph_backend_prefill: breakable
       disable_cuda_graph: false
       sglang_server_config:
         needs_offload: true
@@ -38,6 +38,7 @@ policy:
         pause_generation_mode: retract
         num_gpus: 8
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
       sglang_router_config:
         use_external_router: false
         use_distributed_post: true

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

@@ -33,7 +33,7 @@ policy:
       max_running_requests: null
       mem_fraction_static: 0.6
       skip_server_warmup: true
-      disable_piecewise_cuda_graph: true
+      cuda_graph_backend_prefill: breakable
       disable_cuda_graph: false
       sglang_server_config:
         needs_offload: true
@@ -42,6 +42,7 @@ policy:
         pause_generation_mode: retract
         num_gpus: 16
         num_gpus_per_engine: ${policy.generation.sglang_cfg.tp_size}
+        weight_transfer_mode: ipc
       sglang_router_config:
         use_external_router: false
         use_distributed_post: true

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

@@ -0,0 +1,70 @@
+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-megatrontp1-sglang-noncolocated
+policy:
+  model_name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  tokenizer:
+    name: Qwen/Qwen2.5-Math-1.5B-Instruct
+  optimizer: null
+  megatron_cfg:
+    enabled: true
+    tensor_model_parallel_size: 1
+    pipeline_model_parallel_size: 1
+    scheduler:
+      lr_warmup_iters: 50
+  dtensor_cfg:
+    enabled: false
+  dynamic_batching:
+    enabled: true
+  sequence_packing:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  generation:
+    backend: sglang
+    max_new_tokens: 512
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 1
+    sglang_cfg:
+      model_path: ${policy.model_name}
+      dtype: ${policy.precision}
+      context_length: 512
+      allow_auto_truncate: true
+      tp_size: 2
+      dp_size: 1
+      pp_size: 1
+      ep_size: 1
+      random_seed: 42
+      max_running_requests: null
+      mem_fraction_static: 0.6
+      skip_server_warmup: true
+      cuda_graph_backend_prefill: breakable
+      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}
+        weight_transfer_mode: broadcast
+      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-megatrontp1-sglang-noncolocated
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-qwen2.5-math-1.5b-instruct-2n8g-megatrontp1-sglang-noncolocated
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2

nemo_rl/algorithms/grpo.py

@@ -1266,6 +1266,19 @@ def init_vllm_then_policy():
         if "model_path" not in generation_config["sglang_cfg"]:
             generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"]
 
+        expected_weight_transfer_mode = "ipc" if colocated_inference else "broadcast"
+        weight_transfer_mode = (
+            generation_config["sglang_cfg"]
+            .get("sglang_server_config", {})
+            .get("weight_transfer_mode", expected_weight_transfer_mode)
+        )
+        if weight_transfer_mode != expected_weight_transfer_mode:
+            raise ValueError(
+                f"sglang_server_config.weight_transfer_mode={weight_transfer_mode!r} "
+                f"is inconsistent with colocated.enabled={colocated_inference}: "
+                f"expected {expected_weight_transfer_mode!r}."
+            )
+
         policy_generation, policy = initialize_generation_with_policy(
             init_generation_fn=init_sglang,
             generation_name="SGLang",
@@ -1274,9 +1287,6 @@ 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,
@@ -1337,9 +1347,13 @@ def init_trtllm():
             "https://github.com/NVIDIA-NeMo/RL/issues/3288."
         )
 
-    # if it is not colocated inference, initialize collective communication for update weights
+    # if it is not colocated inference, initialize collective communication for update weights.
+    # SGLang owns its own weight-update process group (set up lazily on the
+    # first refit through ``connect_sglang_rollout_engines_distributed``), so
+    # skip the legacy trainer/vLLM init_collective handshake for SGLang.
     if (
         not colocated_inference
+        and backend != "sglang"
         and remote_transport is None
         and checkpoint_engine_config is None
     ):
@@ -1429,6 +1443,21 @@ def init_trtllm():
             f"Using checkpoint-engine refit backend: {checkpoint_engine_config['backend']}",
             flush=True,
         )
+    elif backend == "sglang":
+        t0 = time.perf_counter()
+        policy_generation.weight_synchronizer = create_weight_synchronizer(
+            policy=policy,
+            generation=policy_generation,
+            generation_backend=backend,
+            colocated=colocated_inference,
+            refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"),
+        )
+        # Only exchanges refit metadata. SGLang's own weight-update group is
+        # established lazily on the first refit.
+        policy_generation.weight_synchronizer.init_communicator()
+        worker_init_timing_metrics["sglang_weight_sync_init_time_s"] = (
+            time.perf_counter() - t0
+        )
     else:
         if not (nccl_reshard_refit_enabled and not colocated_inference):
             state_dict_info = policy.prepare_refit_info()
@@ -2223,10 +2252,22 @@ def refit_policy_generation(
     Returns:
         Scalar metrics reported by the selected weight synchronizer.
     """
+    # Every SGLang deployment reaches its refit through this hook: `setup`
+    # attaches an SGLang synchronizer that owns the whole lifecycle (phase
+    # transitions, engine recovery, pause/flush, transport), so SGLang never
+    # touches the branches below.
     synchronizer = getattr(policy_generation, "weight_synchronizer", None)
     if synchronizer is not None:
         return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
 
+    if isinstance(policy_generation, SGLangGeneration):
+        # Fail loudly rather than falling through to the vLLM branches, which
+        # would call methods the SGLang path does not implement.
+        raise RuntimeError(
+            "SGLang refits require policy_generation.weight_synchronizer to be "
+            "set. Attach one with create_weight_synchronizer(...) during setup."
+        )
+
     # Megatron generation backend needs explicit suspend/resume around refits.
     if isinstance(policy_generation, MegatronGeneration):
         policy_generation.suspend_for_refit()
@@ -2269,36 +2310,21 @@ def refit_policy_generation(
                     policy.get_free_memory_bytes() * float(memory_ratio)
                 )
 
-            if isinstance(policy_generation, SGLangGeneration):
-                # Stream weights to colocated SGLang engines via CUDA IPC over HTTP.
-                futures_train = policy.stream_weights_via_http(
-                    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)
-                update_success = True
-            else:
-                # ZMQ IPC path: shared by vLLM and TRT-LLM colocated. Trainer
-                # streams CUDA IPC handles in chunks; receiver reconstructs
-                # tensors in-place and feeds them into the inference engine's
-                # loader.
-                futures_train = policy.stream_weights_via_ipc_zmq(
-                    buffer_size_bytes=buffer_size_bytes,
-                    kv_scales=kv_scales,
-                )
-                futures_inference = policy_generation.update_weights_via_ipc_zmq()
-                # wait for all futures to complete
-                ray.get(futures_train)
-                results = ray.get(futures_inference)
-                update_success = all(result for result in results if result is not None)
+            # ZMQ IPC path: shared by vLLM and TRT-LLM colocated. Trainer
+            # streams CUDA IPC handles in chunks; receiver reconstructs
+            # tensors in-place and feeds them into the inference engine's
+            # loader.
+            futures_train = policy.stream_weights_via_ipc_zmq(
+                buffer_size_bytes=buffer_size_bytes,
+                kv_scales=kv_scales,
+            )
+            futures_inference = policy_generation.update_weights_via_ipc_zmq()
+            # wait for all futures to complete
+            ray.get(futures_train)
+            results = ray.get(futures_inference)
+            update_success = all(result for result in results if result is not None)
         else:
             # update weights through nccl (vLLM) or megatron reshard
-            # SGLang haven't implemented non-colocated inference mode.
-            if isinstance(policy_generation, SGLangGeneration):
-                raise NotImplementedError(
-                    "SGLang haven't implemented non-colocated inference mode. "
-                )
             if isinstance(policy_generation, MegatronGeneration):
                 futures_train = policy.swap_weights_via_reshard(is_source=True)
             else: