diff --git a/examples/configs/grpo_math_1B_sglang.yaml b/examples/configs/grpo_math_1B_sglang.yaml index 854baa30b6..66ee4f35bc 100644 --- a/examples/configs/grpo_math_1B_sglang.yaml +++ b/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: diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e170954cbc..22f2e44f47 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/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: diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 16e589d0ed..337e6b7d5f 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/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. diff --git a/nemo_rl/models/automodel/setup.py b/nemo_rl/models/automodel/setup.py index 5c91166ce4..c87165a9c1 100644 --- a/nemo_rl/models/automodel/setup.py +++ b/nemo_rl/models/automodel/setup.py @@ -261,11 +261,13 @@ 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"], @@ -273,10 +275,14 @@ def validate_and_prepare_config( 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 diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 025707c700..a61010a215 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/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.""" diff --git a/nemo_rl/models/generation/sglang/config.py b/nemo_rl/models/generation/sglang/config.py index c309f5fa75..d08f5385bd 100644 --- a/nemo_rl/models/generation/sglang/config.py +++ b/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). diff --git a/nemo_rl/models/generation/sglang/fault_tolerance.py b/nemo_rl/models/generation/sglang/fault_tolerance.py new file mode 100644 index 0000000000..0f5fcda4be --- /dev/null +++ b/nemo_rl/models/generation/sglang/fault_tolerance.py @@ -0,0 +1,200 @@ +# 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. + +import logging +import threading + +import ray + +from nemo_rl.models.generation.sglang.config import SGLangConfig + +logger = logging.getLogger(__name__) + + +class RolloutHealthMonitor: + """Health monitor for rollout engines. + + The monitor runs continuously once started, but can be paused/resumed + based on whether the engines are offloaded (cannot health check when offloaded). + + Lifecycle: + - start(): Start the monitor thread (called once during initialization) + - pause(): Pause health checking (called when offloading engines) + - resume(): Resume health checking (called when onloading engines) + - stop(): Stop the monitor thread completely (called during dispose) + """ + + def __init__(self, sglang_generation, sglang_cfg: SGLangConfig): + self._sglang_generation = sglang_generation + + self._thread = None + self._stop_event = None + self._pause_event = None # When set, health checking is paused + self._check_interval = sglang_cfg["sglang_cfg"]["rollout_health_check_interval"] + self._check_timeout = sglang_cfg["sglang_cfg"]["rollout_health_check_timeout"] + self._check_first_wait = sglang_cfg["sglang_cfg"][ + "rollout_health_check_first_wait" + ] + self._need_first_wait = True # Need to wait after each resume + self._is_checking_enabled = False # Track if health checking should be active + + def start(self) -> bool: + """Start the health monitor thread. Called once during initialization. + + Returns: + True if the monitor was started, False if there are no engines to monitor. + """ + if not self._sglang_generation.all_engines: + return False + + if self._thread is not None: + logger.warning("Health monitor thread is already running.") + return True + + logger.info("Starting RolloutHealthMonitor...") + self._stop_event = threading.Event() + self._pause_event = threading.Event() + self._pause_event.set() # Start in paused state until resume() is called + self._thread = threading.Thread( + target=self._health_monitor_loop, + name="RolloutHealthMonitor", + daemon=True, + ) + self._thread.start() + logger.info("RolloutHealthMonitor started (in paused state).") + return True + + def stop(self) -> None: + """Stop the health monitor thread completely. Called during dispose.""" + if not self._thread: + return + + logger.info("Stopping RolloutHealthMonitor...") + assert self._stop_event is not None + self._stop_event.set() + # Also clear pause to let the thread exit + if self._pause_event: + self._pause_event.clear() + timeout = self._check_timeout + self._check_interval + 5 + self._thread.join(timeout=timeout) + if self._thread.is_alive(): + logging.warning( + "Rollout health monitor thread did not terminate within %.1fs", timeout + ) + else: + logger.info("RolloutHealthMonitor stopped.") + + self._thread = None + self._stop_event = None + self._pause_event = None + self._is_checking_enabled = False + + def pause(self) -> None: + """Pause health checking. Called when engines are offloaded.""" + if self._pause_event is None: + return + logger.info("Pausing health monitor...") + self._pause_event.set() + self._is_checking_enabled = False + + def resume(self) -> None: + """Resume health checking. Called when engines are onloaded.""" + if self._pause_event is None: + return + logger.info("Resuming health monitor...") + self._need_first_wait = True # Need to wait after each resume + self._pause_event.clear() + self._is_checking_enabled = True + + def is_checking_enabled(self) -> bool: + """Return whether health checking is currently enabled (not paused).""" + return self._is_checking_enabled + + def _health_monitor_loop(self) -> None: + assert self._stop_event is not None + assert self._pause_event is not None + + while not self._stop_event.is_set(): + # Wait while paused + while self._pause_event.is_set() and not self._stop_event.is_set(): + self._stop_event.wait(timeout=0.5) + + if self._stop_event.is_set(): + break + + # Do first wait after each resume (for large MoE models to be ready) + if self._need_first_wait: + logger.info( + f"Health monitor doing first wait after resume: {self._check_first_wait}s" + ) + if self._stop_event.wait(self._check_first_wait): + logger.info("Health monitor stopped during first wait.") + break + if self._pause_event.is_set(): + # Got paused during first wait, skip this round and wait again next resume + logger.info( + "Health monitor paused during first wait, will wait again next resume." + ) + continue + self._need_first_wait = False + + # Run health checks + if not self._pause_event.is_set() and not self._stop_event.is_set(): + self._run_health_checks() + + # Wait for next check interval + if self._stop_event.wait(self._check_interval): + break + + def _run_health_checks(self) -> None: + for rollout_engine_id, engine in enumerate(self._sglang_generation.engines): + if self._stop_event is not None and self._stop_event.is_set(): + break + if self._pause_event is not None and self._pause_event.is_set(): + break + self._check_engine_health(rollout_engine_id, engine) + + def _check_engine_health(self, rollout_engine_id, engine) -> None: + if engine is None: + logger.info(f"Skipping health check for engine {rollout_engine_id} (None)") + return + + try: + ray.get(engine.health_generate.remote(timeout=self._check_timeout)) + except Exception as e: + logger.error( + f"Health check failed for rollout engine {rollout_engine_id} (ray timeout or error). Killing actor. Exception: {e}" + ) + self._kill_engine(rollout_engine_id=rollout_engine_id) + else: + logger.debug(f"Health check passed for rollout engine {rollout_engine_id}") + + def _kill_engine(self, rollout_engine_id: int): + logger.info(f"Killing server group {rollout_engine_id}...") + for i in range( + rollout_engine_id * self._sglang_generation.nodes_per_engine, + (rollout_engine_id + 1) * self._sglang_generation.nodes_per_engine, + ): + engine = self._sglang_generation.all_engines[i] + if engine: + logger.info(f"Shutting down and killing engine at index {i}") + try: + ray.get(engine.shutdown.remote()) + ray.kill(engine) + logger.info(f"Successfully killed engine at index {i}") + except Exception as e: + logger.warning(f"Fail to kill engine at index {i} (e: {e})") + else: + logger.info(f"Engine at index {i} is already None") + self._sglang_generation.all_engines[i] = None diff --git a/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py b/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py new file mode 100644 index 0000000000..b6f141c640 --- /dev/null +++ b/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py @@ -0,0 +1,225 @@ +# Copyright (c) 2025, 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. + +"""Shared MXFP8 tensor quantization rules for SGLang rollout weight updates. + +Offline conversion (``mxfp8_setup.py``) and online refit (the Megatron SGLang +weight iterator) must call into this module so they make the exact same +quantization decision for any given HF tensor name. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +SKIP_WEIGHT_SUBSTRINGS: tuple[str, ...] = ( + "layernorm", + "embed", + "router", + "mlp.gate.", + "norm", + "lm_head", + "eh_proj", + "weights_proj", +) +SOURCE_FP8_BLOCK_SIZE: list[int] = [128, 128] +TARGET_MXFP8_BLOCK_SIZE: list[int] = [1, 32] +SOURCE_FP8_SCALE_KEY_SUFFIX: str = ".weight_scale_inv" +SOURCE_FP8_DTYPES: tuple[torch.dtype, ...] = (torch.float8_e4m3fn,) + ( + (torch.float8_e4m3fnuz,) if hasattr(torch, "float8_e4m3fnuz") else () +) + +MXFP8_QUANTIZATION_CONFIG: dict[str, Any] = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "mxfp8", + "weight_block_size": TARGET_MXFP8_BLOCK_SIZE, + "scale_fmt": "ue8m0", +} + + +def strip_weight_suffix(weight_key: str) -> str: + if not weight_key.endswith(".weight"): + raise ValueError(f"Expected key ending with '.weight', got: {weight_key}") + return weight_key[: -len(".weight")] + + +def is_mxfp8_quantization_config(config: dict[str, Any] | None) -> bool: + if not isinstance(config, dict): + return False + return ( + config.get("quant_method") == "mxfp8" + and list(config.get("weight_block_size", [])) == TARGET_MXFP8_BLOCK_SIZE + and config.get("scale_fmt") == "ue8m0" + ) + + +def is_source_block_fp8_ue8m0_checkpoint(cfg: dict[str, Any]) -> bool: + qcfg = cfg.get("quantization_config", {}) if isinstance(cfg, dict) else {} + return ( + qcfg.get("quant_method") == "fp8" + and list(qcfg.get("weight_block_size", [])) == SOURCE_FP8_BLOCK_SIZE + and qcfg.get("scale_fmt") == "ue8m0" + ) + + +def is_bf16_source_checkpoint(cfg: dict[str, Any]) -> bool: + qcfg = cfg.get("quantization_config", {}) if isinstance(cfg, dict) else {} + if not isinstance(qcfg, dict) or not qcfg: + return True + return qcfg.get("quant_method") in (None, "", "bf16") + + +def should_quantize( + name: str, + weight: torch.Tensor, + *, + skip_weight_substrings: tuple[str, ...] = SKIP_WEIGHT_SUBSTRINGS, + allow_source_fp8: bool = False, +) -> bool: + allowed_dtypes: tuple[torch.dtype, ...] = ( + torch.float16, + torch.bfloat16, + torch.float32, + ) + if allow_source_fp8: + allowed_dtypes = allowed_dtypes + SOURCE_FP8_DTYPES + if not name.endswith(".weight"): + return False + if any(substr in name for substr in skip_weight_substrings): + return False + if weight.dtype not in allowed_dtypes: + return False + if weight.dim() < 2: + return False + if weight.shape[-1] % 32 != 0: + return False + return True + + +def quantize_mxfp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(qweight, scale)`` in the SGLang MXFP8 layout. + + Uses flashinfer's swizzle-free MXFP8 kernel (``flashinfer.mxfp8_quantize`` + with ``is_sf_swizzled_layout=False``). flashinfer is a hard requirement + here — both the SGLang and Megatron actor environments pin it via + ``pyproject.toml``'s global ``flashinfer-python==0.6.4`` constraint, so a + missing import means the env was built incorrectly. + """ + try: + from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize + except ImportError as e: + raise ImportError( + "flashinfer is required for MXFP8 weight quantization but is not " + "installed in the current actor environment. Install " + "`flashinfer-python==0.6.4` (and `flashinfer-cubin==0.6.4`); " + "in NeMo-RL this is normally provided by the `mcore` or `sglang` " + "extras (see pyproject.toml constraint-dependencies)." + ) from e + + weight = weight.contiguous() + k = weight.shape[-1] + if k % 32 != 0: + raise ValueError(f"Last dim {k} must be divisible by 32 for MXFP8.") + + weight_flat = weight.view(-1, k).contiguous() + qweight, scale = flashinfer_mxfp8_quantize(weight_flat, is_sf_swizzled_layout=False) + qweight = qweight.view_as(weight) + scale = scale.view(*weight.shape[:-1], k // 32).contiguous() + return qweight, scale + + +def source_fp8_to_mxfp8_scale_u8( + weight: torch.Tensor, source_scale_u8: torch.Tensor +) -> torch.Tensor: + n, k = weight.shape[-2], weight.shape[-1] + mxfp8_scale_u8 = source_scale_u8.repeat_interleave( + SOURCE_FP8_BLOCK_SIZE[0], dim=-2 + ).repeat_interleave(SOURCE_FP8_BLOCK_SIZE[1] // TARGET_MXFP8_BLOCK_SIZE[1], dim=-1) + return mxfp8_scale_u8[..., :n, : (k // TARGET_MXFP8_BLOCK_SIZE[1])].contiguous() + + +def build_dynamic_skip_substrings( + *, + quantization_config: dict[str, Any], + num_hidden_layers: int, +) -> tuple[str, ...]: + """Compute the dynamic skip substrings for one HF model. + + Combines the static ``SKIP_WEIGHT_SUBSTRINGS`` list with the user-provided + ``extra_high_precision_layers_hf`` / ``modules_to_not_convert`` lists from + the quantization config, plus per-layer prefixes for the ``head`` / ``tail`` + BF16-band layers. + """ + extra_high_precision_layers_hf = tuple( + quantization_config.get("extra_high_precision_layers_hf", ()) or () + ) + modules_to_not_convert = tuple( + quantization_config.get("modules_to_not_convert", ()) or () + ) + num_layers_at_start_in_bf16 = int( + quantization_config.get("num_layers_at_start_in_bf16", 0) or 0 + ) + num_layers_at_end_in_bf16 = int( + quantization_config.get("num_layers_at_end_in_bf16", 0) or 0 + ) + + head_end_idx = num_layers_at_start_in_bf16 + tail_start_idx = num_hidden_layers - num_layers_at_end_in_bf16 + dynamic_skip_layer_prefixes: set[str] = set() + dynamic_skip_layer_prefixes.update( + f"model.layers.{i}." for i in range(0, head_end_idx) + ) + dynamic_skip_layer_prefixes.update( + f"model.layers.{i}." for i in range(tail_start_idx, num_hidden_layers) + ) + return ( + *SKIP_WEIGHT_SUBSTRINGS, + *extra_high_precision_layers_hf, + *modules_to_not_convert, + *sorted(dynamic_skip_layer_prefixes), + ) + + +def maybe_quantize_hf_weight_mxfp8( + name: str, + tensor: torch.Tensor, + *, + quantization_config: dict[str, Any], + num_hidden_layers: int, +) -> list[tuple[str, torch.Tensor]]: + """Apply the HF-name MXFP8 policy to one finalized HF tensor. + + Returns a list of ``(name, tensor)`` pairs: + + - For an unquantized tensor: ``[(name, tensor)]``. + - For a quantized tensor: ``[(name, qweight), (name_scale_inv, scale)]``. + """ + skip_weight_substrings = build_dynamic_skip_substrings( + quantization_config=quantization_config, + num_hidden_layers=num_hidden_layers, + ) + if not should_quantize( + name, + tensor, + skip_weight_substrings=skip_weight_substrings, + allow_source_fp8=False, + ): + return [(name, tensor)] + + qweight, scale = quantize_mxfp8(tensor) + scale_name = strip_weight_suffix(name) + SOURCE_FP8_SCALE_KEY_SUFFIX + return [(name, qweight), (scale_name, scale)] diff --git a/nemo_rl/models/generation/sglang/mxfp8_setup.py b/nemo_rl/models/generation/sglang/mxfp8_setup.py new file mode 100644 index 0000000000..daad90f788 --- /dev/null +++ b/nemo_rl/models/generation/sglang/mxfp8_setup.py @@ -0,0 +1,451 @@ +# Copyright (c) 2025, 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. + +"""Offline HF -> MXFP8 conversion + startup helper for SGLang. + +Wraps NeMo-RL's quantization core so SGLang can boot from an MXFP8 HF +checkpoint and the online weight-update path can reuse the exact same +per-tensor decisions. +""" + +from __future__ import annotations + +import gc +import hashlib +import json +import logging +import os +import re +import shutil +from typing import Any + +import torch + +from nemo_rl.models.generation.sglang.mxfp8_quantization_core import ( + MXFP8_QUANTIZATION_CONFIG, + SKIP_WEIGHT_SUBSTRINGS, + SOURCE_FP8_BLOCK_SIZE, + SOURCE_FP8_DTYPES, + SOURCE_FP8_SCALE_KEY_SUFFIX, + is_bf16_source_checkpoint, + is_mxfp8_quantization_config, + is_source_block_fp8_ue8m0_checkpoint, + quantize_mxfp8, + should_quantize, + source_fp8_to_mxfp8_scale_u8, + strip_weight_suffix, +) + +logger = logging.getLogger(__name__) + +CONVERTER_VERSION: str = "1" + + +class _ConversionResult: + def __init__(self) -> None: + self.weight_map: dict[str, str] = {} + self.total_size: int = 0 + self.modules_to_not_convert: list[str] = [] + + def add_result( + self, + filename: str, + q_weights: dict[str, torch.Tensor], + module_names: list[str], + ) -> None: + for key, tensor in q_weights.items(): + self.weight_map[key] = filename + self.total_size += tensor.numel() * tensor.element_size() + self.modules_to_not_convert.extend(module_names) + + +def _load_source_scale_u8( + weights: dict[str, torch.Tensor], + weight_key: str, + weight: torch.Tensor, + *, + source_scale_index: dict[str, str], + input_path: str, + device: str, + current_filename: str, +) -> tuple[torch.Tensor, torch.Tensor | None, str]: + import safetensors + + scale_key = strip_weight_suffix(weight_key) + SOURCE_FP8_SCALE_KEY_SUFFIX + scale_file = source_scale_index[scale_key] + if scale_file == current_filename and scale_key in weights: + scale = weights[scale_key] + else: + with safetensors.safe_open( + os.path.join(input_path, scale_file), framework="pt", device=device + ) as f: + scale = f.get_tensor(scale_key) + + if scale.dtype == torch.uint8: + scale_u8: torch.Tensor | None = scale + else: + if scale.dtype != torch.float32: + raise ValueError( + f"Unexpected source FP8 scale dtype {scale.dtype} for {scale_key}" + ) + n, k = weight.shape[-2], weight.shape[-1] + n_tiles = (n + SOURCE_FP8_BLOCK_SIZE[0] - 1) // SOURCE_FP8_BLOCK_SIZE[0] + k_tiles = (k + SOURCE_FP8_BLOCK_SIZE[1] - 1) // SOURCE_FP8_BLOCK_SIZE[1] + scale_fp32 = scale[..., :n_tiles, :k_tiles].contiguous() + bits = scale_fp32.contiguous().view(torch.int32) + mantissa_all_zero = not torch.any((bits & 0x007FFFFF) != 0).item() + non_negative = not torch.any(bits < 0).item() + if mantissa_all_zero and non_negative: + scale_u8 = ((bits >> 23) & 0xFF).to(torch.uint8) + else: + scale_u8 = None + return scale_fp32, scale_u8, scale_key + + n, k = weight.shape[-2], weight.shape[-1] + n_tiles = (n + SOURCE_FP8_BLOCK_SIZE[0] - 1) // SOURCE_FP8_BLOCK_SIZE[0] + k_tiles = (k + SOURCE_FP8_BLOCK_SIZE[1] - 1) // SOURCE_FP8_BLOCK_SIZE[1] + scale_u8 = scale_u8[..., :n_tiles, :k_tiles].contiguous() + scale_fp32 = (scale_u8.to(torch.int32) << 23).view(torch.float32) + return scale_fp32, scale_u8, scale_key + + +def _process_file( + input_path: str, + output_path: str, + filename: str, + *, + result_collector: _ConversionResult, + device: str, + num_hidden_layers: int, + num_layers_at_start_in_bf16: int, + num_layers_at_end_in_bf16: int, + source_is_block_fp8_ue8m0: bool, + extra_high_precision_layers_hf: tuple[str, ...], + source_scale_index: dict[str, str], +) -> None: + import safetensors + import safetensors.torch + from sglang.srt.layers.quantization.fp8_utils import block_quant_dequant + + weights: dict[str, torch.Tensor] = {} + q_weights: dict[str, torch.Tensor] = {} + + with safetensors.safe_open( + os.path.join(input_path, filename), framework="pt", device=device + ) as f: + for key in f.keys(): + weights[key] = f.get_tensor(key) + + modules_to_not_convert: list[str] = [] + head_end_idx = num_layers_at_start_in_bf16 + tail_start_idx = num_hidden_layers - num_layers_at_end_in_bf16 + dynamic_skip_layer_prefixes: set[str] = set() + dynamic_skip_layer_prefixes.update( + f"model.layers.{i}." for i in range(0, head_end_idx) + ) + dynamic_skip_layer_prefixes.update( + f"model.layers.{i}." for i in range(tail_start_idx, num_hidden_layers) + ) + + if num_layers_at_end_in_bf16 > 0 or num_layers_at_start_in_bf16 > 0: + modules_to_not_convert.extend(sorted(dynamic_skip_layer_prefixes)) + + dynamic_skip_substrings = ( + *SKIP_WEIGHT_SUBSTRINGS, + *extra_high_precision_layers_hf, + *sorted(dynamic_skip_layer_prefixes), + ) + + for key, tensor in weights.items(): + if not key.endswith(".weight"): + continue + + should_quant = should_quantize( + key, + tensor, + skip_weight_substrings=dynamic_skip_substrings, + allow_source_fp8=source_is_block_fp8_ue8m0, + ) + + if should_quant: + if source_is_block_fp8_ue8m0 and tensor.dtype in SOURCE_FP8_DTYPES: + source_scale_fp32, source_scale_u8, scale_key = _load_source_scale_u8( + weights, + key, + tensor, + source_scale_index=source_scale_index, + input_path=input_path, + device=device, + current_filename=filename, + ) + if source_scale_u8 is not None: + qweight = tensor.contiguous() + scale = source_fp8_to_mxfp8_scale_u8(tensor, source_scale_u8) + else: + weight_fp32 = block_quant_dequant( + tensor, + source_scale_fp32, + SOURCE_FP8_BLOCK_SIZE, + torch.float32, + ).contiguous() + qweight, scale = quantize_mxfp8(weight_fp32) + q_weights[key] = qweight + q_weights[scale_key] = scale + else: + qweight, scale = quantize_mxfp8(tensor) + q_weights[key] = qweight + q_weights[strip_weight_suffix(key) + SOURCE_FP8_SCALE_KEY_SUFFIX] = ( + scale + ) + else: + if ".experts." not in key: + modules_to_not_convert.append(strip_weight_suffix(key)) + if source_is_block_fp8_ue8m0 and tensor.dtype in SOURCE_FP8_DTYPES: + source_scale_fp32, _, _ = _load_source_scale_u8( + weights, + key, + tensor, + source_scale_index=source_scale_index, + input_path=input_path, + device=device, + current_filename=filename, + ) + q_weights[key] = block_quant_dequant( + tensor, + source_scale_fp32, + SOURCE_FP8_BLOCK_SIZE, + torch.bfloat16, + ).contiguous() + else: + q_weights[key] = tensor + + for key, tensor in weights.items(): + if key.endswith(".weight"): + continue + if source_is_block_fp8_ue8m0 and key.endswith(SOURCE_FP8_SCALE_KEY_SUFFIX): + continue + q_weights[key] = tensor + + safetensors.torch.save_file( + q_weights, os.path.join(output_path, filename), metadata={"format": "pt"} + ) + result_collector.add_result(filename, q_weights, modules_to_not_convert) + + +def convert_mxfp8( + model_dir: str, + save_dir: str, + *, + device: str = "cuda", + num_layers_at_start_in_bf16: int = 0, + num_layers_at_end_in_bf16: int = 0, + extra_high_precision_layers_hf: tuple[str, ...] = (), +) -> None: + """Convert an HF safetensors checkpoint to MXFP8 with UE8M0 scales. + + Uses the shared quantization core in ``mxfp8_quantization_core``. + """ + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available, cannot run MXFP8 quantization.") + + input_path = os.path.abspath(model_dir) + output_path = os.path.abspath(save_dir) + os.makedirs(output_path, exist_ok=True) + config_path = os.path.join(input_path, "config.json") + with open(config_path) as f: + cfg = json.load(f) + num_hidden_layers = int(cfg["num_hidden_layers"]) + if is_source_block_fp8_ue8m0_checkpoint(cfg): + source_is_block_fp8_ue8m0 = True + elif is_bf16_source_checkpoint(cfg): + source_is_block_fp8_ue8m0 = False + else: + raise ValueError( + "Unsupported source quantization_config. " + "Only BF16/FP16/FP32 sources and " + "{quant_method=fp8, weight_block_size=[128, 128], scale_fmt=ue8m0} sources are supported." + ) + + for filename in os.listdir(input_path): + if not filename.endswith(".safetensors") and not os.path.isdir( + os.path.join(input_path, filename) + ): + shutil.copyfile( + os.path.join(input_path, filename), + os.path.join(output_path, filename), + ) + + index_path = os.path.join(input_path, "model.safetensors.index.json") + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + safetensors_files = sorted(set(weight_map.values())) + source_scale_index: dict[str, str] = {} + if source_is_block_fp8_ue8m0: + source_scale_index = { + key: filename + for key, filename in weight_map.items() + if key.endswith(SOURCE_FP8_SCALE_KEY_SUFFIX) + } + + result_collector = _ConversionResult() + for filename in safetensors_files: + logger.info(f"[mxfp8] Processing {filename}") + _process_file( + input_path, + output_path, + filename, + result_collector=result_collector, + device=device, + num_hidden_layers=num_hidden_layers, + num_layers_at_start_in_bf16=num_layers_at_start_in_bf16, + num_layers_at_end_in_bf16=num_layers_at_end_in_bf16, + source_is_block_fp8_ue8m0=source_is_block_fp8_ue8m0, + extra_high_precision_layers_hf=extra_high_precision_layers_hf, + source_scale_index=source_scale_index, + ) + gc.collect() + torch.cuda.empty_cache() + + quantization_config: dict[str, Any] = dict(MXFP8_QUANTIZATION_CONFIG) + if len(result_collector.modules_to_not_convert) > 0: + + def natural_key(s: str) -> list[Any]: + return [int(t) if t.isdigit() else t for t in re.findall(r"\d+|\D+", s)] + + quantization_config["modules_to_not_convert"] = sorted( + list(set(result_collector.modules_to_not_convert)), key=natural_key + ) + + cfg["quantization_config"] = quantization_config + with open(os.path.join(output_path, "config.json"), "w") as f: + json.dump(cfg, f, indent=2) + + index_dict = { + "weight_map": result_collector.weight_map, + "metadata": {"total_size": result_collector.total_size}, + } + with open(os.path.join(output_path, "model.safetensors.index.json"), "w") as f: + json.dump(index_dict, f, indent=2) + + gc.collect() + torch.cuda.empty_cache() + + +def _read_source_config(model_dir: str) -> dict[str, Any]: + config_path = os.path.join(model_dir, "config.json") + if not os.path.isfile(config_path): + return {} + with open(config_path) as f: + return json.load(f) + + +def _quantization_fingerprint(quantization_cfg: dict[str, Any]) -> str: + relevant_keys = ( + "extra_high_precision_layers_hf", + "modules_to_not_convert", + "num_layers_at_start_in_bf16", + "num_layers_at_end_in_bf16", + "weight_block_size", + "scale_fmt", + ) + payload = {k: quantization_cfg.get(k) for k in relevant_keys} + return hashlib.sha1( + json.dumps(payload, sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:12] + + +def _hash_qualified_save_dir( + *, model_dir: str, cache_root: str, quantization_cfg: dict[str, Any] +) -> str: + abs_model = os.path.abspath(model_dir) + src_cfg = _read_source_config(model_dir) + src_fingerprint = hashlib.sha1( + json.dumps(src_cfg, sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:12] + quant_fingerprint = _quantization_fingerprint(quantization_cfg) + payload = f"{abs_model}|{src_fingerprint}|{quant_fingerprint}|v{CONVERTER_VERSION}" + digest = hashlib.sha1(payload.encode("utf-8")).hexdigest()[:16] + base = os.path.basename(os.path.normpath(abs_model)) or "hf" + return os.path.join(os.path.abspath(cache_root), f"{base}-mxfp8-{digest}") + + +def is_existing_mxfp8_checkpoint(path: str) -> bool: + cfg = _read_source_config(path) + qcfg = cfg.get("quantization_config") if isinstance(cfg, dict) else None + return is_mxfp8_quantization_config(qcfg) + + +def ensure_mxfp8_checkpoint( + *, + model_path: str, + quantization_cfg: dict[str, Any], +) -> str: + """Return a path to an MXFP8-loadable HF checkpoint for SGLang. + + - If ``model_path`` is already an MXFP8 checkpoint, return it as-is. + - If ``quantization_cfg.converted_model_path`` is an MXFP8 checkpoint, + return it. + - Otherwise convert ``model_path`` into a hash-qualified subdirectory + under ``quantization_cfg.cache_root`` (or ``$NRL_MXFP8_CACHE`` / + ``~/.cache/nemo_rl/mxfp8`` if not set) and return that path. + + The hash includes absolute model path, source config fingerprint, + quantization config fingerprint and converter version, so different + sources / settings never collide. + """ + if is_existing_mxfp8_checkpoint(model_path): + return model_path + + converted = quantization_cfg.get("converted_model_path") + if converted and is_existing_mxfp8_checkpoint(converted): + return converted + + cache_root = ( + quantization_cfg.get("cache_root") + or os.environ.get("NRL_MXFP8_CACHE") + or os.path.join(os.path.expanduser("~"), ".cache", "nemo_rl", "mxfp8") + ) + save_dir = converted or _hash_qualified_save_dir( + model_dir=model_path, + cache_root=cache_root, + quantization_cfg=quantization_cfg, + ) + + if is_existing_mxfp8_checkpoint(save_dir): + return save_dir + + extra_high_precision_layers_hf = tuple( + quantization_cfg.get("extra_high_precision_layers_hf", ()) or () + ) + num_layers_at_start_in_bf16 = int( + quantization_cfg.get("num_layers_at_start_in_bf16", 0) or 0 + ) + num_layers_at_end_in_bf16 = int( + quantization_cfg.get("num_layers_at_end_in_bf16", 0) or 0 + ) + + logger.info( + f"[mxfp8] Converting {model_path} -> {save_dir} " + f"(start_bf16={num_layers_at_start_in_bf16}, " + f"end_bf16={num_layers_at_end_in_bf16}, " + f"extra_hp={extra_high_precision_layers_hf})" + ) + convert_mxfp8( + model_dir=model_path, + save_dir=save_dir, + num_layers_at_start_in_bf16=num_layers_at_start_in_bf16, + num_layers_at_end_in_bf16=num_layers_at_end_in_bf16, + extra_high_precision_layers_hf=extra_high_precision_layers_hf, + ) + return save_dir diff --git a/nemo_rl/models/generation/sglang/sglang_generation.py b/nemo_rl/models/generation/sglang/sglang_generation.py index bff4b4832c..454985681a 100644 --- a/nemo_rl/models/generation/sglang/sglang_generation.py +++ b/nemo_rl/models/generation/sglang/sglang_generation.py @@ -15,7 +15,7 @@ import asyncio import logging import os -from typing import Any, AsyncGenerator +from typing import Any, AsyncGenerator, Optional import ray import torch @@ -23,6 +23,8 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.virtual_cluster import ( + DEFAULT_GENERATION_PORT_RANGE_HIGH, + DEFAULT_GENERATION_PORT_RANGE_LOW, RayVirtualCluster, get_reordered_bundle, ) @@ -34,6 +36,7 @@ verify_right_padding, ) from nemo_rl.models.generation.sglang.config import SGLangConfig +from nemo_rl.models.generation.sglang.fault_tolerance import RolloutHealthMonitor from nemo_rl.models.generation.sglang.sglang_router import _start_router from nemo_rl.models.generation.sglang.sglang_worker import SGLangGenerationWorker from nemo_rl.models.generation.sglang.utils.async_utils import AsyncLoopThread @@ -43,6 +46,7 @@ ) from nemo_rl.models.generation.sglang.utils.ray_utils import ( NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, + Lock, ) from nemo_rl.utils.nsys import wrap_with_nvtx_name from nemo_rl.utils.venvs import make_actor_runtime_env @@ -101,6 +105,13 @@ def __init__( self.needs_offload: bool = sglang_server_cfg["needs_offload"] self.model_path: str | None = sglang_cfg["sglang_cfg"]["model_path"] + # --- Weight-refit / fault-tolerance state ------------------------ + # Number of engines created by the most recent ``_start_engines`` + # call that the refit dispatch has not connected yet. + self.num_new_engines: int = 0 + self.pause_generation_mode: str = sglang_server_cfg["pause_generation_mode"] + self._health_monitor: RolloutHealthMonitor | None = None + # --- Router bootstrap -------------------------------------------- # Resolved router endpoint is held only on the instance; we don't # mutate the caller's config dict. Workers receive these as explicit @@ -120,6 +131,14 @@ def __init__( if init_handles: ray.get(init_handles) + # Serializes weight refits against engine recovery across processes. + self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() + + if sglang_cfg["sglang_cfg"].get("use_fault_tolerance"): + monitor = RolloutHealthMonitor(self, sglang_cfg) + monitor.start() + self._health_monitor = monitor + # ------------------------------------------------------------------ # Engine topology properties (formerly ``ServerGroup``) # ------------------------------------------------------------------ @@ -163,7 +182,6 @@ def _start_engines( """ if port_cursors is None: port_cursors = {} - num_gpu_per_engine = min(self.num_gpus_per_engine, self.num_gpus_per_node) pg = self.pg reordered_bundle_indices = self.pg_reordered_bundle_indices @@ -211,6 +229,10 @@ def _start_engines( # Explicitly pass CUDA_VISIBLE_DEVICES through to the engine actor so # all engines see the same global value (Ray would otherwise remap it # because we set the NOSET_* flags above). + # Trainer and engine must agree on the NCCL transport; sglang's + # scheduler subprocess defaults to NCCL_CUMEM_ENABLE=0. + env_vars["NCCL_CUMEM_ENABLE"] = "0" + global_cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None) if global_cvd: env_vars["CUDA_VISIBLE_DEVICES"] = global_cvd @@ -246,18 +268,28 @@ def _start_engines( local_all_engines.append((global_rank, engine)) self.all_engines[i] = engine - if len(local_all_engines) == 0: + self.num_new_engines = len(local_all_engines) + + if self.num_new_engines == 0: return [], port_cursors - # SGLang engine ports live in the below-ephemeral-floor engine band - # (7000-8999), shared with vLLM; see virtual_cluster.py port layout. - base_port = max(port_cursors.values()) if port_cursors else 7000 + # SGLang engine server/NCCL/dist_init ports come from the reserved + # generation band (3000-4999 by default), below the ephemeral floor; + # see the port layout in virtual_cluster.py. + gen_port_low = self.sglang_cfg.get("port_range_low") + if gen_port_low is None: + gen_port_low = DEFAULT_GENERATION_PORT_RANGE_LOW + gen_port_high = self.sglang_cfg.get("port_range_high") + if gen_port_high is None: + gen_port_high = DEFAULT_GENERATION_PORT_RANGE_HIGH addr_and_ports, port_cursors = _allocate_rollout_engine_addr_and_ports_normal( gpus_per_node=self.num_gpus_per_node, sglang_cfg=self.sglang_cfg, local_all_engines=local_all_engines, rank_offset=self.rank_offset, - base_port=base_port, + port_range_low=gen_port_low, + port_range_high=gen_port_high, + node_port_cursor=port_cursors, ) init_handles = [ @@ -283,7 +315,132 @@ def check_weights(self, action: str): ] ) + def _recover(self) -> None: + """Recover dead engines, overlapping init.""" + dead_indices = [ + i for i, engine in enumerate(self.all_engines) if engine is None + ] + + port_cursors: dict[int, int] = {} + handles, _ = self._start_engines(port_cursors) + if handles: + ray.get(handles) + + assert self.num_new_engines == len(dead_indices), ( + "num_new_engines does not match dead_indices length" + ) + + if self.needs_offload and dead_indices: + new_engines = [self.all_engines[i] for i in dead_indices] + ray.get( + [ + engine.release_memory_occupation.remote(tags=["weights"]) + for engine in new_engines + ] + ) + ray.get( + [ + engine.release_memory_occupation.remote(tags=["kv_cache"]) + for engine in new_engines + ] + ) + ray.get( + [ + engine.resume_memory_occupation.remote(tags=["weights"]) + for engine in new_engines + ] + ) + + def get_updatable_engines_and_lock(self): + """Return engines eligible for weight updates.""" + return ( + self.engines, + self.rollout_engine_lock, + self.num_new_engines, + self.engine_gpu_counts, + self.engine_gpu_offsets, + ) + + def recover_updatable_engines(self): + """Restart any dead rollout engines and update ``num_new_engines``.""" + self.health_monitoring_pause() + + self._recover() + + return ( + self.engines, + self.rollout_engine_lock, + self.num_new_engines, + self.engine_gpu_counts, + self.engine_gpu_offsets, + ) + + def clear_updatable_num_new_engines(self): + # When fault tolerance is not enabled, num_new_engines must be cleared + # manually after the refit dispatch connects the new engines. + self.num_new_engines = 0 + + def pause_generation(self, mode: Optional[str] = None) -> None: + """Pause generation on every node-0 engine. + + Args: + mode: Pause mode override. When ``None`` (default), the mode + configured in ``sglang_server_config.pause_generation_mode`` + is used. Callers (e.g. the SGLang refit dispatch helpers) + pass an explicit mode when they also need to gate follow-up + steps such as ``invalidate_kv_cache`` on the same value. + """ + engines = [e for e in self.engines if e is not None] + if not engines: + return + if mode is None: + mode = self.pause_generation_mode + ray.get([e.pause_generation.remote(mode=mode) for e in engines]) + + def continue_generation(self) -> None: + """Resume generation on every node-0 engine.""" + engines = [e for e in self.engines if e is not None] + if not engines: + return + ray.get([e.continue_generation.remote() for e in engines]) + + def post_process_weights( + self, + *, + restore_weights_before_load: bool = False, + post_process_quantization: bool = True, + ) -> None: + """Run SGLang's ``/post_process_weights`` RPC on every node-0 engine. + + Called by the refit dispatch helpers after a colocate IPC or + distributed broadcast refit so SGLang finalizes its weight tables + (e.g. materializes quantized scales, swaps in the fresh buffer). + """ + engines = [e for e in self.engines if e is not None] + if not engines: + return + ray.get( + [ + e.post_process_weights.remote( + restore_weights_before_load=restore_weights_before_load, + post_process_quantization=post_process_quantization, + ) + for e in engines + ] + ) + + def health_monitoring_pause(self) -> None: + if self._health_monitor: + self._health_monitor.pause() + + def health_monitoring_resume(self) -> None: + if self._health_monitor: + self._health_monitor.resume() + def shutdown(self) -> bool: + if self._health_monitor: + self._health_monitor.stop() + ok = True engines = [e for e in self.all_engines if e is not None] if engines: diff --git a/nemo_rl/models/generation/sglang/sglang_router.py b/nemo_rl/models/generation/sglang/sglang_router.py index 720aebd894..d523e99d68 100644 --- a/nemo_rl/models/generation/sglang/sglang_router.py +++ b/nemo_rl/models/generation/sglang/sglang_router.py @@ -14,17 +14,20 @@ import logging import multiprocessing -import random import time import ray +from nemo_rl.distributed.virtual_cluster import ( + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_HIGH, + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_LOW, + DEFAULT_SGLANG_ROUTER_PORT_RANGE_HIGH, + DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW, + _get_free_port_local, +) from nemo_rl.models.generation.sglang.config import SGLangRouterConfig from nemo_rl.models.generation.sglang.utils.ip_port_utils import _wrap_ipv6 -from nemo_rl.models.generation.sglang.utils.ray_utils import ( - find_available_port, - get_host_info, -) +from nemo_rl.models.generation.sglang.utils.ray_utils import get_host_info from nemo_rl.utils.venvs import make_actor_runtime_env logger = logging.getLogger(__name__) @@ -51,14 +54,20 @@ def init(self, router_cfg: SGLangRouterConfig) -> tuple[str, int]: router_ip = _wrap_ipv6(get_host_info()[1]) router_port = router_cfg.get("sglang_router_port") if router_port is None: - router_port = find_available_port(random.randint(3000, 4000)) + router_port = _get_free_port_local( + DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW, + DEFAULT_SGLANG_ROUTER_PORT_RANGE_HIGH, + ) router_args = RouterArgs() router_args.host = router_ip router_args.port = router_port if router_cfg.get("router_policy") is not None: router_args.router_policy = router_cfg["router_policy"] - router_args.prometheus_port = find_available_port(random.randint(4000, 5000)) + router_args.prometheus_port = _get_free_port_local( + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_LOW, + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_HIGH, + ) router_args.log_level = "warn" request_timeout_secs = router_cfg.get("sglang_router_request_timeout_secs") if request_timeout_secs is not None: diff --git a/nemo_rl/models/generation/sglang/sglang_worker.py b/nemo_rl/models/generation/sglang/sglang_worker.py index c9af0884eb..80edf3d5ba 100644 --- a/nemo_rl/models/generation/sglang/sglang_worker.py +++ b/nemo_rl/models/generation/sglang/sglang_worker.py @@ -21,12 +21,14 @@ import ray import requests +from nemo_rl.distributed.virtual_cluster import ( + DEFAULT_GENERATION_PORT_RANGE_HIGH, + DEFAULT_GENERATION_PORT_RANGE_LOW, + _get_free_consecutive_ports_local, +) from nemo_rl.models.generation.sglang.utils.ip_port_utils import _format_v6_uri from nemo_rl.models.generation.sglang.utils.patches import _apply_sglang_compat_patches -from nemo_rl.models.generation.sglang.utils.ray_utils import ( - get_current_node_ip, - get_free_port, -) +from nemo_rl.models.generation.sglang.utils.ray_utils import get_current_node_ip logger = logging.getLogger(__name__) @@ -162,11 +164,23 @@ def _make_request(self, endpoint: str, payload: dict | None = None): return response.json() @staticmethod - def _get_current_node_ip_and_free_port(start_port=7000, consecutive=1): - return get_current_node_ip(), get_free_port( - start_port=start_port, consecutive=consecutive + def _get_current_free_port( + port_range_low=DEFAULT_GENERATION_PORT_RANGE_LOW, + port_range_high=DEFAULT_GENERATION_PORT_RANGE_HIGH, + consecutive=1, + start_port=None, + ): + return _get_free_consecutive_ports_local( + port_range_low=port_range_low, + port_range_high=port_range_high, + consecutive=consecutive, + start_port=start_port, ) + @staticmethod + def _get_current_node_ip(): + return get_current_node_ip() + def health_generate(self, timeout: float = 5.0) -> bool: """Run /health_generate on the underlying SGLang HTTP server. @@ -290,6 +304,108 @@ def resume_memory_occupation(self, tags: list[str] | None = None): def check_weights(self, action: str): return self._make_request("weights_checker", {"action": action}) + def get_weight_version(self): + if self.node_rank != 0: + return + # newer sglang moved /get_weight_version into /model_info + for endpoint in ("/model_info", "/get_weight_version"): + response = requests.get(f"{self.server_base_url}{endpoint}") + if response.status_code == 200: + return response.json()["weight_version"] + response.raise_for_status() + + def init_weights_update_group( + self, master_address, master_port, rank_offset, world_size, group_name, backend + ): + return self._make_request( + "init_weights_update_group", + { + "master_address": master_address, + "master_port": master_port, + "rank_offset": rank_offset, + "world_size": world_size, + "group_name": group_name, + "backend": backend, + }, + ) + + def destroy_weights_update_group(self, group_name): + try: + return self._make_request( + "destroy_weights_update_group", + { + "group_name": group_name, + }, + ) + except requests.exceptions.RequestException: + # catch the case where the engine is just created and does not have the group. + pass + + def update_weights_from_distributed( + self, + names, + dtypes, + shapes, + group_name, + flush_cache=False, + weight_version: str | None = None, + ): + payload = { + "names": names, + "dtypes": [str(dtype).replace("torch.", "") for dtype in dtypes], + "shapes": shapes, + "group_name": group_name, + "flush_cache": flush_cache, + } + if weight_version is not None: + payload["weight_version"] = weight_version + return self._make_request( + "update_weights_from_distributed", + payload, + ) + + def pause_generation(self, mode: str = "retract"): + response = requests.post( + f"{self.server_base_url}/pause_generation", + json={"mode": mode}, + ) + response.raise_for_status() + return response + + def continue_generation(self): + response = requests.post(f"{self.server_base_url}/continue_generation", json={}) + response.raise_for_status() + return response + + def post_process_weights( + self, + restore_weights_before_load: bool = False, + post_process_quantization: bool = False, + ): + """Finalize engine-side weights after a distributed/IPC refit. + + The HTTP server only posts metadata; the real weights were already + copied on-GPU by the preceding update path. + """ + return self._make_request( + "post_process_weights", + { + "restore_weights_before_load": restore_weights_before_load, + "post_process_quantization": post_process_quantization, + }, + ) + + def _simulate_crash(self): + """Test-only: tear the engine down to simulate a crash. + + Underscore-prefixed to signal this is **not** part of the public + worker API; production code should never call it. + """ + logger.info( + f"Simulating crash on engine {self.server_host}:{self.server_port}..." + ) + self.shutdown() + def start_profile( self, # The output directory diff --git a/nemo_rl/models/generation/sglang/utils/ip_port_utils.py b/nemo_rl/models/generation/sglang/utils/ip_port_utils.py index 9c3252dccf..8147628faf 100644 --- a/nemo_rl/models/generation/sglang/utils/ip_port_utils.py +++ b/nemo_rl/models/generation/sglang/utils/ip_port_utils.py @@ -17,6 +17,10 @@ import ray +from nemo_rl.distributed.virtual_cluster import ( + DEFAULT_GENERATION_PORT_RANGE_HIGH, + DEFAULT_GENERATION_PORT_RANGE_LOW, +) from nemo_rl.models.generation.sglang.utils.ray_utils import get_host_info logger = logging.getLogger(__name__) @@ -48,14 +52,18 @@ def _allocate_rollout_engine_addr_and_ports_normal( sglang_cfg, local_all_engines, rank_offset=0, - base_port=7000, + port_range_low: int = DEFAULT_GENERATION_PORT_RANGE_LOW, + port_range_high: int = DEFAULT_GENERATION_PORT_RANGE_HIGH, + node_port_cursor: dict[int, int] | None = None, ): # get ports # there are 4 ports we need to allocate # 1. server port # 2. nccl port # 3. dist_init_addr port - # 4. other ports for dp_attention, which is of size 4 + dp_size + # 4. other ports for dp_attention, which is of size 30 + dp_size + if node_port_cursor is None: + node_port_cursor = {} sglang_dp_size = sglang_cfg["sglang_cfg"]["dp_size"] num_gpus_per_engine = sglang_cfg["sglang_cfg"]["sglang_server_config"][ @@ -67,10 +75,6 @@ def _allocate_rollout_engine_addr_and_ports_normal( num_engines_per_node = max(1, num_gpus_per_node // _gpus_per_engine) addr_and_ports: dict[int, dict] = {} - # Track per-node port cursors so that different server groups (called - # sequentially) never race for the same ports on a given node. - node_port_cursor: dict[int, int] = {} - visited_nodes = set() for rank, engine in local_all_engines: local_rank = rank - rank_offset @@ -85,18 +89,19 @@ def _allocate_rollout_engine_addr_and_ports_normal( ) def get_addr_and_ports(engine, node_idx): - # Keep engine ports below the OS ephemeral floor (9000 on some GB200 nodes, - # 32768 on stock Linux) to avoid TOCTOU collisions. SGLang shares - # the vLLM engine rendezvous band (7000-8999); see the port layout - # in ray.sub / nemo_rl/distributed/virtual_cluster.py. - start_port = node_port_cursor.get(node_idx, base_port) + # Allocate from the reserved generation band (below the ephemeral + # floor; see virtual_cluster.py), advancing a per-node cursor so + # blocks never overlap on a given node. + start_port = node_port_cursor.get(node_idx, port_range_low) def port(consecutive=1): nonlocal start_port - _, port = ray.get( - engine._get_current_node_ip_and_free_port.remote( - start_port=start_port, + port = ray.get( + engine._get_current_free_port.remote( + port_range_low=port_range_low, + port_range_high=port_range_high, consecutive=consecutive, + start_port=start_port, ) ) start_port = port + consecutive @@ -104,7 +109,7 @@ def port(consecutive=1): return port def addr(): - addr, _ = ray.get(engine._get_current_node_ip_and_free_port.remote()) + addr = ray.get(engine._get_current_node_ip.remote()) if addr is None: addr = get_host_info()[1] return addr diff --git a/nemo_rl/models/generation/sglang/utils/patches.py b/nemo_rl/models/generation/sglang/utils/patches.py index a565084367..cab29239b2 100644 --- a/nemo_rl/models/generation/sglang/utils/patches.py +++ b/nemo_rl/models/generation/sglang/utils/patches.py @@ -207,70 +207,11 @@ def _patch_sglang_custom_all_reduce_v2_tms_cudagraph() -> None: """Backport sglang#27948 for colocated TMS CUDA graph capture. With ``SGLANG_MEMORY_SAVER_CUDA_GRAPH=true``, custom all-reduce v2 must - avoid registering captured IPC addresses. TMS will otherwise replace those - addresses during capture and custom_all_reduce.cuh can fail at replay time. + not flag the kernel as capturing. TMS replaces the IPC addresses during + capture, so the addresses registered while ``set_cuda_graph_capture`` is + on become stale and custom_all_reduce.cuh can fail at replay time. Passing + ``not self.tms_cudagraph`` keeps the capture flag off in that mode. """ - _patch_sglang_file_replacements( - "jit_kernel/all_reduce.py", - ( - ( - " def set_cuda_graph_register_inputs(self, register_inputs: bool) -> None: ...\n", - " def set_cuda_graph_capture(self, is_capturing: bool) -> None: ...\n", - " def set_cuda_graph_capture(self, is_capturing: bool) -> None: ...\n" - " def set_cuda_graph_register_inputs(self, register_inputs: bool) -> None: ...\n", - ), - ), - "custom all-reduce type stub graph-input registration toggle", - ) - _patch_sglang_file_replacements( - "jit_kernel/csrc/distributed/custom_all_reduce_base.cuh", - ( - ( - ' .def("set_cuda_graph_register_inputs", &Class::set_cuda_graph_register_inputs)\n', - ' .def("set_cuda_graph_capture", &Class::set_cuda_graph_capture)\n', - ' .def("set_cuda_graph_capture", &Class::set_cuda_graph_capture)\n' - ' .def("set_cuda_graph_register_inputs", &Class::set_cuda_graph_register_inputs)\n', - ), - ), - "custom all-reduce C++ binding graph-input registration toggle", - ) - _patch_sglang_file_replacements( - "jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh", - ( - ( - " void set_cuda_graph_register_inputs(bool enabled) {\n", - " void set_cuda_graph_capture(bool enabled) {\n" - " m_is_graph_capturing = enabled;\n" - " }\n\n", - " void set_cuda_graph_capture(bool enabled) {\n" - " m_is_graph_capturing = enabled;\n" - " }\n\n" - " void set_cuda_graph_register_inputs(bool enabled) {\n" - " m_register_graph_inputs = enabled;\n" - " }\n\n", - ), - ( - " bool m_register_graph_inputs = true;\n", - " bool m_is_graph_capturing = false;\n" - " int64_t m_cum_registered_count = 0;\n", - " bool m_is_graph_capturing = false;\n" - " bool m_register_graph_inputs = true;\n" - " int64_t m_cum_registered_count = 0;\n", - ), - ), - "custom all-reduce graph-input registration flag", - ) - _patch_sglang_file_replacements( - "jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh", - ( - ( - " if (check_capturing() && m_register_graph_inputs) {\n", - " if (check_capturing()) {\n", - " if (check_capturing() && m_register_graph_inputs) {\n", - ), - ), - "custom all-reduce pull graph-input registration gate", - ) _patch_sglang_file_replacements( "srt/distributed/device_communicators/custom_all_reduce_v2.py", ( @@ -282,33 +223,24 @@ def _patch_sglang_custom_all_reduce_v2_tms_cudagraph() -> None: ), ( " self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()\n", - " self.override_shot(None) # set default config based on world size\n" " self.override_algo: Optional[AllReduceAlgo] = None\n" " self.obj = get_custom_all_reduce_cls()(\n", - " self.override_shot(None) # set default config based on world size\n" " self.override_algo: Optional[AllReduceAlgo] = None\n" " self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()\n" " self.obj = get_custom_all_reduce_cls()(\n", ), ( - ' log_info_on_rank0(logger, "Registering 0 cuda graph addresses because tms is used")\n', + " self.obj.set_cuda_graph_capture(not self.tms_cudagraph)\n", + " try:\n self.obj.set_cuda_graph_capture(True)\n", " try:\n" - " self.obj.set_cuda_graph_capture(True)\n" - " yield\n" - " finally:\n" - " self.obj.set_cuda_graph_capture(False)\n" - " # cannot call when graph is capturing\n", - " try:\n" - " self.obj.set_cuda_graph_register_inputs(not self.tms_cudagraph)\n" - " self.obj.set_cuda_graph_capture(True)\n" - " yield\n" - " finally:\n" - " self.obj.set_cuda_graph_capture(False)\n" - " self.obj.set_cuda_graph_register_inputs(True)\n" - " if self.tms_cudagraph:\n" - ' log_info_on_rank0(logger, "Registering 0 cuda graph addresses because tms is used")\n' - " return\n" - " # cannot call when graph is capturing\n", + " self.obj.set_cuda_graph_capture(not self.tms_cudagraph)\n", + ), + ( + " self.obj.set_cuda_graph_capture(not self.tms_cudagraph)\n", + " finally:\n" + " self.obj.set_cuda_graph_capture(True)\n", + " finally:\n" + " self.obj.set_cuda_graph_capture(not self.tms_cudagraph)\n", ), ), "custom all-reduce v2 TMS CUDA graph capture path", diff --git a/nemo_rl/models/generation/sglang/utils/ray_utils.py b/nemo_rl/models/generation/sglang/utils/ray_utils.py index c37c579821..ad7360f244 100644 --- a/nemo_rl/models/generation/sglang/utils/ray_utils.py +++ b/nemo_rl/models/generation/sglang/utils/ray_utils.py @@ -13,11 +13,12 @@ # limitations under the License. import os -import random import socket import ray +from nemo_rl.distributed.virtual_cluster import _get_node_ip_local + # Env vars Ray uses to gate its visible-device manipulation. Setting any of # these to "1" tells Ray not to override the corresponding *_VISIBLE_DEVICES # in actor processes — used by sglang workers that want to manage CUDA @@ -33,29 +34,26 @@ ] -def find_available_port(base_port: int): - port = base_port + random.randint(100, 1000) - while True: - if is_port_available(port): - return port - if port < 60000: - port += 42 - else: - port -= 43 +@ray.remote +class Lock: + def __init__(self): + self._locked = False # False: unlocked, True: locked + def acquire(self): + """Try to acquire the lock. -def is_port_available(port): - """Return whether a port is available.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - try: - s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - s.bind(("", port)) - s.listen(1) + Returns True if acquired, False otherwise. Caller should retry until + it returns True. + """ + if not self._locked: + self._locked = True return True - except OSError: - return False - except OverflowError: - return False + return False + + def release(self): + """Release the lock, allowing others to acquire.""" + assert self._locked, "Lock is not acquired, cannot release." + self._locked = False def get_host_info(): @@ -123,15 +121,6 @@ def _resolve_ip(family, test_target_ip): def get_current_node_ip(): - address = ray._private.services.get_node_ip_address() - # strip ipv6 address - address = address.strip("[]") - return address - - -def get_free_port(start_port=7000, consecutive=1): - # find the port where port, port + 1, port + 2, ... port + consecutive - 1 are all available - port = start_port - while not all(is_port_available(port + i) for i in range(consecutive)): - port += 1 - return port + ip = _get_node_ip_local() + # strip ipv6 brackets so callers get a bare ip + return ip.strip("[]") if ip is not None else None diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index d68bf512bc..f0a731f376 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -900,6 +900,12 @@ def finish_generation(self, *args: Any, **kwargs: Any) -> bool: print(f"Error during policy preparation: {e}") return False + def pause_generation(self) -> None: + pass + + def continue_generation(self) -> None: + pass + def shutdown(self) -> bool: """Shut down all vLLM workers and clean up resources.""" try: diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index a88cd33411..5796555fee 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -205,8 +205,16 @@ def destroy_parallel_state(): pass -def setup_distributed() -> None: +def setup_distributed(config: Optional[PolicyConfig] = None) -> None: """Handle NCCL settings, dtype mapping, and basic config setup.""" + if ( + config is not None + and "generation" in config + and config["generation"] is not None + and config["generation"].get("backend") == "sglang" + ): + os.environ["NCCL_CUMEM_ENABLE"] = "0" + # Disable dynamo autotune_local_cache to avoid crash when there's already a cache # with different order of node_bundles configure_dynamo_cache() @@ -226,11 +234,13 @@ def validate_and_set_config( ): # Handle generation configuration 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"], @@ -238,10 +248,14 @@ def validate_and_set_config( 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") # Setup data types dtype_map = { diff --git a/nemo_rl/models/policy/interfaces.py b/nemo_rl/models/policy/interfaces.py index f0c1ad6bb8..fb1818615e 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -191,31 +191,6 @@ def stream_weights_via_ipc_zmq( ) -> list[ray.ObjectRef]: pass - def stream_weights_via_http( - self, - rollout_engine_urls: list[str], - buffer_size_bytes: int, - ) -> list[ray.ObjectRef]: - """Stream model weights to colocated SGLang engines via CUDA IPC over HTTP. - - Args: - rollout_engine_urls: ``http://host:port`` base URLs of each - engine's ``node_rank=0`` SGLang HTTP server. - buffer_size_bytes: Max bucket size in bytes before flushing. - - The rollout TP size (``num_gpus_per_engine``) is captured once via - ``set_rollout_num_gpus_per_engine`` and reused on every refit. - """ - raise NotImplementedError( - "stream_weights_via_http is not implemented for this policy worker" - ) - - def set_rollout_num_gpus_per_engine(self, num_gpus_per_engine: int) -> None: - """Record the rollout engine's TP size for later use in ``stream_weights_via_http``.""" - raise NotImplementedError( - "set_rollout_num_gpus_per_engine is not implemented for this policy worker" - ) - @abstractmethod def broadcast_weights_for_collective( self, kv_scales: Optional[dict[str, float]] = None diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 397b4e086b..f19df5180f 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -1030,38 +1030,77 @@ def stream_weights_via_ipc_zmq( ) return futures - def stream_weights_via_http( + def connect_sglang_rollout_engines( self, - rollout_engine_urls: list[str], - buffer_size_bytes: int, - ) -> list[ray.ObjectRef]: - """Send the weights to colocated SGLang engines via CUDA IPC over HTTP. + *, + engine_gpu_counts: list[int], + engine_gpu_offsets: Optional[list[int]] = None, + ) -> None: + """Set up the colocate Gloo gather topology for SGLang weight refit. - Args: - rollout_engine_urls: ``http://host:port`` base URLs of each - engine's ``node_rank=0`` SGLang HTTP server. The caller - resolves these once (via ``engine.get_base_url``) and passes - them in, so every FSDP rank doesn't redo the Ray RPC. - buffer_size_bytes: Max bucket size in bytes before flushing. - - The rollout TP size is captured once via - ``set_rollout_num_gpus_per_engine`` and reused by each worker. + Called by the SGLang colocated refit drivers (Megatron and FSDP) + whenever engines are added or recovered. """ futures = self.worker_group.run_all_workers_single_data( - "stream_weights_via_http", - rollout_engine_urls=rollout_engine_urls, + "connect_sglang_rollout_engines", + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + ray.get(futures) + + def update_weights_to_sglang_colocated( + self, + *, + rollout_engines: list[ray.actor.ActorHandle], + buffer_size_bytes: int, + target_precision: str = "bf16", + sglang_quantization_cfg: Optional[dict[str, Any]] = None, + ) -> list[ray.ObjectRef]: + """Send Megatron-restored HF tensors to colocated SGLang via Ray IPC.""" + futures = self.worker_group.run_all_workers_single_data( + "update_weights_to_sglang_colocated", + rollout_engines=rollout_engines, buffer_size_bytes=buffer_size_bytes, + target_precision=target_precision, + sglang_quantization_cfg=sglang_quantization_cfg, ) return futures - def set_rollout_num_gpus_per_engine(self, num_gpus_per_engine: int) -> None: - """Broadcast the rollout engine TP size to every policy worker.""" - ray.get( - self.worker_group.run_all_workers_single_data( - "set_rollout_num_gpus_per_engine", - num_gpus_per_engine=num_gpus_per_engine, - ) + def connect_sglang_rollout_engines_distributed( + self, + *, + rollout_engines: list[ray.actor.ActorHandle], + engine_gpu_counts: list[int], + group_name: Optional[str] = None, + ) -> None: + """Bring up the trainer-rank-0 NCCL group for SGLang disaggregate refit.""" + futures = self.worker_group.run_all_workers_single_data( + "connect_sglang_rollout_engines_distributed", + rollout_engines=rollout_engines, + engine_gpu_counts=engine_gpu_counts, + group_name=group_name, + ) + ray.get(futures) + + def update_weights_to_sglang_distributed( + self, + *, + rollout_engines: list[ray.actor.ActorHandle], + rollout_engine_lock: ray.actor.ActorHandle, + buffer_size_bytes: int, + target_precision: str = "bf16", + sglang_quantization_cfg: Optional[dict[str, Any]] = None, + ) -> list[ray.ObjectRef]: + """Broadcast Megatron-restored HF tensors to SGLang via NCCL (rank 0 only).""" + futures = self.worker_group.run_all_workers_single_data( + "update_weights_to_sglang_distributed", + rollout_engines=rollout_engines, + rollout_engine_lock=rollout_engine_lock, + buffer_size_bytes=buffer_size_bytes, + target_precision=target_precision, + sglang_quantization_cfg=sglang_quantization_cfg, ) + return futures def broadcast_weights_for_collective( self, kv_scales: Optional[dict[str, float]] = None diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index aaa90531c8..6d64945b7e 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -15,10 +15,10 @@ import gc import os import traceback +from datetime import timedelta from enum import Enum from typing import Any, Dict, Iterable, Optional -import requests import torch import torch.distributed as dist import zmq @@ -510,230 +510,520 @@ def rebuild_cuda_tensor_from_ipc( return func(*list_args) -def _ensure_ipc_topology( - num_engines: int, - num_gpus_per_engine: int, - worker_state: dict, -) -> None: - """Lazily create a per-engine Gloo subgroup and cache rank-only routing state. +# --------------------------------------------------------------------------- +# SGLang weight-update plumbing (colocate IPC gather + disaggregate broadcast) +# --------------------------------------------------------------------------- +def _derive_engine_gpu_offsets(engine_gpu_counts: list[int]) -> list[int]: + """Cumulative-sum offsets for a dense engine layout.""" + offsets: list[int] = [] + cursor = 0 + for c in engine_gpu_counts: + offsets.append(cursor) + cursor += c + return offsets - Every FSDP rank must call ``dist.new_group`` for every engine's rank range - (collective). Only the ranks inside a given range stash ``gather_src`` and - ``gather_group`` into ``worker_state``. The engine handle itself is resolved - at call time from the caller-provided ``rollout_engines`` list so that - post-recover actor swaps are picked up without cache invalidation. - Note: callers must have already applied ``monkey_patch_torch_reductions`` - once during worker setup; this function no longer applies it. +def connect_colocate_topology( + *, + engine_gpu_counts: list[int], + engine_gpu_offsets: Optional[list[int]] = None, + worker_state: dict, +) -> None: + """Generalized colocate rollout-engine connect for FSDP and Megatron. + + Builds a Gloo gather subgroup for each engine's GPU rank range and stashes + rank-only routing state into ``worker_state``: + + - ``worker_state["_ipc_gather_group"]``: ``ProcessGroup`` covering this + trainer rank's engine, or ``None`` if the rank is a placeholder / + not covered by any engine. + - ``worker_state["_ipc_gather_src"]``: the source rank inside the gather + group (the first GPU index of the covering engine), or ``None``. + - ``worker_state["_ipc_engine_index"]``: index into the caller's engine + list, or ``None``. The caller is responsible for resolving the actor + handle / URL at call time so post-recover actor swaps are picked up. + - ``worker_state["_ipc_layout_key"]``: cached topology signature so + subsequent connects with the same layout are no-ops. + + All trainer ranks must enter this function collectively (each call to + ``dist.new_group`` is collective). When the layout changes (e.g. a + recovered engine resizes the topology) the cached subgroup is destroyed + and rebuilt for the new layout. """ - if worker_state.get("ready"): + if not engine_gpu_counts: + raise ValueError("engine_gpu_counts must be non-empty") + if engine_gpu_offsets is None: + engine_gpu_offsets = _derive_engine_gpu_offsets(engine_gpu_counts) + elif len(engine_gpu_offsets) != len(engine_gpu_counts): + raise ValueError( + "engine_gpu_offsets and engine_gpu_counts must have the same length, " + f"got {len(engine_gpu_offsets)} vs {len(engine_gpu_counts)}" + ) + + layout_key = (tuple(engine_gpu_counts), tuple(engine_gpu_offsets)) + if worker_state.get("_ipc_layout_key") == layout_key: return + old_group = worker_state.get("_ipc_gather_group") + if old_group is not None: + try: + dist.destroy_process_group(old_group) + except Exception: + # Some torch builds raise when the group has no peers; safe to + # ignore — the new group below replaces it. + pass + my_rank = dist.get_rank() - for i in range(num_engines): - start = i * num_gpus_per_engine - group_ranks = list(range(start, start + num_gpus_per_engine)) + new_group = None + new_src: Optional[int] = None + new_engine_idx: Optional[int] = None + for i, (offset, count) in enumerate( + zip(engine_gpu_offsets, engine_gpu_counts, strict=True) + ): + group_ranks = list(range(offset, offset + count)) grp = dist.new_group(ranks=group_ranks, backend="gloo") if my_rank in group_ranks: - worker_state["gather_src"] = start - worker_state["gather_group"] = grp - + new_group = grp + new_src = offset + new_engine_idx = i + + worker_state["_ipc_gather_group"] = new_group + worker_state["_ipc_gather_src"] = new_src + worker_state["_ipc_engine_index"] = new_engine_idx + worker_state["_ipc_layout_key"] = layout_key worker_state.setdefault("weight_version", 0) - worker_state["ready"] = True - - -def _flush_bucket( - named_tensors, - gather_src: int, - gather_group, - engine_url: str, - weight_version: int, - flattened_tensor_bucket_cls, - multiprocessing_serializer_cls, -) -> None: - """Flatten ``named_tensors`` per dtype, gather to ``gather_src``, and POST to the engine.""" - # Wait on any async DTensor redistributes. - named_tensors = [ - (n, (t.wait() if hasattr(t, "wait") else t)) for n, t in named_tensors - ] - by_dtype: dict = {} - for n, t in named_tensors: - by_dtype.setdefault(t.dtype, []).append((n, t)) - - serialized: list[str] = [] - for _dtype, tensors in by_dtype.items(): - bkt = flattened_tensor_bucket_cls(named_tensors=tensors) - payload = { - "flattened_tensor": bkt.get_flattened_tensor(), - "metadata": bkt.get_metadata(), - } - serialized.append( - multiprocessing_serializer_cls.serialize(payload, output_str=True) - ) - my_rank = dist.get_rank() - group_world = dist.get_world_size(gather_group) - gathered = [None] * group_world if my_rank == gather_src else None - dist.gather_object( - serialized, - object_gather_list=gathered, - dst=gather_src, - group=gather_group, - ) +def _check_weight_sync_results(results: list) -> None: + from collections.abc import Mapping - if my_rank != gather_src: - return + for result in results: + if isinstance(result, Mapping): + success = result.get("success") + error_msg = ( + result.get("error_message") or result.get("error") or "unknown error" + ) + elif hasattr(result, "success"): + success = result.success + error_msg = getattr(result, "error_message", "unknown error") + else: + continue - assert gathered is not None - gathered_payloads: list[list[str]] = [] - for item in gathered: - assert item is not None - gathered_payloads.append(item) - - num_dtypes = len(gathered_payloads[0]) - assert num_dtypes > 0 - for i in range(num_dtypes): - body = { - "serialized_named_tensors": [g[i] for g in gathered_payloads], - "load_format": "flattened_bucket", - "flush_cache": False, - "weight_version": str(weight_version), - } - response = requests.post(f"{engine_url}/update_weights_from_tensor", json=body) - try: - response.raise_for_status() - except requests.exceptions.HTTPError as e: - e.add_note(f"{response.text=}") - raise - result = response.json() - success = result.get("success", True) - error_msg = result.get("error_message") or result.get( - "message", "unknown error" - ) - if not success: + if success is False: raise RuntimeError( - f"Weight sync failed on rollout engine: {error_msg}. " - f"Check SGLang version compatibility." + f"SGLang weight sync failed on rollout engine: {error_msg}. " + "Check SGLang version compatibility." ) -def stream_weights_via_http_impl( +def iter_named_tensor_buckets( params_generator: Iterable[tuple[str, torch.Tensor]], - rollout_engine_urls: Iterable[str], - num_gpus_per_engine: int, - rank: int, - world_size: int, - worker_name: str, buffer_size_bytes: int, +) -> "Iterable[list[tuple[str, torch.Tensor]]]": + """Group ``(name, tensor)`` pairs into buckets of at most ``buffer_size_bytes``. + + Waits on async DTensor redistributes (``.wait()``) before sizing, so the + yielded tensors are always materialized and safe to serialize. + """ + if buffer_size_bytes <= 0: + raise ValueError(f"buffer_size_bytes must be positive, got {buffer_size_bytes}") + + bucket: list[tuple[str, torch.Tensor]] = [] + bucket_size = 0 + for name, tensor in params_generator: + if hasattr(tensor, "wait"): + tensor = tensor.wait() + tensor_size = tensor.numel() * tensor.element_size() + if bucket and bucket_size + tensor_size > buffer_size_bytes: + yield bucket + bucket = [] + bucket_size = 0 + bucket.append((name, tensor)) + bucket_size += tensor_size + + if bucket: + yield bucket + + +def send_hf_buckets_via_ipc_actor_impl( + *, + bucket_iterator: Iterable[list[tuple[str, torch.Tensor]]], + rollout_engines: list, worker_state: dict, + weight_version: Optional[int] = None, ) -> None: - """Stream FSDP weights to colocated SGLang engines via CUDA IPC over HTTP. - - Args: - params_generator: Iterable yielding ``(name, tensor)`` pairs to stream. - Caller is responsible for any pre-processing (LoRA merge, HF - adaptation, dtype cast). - rollout_engine_urls: ``http://host:port`` base URLs of each engine's - ``node_rank=0`` SGLang HTTP server. One entry per engine, in TP - rank-range order: engine ``i`` owns global ranks - ``[i * num_gpus_per_engine, (i + 1) * num_gpus_per_engine)``. - num_gpus_per_engine: TP size per SGLang engine. - rank: Global FSDP rank. - world_size: Global FSDP world size. - worker_name: Human label for logs. - buffer_size_bytes: Max bucket size in bytes. - worker_state: Mutable dict on the worker used to cache topology and - weight version across refits. + """Send finalized HF tensor buckets to colocated SGLang engines via Ray IPC. + + Per bucket: group by dtype, serialize a ``FlattenedTensorBucket`` per + dtype, ``dist.gather_object`` to the gather source rank, then on the + source rank call ``ipc_engine.update_weights_from_tensor.remote(...)`` + once per dtype, **block on ``ray.get(refs)`` per chunk**, validate + engine return values, synchronize all trainer ranks, then drop the + trainer-side ``flattened_tensor`` references before moving on. + + The trainer-side topology (``_ipc_gather_group`` / ``_ipc_gather_src`` / + ``_ipc_engine_index``) must already have been set up by + :func:`connect_colocate_topology`. Placeholder ranks (no covering engine) + return immediately — they must not call ``gather_object``. Non-source + trainer ranks participate in the gather and completion broadcast; they + don't issue Ray RPCs and don't ``ray.get``. + + Returns ``None``. Raises ``RuntimeError`` if any chunk fails on the + engine side. """ + import ray + from nemo_rl.models.generation.sglang.utils.train_utils import ( FlattenedTensorBucket, MultiprocessingSerializer, ) - rollout_engine_urls = list(rollout_engine_urls) + gather_group = worker_state.get("_ipc_gather_group") + gather_src = worker_state.get("_ipc_gather_src") + engine_idx = worker_state.get("_ipc_engine_index") - _ensure_ipc_topology( - num_engines=len(rollout_engine_urls), - num_gpus_per_engine=num_gpus_per_engine, - worker_state=worker_state, - ) + if gather_group is None or gather_src is None or engine_idx is None: + # Placeholder rank: must not participate in the per-engine gather. + return None - worker_state["weight_version"] = worker_state.get("weight_version", 0) + 1 - weight_version = worker_state["weight_version"] - gather_src = worker_state["gather_src"] - gather_group = worker_state["gather_group"] - - engine_url = None - for i, candidate in enumerate(rollout_engine_urls): - start = i * num_gpus_per_engine - end = start + num_gpus_per_engine - if start <= rank < end: - engine_url = candidate - break - if engine_url is None: - raise RuntimeError( - f"No rollout engine matched rank={rank} with " - f"num_gpus_per_engine={num_gpus_per_engine} and " - f"{len(rollout_engine_urls)} engine URL(s); " - f"rank must fall within [0, {num_gpus_per_engine * len(rollout_engine_urls)})." - ) + if weight_version is None: + worker_state["weight_version"] = worker_state.get("weight_version", 0) + 1 + weight_version = worker_state["weight_version"] + + ipc_engine = rollout_engines[engine_idx] + my_rank = dist.get_rank() try: - bucket: list = [] - bucket_size = 0 - for name, param in params_generator: - param_size = param.numel() * param.element_size() - if bucket and bucket_size + param_size >= buffer_size_bytes: - _flush_bucket( - bucket, - gather_src=gather_src, - gather_group=gather_group, - engine_url=engine_url, - weight_version=weight_version, - flattened_tensor_bucket_cls=FlattenedTensorBucket, - multiprocessing_serializer_cls=MultiprocessingSerializer, - ) - bucket = [] - bucket_size = 0 - - param = param.cuda() - bucket.append((name, param)) - bucket_size += param_size - - if bucket: - _flush_bucket( - bucket, - gather_src=gather_src, - gather_group=gather_group, - engine_url=engine_url, - weight_version=weight_version, - flattened_tensor_bucket_cls=FlattenedTensorBucket, - multiprocessing_serializer_cls=MultiprocessingSerializer, - ) + for bucket in bucket_iterator: + if not bucket: + continue - if dist.get_rank() == gather_src: - # Mirror SGLangGenerationWorker.invalidate_kv_cache: the endpoint - # returns non-200 while requests are still pending, so retry up to 60s. - import time + # No async-collective ``.wait()`` here — Megatron's AutoBridge + # yields plain ``torch.Tensor``, no DTensor wrapping. - for _ in range(60): - try: - response = requests.get(f"{engine_url}/flush_cache") - if response.status_code == 200: - break - except requests.RequestException: - pass - time.sleep(1) + if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False): + by_dtype: dict = {"dtype": list(bucket)} else: - raise TimeoutError(f"Timeout while flushing cache at {engine_url}.") + by_dtype = {} + for name, tensor in bucket: + by_dtype.setdefault(tensor.dtype, []).append((name, tensor)) + + serialized: list[str] = [] + long_lived_tensors: list[dict] = [] + for _dtype, named_tensors in by_dtype.items(): + bkt = FlattenedTensorBucket(named_tensors=named_tensors) + payload = { + "flattened_tensor": bkt.get_flattened_tensor(), + "metadata": bkt.get_metadata(), + } + long_lived_tensors.append(payload) + serialized.append( + MultiprocessingSerializer.serialize(payload, output_str=True) + ) - except Exception as e: - print( - f"{worker_name} (rank {rank}): Error during HTTP weight streaming: {e}.\n" - f"{traceback.format_exc()}" - ) - raise + group_world = dist.get_world_size(gather_group) + gathered = [None] * group_world if my_rank == gather_src else None + dist.gather_object( + serialized, + object_gather_list=gathered, + dst=gather_src, + group=gather_group, + ) + + refs: list = [] + if my_rank == gather_src: + num_dtypes = len(gathered[0]) + for i in range(num_dtypes): + refs.append( + ipc_engine.update_weights_from_tensor.remote( + serialized_named_tensors=[g[i] for g in gathered], + load_format="flattened_bucket", + weight_version=str(weight_version), + ) + ) + + # The serialized IPC handles gathered on the source may point at + # flattened tensors owned by non-source trainer ranks. Keep every + # rank's tensors alive until the source finishes the engine RPCs. + sync_error: Optional[str] = None + source_exc: Optional[BaseException] = None + if my_rank == gather_src: + try: + results = ray.get(refs) + _check_weight_sync_results(results) + except BaseException as exc: + source_exc = exc + sync_error = repr(exc) + + sync_state = [sync_error] + dist.broadcast_object_list(sync_state, src=gather_src, group=gather_group) + del long_lived_tensors, refs + + if source_exc is not None: + raise source_exc + if sync_state[0] is not None: + raise RuntimeError( + f"SGLang IPC weight update failed on gather src rank " + f"{gather_src}: {sync_state[0]}" + ) finally: gc.collect() torch.cuda.empty_cache() + + return None + + +def init_process_group( + backend: "str | dist.Backend | None" = None, + init_method: Optional[str] = None, + timeout: Optional[timedelta] = None, + world_size: int = -1, + rank: int = -1, + store: "Optional[dist.Store]" = None, + group_name: Optional[str] = None, + pg_options: Any = None, +) -> "torch.distributed.ProcessGroup": + """Create a side-by-side ``ProcessGroup`` without touching the default world. + + ``torch.distributed.init_process_group`` initializes the *default* world + process group. Once the Megatron trainer has stood up its own world during + Policy construction, calling it again to talk to SGLang either errors with + "trying to initialize the default process group twice" or — depending on + torch version — silently hangs in rendezvous against a peer that has + already finished its own custom-group setup. + + Same approach as SGLang's ``sglang.srt.utils.common.init_custom_process_group``: + replay the public API's wiring (rendezvous → ``PrefixStore`` → + ``_new_process_group_helper``) but skip the "set as default PG" step, so + multiple independent groups can coexist in the same process. + + Only one of ``init_method`` and ``store`` may be set; otherwise the + rendezvous source is ambiguous. + """ + from torch.distributed.distributed_c10d import ( + Backend, + PrefixStore, + _get_default_group, + _new_process_group_helper, + _world, + default_pg_timeout, + rendezvous, + ) + + assert (store is None) or (init_method is None), ( + "Cannot specify both init_method and store." + ) + + if store is not None: + assert world_size > 0, "world_size must be positive if using store" + assert rank >= 0, "rank must be non-negative if using store" + elif init_method is None: + init_method = "env://" + + backend = Backend(backend) if backend else Backend("undefined") + if timeout is None: + timeout = default_pg_timeout + + if store is None: + rendezvous_iterator = rendezvous(init_method, rank, world_size, timeout=timeout) + store, rank, world_size = next(rendezvous_iterator) + store.set_timeout(timeout) + # PrefixStore so multiple co-tenant groups don't trample each other's keys. + store = PrefixStore(group_name or "", store) + + # ``pg_options`` was renamed to ``backend_options`` in PyTorch 2.6: + # https://github.com/pytorch/pytorch/commit/a0c7029a75628cd5fa8df83c0de0ea98ee7fd844 + # Use numeric tuple compare — string compare ``"2.10" >= "2.6"`` returns + # False because ``"1"`` sorts before ``"6"`` lexicographically. + _torch_mm = tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:2]) + pg_options_kw = "backend_options" if _torch_mm >= (2, 6) else "pg_options" + + # Disable the ncclCommSplit path (see docstring). Safe to mutate here: + # nothing else reads ``bound_device_id`` during group construction, and + # refit setup does not create process groups from other threads. + default_pg = _get_default_group() if dist.is_initialized() else None + saved_bound_device_id = getattr(default_pg, "bound_device_id", None) + if saved_bound_device_id is not None: + default_pg.bound_device_id = None + try: + pg, _ = _new_process_group_helper( + world_size, + rank, + [], + backend, + store, + group_name=group_name, + **{pg_options_kw: pg_options}, + timeout=timeout, + ) + finally: + if saved_bound_device_id is not None: + default_pg.bound_device_id = saved_bound_device_id + + # Map identity ranks so collective ops can resolve member ranks for ``pg``. + _world.pg_group_ranks[pg] = {i: i for i in range(world_size)} + return pg + + +def connect_rollout_engines_from_distributed( + *, + group_name: str, + rollout_engines: list, + engine_gpu_counts: list[int], +) -> "torch.distributed.ProcessGroup": + """Set up the SGLang NCCL weight-update group with trainer rank 0 as rank 0. + + Only trainer rank 0 broadcasts because the AutoBridge path restores + full HF weights, not per-PP slices. + + The caller (a trainer) must invoke this only on rank 0; other ranks must + not call it. + """ + import ray + + from nemo_rl.distributed.virtual_cluster import _get_free_port_local + + master_address = ray._private.services.get_node_ip_address() + master_port = _get_free_port_local() + world_size = 1 + sum(engine_gpu_counts) + + refs = [] + rank_cursor = 1 + for engine, gpu_count in zip(rollout_engines, engine_gpu_counts, strict=True): + refs.append( + engine.init_weights_update_group.remote( + master_address, + master_port, + rank_cursor, + world_size, + group_name, + "nccl", + ) + ) + rank_cursor += gpu_count + + group = init_process_group( + backend="nccl", + init_method=f"tcp://{master_address}:{master_port}", + world_size=world_size, + rank=0, + group_name=group_name, + ) + ray.get(refs) + return group + + +def disconnect_rollout_engines_from_distributed( + *, + group_name: str, + model_update_group: "torch.distributed.ProcessGroup", + rollout_engines: list, +) -> None: + """Tear down trainer-side and engine-side NCCL state for ``group_name``.""" + import ray + + refs = [ + engine.destroy_weights_update_group.remote(group_name) + for engine in rollout_engines + ] + try: + dist.destroy_process_group(model_update_group) + except Exception: + pass + try: + ray.get(refs) + except Exception: + pass + + +def get_sglang_quantization_cfg(policy_generation: Any) -> dict: + """Read the active SGLang quantization block from the generation handle. + + Returns an empty dict when no quantization config is set, so callers can + treat the result as a stable mapping without ``None`` checks. + """ + return dict(policy_generation.sglang_cfg["sglang_cfg"].get("quantization") or {}) + + +def fetch_updatable_engines_with_recover(policy_generation: Any) -> tuple: + """Run the design-mandated weight-update prelude. + + 1. If ``sglang_cfg.use_fault_tolerance`` is enabled, call + ``rollout_manager.recover_updatable_engines`` which internally pauses + health monitoring, restarts dead engines, and runs + release/resume_memory_occupation on every recovered node-0 engine. + 2. Read the current updatable-engine state via + ``get_updatable_engines_and_lock``. + + Both calls are idempotent — recover is a no-op when no engines have died. + """ + use_ft = bool( + policy_generation.sglang_cfg["sglang_cfg"].get("use_fault_tolerance", False) + ) + if use_ft: + policy_generation.recover_updatable_engines() + return policy_generation.get_updatable_engines_and_lock() + + +def broadcast_hf_buckets_via_distributed_impl( + *, + bucket_iterator: Iterable[list[tuple[str, torch.Tensor]]], + rollout_engines: list, + rollout_engine_lock, + group_name: str, + model_update_group: "torch.distributed.ProcessGroup", + weight_version: int, +) -> None: + """Broadcast finalized HF tensor buckets to SGLang via NCCL (rank 0 only). + + Per-bucket protocol: trainer rank 0 sends per-tensor metadata to every + engine via Ray (``update_weights_from_distributed``), then issues one + ``dist.broadcast`` per tensor over the NCCL group, then waits for the Ray + refs to confirm engines finished loading the bucket. + + The rollout-engine lock wraps each bucket's broadcast so concurrent SGLang + NCCL operations (e.g. health-check pings) cannot collide with the + weight-update broadcast. + """ + import time as _time + + import ray + + bucket_idx = 0 + for bucket in bucket_iterator: + if not bucket: + continue + + bucket_idx += 1 + # No async-collective ``.wait()`` here — AutoBridge yields plain + # ``torch.Tensor`` for the Megatron path (no DTensor wrapping). + + names = [name for name, _ in bucket] + dtypes = [tensor.dtype for _, tensor in bucket] + shapes = [tensor.shape for _, tensor in bucket] + + while not ray.get(rollout_engine_lock.acquire.remote()): + _time.sleep(0.1) + try: + refs = [ + engine.update_weights_from_distributed.remote( + names=names, + dtypes=dtypes, + shapes=shapes, + group_name=group_name, + weight_version=str(weight_version), + ) + for engine in rollout_engines + ] + handles = [] + for i, (_, tensor) in enumerate(bucket): + handles.append( + dist.broadcast( + tensor.data, 0, group=model_update_group, async_op=True + ) + ) + for i, handle in enumerate(handles): + handle.wait() + ray.get(refs) + finally: + ray.get(rollout_engine_lock.release.remote()) diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py index 12a01b2a6b..ad10c73347 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -362,21 +362,20 @@ def __init__( _runtime_is_reward_model, # Duplicate, already set as _is_reward_model ) = runtime_config - # Rollout topology constant for SGLang colocated refit: set once via - # ``set_rollout_num_gpus_per_engine`` after the SGLang generation - # handle exists and consumed by ``stream_weights_via_http`` on each - # refit. Only initialized on the SGLang colocated path since no other - # generation backend uses this attribute. + ## SGLang weight-update state. Populated lazily by + ## ``connect_sglang_rollout_engines`` on the first colocated refit. generation_backend = config.get("generation", {}).get("backend") if generation_backend == "sglang": - from nemo_rl.models.generation.sglang.utils.train_utils import ( - monkey_patch_torch_reductions, - ) - - monkey_patch_torch_reductions() + self._sglang_ipc_state: dict = {} if self.is_generation_colocated: - self._rollout_num_gpus_per_engine: Optional[int] = None - self._ipc_worker_state: dict = {} + # Colocate refit serializes CUDA-IPC tensor handles for + # SGLang; the torch reductions monkey patch must be in place + # before any tensor is serialized. + from nemo_rl.models.generation.sglang.utils.train_utils import ( + monkey_patch_torch_reductions, + ) + + monkey_patch_torch_reductions() def _update_moe_gate_bias_if_supported(self) -> None: """Update the non-gradient MoE routing bias after the optimizer step.""" @@ -384,10 +383,6 @@ def _update_moe_gate_bias_if_supported(self) -> None: if update_moe_gate_bias is not None: update_moe_gate_bias() - def set_rollout_num_gpus_per_engine(self, num_gpus_per_engine: int) -> None: - """Record the rollout engine's TP size for later use in ``stream_weights_via_http``.""" - self._rollout_num_gpus_per_engine = num_gpus_per_engine - @wrap_with_nvtx_name("dtensor_policy_worker_v2/train") def train( self, @@ -1123,45 +1118,70 @@ def stream_weights_via_ipc_zmq( ) @torch.no_grad() - @wrap_with_nvtx_name("dtensor_policy_worker_v2/stream_weights_via_http") - def stream_weights_via_http( + @wrap_with_nvtx_name("dtensor_policy_worker_v2/connect_sglang_rollout_engines") + def connect_sglang_rollout_engines( self, - rollout_engine_urls: list[str], - buffer_size_bytes: int, + *, + engine_gpu_counts: list[int], + engine_gpu_offsets: Optional[list[int]] = None, ) -> None: - """Stream FSDP weights to colocated SGLang engines via CUDA IPC over HTTP. + """Set up the colocate Gloo gather topology for SGLang weight refit. - Args: - rollout_engine_urls: ``http://host:port`` base URLs of each - engine's ``node_rank=0`` SGLang HTTP server. The driver - resolves these once via ``engine.get_base_url`` and passes - them down so every FSDP rank doesn't redo the Ray RPC. - buffer_size_bytes: Max bucket size in bytes before flushing. - - ``num_gpus_per_engine`` is recorded once via - ``set_rollout_num_gpus_per_engine`` after the SGLang generation handle - is created, so the caller doesn't have to pass it on every refit. + Must be called collectively by every FSDP rank when SGLang engines + are added or recovered. Subsequent calls with the same layout are + no-ops. """ - assert self._rollout_num_gpus_per_engine is not None, ( - "stream_weights_via_http called before set_rollout_num_gpus_per_engine; " - "wire the rollout TP size on the policy after SGLangGeneration is built." + from nemo_rl.models.policy.utils import connect_colocate_topology + + connect_colocate_topology( + engine_gpu_counts=list(engine_gpu_counts), + engine_gpu_offsets=( + list(engine_gpu_offsets) if engine_gpu_offsets is not None else None + ), + worker_state=self._sglang_ipc_state, ) + @torch.no_grad() + @wrap_with_nvtx_name("dtensor_policy_worker_v2/update_weights_to_sglang_colocated") + def update_weights_to_sglang_colocated( + self, + *, + rollout_engines: list, + buffer_size_bytes: int, + target_precision: str = "bf16", + sglang_quantization_cfg: Optional[dict[str, Any]] = None, + ) -> None: + """Send FSDP weights to colocated SGLang engines via Ray CUDA IPC. + + Synchronous: each chunk is awaited via ``ray.get`` inside + :func:`send_hf_buckets_via_ipc_actor_impl` before the next chunk is + sent, so trainer-side IPC tensors stay alive until the engine has + copied them and per-chunk engine failures surface immediately. + """ + if target_precision != "bf16": + raise NotImplementedError( + "The FSDP/DTensor policy only supports BF16 SGLang refits; " + f"got target_precision={target_precision!r}." + ) + del sglang_quantization_cfg # accepted for dispatch parity, bf16-only + # Manually move model to cuda for cpu offload case if self.cpu_offload: self.model = self.move_to_cuda(self.model) - from nemo_rl.models.policy.utils import stream_weights_via_http_impl + from nemo_rl.models.policy.utils import ( + iter_named_tensor_buckets, + send_hf_buckets_via_ipc_actor_impl, + ) - stream_weights_via_http_impl( - params_generator=dtensor_params_generator(self.model, self.dtype), - rollout_engine_urls=rollout_engine_urls, - num_gpus_per_engine=self._rollout_num_gpus_per_engine, - rank=self.rank, - world_size=torch.distributed.get_world_size(), - worker_name=str(self), + bucket_iter = iter_named_tensor_buckets( + dtensor_params_generator(self.model, self.dtype), buffer_size_bytes=buffer_size_bytes, - worker_state=self._ipc_worker_state, + ) + send_hf_buckets_via_ipc_actor_impl( + bucket_iterator=bucket_iter, + rollout_engines=list(rollout_engines), + worker_state=self._sglang_ipc_state, ) @torch.no_grad() @@ -1369,3 +1389,78 @@ def _init_checkpoint_manager( ) # pragma: no cover class DTensorPolicyWorkerV2(DTensorPolicyWorkerV2Impl): pass + + +# --------------------------------------------------------------------------- +# Driver-side SGLang weight-update dispatch (FSDP backend) +# --------------------------------------------------------------------------- +def refit_sglang_colocated( + *, + policy: Any, + policy_generation: Any, + buffer_size_bytes: int, +) -> bool: + """Refit colocated SGLang engines from the FSDP/DTensor policy. + + Lifecycle: optional fault-tolerance recover, connect (when new / + recovered engines), pause + KV invalidation, send HF tensor buckets via + Ray IPC, post-process, continue. Mirrors the Megatron colocated driver; + the FSDP path is BF16-only. + """ + from nemo_rl.models.policy.utils import fetch_updatable_engines_with_recover + + ( + rollout_engines, + _rollout_engine_lock, + num_new_engines, + engine_gpu_counts, + engine_gpu_offsets, + ) = fetch_updatable_engines_with_recover(policy_generation) + + if num_new_engines > 0: + policy.connect_sglang_rollout_engines( + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + policy_generation.clear_updatable_num_new_engines() + assert policy_generation.num_new_engines == 0, ( + "clear_updatable_num_new_engines did not zero num_new_engines" + ) + + # Pause with the configured mode, but only invalidate the KV cache when + # the mode actually drops generation state. "in_place" leaves the engine + # paused without dropping its KV cache, so flushing would clobber the + # still-valid in-place state. + pause_mode = policy_generation.pause_generation_mode + policy_generation.pause_generation(mode=pause_mode) + policy_generation.invalidate_kv_cache() + try: + futures = policy.update_weights_to_sglang_colocated( + rollout_engines=rollout_engines, + buffer_size_bytes=buffer_size_bytes, + ) + ray.get(futures) + policy_generation.post_process_weights() + finally: + policy_generation.continue_generation() + return True + + +def refit_sglang_distributed( + *, + policy: Any, # noqa: ARG001 — accepted for dispatch parity + policy_generation: Any, # noqa: ARG001 + buffer_size_bytes: int, # noqa: ARG001 +) -> bool: + """SGLang disaggregate broadcast is not currently supported for FSDP. + + Per the design, only the Megatron backend implements the distributed + refit path (it depends on AutoBridge restoring full HF tensors on + trainer rank 0). FSDP non-colocated refits should keep using the + legacy ``broadcast_weights_for_collective`` flow with a non-SGLang + generation backend. + """ + raise NotImplementedError( + "SGLang weight_transfer_mode='broadcast' is currently only supported " + "for the Megatron policy backend." + ) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index cbf9ce1bfc..12fafec27b 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -90,7 +90,14 @@ LogprobOutputSpec, ReferenceLogprobOutputSpec, ) -from nemo_rl.models.policy.utils import get_runtime_env_for_policy_worker +from nemo_rl.models.policy.utils import ( + broadcast_hf_buckets_via_distributed_impl, + connect_colocate_topology, + connect_rollout_engines_from_distributed, + disconnect_rollout_engines_from_distributed, + get_runtime_env_for_policy_worker, + send_hf_buckets_via_ipc_actor_impl, +) from nemo_rl.models.policy.workers.base_policy_worker import AbstractPolicyWorker from nemo_rl.models.policy.workers.patches import apply_transformer_engine_patch from nemo_rl.utils.grad_norm import warn_if_inf_grad_norm @@ -295,7 +302,7 @@ def __init__( self.timer = Timer(context={"worker": "megatron_policy", "rank": self.rank}) # Step 1: Setup distributed - setup_distributed() + setup_distributed(config) log_gpu_memory_diagnostics( label="after_nccl_init", worker_type="MegatronPolicyWorker" ) @@ -455,6 +462,29 @@ def __init__( ## used for streaming update inference engine weights self._held_gather_buffer = None + ## SGLang weight-update state. Populated lazily by + ## ``connect_sglang_rollout_engines`` (colocate) or + ## ``connect_sglang_rollout_engines_distributed`` (broadcast). + self._sglang_ipc_state: dict = {} + generation_cfg = config.get("generation") + if ( + generation_cfg is not None + and generation_cfg.get("backend") == "sglang" + and generation_cfg["colocated"]["enabled"] + ): + # Colocate refit serializes CUDA-IPC tensor handles for SGLang; + # the torch reductions monkey patch must be in place before any + # tensor is serialized. + from nemo_rl.models.generation.sglang.utils.train_utils import ( + monkey_patch_torch_reductions, + ) + + monkey_patch_torch_reductions() + self._sglang_dist_group: Any = None + self._sglang_dist_group_name: str = "nemo_rl_sglang" + self._sglang_dist_engines: list = [] + self._sglang_weight_version: int = 0 + self._init_inference_engine_state() log_gpu_memory_diagnostics( @@ -1881,6 +1911,183 @@ def _iter_params_with_optional_kv_scales( ).reshape(1) yield param_name, scale_tensor + # ------------------------------------------------------------------ + # SGLang weight update (colocate IPC + disaggregate broadcast) + # ------------------------------------------------------------------ + def _build_sglang_hf_iterator( + self, + *, + target_precision: str, + sglang_quantization_cfg: Optional[dict] = None, + ): + from nemo_rl.models.policy.workers.megatron_sglang_weight_iterator import ( + MegatronSGLangHfWeightIterator, + ) + + if self.refit_conversion_tasks is None: + self.refit_conversion_tasks = self.megatron_bridge.get_conversion_tasks( + [self.model] + ) + + num_hidden_layers = 0 + if target_precision == "mxfp8": + num_hidden_layers = int( + getattr(self.megatron_bridge.transformer_config, "num_layers", 0) + ) + + return MegatronSGLangHfWeightIterator( + megatron_bridge=self.megatron_bridge, + models=[self.model], + conversion_tasks=self.refit_conversion_tasks, + quantization_config=dict(sglang_quantization_cfg or {}), + num_hidden_layers=num_hidden_layers, + ) + + @torch.no_grad() + @wrap_with_nvtx_name("megatron_policy_worker/connect_sglang_rollout_engines") + def connect_sglang_rollout_engines( + self, + *, + engine_gpu_counts: list[int], + engine_gpu_offsets: Optional[list[int]] = None, + ) -> None: + """Set up the colocate Gloo gather topology for SGLang weight refit. + + Must be called collectively by every Megatron rank when SGLang + engines are added or recovered. Subsequent calls with the same + layout are no-ops. + """ + connect_colocate_topology( + engine_gpu_counts=list(engine_gpu_counts), + engine_gpu_offsets=( + list(engine_gpu_offsets) if engine_gpu_offsets is not None else None + ), + worker_state=self._sglang_ipc_state, + ) + + @torch.no_grad() + @wrap_with_nvtx_name("megatron_policy_worker/update_weights_to_sglang_colocated") + def update_weights_to_sglang_colocated( + self, + *, + rollout_engines: list, + buffer_size_bytes: int, + target_precision: str = "bf16", + sglang_quantization_cfg: Optional[dict] = None, + ) -> None: + """Send finalized HF tensor buckets to colocated SGLang engines. + + Synchronous: each chunk is awaited via ``ray.get`` inside + :func:`send_hf_buckets_via_ipc_actor_impl` before the next chunk + is sent, so trainer-side IPC tensors stay alive until the engine + has copied them and per-chunk engine failures surface immediately. + Raises ``RuntimeError`` on any chunk failure. + """ + self._sglang_weight_version += 1 + iterator = self._build_sglang_hf_iterator( + target_precision=target_precision, + sglang_quantization_cfg=sglang_quantization_cfg, + ) + bucket_iter = iterator.iter_hf_weight_buckets( + target_precision=cast(Any, target_precision), + buffer_size_bytes=buffer_size_bytes, + ) + send_hf_buckets_via_ipc_actor_impl( + bucket_iterator=bucket_iter, + rollout_engines=list(rollout_engines), + worker_state=self._sglang_ipc_state, + weight_version=self._sglang_weight_version, + ) + + @torch.no_grad() + @wrap_with_nvtx_name( + "megatron_policy_worker/connect_sglang_rollout_engines_distributed" + ) + def connect_sglang_rollout_engines_distributed( + self, + *, + rollout_engines: list, + engine_gpu_counts: list[int], + group_name: Optional[str] = None, + ) -> None: + """Bring up the trainer-rank-0 NCCL group for SGLang disaggregate refit. + + Only trainer rank 0 broadcasts to SGLang, so only rank 0 owns the + torch process group. Other ranks return immediately. Calling this + again after engines recover destroys the stale group first. + """ + if self.rank != 0: + return + + if group_name is not None: + self._sglang_dist_group_name = group_name + + if self._sglang_dist_group is not None: + disconnect_rollout_engines_from_distributed( + group_name=self._sglang_dist_group_name, + model_update_group=self._sglang_dist_group, + rollout_engines=self._sglang_dist_engines, + ) + self._sglang_dist_group = None + self._sglang_dist_engines = [] + + self._sglang_dist_group = connect_rollout_engines_from_distributed( + group_name=self._sglang_dist_group_name, + rollout_engines=list(rollout_engines), + engine_gpu_counts=list(engine_gpu_counts), + ) + self._sglang_dist_engines = list(rollout_engines) + + @torch.no_grad() + @wrap_with_nvtx_name("megatron_policy_worker/update_weights_to_sglang_distributed") + def update_weights_to_sglang_distributed( + self, + *, + rollout_engines: list, + rollout_engine_lock, + buffer_size_bytes: int, + target_precision: str = "bf16", + sglang_quantization_cfg: Optional[dict] = None, + ) -> None: + """Broadcast finalized HF tensors to SGLang engines from trainer rank 0. + + Non-rank-0 trainers still walk the AutoBridge iterator (Megatron + gather + AutoBridge restoration is a collective), but they do not + participate in the NCCL broadcast. This matches the design's "trainer + rank 0 as the only source" decision. + """ + self._sglang_weight_version += 1 + iterator = self._build_sglang_hf_iterator( + target_precision=target_precision, + sglang_quantization_cfg=sglang_quantization_cfg, + ) + bucket_iter = iterator.iter_hf_weight_buckets( + target_precision=cast(Any, target_precision), + buffer_size_bytes=buffer_size_bytes, + ) + + if self.rank != 0: + # Drain the iterator so AutoBridge collectives complete on every + # rank, but do not broadcast. + for _ in bucket_iter: + pass + return + + if self._sglang_dist_group is None: + raise RuntimeError( + "connect_sglang_rollout_engines_distributed must be called " + "before update_weights_to_sglang_distributed." + ) + + broadcast_hf_buckets_via_distributed_impl( + bucket_iterator=bucket_iter, + rollout_engines=list(rollout_engines), + rollout_engine_lock=rollout_engine_lock, + group_name=self._sglang_dist_group_name, + model_update_group=self._sglang_dist_group, + weight_version=self._sglang_weight_version, + ) + @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/stream_weights_via_ipc_zmq") def stream_weights_via_ipc_zmq( @@ -2541,3 +2748,130 @@ def _percentile(values: list[float], p: float) -> float: ) # pragma: no cover class MegatronPolicyWorker(MegatronPolicyWorkerImpl): pass + + +# --------------------------------------------------------------------------- +# Driver-side SGLang weight-update dispatch (Megatron backend) +# --------------------------------------------------------------------------- +def refit_sglang_colocated( + *, + policy: Any, + policy_generation: Any, + buffer_size_bytes: int, +) -> bool: + """Refit colocated SGLang engines from the Megatron policy. + + Lifecycle: optional fault-tolerance recover, connect (when new / + recovered engines), pause + flush, send HF tensor buckets via Ray + IPC, post-process, continue. + """ + from nemo_rl.models.policy.utils import ( + fetch_updatable_engines_with_recover, + get_sglang_quantization_cfg, + ) + + sglang_quant = get_sglang_quantization_cfg(policy_generation) + target_precision = sglang_quant.get("scheme", "bf16") + + ( + rollout_engines, + _rollout_engine_lock, + num_new_engines, + engine_gpu_counts, + engine_gpu_offsets, + ) = fetch_updatable_engines_with_recover(policy_generation) + + if num_new_engines > 0: + policy.connect_sglang_rollout_engines( + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + policy_generation.clear_updatable_num_new_engines() + assert policy_generation.num_new_engines == 0, ( + "clear_updatable_num_new_engines did not zero num_new_engines" + ) + + # Pause with the configured mode, but only invalidate the KV cache when + # the mode actually drops generation state. "in_place" leaves the engine + # paused without dropping its KV cache, so flushing would clobber the + # still-valid in-place state. + pause_mode = policy_generation.pause_generation_mode + policy_generation.pause_generation(mode=pause_mode) + policy_generation.invalidate_kv_cache() + try: + # Per-worker actor method is now synchronous (per-chunk ray.get + + # lifetime-safe IPC handled inside send_hf_buckets_via_ipc_actor_impl), + # but the policy-group dispatch still returns one Ray future per + # worker; we await those here to wait for all trainer ranks. + futures = policy.update_weights_to_sglang_colocated( + rollout_engines=rollout_engines, + buffer_size_bytes=buffer_size_bytes, + target_precision=target_precision, + sglang_quantization_cfg=sglang_quant, + ) + ray.get(futures) + policy_generation.post_process_weights() + finally: + policy_generation.continue_generation() + return True + + +def refit_sglang_distributed( + *, + policy: Any, + policy_generation: Any, + buffer_size_bytes: int, +) -> bool: + """Broadcast Megatron-restored HF tensors to disaggregate SGLang via NCCL. + + Trainer rank 0 owns the SGLang weight-update group; non-rank-0 ranks still + walk the AutoBridge collective inside ``update_weights_to_sglang_distributed`` + but do not broadcast. Includes optional fault-tolerance recover prelude. + """ + from nemo_rl.models.policy.utils import ( + fetch_updatable_engines_with_recover, + get_sglang_quantization_cfg, + ) + + sglang_quant = get_sglang_quantization_cfg(policy_generation) + target_precision = sglang_quant.get("scheme", "bf16") + + ( + rollout_engines, + rollout_engine_lock, + num_new_engines, + engine_gpu_counts, + _engine_gpu_offsets, + ) = fetch_updatable_engines_with_recover(policy_generation) + + if num_new_engines > 0: + policy.connect_sglang_rollout_engines_distributed( + rollout_engines=rollout_engines, + engine_gpu_counts=engine_gpu_counts, + ) + policy_generation.clear_updatable_num_new_engines() + assert policy_generation.num_new_engines == 0, ( + "clear_updatable_num_new_engines did not zero num_new_engines" + ) + + # Pause with the configured mode, but only invalidate the KV cache when + # the mode actually drops generation state. "in_place" leaves the engine + # paused without dropping its KV cache, so flushing would clobber the + # still-valid in-place state. + pause_mode = policy_generation.pause_generation_mode + policy_generation.pause_generation(mode=pause_mode) + if pause_mode != "in_place": + policy_generation.invalidate_kv_cache() + try: + futures = policy.update_weights_to_sglang_distributed( + rollout_engines=rollout_engines, + rollout_engine_lock=rollout_engine_lock, + buffer_size_bytes=buffer_size_bytes, + target_precision=target_precision, + sglang_quantization_cfg=sglang_quant, + ) + ray.get(futures) + policy_generation.post_process_weights() + finally: + policy_generation.continue_generation() + return True diff --git a/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py b/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py new file mode 100644 index 0000000000..1e14e8799b --- /dev/null +++ b/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py @@ -0,0 +1,141 @@ +# Copyright (c) 2025, 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. + +"""SGLang-only HF weight iterator for the Megatron policy worker. + +Emits buckets of HF-named tensors restored from Megatron via AutoBridge, +with no vLLM-specific KV/Q scale tensors. When +``target_precision == "mxfp8"`` the iterator additionally applies the +offline ``should_quantize`` / ``quantize_mxfp8`` core to each finalized +HF tensor. +""" + +from __future__ import annotations + +from typing import Any, Iterator, Literal + +import torch + +from nemo_rl.models.generation.sglang.mxfp8_quantization_core import ( + SOURCE_FP8_SCALE_KEY_SUFFIX, + build_dynamic_skip_substrings, + quantize_mxfp8, + should_quantize, + strip_weight_suffix, +) + + +class MegatronSGLangHfWeightIterator: + """Yield buckets of finalized HF named tensors for SGLang weight refit. + + The iterator is bound to a Megatron bridge, the local Megatron model(s), + and the conversion-task list precomputed by the policy worker. For each + refit it walks ``bridge.export_hf_weights`` and packs tensors into buckets + sized by the *post-transformation* tensor footprint, so MXFP8 buckets + correctly account for the added ``weight_scale_inv`` tensor. + """ + + def __init__( + self, + *, + megatron_bridge: Any, + models: list[Any], + conversion_tasks: Any, + quantization_config: dict[str, Any] | None = None, + num_hidden_layers: int = 0, + ) -> None: + self._bridge = megatron_bridge + self._models = models + self._conversion_tasks = conversion_tasks + self._quantization_config = dict(quantization_config or {}) + self._num_hidden_layers = num_hidden_layers + + def iter_hf_weight_buckets( + self, + *, + target_precision: Literal["bf16", "mxfp8"] = "bf16", + buffer_size_bytes: int, + ) -> Iterator[list[tuple[str, torch.Tensor]]]: + """Yield finalized HF tensor buckets sized by transmitted bytes.""" + if buffer_size_bytes <= 0: + raise ValueError( + f"buffer_size_bytes must be positive, got {buffer_size_bytes}" + ) + + skip_weight_substrings = ( + build_dynamic_skip_substrings( + quantization_config=self._quantization_config, + num_hidden_layers=self._num_hidden_layers, + ) + if target_precision == "mxfp8" + else None + ) + + bucket: list[tuple[str, torch.Tensor]] = [] + bucket_size = 0 + + for finalized in self._iter_finalized_hf_named_tensors( + target_precision=target_precision, + skip_weight_substrings=skip_weight_substrings, + ): + for name, tensor in finalized: + tensor_size = tensor.numel() * tensor.element_size() + if bucket and bucket_size + tensor_size > buffer_size_bytes: + yield bucket + bucket = [] + bucket_size = 0 + bucket.append((name, tensor)) + bucket_size += tensor_size + + if bucket: + yield bucket + + def _iter_finalized_hf_named_tensors( + self, + *, + target_precision: Literal["bf16", "mxfp8"], + skip_weight_substrings: tuple[str, ...] | None, + ) -> Iterator[list[tuple[str, torch.Tensor]]]: + """Yield finalized HF (name, tensor) groups from one AutoBridge tensor. + + AutoBridge yields one HF named tensor at a time. For BF16 each AutoBridge + item produces exactly one finalized pair; for MXFP8 each item may + expand to a ``(weight, weight_scale_inv)`` pair when the weight is + quantized. + """ + for hf_param_name, tensor in self._bridge.export_hf_weights( + self._models, + show_progress=False, + conversion_tasks=self._conversion_tasks, + ): + # AutoBridge yields plain ``torch.Tensor`` for Megatron (no + # DTensor / async-collective wrapping), so no ``.wait()`` is + # needed here. The previous ``hasattr(tensor, "wait")`` check + # was a copy-from-FSDP residue. + + if target_precision == "mxfp8" and skip_weight_substrings is not None: + if should_quantize( + hf_param_name, + tensor, + skip_weight_substrings=skip_weight_substrings, + allow_source_fp8=False, + ): + qweight, scale = quantize_mxfp8(tensor) + scale_name = ( + strip_weight_suffix(hf_param_name) + SOURCE_FP8_SCALE_KEY_SUFFIX + ) + yield [(hf_param_name, qweight), (scale_name, scale)] + continue + + yield [(hf_param_name, tensor)] diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py index 061022fc9c..89712e2efe 100644 --- a/nemo_rl/weight_sync/factory.py +++ b/nemo_rl/weight_sync/factory.py @@ -16,7 +16,7 @@ Selects the appropriate weight synchronizer based on the deployment topology (colocated vs. non-colocated) and the generation backend -(vLLM uses IPC/ZMQ, SGLang uses HTTP, non-colocated uses NCCL). +(vLLM uses IPC/ZMQ, SGLang uses Ray CUDA-IPC, non-colocated uses NCCL). """ from typing import Any, Optional @@ -89,11 +89,11 @@ def create_weight_synchronizer( ) if generation_backend == SGLANG_BACKEND: - from nemo_rl.weight_sync.http_weight_synchronizer import ( - HTTPWeightSynchronizer, + from nemo_rl.weight_sync.sglang_weight_synchronizer import ( + SGLangColocatedWeightSynchronizer, ) - return HTTPWeightSynchronizer( + return SGLangColocatedWeightSynchronizer( policy=policy, generation=generation, refit_buffer_size_gb=refit_buffer_size_gb, diff --git a/nemo_rl/weight_sync/sglang_weight_synchronizer.py b/nemo_rl/weight_sync/sglang_weight_synchronizer.py new file mode 100644 index 0000000000..56e9e0f3d3 --- /dev/null +++ b/nemo_rl/weight_sync/sglang_weight_synchronizer.py @@ -0,0 +1,141 @@ +# 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. + +"""Weight synchronizer for colocated SGLang generation. + +Handles weight transfer between a colocated policy and the SGLang generation +backend via the backend-specific colocated refit drivers +(``refit_sglang_colocated`` in ``megatron_policy_worker`` / +``dtensor_policy_worker_v2``), which gather HF tensor buckets per engine and +push them over Ray CUDA IPC (``send_hf_buckets_via_ipc_actor_impl``). + +Lifecycle per sync: + 1. policy.offload_before_refit() -- free GPU for weight staging + 2. generation.prepare_for_generation(tags=["weights"]) -- allocate buffers + 3. refit_sglang_colocated() -- pause, send buckets via Ray IPC, + post-process, continue + 4. policy.offload_after_refit() -- restore optimizer state + 5. generation.prepare_for_generation(tags=["kv_cache"]) -- rebuild KV cache +""" + +import os +from contextlib import nullcontext +from typing import Any, Optional + +from nemo_rl.utils.timer import Timer +from nemo_rl.weight_sync.interfaces import WeightSynchronizer + + +class SGLangColocatedWeightSynchronizer(WeightSynchronizer): + """Weight synchronizer for colocated SGLang deployments. + + Both the policy and generation workers run on the same GPUs. Weights are + bucketed on the trainer, gathered per engine, and handed to SGLang via + Ray CUDA-IPC actor RPCs — the same path used by + ``refit_policy_generation`` in the GRPO loop. + + Args: + policy: Policy object implementing ColocatablePolicyInterface. + generation: SGLangGeneration instance. + refit_buffer_size_gb: Fixed buffer size in GB for weight staging. + If None, buffer size is computed dynamically from free GPU memory. + """ + + def __init__( + self, + policy: Any, + generation: Any, + refit_buffer_size_gb: Optional[int] = None, + ): + self._policy = policy + self._generation = generation + self._refit_buffer_size_gb = refit_buffer_size_gb + self._stale = True + + def sync_weights( + self, + *, + timer: Optional[Timer] = None, + kv_scales: Optional[dict[str, float]] = None, + ) -> None: + self._policy.offload_before_refit() + self._generation.prepare_for_generation(tags=["weights"]) + + sync_succeeded = False + try: + timer_context = ( + timer.time("prepare_for_generation/transfer_and_update_weights") + if timer is not None + else nullcontext() + ) + with timer_context: + buffer_size_bytes = self._compute_buffer_size() + sync_succeeded = bool(self._refit_colocated(buffer_size_bytes)) + finally: + self._policy.offload_after_refit() + self._generation.prepare_for_generation(tags=["kv_cache"]) + + self._stale = not sync_succeeded + + def _refit_colocated(self, buffer_size_bytes: int) -> bool: + """Route to the backend-specific colocated SGLang refit driver.""" + use_megatron = bool( + self._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, + ) + + return _backend.refit_sglang_colocated( + policy=self._policy, + policy_generation=self._generation, + buffer_size_bytes=buffer_size_bytes, + ) + + @property + def is_stale(self) -> bool: + return self._stale + + def mark_stale(self) -> None: + self._stale = True + + def init_communicator(self) -> None: + state_dict_info = self._policy.prepare_refit_info() + self._generation.prepare_refit_info(state_dict_info) + + def shutdown(self) -> None: + pass + + def _compute_buffer_size(self) -> int: + if self._refit_buffer_size_gb is not None: + if self._refit_buffer_size_gb <= 0: + raise ValueError("refit_buffer_size_gb must be > 0") + return self._refit_buffer_size_gb * (1024**3) + + memory_ratio_raw = os.getenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0.3") + try: + memory_ratio = float(memory_ratio_raw) + except ValueError as exc: + raise ValueError( + f"NRL_REFIT_BUFFER_MEMORY_RATIO must be a valid float, got {memory_ratio_raw!r}" + ) from exc + if memory_ratio <= 0: + raise ValueError("NRL_REFIT_BUFFER_MEMORY_RATIO must be > 0") + + return int(self._policy.get_free_memory_bytes() * memory_ratio) diff --git a/ray.sub b/ray.sub index d4f44b2570..9e96064c5f 100644 --- a/ray.sub +++ b/ray.sub @@ -140,12 +140,15 @@ fi # 1301-1312 Ray management (node-mgr, obj-mgr, etc.; odd=worker, even=head) # 1400-1999 Master address / TCPStore (cluster.master_port_range_low/high) # 2000-2999 Ray worker gRPC (min/max-worker-port) -# 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 (Gym global config port_range_low/high) # 6000 Sandbox Nginx (NEMO_SKILLS_SANDBOX_PORT) # 6001-6999 Sandbox uWSGI workers (SANDBOX_BASE_PORT) -# 7000-8999 vLLM / SGLang engine rendezvous (VLLM_PORT in vllm_worker.py / SGLang base_port) +# 7000-8999 vLLM engine rendezvous (VLLM_PORT in vllm_worker.py) # 8265 Ray Dashboard (DASHBOARD_PORT; reserved carve-out inside the 7000-8999 band) +# 8600-8799 SGLang router (carve-out inside the 7000-8999 band; only one rollout backend runs) +# 8800-8999 SGLang Prometheus metrics (carve-out inside the 7000-8999 band) MIN_WORKER_PORT=${MIN_WORKER_PORT:-2000} MAX_WORKER_PORT=${MAX_WORKER_PORT:-2999} ######################################################## diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index c21667d047..a4fd4dcdfc 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -1716,9 +1716,6 @@ def init_collective(self, *_args, **_kwargs): def prepare_refit_info(self): return {} - def set_rollout_num_gpus_per_engine(self, _num_gpus_per_engine): - pass - class DummySGLangGeneration: num_gpus_per_engine = 1 diff --git a/tests/unit/distributed/test_virtual_cluster.py b/tests/unit/distributed/test_virtual_cluster.py index f41afec5da..af9ae99618 100644 --- a/tests/unit/distributed/test_virtual_cluster.py +++ b/tests/unit/distributed/test_virtual_cluster.py @@ -27,12 +27,17 @@ DEFAULT_GYM_PORT_RANGE_LOW, DEFAULT_MASTER_PORT_RANGE_HIGH, DEFAULT_MASTER_PORT_RANGE_LOW, + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_HIGH, + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_LOW, + DEFAULT_SGLANG_ROUTER_PORT_RANGE_HIGH, + DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW, DEFAULT_VLLM_PORT_RANGE_LOW, DEFAULT_VLLM_PORTS_PER_ENGINE, PY_EXECUTABLES, RayVirtualCluster, ResourceInsufficientError, _bind_socket_in_range, + _get_free_consecutive_ports_local, _get_free_port_local, _get_node_ip_and_free_port, ) @@ -322,6 +327,80 @@ def test_multiple_calls_return_different_ports(self): assert len(ports) > 1 +class TestGetFreeConsecutivePortsLocal: + """Tests for _get_free_consecutive_ports_local().""" + + def test_single_port_in_range(self): + port = _get_free_consecutive_ports_local(14000, 14100, consecutive=1) + assert 14000 <= port < 14100 + + def test_returns_bindable_contiguous_block(self): + n = 5 + base = _get_free_consecutive_ports_local(14100, 14300, consecutive=n) + assert 14100 <= base + assert base + n - 1 < 14300 + # All n ports are simultaneously bindable. + socks = [] + try: + for offset in range(n): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("", base + offset)) + socks.append(s) + finally: + for s in socks: + s.close() + + def test_start_port_below_low_is_clamped(self): + base = _get_free_consecutive_ports_local( + 14300, 14400, consecutive=1, start_port=1000 + ) + assert base >= 14300 + + def test_cursor_advance_yields_non_overlapping_blocks(self): + n = 3 + first = _get_free_consecutive_ports_local(14400, 14600, consecutive=n) + second = _get_free_consecutive_ports_local( + 14400, 14600, consecutive=n, start_port=first + n + ) + assert second >= first + n + + def test_raises_when_range_exhausted(self): + # A 3-wide range cannot fit a block of 5. + with pytest.raises(RuntimeError, match="consecutive free ports"): + _get_free_consecutive_ports_local(14600, 14603, consecutive=5) + + def test_port_is_reusable_after_return(self): + base = _get_free_consecutive_ports_local(14700, 14800, consecutive=2) + # Sockets are closed before return, so the block can be re-bound. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", base)) + + def test_raises_when_start_port_at_or_above_high(self): + with pytest.raises(RuntimeError, match="consecutive free ports"): + _get_free_consecutive_ports_local( + 14850, 14900, consecutive=1, start_port=15000 + ) + + def test_block_fits_exactly_against_high_boundary(self): + # high is exclusive: a block whose last port is high-1 must still fit. + n = 4 + base = _get_free_consecutive_ports_local(14900, 14900 + n, consecutive=n) + assert base == 14900 + socks = [] + try: + for offset in range(n): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("", base + offset)) + socks.append(s) + finally: + for s in socks: + s.close() + + def test_rejects_non_positive_consecutive(self): + with pytest.raises(AssertionError, match="consecutive must be >= 1"): + _get_free_consecutive_ports_local(14950, 15000, consecutive=0) + + class TestRayVirtualClusterPortRange: """Tests for port range propagation in RayVirtualCluster.""" @@ -397,5 +476,22 @@ def test_default_port_ranges_ordered_and_below_ephemeral_floor(): DEFAULT_VLLM_PORT_RANGE_LOW + 8 * DEFAULT_VLLM_PORTS_PER_ENGINE < EPHEMERAL_FLOOR ) + # SGLang router / Prometheus carve-outs live inside the vLLM band (only one + # rollout backend runs at a time), stay above the 8-engine vLLM high-water + # mark and the Ray dashboard (8265), and sit below the ephemeral floor. + assert ( + DEFAULT_VLLM_PORT_RANGE_LOW + 8 * DEFAULT_VLLM_PORTS_PER_ENGINE + < DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW + ) + assert 8265 < DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW + assert DEFAULT_SGLANG_ROUTER_PORT_RANGE_LOW < DEFAULT_SGLANG_ROUTER_PORT_RANGE_HIGH + assert ( + DEFAULT_SGLANG_ROUTER_PORT_RANGE_HIGH < DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_LOW + ) + assert ( + DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_LOW + < DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_HIGH + ) + assert DEFAULT_SGLANG_PROMETHEUS_PORT_RANGE_HIGH < EPHEMERAL_FLOOR # Avoid privileged ports (<1024). assert DEFAULT_MASTER_PORT_RANGE_LOW > 1024 diff --git a/tests/unit/models/generation/sglang/helpers.py b/tests/unit/models/generation/sglang/helpers.py index 150d5c122d..fda12e1daf 100644 --- a/tests/unit/models/generation/sglang/helpers.py +++ b/tests/unit/models/generation/sglang/helpers.py @@ -34,10 +34,13 @@ import ray +from nemo_rl.distributed.virtual_cluster import ( + DEFAULT_GENERATION_PORT_RANGE_LOW, + _get_free_port_local, +) from nemo_rl.models.generation.sglang.sglang_worker import SGLangGenerationWorker from nemo_rl.models.generation.sglang.utils.ray_utils import ( NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, - find_available_port, get_host_info, ) from nemo_rl.utils.venvs import make_actor_runtime_env @@ -141,9 +144,10 @@ def create_worker(router_info, base_gpu_id=0, tp_size=1, rank=0): ) host_ip = get_host_info()[1] - port = find_available_port(30000 + rank * 1000) - nccl_port = find_available_port(40000 + rank * 1000) - dist_init_port = find_available_port(50000 + rank * 1000) + band_low = DEFAULT_GENERATION_PORT_RANGE_LOW + rank * 1000 + port = _get_free_port_local(band_low, band_low + 300) + nccl_port = _get_free_port_local(band_low + 300, band_low + 600) + dist_init_port = _get_free_port_local(band_low + 600, band_low + 1000) ray.get( worker.init.remote( diff --git a/tests/unit/models/generation/sglang/test_sglang_router.py b/tests/unit/models/generation/sglang/test_sglang_router.py index 479ccb9f5d..7da4ecaa6a 100644 --- a/tests/unit/models/generation/sglang/test_sglang_router.py +++ b/tests/unit/models/generation/sglang/test_sglang_router.py @@ -22,8 +22,8 @@ import ray import requests +from nemo_rl.distributed.virtual_cluster import _get_free_port_local from nemo_rl.models.generation.sglang.sglang_router import RouterActor, _start_router -from nemo_rl.models.generation.sglang.utils.ray_utils import find_available_port from . import ( helpers, # noqa: F401 — installs env vars + module stubs before nemo_rl imports @@ -68,7 +68,7 @@ def test_start_returns_ip_and_port(ray_cluster): def test_start_uses_configured_port(ray_cluster): """When sglang_router_port is set, the router uses that exact port.""" - configured_port = find_available_port(9000) + configured_port = _get_free_port_local(9000, 10000) actor = RouterActor.remote() try: ip, port = _start_and_cleanup(actor, {"sglang_router_port": configured_port}) diff --git a/tests/unit/models/generation/sglang/test_utils_smoke.py b/tests/unit/models/generation/sglang/test_utils_smoke.py index c7db029e62..b186ec40fe 100644 --- a/tests/unit/models/generation/sglang/test_utils_smoke.py +++ b/tests/unit/models/generation/sglang/test_utils_smoke.py @@ -27,11 +27,7 @@ pytestmark = pytest.mark.sglang from nemo_rl.models.generation.sglang.utils.ip_port_utils import _wrap_ipv6 -from nemo_rl.models.generation.sglang.utils.ray_utils import ( - find_available_port, - get_host_info, - is_port_available, -) +from nemo_rl.models.generation.sglang.utils.ray_utils import get_host_info from nemo_rl.models.generation.sglang.utils.train_utils import ( MultiprocessingSerializer, ) @@ -40,14 +36,6 @@ # --------------------------------------------------------------------------- # ray_utils # --------------------------------------------------------------------------- -def test_find_available_port(): - """find_available_port returns a port that passes is_port_available.""" - port = find_available_port(20000) - assert isinstance(port, int) - assert port > 0 - assert is_port_available(port) - - def test_wrap_ipv6_noop_for_ipv4(): """IPv4 addresses are returned unchanged by _wrap_ipv6.""" assert _wrap_ipv6("192.168.1.1") == "192.168.1.1" diff --git a/tests/unit/models/generation/sglang/test_weight_update_real.py b/tests/unit/models/generation/sglang/test_weight_update_real.py index 05dadb7b07..baf935ccd0 100644 --- a/tests/unit/models/generation/sglang/test_weight_update_real.py +++ b/tests/unit/models/generation/sglang/test_weight_update_real.py @@ -35,12 +35,14 @@ import torch from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from nemo_rl.distributed.virtual_cluster import RayVirtualCluster -from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration -from nemo_rl.models.generation.sglang.utils.ray_utils import ( - find_available_port, - get_host_info, +from nemo_rl.distributed.virtual_cluster import ( + DEFAULT_MASTER_PORT_RANGE_HIGH, + DEFAULT_MASTER_PORT_RANGE_LOW, + RayVirtualCluster, + _get_free_port_local, ) +from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration +from nemo_rl.models.generation.sglang.utils.ray_utils import get_host_info from tests.unit.models.generation.sglang.weight_update_actor import MockFSDPWorker from .helpers import make_actor_env_vars, post_and_assert_200 @@ -160,7 +162,9 @@ def mock_trainer(ray_cluster, sglang_gen): fit both worker groups. """ host_ip = get_host_info()[1] - master_port = find_available_port(29500) + master_port = _get_free_port_local( + DEFAULT_MASTER_PORT_RANGE_LOW, DEFAULT_MASTER_PORT_RANGE_HIGH + ) env_vars = make_actor_env_vars() pg = sglang_gen.cluster.get_placement_groups()[0] diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py index edec20b6c0..64bd3a54f5 100644 --- a/tests/unit/weight_sync/test_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_weight_synchronizer.py @@ -27,13 +27,13 @@ CollectiveWeightSynchronizer, ) from nemo_rl.weight_sync.factory import create_weight_synchronizer -from nemo_rl.weight_sync.http_weight_synchronizer import ( - HTTPWeightSynchronizer, -) from nemo_rl.weight_sync.interfaces import WeightSynchronizer from nemo_rl.weight_sync.ipc_weight_synchronizer import ( IPCWeightSynchronizer, ) +from nemo_rl.weight_sync.sglang_weight_synchronizer import ( + SGLangColocatedWeightSynchronizer, +) # --------------------------------------------------------------------------- # Helpers @@ -46,7 +46,7 @@ def _mock_policy(**overrides): policy.offload_after_refit.return_value = None policy.prepare_refit_info.return_value = {"layer_0": {"shape": [4096, 4096]}} policy.stream_weights_via_ipc_zmq.return_value = [MagicMock()] - policy.stream_weights_via_http.return_value = [MagicMock()] + policy.cfg = {"megatron_cfg": {"enabled": False}} policy.broadcast_weights_for_collective.return_value = [MagicMock()] policy.init_collective.return_value = [MagicMock()] policy.get_free_memory_bytes.return_value = 1024**3 # 1 GB @@ -62,7 +62,6 @@ def _mock_generation(**overrides): gen.prepare_refit_info.return_value = None gen.update_weights_via_ipc_zmq.return_value = [MagicMock()] gen.update_weights_from_collective.return_value = [MagicMock()] - gen.get_rollout_engine_urls.return_value = ["http://localhost:30000"] gen.init_collective.return_value = [MagicMock()] for k, v in overrides.items(): setattr(gen, k, v) @@ -217,17 +216,21 @@ def test_zero_env_ratio_raises(self, mock_ray, monkeypatch): # --------------------------------------------------------------------------- -# HTTPWeightSynchronizer +# SGLangColocatedWeightSynchronizer # --------------------------------------------------------------------------- +_SGLANG_REFIT_DRIVER = ( + "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.refit_sglang_colocated" +) -class TestHTTPWeightSynchronizer: - @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray") - def test_sync_weights_calls_full_lifecycle(self, mock_ray): - mock_ray.get.return_value = [True] + +class TestSGLangColocatedWeightSynchronizer: + @patch(_SGLANG_REFIT_DRIVER) + def test_sync_weights_calls_full_lifecycle(self, mock_refit): + mock_refit.return_value = True policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen) + sync = SGLangColocatedWeightSynchronizer(policy, gen) assert sync.is_stale sync.sync_weights() @@ -235,30 +238,28 @@ def test_sync_weights_calls_full_lifecycle(self, mock_ray): policy.offload_before_refit.assert_called_once() gen.prepare_for_generation.assert_any_call(tags=["weights"]) - policy.stream_weights_via_http.assert_called_once() - gen.get_rollout_engine_urls.assert_called_once() - call_kwargs = policy.stream_weights_via_http.call_args - assert call_kwargs.kwargs["rollout_engine_urls"] == ["http://localhost:30000"] - assert call_kwargs.kwargs["buffer_size_bytes"] == int((1024**3) * 0.3) + mock_refit.assert_called_once() + call_kwargs = mock_refit.call_args.kwargs + assert call_kwargs["policy"] is policy + assert call_kwargs["policy_generation"] is gen + assert call_kwargs["buffer_size_bytes"] == int((1024**3) * 0.3) policy.offload_after_refit.assert_called_once() gen.prepare_for_generation.assert_any_call(tags=["kv_cache"]) - @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray") - def test_fixed_buffer_size(self, mock_ray): - mock_ray.get.return_value = [True] + @patch(_SGLANG_REFIT_DRIVER) + def test_fixed_buffer_size(self, mock_refit): + mock_refit.return_value = True policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen, refit_buffer_size_gb=2) + sync = SGLangColocatedWeightSynchronizer(policy, gen, refit_buffer_size_gb=2) sync.sync_weights() - call_kwargs = policy.stream_weights_via_http.call_args - assert call_kwargs.kwargs["rollout_engine_urls"] == ["http://localhost:30000"] - assert call_kwargs.kwargs["buffer_size_bytes"] == 2 * (1024**3) + assert mock_refit.call_args.kwargs["buffer_size_bytes"] == 2 * (1024**3) def test_mark_stale(self): policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen) + sync = SGLangColocatedWeightSynchronizer(policy, gen) sync._stale = False assert not sync.is_stale @@ -268,21 +269,21 @@ def test_mark_stale(self): def test_init_communicator(self): policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen) + sync = SGLangColocatedWeightSynchronizer(policy, gen) sync.init_communicator() policy.prepare_refit_info.assert_called_once() gen.prepare_refit_info.assert_called_once() - @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray") - def test_phase_restoration_on_transfer_failure(self, mock_ray): + @patch(_SGLANG_REFIT_DRIVER) + def test_phase_restoration_on_transfer_failure(self, mock_refit): """offload_after_refit and kv_cache prep run even when transfer raises.""" - mock_ray.get.side_effect = RuntimeError("HTTP transfer exploded") + mock_refit.side_effect = RuntimeError("IPC transfer exploded") policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen) + sync = SGLangColocatedWeightSynchronizer(policy, gen) - with pytest.raises(RuntimeError, match="HTTP transfer exploded"): + with pytest.raises(RuntimeError, match="IPC transfer exploded"): sync.sync_weights() policy.offload_after_refit.assert_called_once() @@ -292,25 +293,23 @@ def test_phase_restoration_on_transfer_failure(self, mock_ray): def test_negative_buffer_size_raises(self): policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen, refit_buffer_size_gb=-1) + sync = SGLangColocatedWeightSynchronizer(policy, gen, refit_buffer_size_gb=-1) with pytest.raises(ValueError, match="refit_buffer_size_gb must be > 0"): sync._compute_buffer_size() - @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray") - def test_invalid_env_ratio_raises(self, mock_ray, monkeypatch): + def test_invalid_env_ratio_raises(self, monkeypatch): monkeypatch.setenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "not_a_number") policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen) + sync = SGLangColocatedWeightSynchronizer(policy, gen) with pytest.raises(ValueError, match="must be a valid float"): sync._compute_buffer_size() - @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray") - def test_zero_env_ratio_raises(self, mock_ray, monkeypatch): + def test_zero_env_ratio_raises(self, monkeypatch): monkeypatch.setenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0") policy = _mock_policy() gen = _mock_generation() - sync = HTTPWeightSynchronizer(policy, gen) + sync = SGLangColocatedWeightSynchronizer(policy, gen) with pytest.raises(ValueError, match="must be > 0"): sync._compute_buffer_size() @@ -408,7 +407,7 @@ def test_colocated_vllm_returns_ipc(self): ) assert isinstance(sync, IPCWeightSynchronizer) - def test_colocated_sglang_returns_http(self): + def test_colocated_sglang_returns_sglang_colocated(self): policy = _mock_policy() gen = _mock_generation() sync = create_weight_synchronizer( @@ -417,7 +416,7 @@ def test_colocated_sglang_returns_http(self): generation_backend=SGLANG_BACKEND, colocated=True, ) - assert isinstance(sync, HTTPWeightSynchronizer) + assert isinstance(sync, SGLangColocatedWeightSynchronizer) def test_colocated_megatron_returns_ipc(self): policy = _mock_policy()