fix(sglang): avoid ncclCommSplit hang in the refit weight-update group

examples/configs/grpo_math_1B_sglang.yaml

@@ -24,6 +24,17 @@ policy:
       # (likely torch 2.10 + sglang incompatibility). Keep disabled until upstream fix.
       disable_piecewise_cuda_graph: true
       disable_cuda_graph: false
+      # Fault tolerance (RolloutHealthMonitor). Off by default; when enabled,
+      # a daemon thread health-checks each engine and restarts hung/dead actors.
+      use_fault_tolerance: false
+      rollout_health_check_interval: 60
+      rollout_health_check_timeout: 60
+      rollout_health_check_first_wait: 60
+      # Weight precision for rollout/refit. scheme=bf16 (default) sends BF16
+      # HF tensors; scheme=mxfp8 boots SGLang from an MXFP8 HF checkpoint and
+      # quantizes refit tensors online (see SglangQuantizationConfig).
+      quantization:
+        scheme: bf16
       sglang_server_config:
         needs_offload: true
         cpu_weight_backup: true
@@ -32,6 +43,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:

nemo_rl/algorithms/grpo.py

@@ -1201,6 +1201,23 @@ def init_vllm_then_policy():
         if "model_path" not in generation_config["sglang_cfg"]:
             generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"]
 
+        # If MXFP8 is requested, ensure SGLang boots from an MXFP8 HF
+        # checkpoint. This must happen before ``init_sglang`` so the engine
+        # loads quantized weights.
+        sglang_quantization_cfg = (
+            generation_config["sglang_cfg"].get("quantization") or {}
+        )
+        if sglang_quantization_cfg.get("scheme", "bf16") == "mxfp8":
+            from nemo_rl.models.generation.sglang.mxfp8_setup import (
+                ensure_mxfp8_checkpoint,
+            )
+
+            mxfp8_path = ensure_mxfp8_checkpoint(
+                model_path=generation_config["sglang_cfg"]["model_path"],
+                quantization_cfg=sglang_quantization_cfg,
+            )
+            generation_config["sglang_cfg"]["model_path"] = mxfp8_path
+
         policy_generation, policy = initialize_generation_with_policy(
             init_generation_fn=init_sglang,
             generation_name="SGLang",
@@ -1209,9 +1226,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,
@@ -1223,8 +1237,11 @@ def init_vllm_then_policy():
     # print the node IP and GPU ID of the policy workers for debugging
     policy.print_node_ip_and_gpu_id()
 
-    # if it is not colocated inference, initialize collective communication for update weights
-    if not colocated_inference:
+    # 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 not isinstance(policy_generation, SGLangGeneration):
         t0 = time.perf_counter()
         ip, port = train_cluster.get_master_address_and_port()
         print(f"Using ip: {ip}, port: {port} for collective communication", flush=True)
@@ -1263,6 +1280,19 @@ def init_vllm_then_policy():
             ray.get(futures_train + futures_inference)
         worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0
 
+    if backend == "sglang" and isinstance(policy_generation, SGLangGeneration):
+        weight_transfer_mode = generation_config["sglang_cfg"][
+            "sglang_server_config"
+        ].get("weight_transfer_mode", "ipc" if colocated_inference else "broadcast")
+        expected = "ipc" if colocated_inference else "broadcast"
+        if weight_transfer_mode != expected:
+            raise ValueError(
+                f"sglang_server_config.weight_transfer_mode={weight_transfer_mode!r} "
+                f"is inconsistent with colocated.enabled={colocated_inference}: "
+                f"expected {expected!r}."
+            )
+
+    # prepare refit info
     state_dict_info = policy.prepare_refit_info()
     if policy_generation is not None:
         policy_generation.prepare_refit_info(state_dict_info)
@@ -2018,6 +2048,44 @@ def _clip_grpo_advantages(
     return advantages
 
 
+def _refit_sglang_dispatch(
+    *,
+    policy: ColocatablePolicyInterface,
+    policy_generation: SGLangGeneration,
+    buffer_size_bytes: int,
+    mode: str,
+) -> bool:
+    """Route an SGLang refit to the backend-specific helper.
+
+    Backend-specific lifecycle (lock + pause/flush + send + post_process +
+    continue) lives in the corresponding worker module:
+
+    - ``megatron_policy_worker.refit_sglang_{colocated,distributed}``
+    - ``dtensor_policy_worker_v2.refit_sglang_{colocated,distributed}``
+
+    so this function only picks the right module by trainer backend and
+    transfer mode.
+    """
+    use_megatron = bool(policy.cfg.get("megatron_cfg", {}).get("enabled", False))
+    if use_megatron:
+        from nemo_rl.models.policy.workers import megatron_policy_worker as _backend
+    else:
+        from nemo_rl.models.policy.workers import dtensor_policy_worker_v2 as _backend
+
+    if mode == "ipc":
+        helper = _backend.refit_sglang_colocated
+    elif mode == "broadcast":
+        helper = _backend.refit_sglang_distributed
+    else:
+        raise ValueError(f"unknown SGLang weight_transfer_mode: {mode!r}")
+
+    return helper(
+        policy=policy,
+        policy_generation=policy_generation,
+        buffer_size_bytes=buffer_size_bytes,
+    )
+
+
 def refit_policy_generation(
     policy: ColocatablePolicyInterface,
     policy_generation: GenerationInterface,
@@ -2066,8 +2134,9 @@ def refit_policy_generation(
     with timer_context:
         # update weights
         update_success = False
-        if colocated_inference:
-            # get model param keys, which is grouped by size
+        # Bucket size for streamed refits: every colocated path and the SGLang
+        # broadcast dispatch group parameters into buffers of this size.
+        if colocated_inference or isinstance(policy_generation, SGLangGeneration):
             if _refit_buffer_size_gb is not None:
                 buffer_size_bytes = int(_refit_buffer_size_gb * (1024**3))
             else:
@@ -2078,15 +2147,14 @@ def refit_policy_generation(
                     policy.get_free_memory_bytes() * float(memory_ratio)
                 )
 
+        if colocated_inference:
             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(),
+                update_success = _refit_sglang_dispatch(
+                    policy=policy,
+                    policy_generation=policy_generation,
                     buffer_size_bytes=buffer_size_bytes,
+                    mode="ipc",
                 )
-                # Wait for all workers to complete
-                ray.get(futures_train)
-                update_success = True
             else:
                 # Original ZMQ IPC path for vLLM
                 futures_train = policy.stream_weights_via_ipc_zmq(
@@ -2098,23 +2166,27 @@ def refit_policy_generation(
                 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.
+            # update weights through nccl (vLLM), megatron reshard, or the
+            # SGLang broadcast dispatch
             if isinstance(policy_generation, SGLangGeneration):
-                raise NotImplementedError(
-                    "SGLang haven't implemented non-colocated inference mode. "
+                update_success = _refit_sglang_dispatch(
+                    policy=policy,
+                    policy_generation=policy_generation,
+                    buffer_size_bytes=buffer_size_bytes,
+                    mode="broadcast",
                 )
-            if isinstance(policy_generation, MegatronGeneration):
-                futures_train = policy.swap_weights_via_reshard(is_source=True)
             else:
-                futures_train = policy.broadcast_weights_for_collective(
-                    kv_scales=kv_scales
-                )
-            futures_inference = policy_generation.update_weights_from_collective()
-            # 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)
+                if isinstance(policy_generation, MegatronGeneration):
+                    futures_train = policy.swap_weights_via_reshard(is_source=True)
+                else:
+                    futures_train = policy.broadcast_weights_for_collective(
+                        kv_scales=kv_scales
+                    )
+                futures_inference = policy_generation.update_weights_from_collective()
+                # 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)
 
         # check if update is successful
         if not update_success:

nemo_rl/distributed/virtual_cluster.py

@@ -83,9 +83,14 @@ class PY_EXECUTABLES:
 # the full layout including Ray's own GCS / worker gRPC ports.
 #
 #   1400-1999    Master address / TCPStore       (cluster.master_port_range_low/high)
-#   3000-4999    NeMo RL generation HTTP servers (policy.generation.port_range_low/high)
+#   3000-4999    NeMo RL generation HTTP servers + SGLang engine NCCL/dist_init
+#                                                 (policy.generation.port_range_low/high)
 #   5000-5999    NeMo Gym HTTP servers           (env.nemo_gym.port_range_low/high)
-#   7000-8999    vLLM / SGLang engine rendezvous (VLLM_PORT env var / SGLang base_port)
+#   7000-8999    vLLM engine rendezvous          (VLLM_PORT env var, 100-port spacing)
+#   8600-8799    SGLang router                   (DEFAULT_SGLANG_ROUTER_PORT_RANGE_*, hard-coded;
+#                                                 carved out of the vLLM band — only one rollout
+#                                                 backend runs at a time)
+#   8800-8999    SGLang Prometheus metrics       (DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_*, hard-coded)
 DEFAULT_GENERATION_PORT_RANGE_LOW = 3000
 DEFAULT_GENERATION_PORT_RANGE_HIGH = 4999
 DEFAULT_GYM_PORT_RANGE_LOW = 5000
@@ -95,6 +100,14 @@ class PY_EXECUTABLES:
 # 7000 + 8*100 = 7800, still below the 9000 ephemeral floor.
 DEFAULT_VLLM_PORT_RANGE_LOW = 7000
 DEFAULT_VLLM_PORTS_PER_ENGINE = 100
+# SGLang control-plane ports, carved out of the top of the vLLM rendezvous band —
+# safe because only one rollout backend runs at a time, and a vLLM run only
+# climbs past 8600 with >=16 engines on a single node.  Both bands also steer
+# clear of the Ray dashboard carve-out at 8265 (see ray.sub).
+DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW = 8600
+DEFAULT_SGLANG_ROUTER_PORT_RANGE_HIGH = 8799
+DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_LOW = 8800
+DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_HIGH = 8999
 # Master address / TCPStore range, tucked below the Ray worker-gRPC band (2000+).
 DEFAULT_MASTER_PORT_RANGE_LOW = 1400
 DEFAULT_MASTER_PORT_RANGE_HIGH = 1999
@@ -183,6 +196,40 @@ def _get_free_port_local(
     return port
 
 
+def _get_free_consecutive_ports_local(
+    port_range_low: int,
+    port_range_high: int,
+    consecutive: int = 1,
+    start_port: Optional[int] = None,
+) -> int:
+    """Find ``consecutive`` contiguous bindable ports and return the base.
+
+    Scans upward from *start_port* within [port_range_low, port_range_high).
+    *start_port* lets a caller thread a per-node cursor so successive blocks do
+    not overlap. Raises ``RuntimeError`` if no such block exists in the range.
+    """
+    assert consecutive >= 1, f"consecutive must be >= 1, got {consecutive}"
+    base = port_range_low if start_port is None else max(start_port, port_range_low)
+    while base + consecutive - 1 < port_range_high:
+        socks: list[socket.socket] = []
+        try:
+            for offset in range(consecutive):
+                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                s.bind(("", base + offset))
+                s.listen(1)
+                socks.append(s)
+            return base
+        except OSError:
+            base += 1
+        finally:
+            for s in socks:
+                s.close()
+    raise RuntimeError(
+        f"Could not find {consecutive} consecutive free ports in "
+        f"[{port_range_low}, {port_range_high})."
+    )
+
+
 def init_ray(log_dir: Optional[str] = None) -> None:
     """Initialise Ray.
 

nemo_rl/models/automodel/setup.py

@@ -261,22 +261,28 @@ def validate_and_prepare_config(
     # Set basic configuration
     is_vlm = processor is not None
     is_generation_colocated = None
+    rollout_backend = None
     sampling_params = None
     if "generation" in config and config["generation"] is not None:
         generation_cfg = config["generation"]
         # set generation colocated
         is_generation_colocated = generation_cfg["colocated"]["enabled"]
+        rollout_backend = generation_cfg.get("backend")
         # set sampling params
         sampling_params = TrainingSamplingParams(
             top_k=generation_cfg["top_k"],
             top_p=generation_cfg["top_p"],
             temperature=generation_cfg["temperature"],
         )
 
-    # Explicitly set NCCL_CUMEM_ENABLE to 1 to avoid the P2P initialization error for PyNCCLCommunicator.
-    # See https://github.com/NVIDIA-NeMo/RL/issues/564 for more details.
-    if not is_generation_colocated:
-        os.environ["NCCL_CUMEM_ENABLE"] = "1"
+    # SGLang's scheduler subprocess defaults to NCCL_CUMEM_ENABLE=0, and the
+    # trainer / engine must agree on the transport selection.
+    if rollout_backend == "sglang":
+        os.environ["NCCL_CUMEM_ENABLE"] = "0"
+    # Explicitly set NCCL_CUMEM_ENABLE to 1 to avoid the P2P initialization error
+    # for PyNCCLCommunicator (see https://github.com/NVIDIA-NeMo/RL/issues/564).
+    elif not is_generation_colocated:
+        os.environ.setdefault("NCCL_CUMEM_ENABLE", "1")
 
     # Disable dynamo autotune_local_cache to avoid crash when there's already a cache
     # with different order of node_bundles

nemo_rl/models/generation/interfaces.py

@@ -252,6 +252,14 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool:
     def finish_generation(self, *args: Any, **kwargs: Any) -> bool:
         pass
 
+    def pause_generation(self) -> None:
+        """Pause in-flight generation on the backend."""
+        raise NotImplementedError
+
+    def continue_generation(self) -> None:
+        """Resume previously paused generation on the backend."""
+        raise NotImplementedError
+
     @property
     def requires_kv_scale_sync(self) -> bool:
         """Whether the generation backend requires KV cache scales synchronization."""

nemo_rl/models/generation/sglang/config.py

@@ -17,6 +17,25 @@
 from nemo_rl.models.generation.interfaces import GenerationConfig
 
 
+class SglangQuantizationConfig(TypedDict, total=False):
+    """SGLang weight-precision config.
+
+    ``scheme="bf16"`` (or omitting the block) means BF16 rollout/refit. Set
+    ``scheme="mxfp8"`` to boot SGLang from an MXFP8 HF checkpoint and to send
+    MXFP8 HF tensors during online refit.
+    """
+
+    scheme: str  # "bf16" | "mxfp8"
+    weight_block_size: list[int]
+    scale_fmt: str
+    modules_to_not_convert: list[str]
+    extra_high_precision_layers_hf: list[str]
+    num_layers_at_start_in_bf16: int
+    num_layers_at_end_in_bf16: int
+    converted_model_path: str
+    cache_root: str
+
+
 class SGLangServerConfig(TypedDict):
     # When True, sets SGLang `enable_memory_saver=True` so weights/KV can be released
     # during training and re-acquired before generation.
@@ -31,6 +50,10 @@ class SGLangServerConfig(TypedDict):
     pause_generation_mode: str
     # Total number of GPUs allocated to inference across all engines.
     num_gpus: NotRequired[int]
+    # "ipc" -> CUDA-IPC to the colocated SGLang HTTP server (default for
+    # colocated inference). "broadcast" -> NCCL broadcast over a shared
+    # weight-update group (used when SGLang engines run on disaggregate GPUs).
+    weight_transfer_mode: NotRequired[str]
     # GPUs per SGLang engine
     # num_gpus_per_engine = tp_size * pp_size; set ep, dp-attn are not orthgonal to those
     # nodes_per_engine: max(1, num_gpus_per_engine // num_gpus_per_node)
@@ -67,6 +90,15 @@ class SglangSpecificArgs(TypedDict):
     # Nested server/router configs. Kept under ``sglang_cfg`` so YAML and call
     # sites have a single sglang namespace instead of three sibling fields.
     sglang_server_config: SGLangServerConfig
+
+    # Fault tolerance (RolloutHealthMonitor). Off by default; when enabled, a
+    # daemon thread health-checks each engine and restarts hung/dead actors.
+    use_fault_tolerance: NotRequired[bool]
+    rollout_health_check_interval: NotRequired[int]
+    rollout_health_check_timeout: NotRequired[int]
+    rollout_health_check_first_wait: NotRequired[int]
+    # Weight precision and (when scheme=mxfp8) offline-conversion knobs.
+    quantization: NotRequired[SglangQuantizationConfig]
     sglang_router_config: SGLangRouterConfig
 
     # Path to model weights (local folder or HF repo id).