diff --git a/.gitmodules b/.gitmodules index ab84439b11..66c4435fe6 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,8 +1,3 @@ -[submodule "3rdparty/Megatron-LM"] - path = 3rdparty/Megatron-LM-workspace/Megatron-LM - url = https://github.com/NVIDIA/Megatron-LM.git - branch = main - shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge url = https://github.com/NVIDIA-NeMo/Megatron-Bridge.git diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 23cdeb3807..1a4912e231 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 23cdeb38077d7b72a5fbae0927a2e1a74bfc15f7 +Subproject commit 1a4912e231bb2795b062f7de97496caaf382c7f6 diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index a2bb70b91b..95e5f38f87 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit a2bb70b91b827bd6b085a77442c7cf60cfdb59fe +Subproject commit 95e5f38f8727c4ab30830559c68939f35f4e52f6 diff --git a/3rdparty/Megatron-Bridge-workspace/setup.py b/3rdparty/Megatron-Bridge-workspace/setup.py index a1fa9a77e6..55fabf9a54 100644 --- a/3rdparty/Megatron-Bridge-workspace/setup.py +++ b/3rdparty/Megatron-Bridge-workspace/setup.py @@ -25,6 +25,7 @@ bridge_src_dir = "Megatron-Bridge/src/megatron/bridge" bridge_package_name = "megatron.bridge" +# Default dependencies from pyproject.toml CACHED_DEPENDENCIES = [ "transformers>=5.0.0,<=5.3.0", "peft>=0.18.1", @@ -48,14 +49,14 @@ "megatron-core[dev,mlm]", "qwen-vl-utils", # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13 base container - "transformer-engine[pytorch,core_cu12]", + "transformer-engine[pytorch,core_cu13]", "mamba-ssm", - "nvidia-resiliency-ext~=0.5.0", + "nvidia-resiliency-ext", "causal-conv1d", "flash-linear-attention", "timm", "open-clip-torch>=3.2.0", - "mlflow>=3.5.0", + "mlflow>=3.9.0", "comet-ml>=3.50.0", "torch>=2.6.0", ] @@ -73,20 +74,9 @@ project = data["project"] deps_list = project["dependencies"] submodule_deps = set(str(d).strip() for d in deps_list) - - # Normalize the transformer-engine CUDA variant extra (core_cu12 vs core_cu13) - # so our CUDA 12 override doesn't trip the consistency check against the - # submodule's CUDA 13 default. - # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): remove this when we upgrade to CUDA 13 - def _normalize_te_cuda(dep): - if dep.startswith("transformer-engine") or dep.startswith("transformer_engine"): - return dep.replace("core_cu13", "core_cu12") - return dep - - normalized_submodule = set(_normalize_te_cuda(d) for d in submodule_deps) - normalized_cached = set(_normalize_te_cuda(d) for d in CACHED_DEPENDENCIES) - missing_in_cached = normalized_submodule - normalized_cached - extra_in_cached = normalized_cached - normalized_submodule + cached_deps = set(CACHED_DEPENDENCIES) + missing_in_cached = submodule_deps - cached_deps + extra_in_cached = cached_deps - submodule_deps if missing_in_cached or extra_in_cached: print( diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM deleted file mode 160000 index 9e28104173..0000000000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 9e2810417315a7ee93b41d4e234454abd3c16af5 diff --git a/3rdparty/Megatron-LM-workspace/is_megatron_installed.py b/3rdparty/Megatron-LM-workspace/is_megatron_installed.py deleted file mode 100644 index 9a88db404f..0000000000 --- a/3rdparty/Megatron-LM-workspace/is_megatron_installed.py +++ /dev/null @@ -1,21 +0,0 @@ -# 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. -try: - from megatron.core import parallel_state # noqa: F401 - - INSTALLED = True -except ImportError: - INSTALLED = False - -print(f"Megatron {INSTALLED=}") diff --git a/3rdparty/Megatron-LM-workspace/pyproject.toml b/3rdparty/Megatron-LM-workspace/pyproject.toml deleted file mode 100644 index 4537293a9d..0000000000 --- a/3rdparty/Megatron-LM-workspace/pyproject.toml +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. - -[build-system] -requires = ["setuptools", "pybind11"] -build-backend = "setuptools.build_meta" - -[project] -name = "megatron-core" -dynamic = ["dependencies", "version"] -description = "Megatron Core - a library for efficient and scalable training of transformer based models" -authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }] -maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }] diff --git a/3rdparty/Megatron-LM-workspace/setup.py b/3rdparty/Megatron-LM-workspace/setup.py deleted file mode 100644 index 8a7b438bb7..0000000000 --- a/3rdparty/Megatron-LM-workspace/setup.py +++ /dev/null @@ -1,206 +0,0 @@ -# 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. -"""Setup for pip package.""" - -import os -import subprocess -import sys -import tomllib - -import setuptools -from setuptools import Extension - -############################################################################### -# Extension Making # -# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # - -# --- Configuration Start --- -# These will be populated conditionally or with defaults -final_packages = [] -final_package_dir = {} -final_ext_modules = [] - -# --- megatron.core conditional section --- -# Directory for the megatron.core Python package source -megatron_core_python_package_source_dir = "Megatron-LM/megatron/core" -megatron_core_package_name = "megatron.core" - -# Path for the C++ extension's source file, relative to setup.py -megatron_core_cpp_extension_source_file = "megatron/core/datasets/helpers.cpp" - -# Cached dependencies: default + dev from pyproject.toml -# VCS dependencies use full "pkg @ git+URL@rev" format matching pyproject.toml [tool.uv.sources] -CACHED_DEPENDENCIES = [ - # Default dependencies from pyproject.toml - "torch>=2.6.0", - "numpy", - "packaging>=24.2", - # Dev dependencies from pyproject.toml - "nvidia-modelopt[torch]; sys_platform != 'darwin'", - # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13 base container - "transformer-engine[pytorch,core_cu12]", - "nvidia-resiliency-ext @ git+https://github.com/NVIDIA/nvidia-resiliency-ext.git@v0.5.0", - "tqdm", - "einops~=0.8", - "tensorstore~=0.1,!=0.1.46,!=0.1.72", - "nvtx~=0.2", - "multi-storage-client~=0.27", - "opentelemetry-api~=1.33.1", - "mamba-ssm~=2.2", - "causal-conv1d~=1.5", - "flash-linear-attention~=0.4.0", - "megatron-energon[av_decode]~=6.0", - "av", - "flashinfer-python~=0.5.0", - "wget", - "onnxscript", - # VCS dependency - must match pyproject.toml [tool.uv.sources] - "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1.0", - "datasets", - "fastapi~=0.50", - "hypercorn", - "quart", - "openai[aiohttp]", - "orjson", -] - - -def build_vcs_dependency(pkg_name: str, source_info: dict) -> str: - """Build a PEP 440 VCS dependency string from pyproject.toml [tool.uv.sources] entry.""" - git_url = source_info.get("git") - rev = source_info.get("rev") - if not git_url: - raise ValueError(f"No git URL found for VCS dependency: {pkg_name}") - if not rev: - raise ValueError(f"No rev/commit found for VCS dependency: {pkg_name}") - return f"{pkg_name} @ git+{git_url}@{rev}" - - -# Read pyproject.toml to validate dependencies -pyproject_path = os.path.join("Megatron-LM", "pyproject.toml") - -if os.path.exists(megatron_core_python_package_source_dir): - if not os.path.exists(pyproject_path): - raise FileNotFoundError( - f"[megatron-core][setup] {pyproject_path} not found; skipping dependency consistency check." - ) - - with open(pyproject_path, "rb") as f: - data = tomllib.load(f) - - # Extract [tool.uv.sources] for VCS dependencies - uv_sources = data.get("tool", {}).get("uv", {}).get("sources", {}) - - # Combine default dependencies + dev optional-dependencies - project = data["project"] - default_deps = project.get("dependencies", []) - optional_deps = project.get("optional-dependencies", {}) - dev_deps = optional_deps.get("dev", []) - - submodule_deps = set(str(d).strip() for d in default_deps + dev_deps) - - # Build expected dependencies, converting any in [tool.uv.sources] to full VCS strings - submodule_deps_with_vcs = set() - for dep in submodule_deps: - if dep in uv_sources: - # Replace with full VCS string constructed from [tool.uv.sources] - vcs_dep = build_vcs_dependency(dep, uv_sources[dep]) - submodule_deps_with_vcs.add(vcs_dep) - else: - submodule_deps_with_vcs.add(dep) - - cached_deps_set = set(CACHED_DEPENDENCIES) - - # Normalize the transformer-engine CUDA variant extra (core_cu12 vs core_cu13) - # so our CUDA 12 override doesn't trip the consistency check against the - # submodule's CUDA 13 default. - # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): remove this when we upgrade to CUDA 13 - def _normalize_te_cuda(dep): - if dep.startswith("transformer-engine") or dep.startswith("transformer_engine"): - return dep.replace("core_cu13", "core_cu12") - return dep - - normalized_submodule = set(_normalize_te_cuda(d) for d in submodule_deps_with_vcs) - normalized_cached = set(_normalize_te_cuda(d) for d in cached_deps_set) - missing_in_cached = normalized_submodule - normalized_cached - extra_in_cached = normalized_cached - normalized_submodule - - if missing_in_cached or extra_in_cached: - print( - "[megatron-core][setup] Dependency mismatch between Megatron-LM-workspace/Megatron-LM/pyproject.toml vs Megatron-LM-workspace/setup.py::CACHED_DEPENDENCIES.", - file=sys.stderr, - ) - if missing_in_cached: - print( - " - Present in Megatron-LM/pyproject.toml (default+dev) but missing from CACHED_DEPENDENCIES:", - file=sys.stderr, - ) - for dep in sorted(missing_in_cached): - print(f" * {dep}", file=sys.stderr) - if extra_in_cached: - print( - " - Present in CACHED_DEPENDENCIES but not in Megatron-LM/pyproject.toml (default+dev):", - file=sys.stderr, - ) - for dep in sorted(extra_in_cached): - print(f" * {dep}", file=sys.stderr) - print( - " Please update CACHED_DEPENDENCIES or the submodule pyproject to keep them in sync.", - file=sys.stderr, - ) - sys.exit(1) - else: - print( - "[megatron-core][setup] Dependency sets are consistent with the submodule pyproject (default+dev).", - file=sys.stderr, - ) - -# Check if the main directory for the megatron.core Python package exists -if os.path.exists(megatron_core_python_package_source_dir): - # Add Python package 'megatron.core' - final_packages.append(megatron_core_package_name) - final_package_dir[megatron_core_package_name] = ( - megatron_core_python_package_source_dir - ) - - # If the Python package is being added, then check if its C++ extension can also be added - # This requires the specific C++ source file to exist - if os.path.exists(megatron_core_cpp_extension_source_file): - megatron_extension = Extension( - "megatron.core.datasets.helpers_cpp", # Name of the extension - sources=[megatron_core_cpp_extension_source_file], # Path to C++ source - language="c++", - extra_compile_args=( - subprocess.check_output(["python3", "-m", "pybind11", "--includes"]) - .decode("utf-8") - .strip() - .split() - ) - + ["-O3", "-Wall", "-std=c++17"], - optional=True, # As in your original setup - ) - final_ext_modules.append(megatron_extension) -# --- End of megatron.core conditional section --- - -setuptools.setup( - name="megatron-core", - version="0.0.0", - packages=final_packages, - package_dir=final_package_dir, - py_modules=["is_megatron_installed"], - ext_modules=final_ext_modules, - # Add in any packaged data. - include_package_data=True, - install_requires=CACHED_DEPENDENCIES, -) diff --git a/examples/configs/grpo_math_1B_sglang.yaml b/examples/configs/grpo_math_1B_sglang.yaml deleted file mode 100644 index 72a645f1ff..0000000000 --- a/examples/configs/grpo_math_1B_sglang.yaml +++ /dev/null @@ -1,28 +0,0 @@ -defaults: grpo_math_1B.yaml - -grpo: - val_batch_size: 128 - -policy: - generation: - backend: "sglang" - sglang_cfg: - # SGLang specific configuration - model_path: ${policy.model_name} - gpus_per_server: 1 - dtype: ${policy.precision} - context_length: 512 # Maximum context length - allow_auto_truncate: true - enable_memory_saver: false - dp_size: 1 - pp_size: 1 - ep_size: 1 - max_running_requests: null - mem_fraction_static: 0.7 - skip_server_warmup: true - # Piecewise CUDA graph currently crashes with "illegal memory access" - # (likely torch 2.10 + sglang incompatibility). Keep disabled until upstream fix. - disable_piecewise_cuda_graph: true - -logger: - wandb_enabled: true diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.yaml new file mode 100644 index 0000000000..b7ba75a475 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.yaml @@ -0,0 +1,66 @@ +defaults: ../../grpo_math_1B.yaml + +grpo: + max_num_steps: 450 + val_batch_size: 128 + +checkpointing: + checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang + +policy: + model_name: Qwen/Qwen2.5-Math-1.5B-Instruct + tokenizer: + name: Qwen/Qwen2.5-Math-1.5B-Instruct + dynamic_batching: + enabled: true + sequence_packing: + enabled: false + make_sequence_length_divisible_by: 1 + generation: + backend: "sglang" + max_new_tokens: 512 + sglang_cfg: + model_path: ${policy.model_name} + dtype: ${policy.precision} + context_length: 512 + allow_auto_truncate: true + dp_size: 1 + pp_size: 1 + ep_size: 1 + random_seed: 42 + max_running_requests: null + mem_fraction_static: 0.6 + skip_server_warmup: true + # Piecewise CUDA graph currently crashes with "illegal memory access" + # (likely torch 2.10 + sglang incompatibility). Keep disabled until upstream fix. + disable_piecewise_cuda_graph: true + disable_cuda_graph: false + # Fault tolerance (RolloutHealthMonitor). Off by default; when enabled, + # a daemon thread health-checks each engine and kills hung actors. + use_fault_tolerance: false + rollout_health_check_interval: 60 + rollout_health_check_timeout: 60 + rollout_health_check_first_wait: 60 + sglang_server: + needs_offload: true + cpu_weight_backup: true + sglang_server_concurrency: 1024 + pause_generation_mode: retract + num_gpus: 2 + num_gpus_per_engine: 2 + sglang_router: {} + +data: + max_input_seq_length: 512 + +logger: + log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang + +cluster: + gpus_per_node: 2 + num_nodes: 1 diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm.yaml new file mode 100644 index 0000000000..14219f0508 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm.yaml @@ -0,0 +1,38 @@ +defaults: ../../grpo_math_1B.yaml + +grpo: + max_num_steps: 450 + val_batch_size: 128 + +checkpointing: + checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm + +policy: + model_name: Qwen/Qwen2.5-Math-1.5B-Instruct + tokenizer: + name: Qwen/Qwen2.5-Math-1.5B-Instruct + dynamic_batching: + enabled: true + sequence_packing: + enabled: false + make_sequence_length_divisible_by: 1 + generation: + max_new_tokens: 512 + vllm_cfg: + max_model_len: 512 + tensor_parallel_size: 2 + +data: + max_input_seq_length: 512 + +logger: + log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm + +cluster: + gpus_per_node: 2 + num_nodes: 1 diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml deleted file mode 100644 index 7d5a30e998..0000000000 --- a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.yaml +++ /dev/null @@ -1,51 +0,0 @@ -defaults: ../../grpo_math_1B.yaml - -grpo: - max_num_steps: 450 - -checkpointing: - checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang - -policy: - model_name: Qwen/Qwen2.5-Math-1.5B-Instruct - tokenizer: - name: Qwen/Qwen2.5-Math-1.5B-Instruct - dynamic_batching: - enabled: true - sequence_packing: - enabled: false - make_sequence_length_divisible_by: 1 - generation: - backend: "sglang" - max_new_tokens: 512 - sglang_cfg: - model_path: ${policy.model_name} - gpus_per_server: 1 - dtype: ${policy.precision} - context_length: 512 - allow_auto_truncate: true - enable_memory_saver: false - dp_size: 1 - pp_size: 1 - ep_size: 1 - max_running_requests: null - mem_fraction_static: 0.5 - skip_server_warmup: true - # Piecewise CUDA graphs fail with CUBLAS_STATUS_EXECUTION_FAILED - # inside Ray worker forks. See unit test Err 4 for details. - disable_piecewise_cuda_graph: true - -data: - max_input_seq_length: 512 - -logger: - log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang - -cluster: - gpus_per_node: 8 - diff --git a/examples/configs/recipes/llm/grpo-qwen3-0.6b-1n8g-sglang.yaml b/examples/configs/recipes/llm/grpo-qwen3-0.6b-1n8g-sglang.yaml deleted file mode 100644 index cc0c447e3f..0000000000 --- a/examples/configs/recipes/llm/grpo-qwen3-0.6b-1n8g-sglang.yaml +++ /dev/null @@ -1,52 +0,0 @@ -defaults: ../../grpo_math_1B.yaml - -grpo: - max_num_steps: 500 - val_batch_size: 128 - -checkpointing: - checkpoint_dir: results/grpo-qwen3-0.6b-1n8g-sglang - -policy: - model_name: Qwen/Qwen3-0.6B - tokenizer: - name: Qwen/Qwen3-0.6B - dynamic_batching: - enabled: true - sequence_packing: - enabled: false - make_sequence_length_divisible_by: 1 - generation: - backend: "sglang" - max_new_tokens: 512 - sglang_cfg: - model_path: ${policy.model_name} - gpus_per_server: 8 - dtype: ${policy.precision} - context_length: 512 - allow_auto_truncate: true - enable_memory_saver: false - dp_size: 1 - pp_size: 1 - ep_size: 1 - max_running_requests: null - mem_fraction_static: 0.7 - skip_server_warmup: true - # Piecewise CUDA graphs fail with CUBLAS_STATUS_EXECUTION_FAILED - # inside Ray worker forks. See unit test Err 4 for details. - disable_piecewise_cuda_graph: true - -data: - max_input_seq_length: 512 - -logger: - log_dir: logs/grpo-qwen3-0.6b-1n8g-sglang - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-qwen3-0.6b-1n8g-sglang - -cluster: - gpus_per_node: 8 - diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e550429ce2..00c817685f 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -68,7 +68,8 @@ run_multi_turn_rollout, ) from nemo_rl.models.generation.interfaces import GenerationInterface -from nemo_rl.models.generation.sglang import SGLangConfig, SGLangGeneration +from nemo_rl.models.generation.sglang.config import SGLangConfig +from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface @@ -578,7 +579,13 @@ def init_vllm(): def init_sglang(): """Initialize SGLang generation workers.""" t0 = time.perf_counter() - pg = SGLangGeneration(cluster=inference_cluster, config=generation_config) + pg = SGLangGeneration( + cluster=inference_cluster, + sglang_cfg=generation_config, + ) + if generation_config["sglang_server"].get("check_weight_update_equal", False): + pg.check_weights(action="snapshot") + pg.check_weights(action="reset") pg.finish_generation() return pg, time.perf_counter() - t0 @@ -704,6 +711,23 @@ def initialize_generation_with_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", @@ -723,8 +747,11 @@ def initialize_generation_with_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) @@ -743,11 +770,39 @@ def initialize_generation_with_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_server"].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.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) + if backend == "sglang" and isinstance(policy_generation, SGLangGeneration): + sglang_cfg_typed = cast(SGLangConfig, generation_config) + check_equal = sglang_cfg_typed["sglang_server"].get( + "check_weight_update_equal", False + ) + if check_equal: + refit_policy_generation( + policy=policy, + policy_generation=policy_generation, + colocated_inference=colocated_inference, + ) + policy_generation.check_weights(action="compare") + policy_generation.finish_generation() + + policy.prepare_for_training() + # Calculate total setup time total_setup_time = time.perf_counter() - setup_start_time worker_init_timing_metrics["total_setup_time_s"] = total_setup_time @@ -1094,6 +1149,44 @@ def _extract_prompt_only_messages(message_logs: list) -> list: return prompt_only_message_logs +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, @@ -1139,19 +1232,12 @@ def refit_policy_generation( ) if isinstance(policy_generation, SGLangGeneration): - sglang_url_to_gpu_uuids = ( - policy_generation.get_sglang_url_to_gpu_uuids() - ) - # Stream weights via HTTP - flush_success = policy_generation.invalidate_kv_cache() - if not flush_success: - print("SGLang KV cache invalidation failed before weight update. ") - futures_train = policy.stream_weights_via_http( - sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids, + 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( @@ -1164,17 +1250,22 @@ def refit_policy_generation( update_success = all(result for result in results if result is not None) else: # update weights through nccl - # SGLang haven't implemented non-colocated inference mode. 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", + ) + else: + futures_train = policy.broadcast_weights_for_collective( + kv_scales=kv_scales ) - 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) + 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 96282ad623..f0acfdb085 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -185,6 +185,54 @@ def get_gpu_id(self): return ray.get_gpu_ids()[0] +def get_reordered_bundle( + pg: PlacementGroup, +) -> tuple[list[int], list[int]]: + """Return bundle indices and GPU IDs sorted by (node_id, gpu_id). + + Uses ``GetGPUIDActor`` to discover the physical GPU ID assigned to each + bundle. + + Returns: + (reordered_bundle_indices, reordered_gpu_ids) + """ + pg_data = placement_group_table(pg) + num_bundles = len(pg_data["bundles"]) + bundle_to_node_ids = pg_data["bundles_to_node_id"] + + info_actors = [] + for i in range(num_bundles): + info_actors.append( + GetGPUIDActor.options( + num_cpus=0.01, + num_gpus=0.01, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=i, + ), + ).remote() + ) + + gpu_ids = ray.get([actor.get_gpu_id.remote() for actor in info_actors]) + for actor in info_actors: + ray.kill(actor) + + bundle_infos = [(i, bundle_to_node_ids[i], gpu_ids[i]) for i in range(num_bundles)] + sorted_infos = sorted(bundle_infos, key=lambda x: (x[1], x[2])) + + reordered_bundle_indices = [info[0] for info in sorted_infos] + reordered_gpu_ids = [gpu_ids[info[0]] for info in sorted_infos] + + for i, info in enumerate(sorted_infos): + actual_idx = info[0] + logger.info( + f" bundle {i:4}, actual_bundle_index: {actual_idx:4}, " + f"node: {info[1]}, gpu: {gpu_ids[actual_idx]}" + ) + + return reordered_bundle_indices, reordered_gpu_ids + + class ResourceInsufficientError(Exception): """Exception raised when the cluster does not have enough resources to satisfy the requested configuration.""" @@ -446,39 +494,10 @@ def _get_sorted_bundle_indices(self) -> Optional[list[int]]: if len(self._node_placement_groups) != 1: return None - pg = self._node_placement_groups[0] - pg_data = placement_group_table(pg) - num_bundles = len(pg_data["bundles"]) - bundle_to_node_ids = pg_data["bundles_to_node_id"] - - # use info actor to get the GPU id - info_actors = [] - for i in range(num_bundles): - info_actors.append( - GetGPUIDActor.options( - num_cpus=0.01, # set both num_cpus and num_gpus to be small values to enable assignment in colocated case - num_gpus=0.01, - resources=None, - scheduling_strategy=PlacementGroupSchedulingStrategy( - placement_group=pg, - placement_group_bundle_index=i, - ), - ).remote() - ) - - gpu_ids = ray.get([actor.get_gpu_id.remote() for actor in info_actors]) - for actor in info_actors: - ray.kill(actor) - - # original index, node_id, gpu_id - bundle_infos = [ - (i, bundle_to_node_ids[i], gpu_ids[i]) for i in range(num_bundles) - ] - pg_reordered_bundle_indices = [ - bundle_info[0] - for bundle_info in sorted(bundle_infos, key=lambda x: (x[1], x[2])) - ] # sort by node_id, then gpu_id - return pg_reordered_bundle_indices + reordered_bundle_indices, _ = get_reordered_bundle( + self._node_placement_groups[0] + ) + return reordered_bundle_indices def shutdown(self) -> bool: """Cleans up and releases all resources associated with this virtual cluster. diff --git a/nemo_rl/models/automodel/setup.py b/nemo_rl/models/automodel/setup.py index e754bd1e74..61238a6b83 100644 --- a/nemo_rl/models/automodel/setup.py +++ b/nemo_rl/models/automodel/setup.py @@ -236,11 +236,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"], @@ -248,10 +250,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 037b4880f5..ea6b62342f 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -243,6 +243,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/__init__.py b/nemo_rl/models/generation/sglang/__init__.py deleted file mode 100644 index 4073c3884b..0000000000 --- a/nemo_rl/models/generation/sglang/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# 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. -from nemo_rl.models.generation.sglang.config import SGLangConfig -from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration - -__all__ = [ - "SGLangConfig", - "SGLangGeneration", -] diff --git a/nemo_rl/models/generation/sglang/config.py b/nemo_rl/models/generation/sglang/config.py index d5baa00022..73227378e2 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 SglangSpecificArgs(TypedDict): """SGLang-specific configuration arguments. @@ -25,7 +44,7 @@ class SglangSpecificArgs(TypedDict): """ model_path: NotRequired[str] - gpus_per_server: NotRequired[int] + # Total number of gpus for rollout random_seed: NotRequired[int] skip_tokenizer_init: NotRequired[bool] disable_cuda_graph: NotRequired[bool] @@ -42,6 +61,17 @@ class SglangSpecificArgs(TypedDict): enable_mixed_chunk: NotRequired[bool] enable_dp_attention: NotRequired[bool] enable_ep_moe: NotRequired[bool] + # FP8 GEMM backend, e.g. {"auto", "cutlass", "triton", "flashinfer_trtllm"}. + fp8_gemm_runner_backend: NotRequired[str] + # MoE runner backend, e.g. {"auto", "deep_gemm", "triton", "cutlass"}. + moe_runner_backend: NotRequired[str] + # MoE all-to-all backend. Newer sglang forks replaced ``enable_ep_moe`` + # with this single switch: one of {"none", "deepep", "mooncake", "mori", + # "ascend_fuseep", "flashinfer"}. Defaults to "none" upstream. + moe_a2a_backend: NotRequired[str] + # DeepEP routing mode (only meaningful when ``moe_a2a_backend == "deepep"``): + # one of {"auto", "normal", "low_latency"}. + deepep_mode: NotRequired[str] enable_torch_compile: NotRequired[bool] torch_compile_max_bs: NotRequired[int] cuda_graph_max_bs: NotRequired[int | None] @@ -52,7 +82,6 @@ class SglangSpecificArgs(TypedDict): triton_attention_reduce_in_fp32: NotRequired[bool] triton_attention_num_kv_splits: NotRequired[int] num_continuous_decode_steps: NotRequired[int] - enable_memory_saver: NotRequired[bool] allow_auto_truncate: NotRequired[bool] attention_backend: NotRequired[str | None] enable_multimodal: NotRequired[bool] @@ -93,10 +122,48 @@ class SglangSpecificArgs(TypedDict): enable_fast_load: NotRequired[bool] # Server warmup skip_server_warmup: NotRequired[bool] + # Fault tolerance + 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] + + +class SGLangServer(TypedDict): + # needs_offload true --> enable_memory_saver true + needs_offload: bool + # for testing purpose. memory_saver + cpu_weight_backup: bool + sglang_server_concurrency: int + # for pause/continue gen ("retract" or "kill"); required in YAML. + pause_generation_mode: str + # When true, refit bookends every weight update with a check_weights + # snapshot + reset (pre-refit) and compare (post-refit) to assert that + # streamed weights actually land on the engine. + check_weight_update_equal: NotRequired[bool] + # total num gpus for inference + num_gpus: NotRequired[int] + num_gpus_per_engine: NotRequired[int] + # "ipc" -> CUDA-IPC over 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] + + +class SGLangRouter(TypedDict): + sglang_router_ip: NotRequired[str] + sglang_router_port: NotRequired[int] + router_policy: NotRequired[str] + use_distributed_post: NotRequired[bool] + sglang_router_request_timeout_secs: NotRequired[int] class SGLangConfig(GenerationConfig): """Configuration for SGLang runtime.""" sglang_cfg: SglangSpecificArgs + sglang_server: SGLangServer + sglang_router: SGLangRouter sglang_kwargs: NotRequired[dict[str, Any]] 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..30fb49f0e9 --- /dev/null +++ b/nemo_rl/models/generation/sglang/fault_tolerance.py @@ -0,0 +1,186 @@ +import logging +import threading + +import ray + +logger = logging.getLogger(__name__) + +from nemo_rl.models.generation.sglang.config import SGLangConfig + + +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..5175efa4a3 --- /dev/null +++ b/nemo_rl/models/generation/sglang/mxfp8_setup.py @@ -0,0 +1,452 @@ +# 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, + TARGET_MXFP8_BLOCK_SIZE, + 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_copied_utils.py b/nemo_rl/models/generation/sglang/sglang_copied_utils.py deleted file mode 100644 index aa9eafea01..0000000000 --- a/nemo_rl/models/generation/sglang/sglang_copied_utils.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# 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. -# ============================================================================== -"""Standalone utility functions copied from the SGLang project. - -This module contains utility functions that were originally part of the SGLang -repository (https://github.com/sgl-project/sglang). They have been copied here -to avoid requiring sglang as a runtime dependency for weight refitting functionality. - -IMPORTANT: This module should NOT contain any imports from the sglang package. -All functions are standalone and self-contained. - -Each function includes a permalink to its original source in the SGLang repository. -These functions were copied from sglang version 0.5.2. -""" - -import io -from multiprocessing.reduction import ForkingPickler -from typing import Callable, Union - -import pybase64 -import torch -from torch.multiprocessing import reductions - - -class MultiprocessingSerializer: # pragma: no cover - """Serialize/deserialize Python objects using ForkingPickler for IPC. - - This class enables serialization of objects (including CUDA tensors with IPC - handles) for transfer between processes via HTTP or other mechanisms. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/utils.py#L589-L623 - """ - - @staticmethod - def serialize(obj, output_str: bool = False): - """Serialize a Python object using ForkingPickler. - - Args: - obj: The object to serialize. - output_str (bool): If True, return a base64-encoded string instead of raw bytes. - - Returns: - bytes or str: The serialized object. - """ - buf = io.BytesIO() - ForkingPickler(buf).dump(obj) - buf.seek(0) - output = buf.read() - - if output_str: - # Convert bytes to base64-encoded string - output = pybase64.b64encode(output).decode("utf-8") - - return output - - @staticmethod - def deserialize(data): - """Deserialize a previously serialized object. - - Args: - data (bytes or str): The serialized data, optionally base64-encoded. - - Returns: - The deserialized Python object. - """ - if isinstance(data, str): - # Decode base64 string to bytes - data = pybase64.b64decode(data, validate=True) - - return ForkingPickler.loads(data) - - -def monkey_patch_torch_reductions(): # pragma: no cover - """Monkey patch torch multiprocessing reductions to use GPU UUIDs. - - This patch modifies PyTorch's CUDA tensor IPC mechanism to use GPU UUIDs - instead of device indices. This enables proper weight transfer between - processes that may have different CUDA_VISIBLE_DEVICES configurations. - - The patch is idempotent - calling it multiple times is safe. - - This is a workaround before PyTorch https://github.com/pytorch/pytorch/pull/149248 - is merged and released. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L20-L33 - """ - if hasattr(reductions, "_reduce_tensor_original"): - return - - reductions._reduce_tensor_original = reductions.reduce_tensor - reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor - - reductions.reduce_tensor = _reduce_tensor_modified - reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified - - reductions.init_reductions() - - -# The signature has not been changed for years, and we will not need this when -# the next version is released, so it looks safe to use a constant. -# Original source (sglang v0.5.2): -# https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L36 -_REDUCE_TENSOR_ARG_DEVICE_INDEX = 6 - - -def _reduce_tensor_modified(*args, **kwargs): # pragma: no cover - """Modified reduce_tensor that stores GPU UUID instead of device index. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L39-L43 - """ - output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs) - output_args = _modify_tuple( - output_args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_to_uuid - ) - return output_fn, output_args - - -def _rebuild_cuda_tensor_modified(*args): # pragma: no cover - """Modified rebuild_cuda_tensor that accepts GPU UUID or device index. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L46-L48 - """ - args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_from_maybe_uuid) - return reductions._rebuild_cuda_tensor_original(*args) - - -def _device_to_uuid(device: int) -> str: # pragma: no cover - """Convert a device index to its UUID string. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L51-L52 - """ - return str(torch.cuda.get_device_properties(device).uuid) - - -def _device_from_maybe_uuid( - device_maybe_uuid: Union[int, str], -) -> int: # pragma: no cover - """Convert a device UUID string or index to a device index. - - Args: - device_maybe_uuid: Either an integer device index or a UUID string. - - Returns: - The integer device index. - - Raises: - Exception: If the UUID doesn't match any available device. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L55-L65 - """ - if isinstance(device_maybe_uuid, int): - return device_maybe_uuid - - if isinstance(device_maybe_uuid, str): - for device in range(torch.cuda.device_count()): - if str(torch.cuda.get_device_properties(device).uuid) == device_maybe_uuid: - return device - raise Exception("Invalid device_uuid=" + device_maybe_uuid) - - raise Exception(f"Unknown type: {device_maybe_uuid=}") - - -def _modify_tuple(t, index: int, modifier: Callable): # pragma: no cover - """Create a new tuple with one element modified by a function. - - Original source (sglang v0.5.2): - https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/patch_torch.py#L68-L69 - """ - return *t[:index], modifier(t[index]), *t[index + 1 :] diff --git a/nemo_rl/models/generation/sglang/sglang_generation.py b/nemo_rl/models/generation/sglang/sglang_generation.py index 85122779ee..bf6ec427bd 100644 --- a/nemo_rl/models/generation/sglang/sglang_generation.py +++ b/nemo_rl/models/generation/sglang/sglang_generation.py @@ -1,384 +1,1122 @@ -# 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. - +import asyncio import logging import os -from typing import ( - Any, - Optional, - Union, -) +from typing import Any, AsyncGenerator, Optional -import numpy as np import ray - -from nemo_rl.distributed.batched_data_dict import BatchedDataDict, SlicedDataDict -from nemo_rl.distributed.named_sharding import NamedSharding -from nemo_rl.distributed.virtual_cluster import RayVirtualCluster -from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup +import torch +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.ray_actor_environment_registry import SGLANG_EXECUTABLE +from nemo_rl.distributed.virtual_cluster import ( + RayVirtualCluster, + get_reordered_bundle, +) +from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches from nemo_rl.models.generation.interfaces import ( GenerationDatumSpec, GenerationInterface, GenerationOutputSpec, + 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 RouterActor +from nemo_rl.models.generation.sglang.sglang_worker import SGLangGenerationWorker +from nemo_rl.models.generation.sglang.utils.async_utils import AsyncLoopThread +from nemo_rl.models.generation.sglang.utils.http_utils import ( + HttpClient, + init_http_client, +) +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 -# Global thresholds for top_k and top_p validation. -# While top-k/p are not supported, these values allow for token filtering while the logprobs should be compatible. -# See https://github.com/NVIDIA-NeMo/RL/issues/69 and https://github.com/NVIDIA-NeMo/RL/issues/237 for more details. -TOP_K_THRESHOLD = 8000 # Allow top_k >= 8000 (effectively no filtering) -TOP_P_THRESHOLD = 0.99 # Allow top_p >= 0.99 (close to 1.0) +logging.getLogger("httpx").setLevel(logging.WARNING) +logging.getLogger("httpcore").setLevel(logging.WARNING) logger = logging.getLogger(__name__) class SGLangGeneration(GenerationInterface): + """The class to run rollout and convert rollout data to training data. + + This class owns the full rollout server topology: the placement group, + the router subprocess, and every ``SGLangGenerationWorker`` Ray actor. + The former ``ServerGroup`` dataclass has been folded in so there is a + single source of truth for engine state. + + TODO: one sglang router(router ip, router port) --> different server group(eg: PD, different tp size, ...; each server group multiple engines/servsers with same settings) + router + [[p, ..., p] [d, ..., d]] or router + [[tp = 2, ... tp = 2], ..., [tp = 8, ..., tp = 8]] + """ + def __init__( self, cluster: RayVirtualCluster, - config: SGLangConfig, - name_prefix: str = "sglang_policy", - workers_per_node: Optional[Union[int, list[int]]] = None, + sglang_cfg: SGLangConfig, ): - """Initialize a SGLang policy with distributed workers. + self.cluster = cluster + self.sglang_cfg = sglang_cfg + self._health_monitor = None + self._async_loop: AsyncLoopThread | None = AsyncLoopThread() + self._http_client: HttpClient | None = None + + pgs = cluster._init_placement_groups( + strategy="PACK", + use_unified_pg=True, + ) + self.pg = pgs[0] + self.pg_reordered_bundle_indices, self.pg_reordered_gpu_ids = ( + get_reordered_bundle(self.pg) + ) + self._http_client = init_http_client(sglang_cfg) + + # --- Engine topology (formerly ``ServerGroup``) ------------------ + gpus_per_engine = sglang_cfg["sglang_server"]["num_gpus_per_engine"] + num_gpus_per_node = cluster.num_gpus_per_node + num_gpu_per_engine_local = min(gpus_per_engine, num_gpus_per_node) + num_engines = ( + sglang_cfg["sglang_server"]["num_gpus"] // num_gpu_per_engine_local + ) + + self.num_gpus_per_engine: int = gpus_per_engine + self.num_gpus_per_node: int = num_gpus_per_node + self.all_engines: list = [None] * num_engines + self.num_new_engines: int = 0 + self.rank_offset: int = 0 + self.gpu_offset: int = 0 + self.needs_offload: bool = sglang_cfg["sglang_server"]["needs_offload"] + self.pause_generation_mode: str = sglang_cfg["sglang_server"][ + "pause_generation_mode" + ] + self.model_path: str | None = sglang_cfg["sglang_cfg"]["model_path"] + + # --- 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 + # ``router_ip`` / ``router_port`` kwargs in ``init.remote(...)``. + router_ip, router_port, router_actor = _start_router(sglang_cfg) + self.router_ip: str = router_ip + self.router_port: int = router_port + # Only set when ``_start_router`` actually spawned the router (i.e. + # sglang_router_ip was not already configured). Kept so ``shutdown`` + # can terminate it cleanly. + self._router_actor: ray.actor.ActorHandle | None = router_actor + + # --- Start engines ----------------------------------------------- + init_handles, _ = self._start_engines({}) + if init_handles: + ray.get(init_handles) + + 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``) + # ------------------------------------------------------------------ + @property + def nodes_per_engine(self) -> int: + return max(1, self.num_gpus_per_engine // self.num_gpus_per_node) + + @property + def engines(self) -> list: + """Node-0 engines only (one entry per logical engine). + + For multi-node TP, ``all_engines`` contains ``nodes_per_engine`` + consecutive actors per logical engine; this slice returns just the + node-0 representative for each. + """ + return self.all_engines[:: self.nodes_per_engine] + + @property + def rollout_engines(self) -> list: + """Alias for ``engines`` — node-0 engines across all servers / models.""" + return self.engines + + @property + def engine_gpu_counts(self) -> list[int]: + """Per-engine GPU count, parallel to ``engines``.""" + return [self.num_gpus_per_engine for _ in self.engines] + + @property + def engine_gpu_offsets(self) -> list[int]: + return [ + self.gpu_offset + j * self.num_gpus_per_engine + for j in range(len(self.engines)) + ] - SGLang server manages TP/PP internally, but we still need to: - 1. Manage data parallel distribution across multiple servers - 2. Assign GPU bundles to each server + # ------------------------------------------------------------------ + # Engine lifecycle (formerly ``ServerGroup.start_engines`` / ``recover``) + # ------------------------------------------------------------------ + def _start_engines( + self, port_cursors: dict[int, int] | None = None + ) -> tuple[list, dict[int, int]]: + """Create Ray actors, allocate ports, and fire ``engine.init()`` without waiting. - Each server will see logical GPUs 0-N (via CUDA_VISIBLE_DEVICES set by Ray), - so we just need to tell SGLang how many GPUs to use (tp_size). + Returns ``(init_handles, port_cursors)`` where *init_handles* is a list + of Ray ObjectRefs and *port_cursors* maps node index -> next free port. """ - # Store config - self.cfg = config - self.sglang_cfg = config["sglang_cfg"] - - gpus_per_server = self.sglang_cfg.get("gpus_per_server", None) - if gpus_per_server is None: - raise ValueError("gpus_per_server must be set in SGLangConfig.sglang_cfg.") - - # Calculate number of servers based on available resources - total_gpus = cluster.world_size() - num_servers = total_gpus // gpus_per_server - - if num_servers == 0: - raise ValueError( - f"Not enough GPUs. Need at least {gpus_per_server} GPUs per server, " - f"but only have {total_gpus} GPUs total." + 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 + reordered_gpu_ids = self.pg_reordered_gpu_ids + + local_all_engines = [] + for i in range(len(self.all_engines)): + if self.all_engines[i] is not None: + continue + + global_rank = self.rank_offset + i + num_gpus = 0.2 + num_cpus = num_gpus + + gpu_index = self.gpu_offset + i * num_gpu_per_engine + base_gpu_id = int(reordered_gpu_ids[gpu_index]) + + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[gpu_index], ) - if total_gpus % gpus_per_server != 0: - logger.warning( - f"[WARNING] Total GPUs ({total_gpus}) is not divisible by GPUs per server ({gpus_per_server}). " - f"Will use {num_servers} servers, leaving {total_gpus % gpus_per_server} GPUs unused." + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { + key: os.environ.get(key, default_val) + for key, default_val in { + "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "false", + "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true", + "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT": "true", + "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false", + "SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "false", + }.items() + } + env_vars["NCCL_CUMEM_ENABLE"] = "0" + + # 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). + global_cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None) + if global_cvd: + env_vars["CUDA_VISIBLE_DEVICES"] = global_cvd + + actor_options = { + "num_cpus": num_cpus, + "num_gpus": num_gpus, + "scheduling_strategy": scheduling_strategy, + "runtime_env": { + "py_executable": SGLANG_EXECUTABLE, + "env_vars": env_vars, + **get_nsight_config_if_pattern_matches("sglang_generation_worker"), + }, + } + init_args = (self.num_gpus_per_node, self.sglang_cfg) + init_kwargs = { + "rank": global_rank, + "base_gpu_id": base_gpu_id, + "num_gpus_per_engine": self.num_gpus_per_engine, + } + + # Create worker actor directly — sglang_worker.py uses lazy imports + # so it's importable in SYSTEM env; the actor runs in sglang env. + engine = SGLangGenerationWorker.options(**actor_options).remote( + *init_args, **init_kwargs ) - self.dp_size = num_servers - self.gpus_per_server = gpus_per_server - - # Create sharding annotations - # Even though SGLang manages TP internally, we include it in the layout to support - # RayWorkerGroup's worker management (which creates one worker per GPU bundle). - # The TP dimension becomes a "free axis" in run_all_workers_sharded_data, ensuring - # only the primary workers (TP rank 0) are called. - total_workers = num_servers * gpus_per_server - self.sharding_annotations = NamedSharding( - layout=np.arange(total_workers).reshape(num_servers, gpus_per_server), - names=["data_parallel", "tensor_parallel"], - ) + local_all_engines.append((global_rank, engine)) + self.all_engines[i] = engine - # Initialize placement groups - # For SGLang, we use PACK strategy to keep bundles together - # colocated is always at top level, not in sglang_cfg - strategy = None if self.cfg["colocated"]["enabled"] else "PACK" - cluster._init_placement_groups( - strategy=strategy, - use_unified_pg=False, # SGLang servers don't need cross-node model parallelism + self.num_new_engines = len(local_all_engines) + + if self.num_new_engines == 0: + return [], port_cursors + + base_port = max(port_cursors.values()) if port_cursors else 15000 + 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, ) - # Create worker builder for SGLangGenerationWorker - worker_cls = ( - "nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker" + init_handles = [ + engine.init.remote( + **(addr_and_ports[rank]), + router_ip=self.router_ip, + router_port=self.router_port, + ) + for rank, engine in local_all_engines + ] + return init_handles, port_cursors + + 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" ) - worker_builder = RayWorkerBuilder(worker_cls, config) - - env_vars = {} - global_cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None) - if global_cvd: - # Explicitly pass CUDA_VISIBLE_DEVICES to workers via env_vars - # This ensures all workers see the same global value, even though - env_vars["CUDA_VISIBLE_DEVICES"] = global_cvd - - # Allocate bundles for each server - # Each server gets consecutive bundles - bundle_indices_list = self._allocate_bundles_for_servers( - cluster, num_servers, gpus_per_server + + if self.needs_offload and dead_indices: + new_engines = [self.all_engines[i] for i in dead_indices] + ray.get([engine.release_memory_weights.remote() for engine in new_engines]) + ray.get( + [ + engine.release_memory_kv_cache_and_cuda_graph.remote() + for engine in new_engines + ] + ) + ray.get([engine.resume_memory_weights.remote() for engine in new_engines]) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + 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, ) - # Create worker group with explicit bundle allocation - self.worker_group = RayWorkerGroup( - cluster, - worker_builder, - name_prefix=name_prefix, - bundle_indices_list=bundle_indices_list, - sharding_annotations=self.sharding_annotations, - env_vars=env_vars, + def offload_weights(self): + if not self.needs_offload: + return + + handles = [ + engine.release_memory_weights.remote() + for engine in self.engines + if engine is not None + ] + if handles: + ray.get(handles) + + def offload_kv(self): + if not self.needs_offload: + return + + handles = [ + engine.release_memory_kv_cache_and_cuda_graph.remote() + for engine in self.engines + if engine is not None + ] + if handles: + ray.get(handles) + + def onload_weights(self): + if not self.needs_offload: + return + + handles = [ + engine.resume_memory_weights.remote() + for engine in self.engines + if engine is not None + ] + if handles: + ray.get(handles) + + def onload_kv(self): + if not self.needs_offload: + return + + handles = [ + engine.resume_memory_kv_cache_and_cuda_graph.remote() + for engine in self.engines + if engine is not None + ] + if handles: + ray.get(handles) + + def recover_updatable_engines(self): + """Restart any dead rollout engines and update ``num_new_engines``. + + for weight-update detection. + """ + 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, ) - # Verify data parallel size matches - assert self.dp_size == self.worker_group.dp_size, ( - f"Data parallel size mismatch. Expected {self.dp_size}, got {self.worker_group.dp_size}" + def clear_updatable_num_new_engines(self): + # when fault tolerance is not enabled, we need to manually clear num_new_engines after update_weights + self.num_new_engines = 0 + + def check_weights(self, action: str): + """All node-0 engines across all servers / models.""" + return ray.get( + [ + engine.check_weights.remote(action=action) + for engine in self.engines + if engine is not None + ] ) - # Used to track the round-robin selection of worker groups for generate_async - self.current_generate_dp_shard_idx = 0 + def pause_generation(self, mode: Optional[str] = None) -> None: + """Pause generation on every node-0 engine. - def _allocate_bundles_for_servers( + Args: + mode: Pause mode override. When ``None`` (default), the mode + configured in ``sglang_server.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 ``flush_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, - cluster: RayVirtualCluster, - num_servers: int, - gpus_per_server: int, - ) -> list[tuple[int, list[int]]]: - """Allocate GPU bundles to each SGLang server. + *, + 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 Megatron-side 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 freshly + loaded 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: + try: + ray.get([e.shutdown.remote() for e in engines]) + except Exception as e: + logger.warning(f"Engine shutdown failed: {e}") + ok = False + self.all_engines = [None] * len(self.all_engines) + + if self._router_actor is not None: + try: + ray.get(self._router_actor.stop.remote()) + ray.kill(self._router_actor) + except Exception as e: + logger.warning(f"Router terminate failed: {e}") + ok = False + self._router_actor = None + + if self._http_client is not None: + try: + if self._async_loop is not None: + self._async_loop.run(self._http_client.aclose()) + else: + self._http_client.shutdown() + except Exception as e: + logger.warning(f"HTTP client shutdown failed: {e}") + ok = False + self._http_client = None + + if self._async_loop is not None: + try: + self._async_loop.close() + except Exception as e: + logger.warning(f"AsyncLoopThread close failed: {e}") + ok = False + self._async_loop = None + + return ok + + def __del__(self) -> None: + self.shutdown() - Each server gets consecutive bundles within the same placement group (node). - Ray will automatically set CUDA_VISIBLE_DEVICES so each server sees logical GPUs 0, 1, 2, ..., gpus_per_server-1. + def _merge_stop_strings(self, batch_stop_strings) -> list[list[str]]: + """Merge stop strings from config and batch. Args: - cluster: The Ray virtual cluster - num_servers: Total number of SGLang servers to create - gpus_per_server: Number of GPUs each server needs + batch_stop_strings: List of stop strings from batch (one per sample) Returns: - List of (node_idx, [bundle_indices]) tuples for each server + List of merged stop strings (one per sample) """ - placement_groups = cluster.get_placement_groups() + stop_set: set[str] = set() - if not placement_groups: - raise ValueError("No placement groups available in the cluster") + # Add stop strings from config + if self.sglang_cfg.get("stop_strings"): + stop_set.update(self.sglang_cfg["stop_strings"]) - bundle_indices_list = [] + # Merge stop strings from batch + merged_stop_strings = [] + for sample_ss in batch_stop_strings: + sample_stop_set = stop_set.copy() + if sample_ss: + if isinstance(sample_ss, str): + sample_stop_set.add(sample_ss) + elif isinstance(sample_ss, list): + sample_stop_set.update(sample_ss) - # Each server's bundles must be within the same placement group (node) - server_idx = 0 - for pg_idx, pg in enumerate(placement_groups): - if pg.bundle_count == 0: - continue + merged_stop_strings.append(list(sample_stop_set)) - # Calculate how many servers can fit in this placement group - num_servers_in_pg = pg.bundle_count // gpus_per_server + return merged_stop_strings - # Allocate servers within this placement group - for local_server_idx in range(num_servers_in_pg): - if server_idx >= num_servers: - break + def _build_sampling_params( + self, + *, + greedy: bool, + max_new_tokens: int, + stop_strings: list[str] | None = None, + ) -> dict[str, Any]: + """Build sampling parameters dictionary for SGLang API. - # Calculate which bundles this server gets (consecutive within the PG) - start_bundle = local_server_idx * gpus_per_server - server_bundles = list( - range(start_bundle, start_bundle + gpus_per_server) - ) + Args: + greedy: Whether to use greedy decoding (temperature=0.0) + max_new_tokens: Max new tokens for this sample (already clamped by caller + against ``context_length - input_length``). + stop_strings: Merged stop strings for this sample. - # Each server gets a tuple of (node_idx, [local_bundle_indices]) - bundle_indices_list.append((pg_idx, server_bundles)) - server_idx += 1 + Returns: + Dictionary of sampling parameters compatible with SGLang API. + """ + temperature = 0.0 if greedy else self.sglang_cfg["temperature"] + top_k_cfg = self.sglang_cfg.get("top_k") + top_k_val = 1 if greedy else (top_k_cfg if top_k_cfg is not None else -1) - if server_idx >= num_servers: - break + # Build sampling params dict first, then patch in optional fields so we + # never reference ``sampling_params`` before it's bound. + sampling_params: dict[str, Any] = { + "temperature": temperature, + "top_p": self.sglang_cfg.get("top_p", 1.0), + "max_new_tokens": max_new_tokens, + "no_stop_trim": True, + "spaces_between_special_tokens": False, + } - if len(bundle_indices_list) < num_servers: - total_available = sum( - pg.bundle_count // gpus_per_server - for pg in placement_groups - if pg.bundle_count > 0 - ) - raise ValueError( - f"Not enough bundles to allocate all {num_servers} servers. " - f"Only {total_available} servers can be allocated " - f"(each server needs {gpus_per_server} GPUs)." - ) + if top_k_val != -1: + sampling_params["top_k"] = top_k_val - return bundle_indices_list + stop_token_ids = self.sglang_cfg.get("stop_token_ids") + if stop_token_ids is not None: + sampling_params["stop_token_ids"] = stop_token_ids - def init_collective( - self, ip: str, port: int, world_size: int, *, train_world_size: int - ) -> list[ray.ObjectRef]: - """Initialize the collective communication. + if stop_strings is not None and len(stop_strings) > 0: + sampling_params["stop"] = stop_strings - TODO: if weight updates via NCCL are needed in the future. - """ - return [] + return sampling_params + @wrap_with_nvtx_name("sglang_genertion/generate") def generate( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False ) -> BatchedDataDict[GenerationOutputSpec]: - """Generate a batch of data using SGLang.""" - assert isinstance(data, BatchedDataDict), ( - f"data must be a BatchedDataDict, got type: {type(data)}" - ) - assert "input_ids" in data and "input_lengths" in data, ( - "input_ids and input_lengths are required in data for SGLang generation" - ) + """Generate a batch of data using Sglang generation. - # Shard the data across the data parallel servers - dp_size = self.sharding_annotations.get_axis_size("data_parallel") - sharded_data: list[SlicedDataDict] = data.shard_by_batch_size( - dp_size, allow_uneven_shards=True - ) - future_bundle = self.worker_group.run_all_workers_sharded_data( - "generate", - data=sharded_data, - in_sharded_axes=["data_parallel"], - replicate_on_axes=None, - output_is_replicated=None, - common_kwargs={"greedy": greedy}, - ) + Args: + data: BatchedDataDict containing input_ids and input_lengths tensors + greedy: Whether to use greedy decoding instead of sampling + + Returns: + BatchedDataDict conforming to GenerationOutputSpec: + - output_ids: input + generated token IDs with proper padding + - logprobs: Log probabilities for tokens + - generation_lengths: Lengths of each response + - unpadded_sequence_lengths: Lengths of each input + generated sequence + """ + # Handle empty input case + if len(data["input_ids"]) == 0: + # Return empty BatchedDataDict with all required fields + return BatchedDataDict[GenerationOutputSpec]( + { + "output_ids": torch.zeros((0, 0), dtype=torch.long), + "logprobs": torch.zeros((0, 0), dtype=torch.float), + "generation_lengths": torch.zeros(0, dtype=torch.long), + "unpadded_sequence_lengths": torch.zeros(0, dtype=torch.long), + "truncated": torch.zeros(0, dtype=torch.bool), + } + ) + + input_ids = data["input_ids"] + input_lengths = data["input_lengths"] + batch_stop_strings: list[list[str]] = data.get("stop_strings", []) + stop_strings = self._merge_stop_strings(batch_stop_strings) + + batch_size = len(input_lengths) + padded_input_length = input_ids.size(1) + context_length = self.sglang_cfg["sglang_cfg"]["context_length"] + + # verify inputs have correct padding + verify_right_padding(data, pad_value=self.sglang_cfg["_pad_token_id"]) + + # Build per-sample requests (each sample gets its own sampling params because + # max_new_tokens is adjusted against the per-sample input length). + sample_requests: list[tuple[int, dict[str, Any], list[int]]] = [] + skip_results: set[int] = set() + skip_max_length = 0 + for i in range(batch_size): + input_length = input_lengths[i].item() + valid_input_ids = input_ids[i, :input_length].tolist() + + if context_length is not None: + max_new_tokens = min( + self.sglang_cfg["max_new_tokens"], context_length - input_length + ) + else: + max_new_tokens = self.sglang_cfg["max_new_tokens"] + max_new_tokens = max(0, max_new_tokens) - # Get results from the workers - results = self.worker_group.get_all_worker_results(future_bundle) + if max_new_tokens == 0: + skip_results.add(i) + skip_max_length = max(skip_max_length, input_length) + continue - # Combine results from all servers - combined: BatchedDataDict[GenerationOutputSpec] = BatchedDataDict.from_batches( - results, pad_value_dict={"output_ids": self.cfg["_pad_token_id"]} + sample_sampling_params = self._build_sampling_params( + greedy=greedy, + max_new_tokens=max_new_tokens, + stop_strings=stop_strings[i] if i < len(stop_strings) else None, + ) + sample_requests.append((i, sample_sampling_params, valid_input_ids)) + + # Dispatch concurrently to the SGLang router with bounded concurrency. + # Max concurrency = per-engine concurrency * number of engines. + sglang_server_cfg = self.sglang_cfg["sglang_server"] + max_concurrency = ( + sglang_server_cfg["sglang_server_concurrency"] + * sglang_server_cfg["num_gpus"] + // sglang_server_cfg["num_gpus_per_engine"] ) - # Verify the output has all required fields - required_keys = [ - "output_ids", - "generation_lengths", - "unpadded_sequence_lengths", - "logprobs", - ] - missing_keys = [key for key in required_keys if key not in combined] - if missing_keys: - raise ValueError( - f"Missing required keys for GenerationOutputSpec: {missing_keys}" + router_ip = self.router_ip + router_port = self.router_port + + semaphore = asyncio.Semaphore(max_concurrency) + + async def _bounded_generate_one_sample( + idx: int, sp: dict[str, Any], ids: list[int] + ): + async with semaphore: + return await generate_one_sample( + router_ip, + router_port, + sp, + ids, + idx, + http_client=self._http_client, + ) + + async def _dispatch_all() -> dict[int, tuple[list[int], list[float], bool]]: + gathered = await asyncio.gather( + *( + _bounded_generate_one_sample(idx, sp, ids) + for idx, sp, ids in sample_requests + ) ) + # generate_one_sample returns (index, tokens, logprobs, truncated). + # Re-key by the original sample index so downstream code can look up + # results directly without sorting. + return { + returned_idx: (new_tokens, new_logprobs, is_truncated) + for returned_idx, new_tokens, new_logprobs, is_truncated in gathered + } + + router_results: dict[int, tuple[list[int], list[float], bool]] = ( + self._async_loop.run(_dispatch_all()) if sample_requests else {} + ) - return combined + # Process the outputs - preserve the original input padding structure. + pad_token_id = self.sglang_cfg["_pad_token_id"] + output_ids_list: list[torch.Tensor] = [] + logprobs_list: list[torch.Tensor] = [] + generation_lengths_list: list[int] = [] + unpadded_sequence_lengths_list: list[int] = [] + truncated_list: list[bool] = [] + + # First pass: compute total_length as the max over all samples of + # (input_length + generation_length). Skipped samples contribute only + # their input_length (already tracked in ``skip_max_length``). + max_length = skip_max_length + for returned_idx, (returned_tokens, _, _) in router_results.items(): + sample_input_length = input_lengths[returned_idx].item() + max_length = max(max_length, sample_input_length + len(returned_tokens)) + total_length = max(max_length, padded_input_length) + + # Second pass: materialize the output tensors, using a single set of + # local variable names (``generation_length`` / ``unpadded_length`` are + # always Python ints; tensor promotion happens only at the final stack). + for i in range(batch_size): + input_length = input_lengths[i].item() + full_output = torch.full( + (total_length,), pad_token_id, dtype=input_ids.dtype + ) + full_logprobs = torch.zeros(total_length, dtype=torch.float32) + full_output[:input_length] = input_ids[i][:input_length] - def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: - pass + if i in skip_results: + generation_length = 0 + is_truncated = False + else: + new_tokens, new_logprobs, is_truncated = router_results[i] + generation_length = len(new_tokens) + if new_tokens: + full_output[input_length : input_length + generation_length] = ( + torch.tensor(new_tokens, dtype=input_ids.dtype) + ) + if new_logprobs: + full_logprobs[input_length : input_length + len(new_logprobs)] = ( + torch.tensor(new_logprobs, dtype=torch.float32) + ) + + unpadded_length = input_length + generation_length + output_ids_list.append(full_output) + logprobs_list.append(full_logprobs) + generation_lengths_list.append(generation_length) + unpadded_sequence_lengths_list.append(unpadded_length) + truncated_list.append(bool(is_truncated)) + + return_data = BatchedDataDict[GenerationOutputSpec]( + { + "output_ids": torch.stack(output_ids_list), + "logprobs": torch.stack(logprobs_list), + "generation_lengths": torch.tensor( + generation_lengths_list, dtype=torch.long + ), + "unpadded_sequence_lengths": torch.tensor( + unpadded_sequence_lengths_list, dtype=torch.long + ), + "truncated": torch.tensor(truncated_list, dtype=torch.bool), + } + ) - def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]: - return [] + return return_data - def update_weights_from_collective(self) -> list[ray.ObjectRef]: - return [] + async def generate_async( + self, + data: BatchedDataDict[GenerationDatumSpec], + greedy: bool = False, + ) -> AsyncGenerator[tuple[int, BatchedDataDict[GenerationOutputSpec]], None]: + """Generate a single sample using SGLang, yielding the result when ready. - def get_sglang_server_urls(self) -> list[str]: - """Get base URLs of all SGLang servers. + Args: + data: BatchedDataDict with input_ids and input_lengths (batch_size must be 1) + greedy: Whether to use greedy decoding instead of sampling - Returns: - List of base URLs (e.g., ["http://localhost:30000", "http://localhost:30001"]) + Yields: + Tuple of (original_index, BatchedDataDict conforming to GenerationOutputSpec) """ - if not self.worker_group or not self.worker_group.workers: - raise RuntimeError("Worker group is not initialized") - - # Get base URLs from all workers (only primary workers, TP rank 0) - # Use run_rank_0_only_axes to only get URLs from primary workers - futures = self.worker_group.run_all_workers_single_data( - "get_base_url", - run_rank_0_only_axes=["tensor_parallel"], + # Handle empty input case + if len(data["input_ids"]) == 0: + return + + verify_right_padding(data, pad_value=self.sglang_cfg["_pad_token_id"]) + + input_ids_batch = data["input_ids"] + input_lengths_batch = data["input_lengths"] + batch_size = input_ids_batch.shape[0] + + # Restrict to single-sample batches, matching the vLLM async contract. + assert batch_size == 1, ( + f"generate_async is restricted to handle only single samples, " + f"but received batch_size={batch_size}. Please handle batching outside this method." ) - urls = ray.get(futures) - # Filter out None values and return unique URLs - return list(set(url for url in urls if url is not None)) - def get_sglang_url_to_gpu_uuids(self) -> dict[str, list[str]]: - """Get mapping from SGLang server URL to list of GPU UUIDs it uses. + sample_idx = 0 + input_length = input_lengths_batch[sample_idx].item() + original_input_ids_single_row = input_ids_batch[sample_idx] + device = original_input_ids_single_row.device + dtype = original_input_ids_single_row.dtype + pad_token_id = self.sglang_cfg["_pad_token_id"] + + # Clamp max_new_tokens against the per-sample remaining context window, + # mirroring the logic in ``generate``. + context_length = self.sglang_cfg["sglang_cfg"].get("context_length") + if context_length is not None: + max_new_tokens = min( + self.sglang_cfg["max_new_tokens"], context_length - input_length + ) + else: + max_new_tokens = self.sglang_cfg["max_new_tokens"] + max_new_tokens = max(0, max_new_tokens) + + # Short-circuit when there is no room left in the context window. Yield + # a pure-input row (generation_length=0, truncated=False) without + # touching the SGLang router. + if max_new_tokens == 0: + output_ids_single_item_batched = original_input_ids_single_row[ + :input_length + ].unsqueeze(0) + logprobs_single_item = torch.zeros( + (1, input_length), dtype=torch.float32, device=device + ) + empty_result = BatchedDataDict[GenerationOutputSpec]( + { + "output_ids": output_ids_single_item_batched, + "logprobs": logprobs_single_item, + "generation_lengths": torch.tensor( + [0], dtype=torch.long, device=device + ), + "unpadded_sequence_lengths": torch.tensor( + [input_length], dtype=torch.long, device=device + ), + "truncated": torch.tensor([False], dtype=torch.bool, device=device), + } + ) + yield (sample_idx, empty_result) + return + + # Merge stop strings for this single sample. + batch_stop_strings: list[list[str]] = data.get("stop_strings", []) + stop_strings = self._merge_stop_strings(batch_stop_strings) + per_sample_stop_strings = ( + stop_strings[sample_idx] if sample_idx < len(stop_strings) else None + ) - Returns: - Dict mapping server URL to list of GPU UUIDs - e.g., {"http://localhost:30000": ["GPU-aaa", "GPU-bbb"], ...} - """ - if not self.worker_group or not self.worker_group.workers: - raise RuntimeError("Worker group is not initialized") + sampling_params = self._build_sampling_params( + greedy=greedy, + max_new_tokens=max_new_tokens, + stop_strings=per_sample_stop_strings, + ) + + router_ip = self.router_ip + router_port = self.router_port + valid_input_ids = original_input_ids_single_row[:input_length].tolist() + + # batch_size == 1, so no task fan-out / as_completed is needed. Just + # await the single coroutine directly. + _, new_tokens, new_logprobs, is_truncated = await generate_one_sample( + router_ip, + router_port, + sampling_params, + valid_input_ids, + sample_idx, + http_client=self._http_client, + ) - # Get base URLs and GPU UUIDs from all primary workers (TP rank 0) - futures_url = self.worker_group.run_all_workers_single_data( - "get_base_url", - run_rank_0_only_axes=["tensor_parallel"], + # Build the single-sample output tensor: [input | generated]. + generation_length = len(new_tokens) + unpadded_length = input_length + generation_length + + output_ids_single_item = torch.full( + (unpadded_length,), pad_token_id, dtype=dtype, device=device + ) + output_ids_single_item[:input_length] = original_input_ids_single_row[ + :input_length + ] + # Logprobs: zeros for input tokens, raw floats at generated positions. + logprobs_single_item = torch.zeros( + (1, unpadded_length), dtype=torch.float32, device=device ) - futures_uuids = self.worker_group.run_all_workers_single_data( - "get_gpu_uuids", - run_rank_0_only_axes=["tensor_parallel"], + + if new_tokens: + output_ids_single_item[input_length:unpadded_length] = torch.tensor( + new_tokens, dtype=dtype, device=device + ) + if new_logprobs: + logprobs_single_item[0, input_length : input_length + len(new_logprobs)] = ( + torch.tensor(new_logprobs, dtype=torch.float32, device=device) + ) + + result_batch = BatchedDataDict[GenerationOutputSpec]( + { + "output_ids": output_ids_single_item.unsqueeze(0), + "logprobs": logprobs_single_item, + "generation_lengths": torch.tensor( + [generation_length], dtype=torch.long, device=device + ), + "unpadded_sequence_lengths": torch.tensor( + [unpadded_length], dtype=torch.long, device=device + ), + "truncated": torch.tensor( + [bool(is_truncated)], dtype=torch.bool, device=device + ), + } ) - urls = ray.get(futures_url) - uuids_list = ray.get(futures_uuids) + yield (sample_idx, result_batch) - # Create mapping - url_to_uuids = {} - for url, uuids in zip(urls, uuids_list): - if url is not None and uuids is not None: - url_to_uuids[url] = uuids + # --------------------------------------------------------------------------- + # Compatible with parent class or old interfaces + # --------------------------------------------------------------------------- + def init_collective( + self, ip: str, port: int, world_size: int, *, train_world_size: int + ) -> list[ray.ObjectRef]: + return [] - return url_to_uuids + def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + pass + + def update_weights_via_ipc_zmq(self) -> list[ray.ObjectRef]: + return [] + + def update_weights_from_collective(self) -> list[ray.ObjectRef]: + return [] def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: """Wake workers up for colocated inference.""" - pass + tags = kwargs.get("tags", None) + if self.needs_offload: + if tags is None: + self.onload_weights() + self.onload_kv() + else: + if "weights" in tags: + self.onload_weights() + if "kv_cache" in tags: + self.onload_kv() + + self.health_monitoring_resume() def finish_generation(self, *args: Any, **kwargs: Any) -> bool: """Sleep workers and reset prefix cache.""" - pass - - def shutdown(self) -> bool: - """Shut down all SGLang workers and clean up resources.""" - try: - # Use the worker group's shutdown method with the worker's cleanup method - return self.worker_group.shutdown(cleanup_method="shutdown") - except Exception as e: - logger.error(f"Error during SGLang policy shutdown: {e}") - return False + self.health_monitoring_pause() - def __del__(self) -> None: - """Shuts down the worker groups when the object is deleted or is garbage collected. - - This is an extra safety net in case the user forgets to call shutdown() and the pointer to - the object is lost due to leaving a function scope. It's always recommended that the - user calls shutdown(). - """ - self.shutdown() + tags = kwargs.get("tags", None) + if self.needs_offload: + if tags is None: + self.offload_weights() + self.offload_kv() + else: + if "weights" in tags: + self.offload_weights() + if "kv_cache" in tags: + self.offload_kv() def invalidate_kv_cache(self) -> bool: """Invalidate KV cache before weight updates (Megatron-style). - This flushes the cache before weight updates to clear stale cache. - Only primary workers (TP rank 0, model owners) will flush their cache. - - Returns: - bool: True if all caches were flushed successfully, False otherwise + Flushes the cache on every node-0 engine so stale KV entries are + discarded before new weights land. Returns ``True`` iff every engine + reports success. """ + engines = [e for e in self.engines if e is not None] + if not engines: + return True try: - futures = self.worker_group.run_all_workers_single_data( - "invalidate_kv_cache", - run_rank_0_only_axes=["tensor_parallel"], - ) - results = ray.get(futures) - results = [r for r in results if r is not None] - success = all(result for result in results) if results else True - if success: - logger.info( - "[sglang refit] All SGLang server caches flushed successfully" - ) - else: - logger.warning( - "[sglang refit] WARNING - Some SGLang server caches failed to flush" - ) - return success + results = ray.get([e.invalidate_kv_cache.remote() for e in engines]) except Exception as e: logger.error(f"[sglang refit] Error flushing SGLang caches: {e}") return False + + success = all(results) + if success: + logger.info("[sglang refit] All SGLang server caches flushed successfully") + else: + logger.warning( + "[sglang refit] WARNING - Some SGLang server caches failed to flush" + ) + return success + + +# --------------------------------------------------------------------------- +# Generate one sample helper +# --------------------------------------------------------------------------- +async def generate_one_sample( + sglang_router_ip, + sglang_router_port, + sampling_params, + input_ids, + index: int, + http_client: HttpClient | None = None, +): + """Generate using traditional SGLang router with token-based workflow.""" + url = f"http://{sglang_router_ip}:{sglang_router_port}/generate" + + # Prepare payload for sglang server + payload = { + "sampling_params": sampling_params, + "return_logprob": True, + "input_ids": input_ids, + } + + owns_client = http_client is None + if http_client is None: + http_client = HttpClient() + + try: + output = await http_client.post(url, payload) + finally: + if owns_client: + await http_client.aclose() + + if "output_token_logprobs" in output["meta_info"]: + response_tokens = [ + item[1] for item in output["meta_info"]["output_token_logprobs"] + ] + response_log_probs = [ + item[0] for item in output["meta_info"]["output_token_logprobs"] + ] + else: + response_tokens, response_log_probs = [], [] + + # SGLang reports the termination reason under meta_info.finish_reason.type; + # "length" means the decoder hit max_new_tokens before EOS. + finish_reason = output["meta_info"].get("finish_reason") or {} + response_truncated = finish_reason.get("type") == "length" + + return index, response_tokens, response_log_probs, response_truncated + + +# --------------------------------------------------------------------------- +# Port allocation helpers +# --------------------------------------------------------------------------- +def _allocate_rollout_engine_addr_and_ports_normal( + *, + gpus_per_node: int, + sglang_cfg, + local_all_engines, + rank_offset=0, + base_port=15000, +): + # 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 + + sglang_dp_size = sglang_cfg["sglang_cfg"]["dp_size"] + num_gpus_per_engine = sglang_cfg["sglang_server"]["num_gpus_per_engine"] + num_gpus_per_node = gpus_per_node + + _gpus_per_engine = num_gpus_per_engine + 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 + node_index = local_rank // num_engines_per_node + if node_index in visited_nodes: + continue + visited_nodes.add(node_index) + # TODO: currently when restarting engines, we will set port for all engines on this node starting with this rank. + # e.g. for 8 gpus, if we are restarting engine on gpu 3, we will set port for engine 3,4,5,6,7 on this node. + num_engines_on_this_node = num_engines_per_node - ( + local_rank % num_engines_per_node + ) + + def get_addr_and_ports(engine, node_idx): + # use small ports to prevent ephemeral port between 32768 and 65536. + # also, ray uses port 10002-19999, thus we avoid near-10002 to avoid racing condition + start_port = node_port_cursor.get(node_idx, base_port) + + def port(consecutive=1): + nonlocal start_port + _, port = ray.get( + engine._get_current_node_ip_and_free_port.remote( + start_port=start_port, + consecutive=consecutive, + ) + ) + start_port = port + consecutive + node_port_cursor[node_idx] = start_port + return port + + def addr(): + addr, _ = ray.get(engine._get_current_node_ip_and_free_port.remote()) + return addr + + return addr, port + + get_addr, get_port = get_addr_and_ports(engine, node_index) + + for i in range(num_engines_on_this_node): + current_rank = rank + i + addr_and_ports.setdefault(current_rank, {}) + addr_and_ports[current_rank]["host"] = get_addr() + addr_and_ports[current_rank]["port"] = get_port() + addr_and_ports[current_rank]["nccl_port"] = get_port() + + if _gpus_per_engine > num_gpus_per_node: + num_node_per_engine = _gpus_per_engine // num_gpus_per_node + if local_rank % num_node_per_engine == 0: + dist_init_addr = f"{get_addr()}:{get_port(30 + sglang_dp_size)}" + for i in range(num_node_per_engine): + addr_and_ports.setdefault(rank + i, {}) + addr_and_ports[rank + i]["dist_init_addr"] = dist_init_addr + else: + for i in range(num_engines_on_this_node): + addr_and_ports[rank + i]["dist_init_addr"] = ( + f"{get_addr()}:{get_port(30 + sglang_dp_size)}" + ) + + for i, _ in local_all_engines: + for key in ["port", "nccl_port", "dist_init_addr"]: + assert key in addr_and_ports[i], f"Engine {i} {key} is not set." + logger.info(f"Ports for engine {i}: {addr_and_ports[i]}") + + return addr_and_ports, node_port_cursor + + +def _start_router( + sglang_cfg: SGLangConfig, +) -> tuple[str, int, ray.actor.ActorHandle | None]: + """Start sgl router, returning ``(router_ip, router_port, actor_handle)``. + + If ``sglang_router.sglang_router_ip`` is already set, reuse it and return + ``actor_handle=None`` (we do not own that router and must not terminate it). + Otherwise spawn a ``RouterActor`` in sglang env to own the router process. + """ + router_cfg = sglang_cfg.get("sglang_router") or {} + if router_cfg.get("sglang_router_ip") is not None: + return router_cfg["sglang_router_ip"], router_cfg["sglang_router_port"], None + + router_actor = RouterActor.options( + runtime_env={"py_executable": SGLANG_EXECUTABLE}, + ).remote() + router_ip, router_port = ray.get(router_actor.start.remote(dict(router_cfg))) + logger.info(f"Router launched at {router_ip}:{router_port}") + return router_ip, router_port, router_actor diff --git a/nemo_rl/models/generation/sglang/sglang_router.py b/nemo_rl/models/generation/sglang/sglang_router.py new file mode 100644 index 0000000000..285849298e --- /dev/null +++ b/nemo_rl/models/generation/sglang/sglang_router.py @@ -0,0 +1,63 @@ +import logging + +import ray + +from nemo_rl.models.generation.sglang.config import SGLangRouter + +logger = logging.getLogger(__name__) + + +@ray.remote(num_cpus=1, num_gpus=0) +class RouterActor: + """Starts and owns the sglang router subprocess. + + Runs under SGLANG_EXECUTABLE so it can import sglang_router. + The driver (SYSTEM env) holds a handle to this actor and retrieves + (router_ip, router_port) without ever importing sglang_router itself. + """ + + def start(self, router_cfg: SGLangRouter) -> tuple[str, int]: + import multiprocessing + import random + + from sglang_router.launch_router import RouterArgs + + from nemo_rl.models.generation.sglang.utils.ray_utils import ( + _wrap_ipv6, + find_available_port, + get_host_info, + ) + from nemo_rl.models.generation.sglang.utils.router_utils import run_router + + 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_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.log_level = "warn" + request_timeout_secs = router_cfg.get("sglang_router_request_timeout_secs") + if request_timeout_secs is not None: + router_args.request_timeout_secs = request_timeout_secs + + self._process = multiprocessing.Process(target=run_router, args=(router_args,)) + self._process.daemon = True + self._process.start() + import time + + time.sleep(3) + assert self._process.is_alive(), "Router process died on startup" + return router_ip, router_port + + def stop(self): + from nemo_rl.models.generation.sglang.utils.router_utils import ( + terminate_process, + ) + + if hasattr(self, "_process"): + terminate_process(self._process) diff --git a/nemo_rl/models/generation/sglang/sglang_worker.py b/nemo_rl/models/generation/sglang/sglang_worker.py index 94e7f5df72..2054d7b894 100644 --- a/nemo_rl/models/generation/sglang/sglang_worker.py +++ b/nemo_rl/models/generation/sglang/sglang_worker.py @@ -1,805 +1,807 @@ -# 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. - -import asyncio +import ipaddress import logging import multiprocessing import os import time -from typing import Any, Optional +from collections.abc import Callable -import aiohttp import ray import requests -import torch - -from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.virtual_cluster import _get_free_port_local, _get_node_ip_local -from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches -from nemo_rl.models.generation.interfaces import ( - GenerationDatumSpec, - GenerationOutputSpec, - verify_right_padding, +from urllib3.exceptions import NewConnectionError + +from nemo_rl.models.generation.sglang.utils.ray_utils import ( + get_current_node_ip, + get_free_port, + get_host_info, ) -from nemo_rl.models.generation.sglang.config import SGLangConfig -from nemo_rl.models.generation.sglang.utils import AsyncLoopThread -from nemo_rl.utils.nsys import wrap_with_nvtx_name logger = logging.getLogger(__name__) -def _require_sglang(): - """Import `sglang` lazily so test collection works without the optional extra.""" - try: - from sglang.srt.entrypoints.http_server import launch_server - from sglang.srt.server_args import ServerArgs - from sglang.srt.utils import kill_process_tree - except ModuleNotFoundError as e: # pragma: no cover - raise ModuleNotFoundError( - "Optional dependency `sglang` is required for the SGLang generation backend.\n" - "Install it via the project extra (e.g. `uv run --extra sglang ...`) to use " - "`SGLangGenerationWorker`." - ) from e - - return launch_server, ServerArgs, kill_process_tree +def _get_sglang_file(relative_path: str) -> str: + from importlib.util import find_spec + spec = find_spec("sglang") + if spec is None or not spec.submodule_search_locations: + raise RuntimeError( + f"sglang package not found while attempting to patch '{relative_path}'. " + ) -@ray.remote( - runtime_env={**get_nsight_config_if_pattern_matches("sglang_generation_worker")} -) # pragma: no cover -class SGLangGenerationWorker: - def __repr__(self) -> str: - """Customizes the actor's prefix in the Ray logs. + base_dir = next(iter(spec.submodule_search_locations)) + file_path = os.path.join(base_dir, *relative_path.split("/")) + if not os.path.exists(file_path): + raise RuntimeError( + f"Expected sglang file '{relative_path}' not found at '{file_path}'. " + "The sglang version may have moved this file; compat patch cannot be applied." + ) + return file_path + + +def _write_and_verify(file_path: str, content: str, sentinel: str) -> None: + tmp_path = f"{file_path}.nemo_rl_compat.{os.getpid()}.tmp" + with open(tmp_path, "w") as f: + f.write(content) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, file_path) + + with open(file_path, "r") as f: + verify = f.read() + if sentinel not in verify: + raise RuntimeError( + f"Compat patch verification failed for {file_path}: " + f"sentinel '{sentinel}' not present after write. " + "The write may have been silently dropped by the filesystem." + ) - This makes it easier to identify which worker is producing specific log messages. - """ - return f"{self.__class__.__name__}" - @staticmethod - def configure_worker( - num_gpus: int | float, bundle_indices: Optional[tuple[int, list[int]]] = None - ) -> tuple[dict[str, Any], dict[str, str], dict[str, Any]]: - """Provides complete worker configuration for SGLang server. +def _patch_sglang_safe_unpickler() -> None: + file_to_patch = _get_sglang_file("srt/utils/common.py") - This method configures the worker based on bundle_indices which tells us - how many GPUs this server should use. + with open(file_to_patch, "r") as f: + content = f.read() - Args: - num_gpus: Original GPU allocation for this worker based on the placement group - bundle_indices: Tuple of (node_idx, local_bundle_indices) for this server + sentinel = '"nemo_rl.models.policy.torch_reductions_utils."' + if sentinel in content: + return - Returns: - tuple with complete worker configuration: - - 'resources': Resource allocation (e.g., num_gpus) - - 'env_vars': Environment variables for this worker - - 'init_kwargs': Parameters to pass to __init__ of the worker - """ - # Initialize configuration - resources: dict[str, Any] = {"num_gpus": num_gpus} - init_kwargs: dict[str, Any] = {} - env_vars: dict[str, str] = {} - - local_bundle_indices = None - if bundle_indices is not None: - node_idx = bundle_indices[0] - local_bundle_indices = bundle_indices[1] - init_kwargs["bundle_indices"] = local_bundle_indices - - # Calculate a unique seed from node_idx and bundle_indices - if len(local_bundle_indices) == 1: - seed = node_idx * 1024 + local_bundle_indices[0] - else: - bundle_id = local_bundle_indices[0] // len(local_bundle_indices) - seed = node_idx * 1024 + bundle_id - - init_kwargs["seed"] = seed - - # Check if this worker is part of a parallel group (multiple GPUs per server). - # A worker with local rank =0 owns the server(local_bundle_indices is not None ) - # otherwise it is a placeholder for Ray's resource management (local_bundle_indices is None). - is_part_of_parallel_workers = ( - local_bundle_indices is not None and len(local_bundle_indices) > 1 - ) or local_bundle_indices is None - - if is_part_of_parallel_workers: - # For parallel workers, we manage GPU assignment via base_gpu_id - # All workers see the same global CUDA_VISIBLE_DEVICES, but use different - # logical GPU ranges via base_gpu_id - resources["num_gpus"] = 0 - env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] = "1" - init_kwargs["fraction_of_gpus"] = num_gpus - else: - env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] = "1" + anchor = ' "torch.nn.parameter.",\n' + insertion = anchor + ' "nemo_rl.models.policy.torch_reductions_utils.",\n' + if anchor not in content: + raise RuntimeError( + f"SafeUnpickler allowlist anchor '{anchor.strip()}' not found in " + f"{file_to_patch}." + ) - return resources, env_vars, init_kwargs + content = content.replace(anchor, insertion, 1) + _write_and_verify(file_to_patch, content, sentinel) + logger.info("Patched SafeUnpickler allowlist in %s.", file_to_patch) - def __init__( - self, - config: SGLangConfig, - bundle_indices: Optional[list[int]] = None, - fraction_of_gpus: float = 1.0, - seed: Optional[int] = None, - ): - """Initialize a SGLang worker for distributed inference. - Args: - config: Configuration dictionary for the policy - bundle_indices: List of local bundle indices for this server. - The length of this list determines tp_size (number of GPUs per server). - Only needed for the first worker in each server group (model owner). - fraction_of_gpus: Fraction of GPUs to use for this worker - seed: Random seed for initialization, if None, then defaults to the config's seed - """ - self.cfg = config - self.is_model_owner = bundle_indices is not None - self.global_rank = int(os.environ.get("RANK", "0")) - self.sglang_cfg = config["sglang_cfg"] - - # Create a dedicated event loop thread for async operations - # there will be issues if we use the event loop in the main thread - self.async_loop_thread = AsyncLoopThread() - - # temp: Maximum concurrent requests per server - # we may remove this limit in the future - self.max_concurrent_requests = config.get("max_concurrent_requests", 999999) - - # Only the primary worker (local_rank=0) in each server group starts the SGLang server - # Secondary workers (local_rank!=0) just returns - if not self.is_model_owner: - return +def _override_sglang_imbalance_check_env() -> None: + """Force-disable sglang's per-GPU memory imbalance check. - # `sglang` is an optional dependency; import only when we actually start a server. - _, ServerArgs, _ = _require_sglang() + Pop the legacy names so the shim has nothing to copy, then set + ``ENABLE=false`` directly. Inherited env reaches the subprocesses + cleaned, so the shim no longer overwrites our ENABLE on re-import. + """ + for legacy in ( + "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK", + "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK", + ): + os.environ.pop(legacy, None) + os.environ["SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK"] = "false" - # Determine tp_size from bundle_indices length - tp_size = len(bundle_indices) - base_gpu_id = bundle_indices[0] if bundle_indices else 0 +def _get_megatron_file(subpackage: str, relative_path: str) -> str | None: + """Locate a file inside ``megatron.`` (e.g. ``core``, ``training``). - # Get the global CUDA_VISIBLE_DEVICES (all engines see the same global value) - global_cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None) + Returns ``None`` if megatron isn't importable so callers can treat that + as "nothing to patch". Raises if the package is present but the + expected file is missing (signals a megatron version mismatch). + """ + from importlib.util import find_spec - logger.info( - f"[SGLang Server] Rank {self.global_rank}: " - f"base_gpu_id={base_gpu_id}, tp_size={tp_size}, " - f"bundle_indices={bundle_indices}, global_cvd={global_cvd}" + full_pkg = f"megatron.{subpackage}" + try: + spec = find_spec(full_pkg) + except (ImportError, ValueError): + return None + if spec is None or not spec.submodule_search_locations: + return None + + base_dir = next(iter(spec.submodule_search_locations)) + file_path = os.path.join(base_dir, *relative_path.split("/")) + if not os.path.exists(file_path): + raise RuntimeError( + f"Expected megatron file '{full_pkg}/{relative_path}' not found at " + f"'{file_path}'. The megatron version may have moved this file; " + "compat patch cannot be applied." ) - - # Get current node IP and a free port for the server - node_ip = _get_node_ip_local() - free_port = _get_free_port_local() - - # Build SGLang server arguments - kwargs = { - "model_path": self.sglang_cfg["model_path"], - "trust_remote_code": True, - "random_seed": seed - if seed is not None - else self.sglang_cfg.get("random_seed", 1), - # Memory settings - "enable_memory_saver": self.sglang_cfg["enable_memory_saver"], - "gpu_id_step": 1, - "base_gpu_id": base_gpu_id, - # Parallel settings - "tp_size": tp_size, - "dp_size": self.sglang_cfg["dp_size"], - "pp_size": self.sglang_cfg["pp_size"], - "ep_size": self.sglang_cfg["ep_size"], - # Always skip warmup to prevent warmup timeout - "skip_server_warmup": self.sglang_cfg.get("skip_server_warmup", True), - # Server network settings - listen on all interfaces, use the free port we found - "host": "0.0.0.0", - "port": free_port, - "torchao_config": "", - } - - for key in [ - "dtype", - "kv_cache_dtype", - "context_length", - "max_running_requests", - "chunked_prefill_size", - "max_prefill_tokens", - "schedule_policy", - "schedule_conservativeness", - "cpu_offload_gb", - "log_level", - "mem_fraction_static", - "allow_auto_truncate", - "disable_piecewise_cuda_graph", - ]: - if key in self.sglang_cfg: - kwargs[key] = self.sglang_cfg[key] - - server_args = ServerArgs(**kwargs) - # Save server_args and base_url for use in generate() and _make_request() - self.server_args = server_args - self.base_url = f"http://{node_ip}:{free_port}" - - logger.info( - f"[SGLang Worker] Rank {self.global_rank} Starting on {self.base_url}, CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', None)}, base_gpu_id: {base_gpu_id}" + return file_path + + +def _patch_megatron_hook_mode_in(file_path: str) -> None: + """Comment out ``torch_memory_saver.hook_mode = "torch"`` in a megatron file. + + Megatron sets ``tms.hook_mode = "torch"`` at module import time on the + global ``torch_memory_saver`` singleton. That mutation breaks sglang's + pauseable CUDA graph path, which asserts ``_hook_mode == "preload"`` + inside ``TorchMemorySaver.cuda_graph(...)``. Commenting the line out + leaves the singleton at its default ``"preload"`` mode that sglang + expects. + """ + with open(file_path, "r") as f: + content = f.read() + + sentinel = '# torch_memory_saver.hook_mode = "torch"' + if sentinel in content: + return + + anchor = ' torch_memory_saver.hook_mode = "torch"\n' + if anchor not in content: + raise RuntimeError( + f"Megatron hook_mode anchor '{anchor.strip()}' not found in " + f"{file_path}; the megatron version may have moved or removed it." ) - self.session = None - self.connector = None - - self.server_process = self._launch_server_process(server_args) - - def get_base_url(self) -> str: - """Get the base URL of this SGLang server.""" - return self.base_url + replacement = ( + ' # torch_memory_saver.hook_mode = "torch" ' + "# patched by nemo_rl: conflicts with sglang pauseable CUDA Graph\n" + ) + content = content.replace(anchor, replacement, 1) + _write_and_verify(file_path, content, sentinel) + logger.info("Patched megatron tms.hook_mode mutation in %s.", file_path) + + +def _patch_megatron_dynamic_context_hook_mode() -> None: + file_path = _get_megatron_file("core", "inference/contexts/dynamic_context.py") + if file_path is None: + return + _patch_megatron_hook_mode_in(file_path) + + +def _patch_megatron_training_hook_mode() -> None: + file_path = _get_megatron_file("training", "training.py") + if file_path is None: + return + _patch_megatron_hook_mode_in(file_path) + + +def _apply_sglang_compat_patches() -> None: + _patch_sglang_safe_unpickler() + _override_sglang_imbalance_check_env() + _patch_megatron_dynamic_context_hook_mode() + _patch_megatron_training_hook_mode() + + +def get_base_gpu_id(gpus_per_node: int, sglang_cfg, rank): + num_gpus = min(gpus_per_node, sglang_cfg["sglang_server"]["num_gpus_per_engine"]) + start_index = (rank * num_gpus) % gpus_per_node + return start_index + + +def _to_local_gpu_id(physical_gpu_id: int) -> int: + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if not cvd: + return physical_gpu_id # no remapping + # CUDA_VISIBLE_DEVICES can be like "4,5,6,7" + visible = [int(x) for x in cvd.split(",") if x.strip() != ""] + # In a remapped process, valid torch device indices are 0..len(visible)-1 + if physical_gpu_id in visible: + return visible.index(physical_gpu_id) + # If we're already getting local IDs, allow them + if 0 <= physical_gpu_id < len(visible): + return physical_gpu_id + raise RuntimeError( + f"GPU id {physical_gpu_id} is not valid under CUDA_VISIBLE_DEVICES={cvd}. " + f"Expected one of {visible} (physical) or 0..{len(visible) - 1} (local)." + ) + + +def launch_server_process(server_args) -> multiprocessing.Process: + from sglang.srt.entrypoints.http_server import launch_server + + multiprocessing.set_start_method("spawn", force=True) + server_args.host = server_args.host.strip("[]") + p = multiprocessing.Process(target=launch_server, args=(server_args,)) + p.start() + + if server_args.node_rank != 0: + return p - def invalidate_kv_cache(self) -> bool: - """Invalidate KV cache before weight updates (Megatron-style). + _wait_server_healthy( + base_url=server_args.url(), + api_key=server_args.api_key, + process_alive_fn=lambda: p.is_alive(), + ) - This flushes the cache before weight updates to clear stale cache. - Uses retry logic to handle cases where there are pending requests. + return p - Returns: - bool: True if flush was successful, False otherwise - """ - if not self.is_model_owner: - return True - url = f"{self.base_url}/flush_cache" - max_attempts = 60 - connection_retry_limit = 5 +def _wait_server_healthy( + base_url: str, + api_key: str | None, + process_alive_fn: Callable[[], bool], +) -> None: + headers = { + "Content-Type": "application/json; charset=utf-8", + "Authorization": f"Bearer {api_key}", + } - # flush_cache will not return status_code 200 when there are pending requests - for attempt in range(max_attempts): + with requests.Session() as session: + while True: try: - response = requests.get(url, timeout=10) + response = session.get(f"{base_url}/health_generate", headers=headers) if response.status_code == 200: - if attempt > 0: - logger.info( - f"[SGLang Worker] Rank {self.global_rank} Cache flushed successfully " - f"(attempt {attempt + 1})" - ) - return True - except requests.exceptions.ConnectionError: - # Server might not be ready yet - only retry for first few attempts - if attempt >= connection_retry_limit: - logger.warning( - f"[SGLang Worker] Rank {self.global_rank} Connection failed after " - f"{connection_retry_limit} attempts" - ) - return False - except Exception as e: - # For other errors, log and retry (except on last attempt) - if attempt == max_attempts - 1: - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Failed to flush cache after " - f"{max_attempts} attempts: {e}" - ) - return False + break + except requests.RequestException: + pass - time.sleep(1) - - # All attempts exhausted without success - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Timeout: Cache flush failed after " - f"{max_attempts} attempts. Server may have pending requests." - ) - return False - - def get_gpu_uuids(self) -> list[str]: - """Get list of GPU UUIDs used by this SGLang server. - - Returns: - List of GPU UUIDs (e.g., ["GPU-xxxxx", "GPU-yyyyy"]) - """ - from nemo_rl.utils.nvml import get_device_uuid + if not process_alive_fn(): + raise Exception("Server process terminated unexpectedly.") - # Get all GPU UUIDs used by this server - # SGLang server uses GPUs starting from base_gpu_id with tp_size GPUs - gpu_uuids = [] - for i in range(self.server_args.tp_size): - gpu_id = self.server_args.base_gpu_id + i - uuid = get_device_uuid(gpu_id) - gpu_uuids.append(uuid) + time.sleep(2) - return gpu_uuids + # use flush_cache to make sure the working queue is empty, so that we can do offload + while True: + try: + response = session.get(f"{base_url}/flush_cache", headers=headers) + if response.status_code == 200: + break - def _merge_stop_strings(self, batch_stop_strings): - """Merge stop strings from config and batch. + except requests.RequestException: + pass - Args: - batch_stop_strings: List of stop strings from batch (one per sample) + if not process_alive_fn(): + raise Exception("Server process terminated unexpectedly.") - Returns: - List of merged stop strings (one per sample) - """ - stop_set: set[str] = set() - - # Add stop strings from config - if self.cfg.get("stop_strings"): - stop_set.update(self.cfg["stop_strings"]) - - # Merge stop strings from batch - merged_stop_strings = [] - for sample_ss in batch_stop_strings: - sample_stop_set = stop_set.copy() - if sample_ss: - if isinstance(sample_ss, str): - sample_stop_set.add(sample_ss) - elif isinstance(sample_ss, list): - sample_stop_set.update(sample_ss) - - merged_stop_strings.append( - list(sample_stop_set) if sample_stop_set else None - ) + time.sleep(2) - return merged_stop_strings - def _build_sampling_params( +@ray.remote # pragma: no cover +class SGLangGenerationWorker: + def __init__( + self, + gpus_per_node: int, + sglang_cfg, + rank: int, + base_gpu_id: int | None = None, + num_gpus_per_engine: int | None = None, + ): + _apply_sglang_compat_patches() + self.gpus_per_node = gpus_per_node + self.sglang_cfg = sglang_cfg + self.rank = rank + self.base_gpu_id = base_gpu_id + self.num_gpus_per_engine = num_gpus_per_engine + + def init( self, - *, - greedy: bool, - stop_strings, - max_new_tokens: Optional[int] = None, - input_len: Optional[int] = None, - context_length: Optional[int] = None, - sample_index: Optional[int] = None, - ) -> dict[str, Any]: - """Build sampling parameters dictionary for SGLang API. + dist_init_addr, + port, + nccl_port, + host=None, + router_ip=None, + router_port=None, + ): - Args: - greedy: Whether to use greedy decoding (temperature=0.0) - stop_strings: Merged stop strings (not used here, handled per sample) - max_new_tokens: Override max_new_tokens from config if provided - input_len: Input length for this sample (used for context_length adjustment) - context_length: Maximum context length (if provided, adjusts max_new_tokens) - sample_index: Sample index (used for warning messages, 0-indexed) + self.router_ip = ( + router_ip + if router_ip is not None + else self.sglang_cfg["sglang_router"]["sglang_router_ip"] + ) + self.router_port = ( + router_port + if router_port is not None + else self.sglang_cfg["sglang_router"]["sglang_router_port"] + ) - Returns: - Dictionary of sampling parameters compatible with SGLang API - """ - top_k_cfg = self.cfg.get("top_k") - top_k_val = 1 if greedy else (top_k_cfg if top_k_cfg is not None else -1) - temperature = 0.0 if greedy else self.cfg["temperature"] + host = host or get_host_info()[1] - base_max_tokens = ( - max_new_tokens if max_new_tokens is not None else self.cfg["max_new_tokens"] + def _format_v6_uri(addr): + if not addr or addr.startswith("["): + return addr + try: + if ipaddress.ip_address(addr).version == 6: + return f"[{addr}]" + except ValueError: + pass + return addr + + host = _format_v6_uri(host) + ip_part, port_part = dist_init_addr.rsplit(":", 1) + dist_init_addr = f"{_format_v6_uri(ip_part)}:{port_part}" + + server_args_dict = _compute_server_args( + self.gpus_per_node, + self.sglang_cfg, + self.rank, + dist_init_addr, + nccl_port, + host, + port, + base_gpu_id=self.base_gpu_id, + num_gpus_per_engine=self.num_gpus_per_engine, ) - # TODO: check if this is needed - final_max_tokens = base_max_tokens - if context_length is not None and input_len is not None: - max_allowed_new_tokens = max(0, context_length - input_len - 1) - if base_max_tokens > max_allowed_new_tokens: - final_max_tokens = max_allowed_new_tokens - if sample_index == 0: - logger.warning( - f"[SGLang Worker] Rank {self.global_rank} Warning: " - f"Sample {sample_index} input length ({input_len}) + max_new_tokens ({base_max_tokens}) " - f"would exceed context_length ({context_length}). " - f"Reducing max_new_tokens to {final_max_tokens} for this sample." - ) - - # Build sampling params dict - sampling_params = { - "temperature": temperature, - "top_p": self.cfg.get("top_p", 1.0), - "max_new_tokens": final_max_tokens, - } + self.node_rank = server_args_dict["node_rank"] + self.server_host = server_args_dict["host"] # with [] if ipv6 + self.server_port = server_args_dict["port"] + self.server_base_url = f"http://{self.server_host}:{self.server_port}" - if top_k_val != -1: - sampling_params["top_k"] = top_k_val + self._init_normal(server_args_dict) - stop_token_ids = self.cfg.get("stop_token_ids") - if stop_token_ids is not None: - sampling_params["stop_token_ids"] = stop_token_ids + def _init_normal(self, server_args_dict): + from sglang.srt.server_args import ServerArgs - return sampling_params + logger.info( + f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}" + ) + self.process = launch_server_process(ServerArgs(**server_args_dict)) - async def _ensure_session(self): - if self.session is None: - # Create connector with connection pool limit - self.connector = aiohttp.TCPConnector(limit=512, limit_per_host=512) - # Create session with timeout - timeout = aiohttp.ClientTimeout(total=300) # 5 minutes timeout - self.session = aiohttp.ClientSession( - connector=self.connector, timeout=timeout + if self.node_rank == 0 and self.router_ip and self.router_port: + payload = { + "url": self.server_base_url, + "worker_type": "regular", + } + response = requests.post( + f"http://{self.router_ip}:{self.router_port}/workers", + json=payload, ) - return self.session + response.raise_for_status() - async def _generate_single_sample( - self, - input_ids: list[int], - sampling_params: dict[str, Any], - stop_string: Optional[str] = None, - ) -> tuple[list[int], list[float]]: - """Generate a single sample using SGLang API (async function). + def _make_request(self, endpoint: str, payload: dict | None = None): + """Make a POST request to the specified endpoint with the given payload. Args: - input_ids: List of input token IDs (without padding) - sampling_params: Dictionary of sampling parameters (temperature, top_p, max_new_tokens, etc.) - stop_string: Optional stop string for this sample + endpoint: The API endpoint to call + payload: The JSON payload to send (default: empty dict) Returns: - Tuple of (generated_tokens, logprobs): - - generated_tokens: List of generated token IDs - - logprobs: List of log probabilities for generated tokens + The JSON response from the server """ - # Prepare payload for SGLang API - # Note: stop should be in sampling_params, not in payload top level - # TODO: double check this - if stop_string is not None: - # stop can be a string or list of strings - sampling_params = sampling_params.copy() # Don't modify the original - sampling_params["stop"] = stop_string + if self.node_rank != 0: + return - payload = { - "sampling_params": sampling_params, - "return_logprob": True, - "input_ids": input_ids, - } + url = f"{self.server_base_url}/{endpoint}" + response = requests.post(url, json=payload or {}) + try: + response.raise_for_status() + except requests.exceptions.HTTPError as e: + e.add_note(f"{response.text=}") + raise + return response.json() - url = f"{self.base_url}/generate" - headers = { - "Content-Type": "application/json; charset=utf-8", - } + @staticmethod + def _get_current_node_ip_and_free_port(start_port=10000, consecutive=1): + return get_current_node_ip(), get_free_port( + start_port=start_port, consecutive=consecutive + ) - session = await self._ensure_session() + def health_generate(self, timeout: float = 5.0) -> bool: + """Run /health_generate on the underlying SGLang HTTP server. - try: - async with session.post(url, json=payload, headers=headers) as response: - response.raise_for_status() - result = await response.json() - except Exception as e: - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Request failed for input_len={len(input_ids)}: {e}" - ) - raise + Args: + timeout: Timeout for the health request in seconds. - # Extract generated tokens and logprobs - meta_info = result.get("meta_info", {}) - output_token_logprobs = meta_info.get("output_token_logprobs", []) + Returns: + True if the server responds with HTTP 200. - if output_token_logprobs: - new_tokens = [item[1] for item in output_token_logprobs] - new_logprobs = [item[0] for item in output_token_logprobs] - else: - # Fallback: empty if token logprobs not available - new_tokens = [] - new_logprobs = [] + Raises: + requests.RequestException: If the request fails for any reason, including timeout. + """ + if self.node_rank != 0: + return True - return new_tokens, new_logprobs + response = requests.get( + f"{self.server_base_url}/health_generate", + timeout=timeout, + ) + response.raise_for_status() + return True - async def _generate_async(self, tasks): - """Execute generation tasks with concurrency control. + def update_weights_from_tensor( + self, + serialized_named_tensors: list[str], + load_format: str | None = None, + flush_cache: bool = False, + weight_version: str | None = None, + ): + """Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs. - TEMP: Uses a semaphore to limit the number of concurrent requests per server, preventing server overload. - A router based solution is preffered in the future. + Note: The model should be on GPUs rather than CPU for this functionality to work properly. + If you encounter issues, ensure your model is loaded on GPU devices rather than CPU. """ - semaphore = asyncio.Semaphore(self.max_concurrent_requests) - - async def wrap(idx, coro): - async with semaphore: - try: - result = await coro - return idx, result - except Exception as e: - raise - - wrapped = [wrap(i, t) for i, t in enumerate(tasks)] - results = [None] * len(tasks) - count = 0 - - for fut in asyncio.as_completed(wrapped): - idx, value = await fut - results[idx] = value - count += 1 - if count % 50 == 0 or count == len(tasks): - logger.debug( - f"[SGLang Worker] Rank {self.global_rank} Completed {count}/{len(tasks)} tasks" - ) - - return results - - def _launch_server_process(self, server_args: Any) -> multiprocessing.Process: - """Launch the SGLang server process and wait for it to be ready.""" - # Ensure `sglang` is importable when we actually start a server. - launch_server, _, kill_process_tree = _require_sglang() - p = multiprocessing.Process(target=launch_server, args=(server_args,)) - p.start() - - # Wait for server to be ready by checking health endpoint - # Use the base_url we stored earlier - headers = { - "Content-Type": "application/json; charset=utf-8", + payload = { + "serialized_named_tensors": serialized_named_tensors, + "load_format": load_format, + "flush_cache": flush_cache, } + if weight_version is not None: + payload["weight_version"] = weight_version + return self._make_request( + "update_weights_from_tensor", + payload, + ) - max_wait_time = 300 # 5 minutes timeout - start_time = time.time() - with requests.Session() as session: - while True: - if time.time() - start_time > max_wait_time: - kill_process_tree(p.pid) - raise TimeoutError( - f"[SGLang Server] Rank {self.global_rank} Server failed to start within {max_wait_time}s" - ) - try: - response = session.get( - f"{self.base_url}/health_generate", headers=headers, timeout=10 - ) - if response.status_code == 200: - logger.info( - f"[SGLang Server] Rank {self.global_rank} Server is ready at {self.base_url}" + def flush_cache(self): + """Flush the cache of the server.""" + if self.node_rank != 0: + return + # flush cache will not return status_code 200 when there are pending requests + for _ in range(60): + try: + response = requests.get(f"{self.server_base_url}/flush_cache") + if response.status_code == 200: + break + except NewConnectionError as e: + raise e + except Exception as e: + logger.info(f"Error flushing cache: {e}") + # Pace retries on both non-200 and exception paths; otherwise the + # 60 iterations fly by in milliseconds and timeout before sglang + # has a chance to drain its queue. + time.sleep(1) + else: + raise TimeoutError("Timeout while flushing cache.") + + def shutdown(self): + from sglang.srt.utils import kill_process_tree + + logger.info(f"Shutdown engine {self.server_host}:{self.server_port}...") + if self.node_rank == 0: + worker_url = self.server_base_url + response = None + try: + all_workers = requests.get( + f"http://{self.router_ip}:{self.router_port}/workers" + ).json()["workers"] + for worker in all_workers: + if worker["url"] == worker_url: + worker_id = worker["id"] + response = requests.delete( + f"http://{self.router_ip}:{self.router_port}/workers/{worker_id}" ) break - except requests.RequestException: - pass - - if not p.is_alive(): - raise RuntimeError( - f"[SGLang Server] Rank {self.global_rank} Server process terminated unexpectedly." + else: + logger.warning( + f"Worker {worker_url} not found in router during shutdown." ) + except Exception as e: + logger.warning(f"Failed to fetch workers list or remove worker: {e}") - time.sleep(2) - return p - - @wrap_with_nvtx_name("sglang_genertion_worker/generate") - def generate( - self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False - ) -> BatchedDataDict[GenerationOutputSpec]: - """Generate a batch of data using SGLang generation. + if response is not None: + response.raise_for_status() + kill_process_tree(self.process.pid) - Args: - data: BatchedDataDict containing input_ids and input_lengths tensors - greedy: Whether to use greedy decoding instead of sampling + def get_weight_version(self): + if self.node_rank != 0: + return + # new sglang change api from /get_weight_version to /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() - Returns: - BatchedDataDict conforming to GenerationOutputSpec: - - output_ids: input + generated token IDs with proper padding - - logprobs: Log probabilities for tokens - - generation_lengths: Lengths of each response - - unpadded_sequence_lengths: Lengths of each input + generated sequence - """ - # Handle empty input case - if len(data["input_ids"]) == 0: - return BatchedDataDict[GenerationOutputSpec]( - { - "output_ids": torch.zeros((0, 0), dtype=torch.long), - "logprobs": torch.zeros((0, 0), dtype=torch.float), - "generation_lengths": torch.zeros(0, dtype=torch.long), - "unpadded_sequence_lengths": torch.zeros(0, dtype=torch.long), - } - ) + def release_memory_occupation(self, tags: list[str] | None = None): + """Release memory occupation. Available tags: weights, kv_cache.""" + self.flush_cache() + return self._make_request( + "release_memory_occupation", + {"tags": tags}, + ) - input_ids = data["input_ids"] - input_lengths = data["input_lengths"] - batch_stop_strings = data.get("stop_strings", [None] * len(input_lengths)) - stop_strings = self._merge_stop_strings(batch_stop_strings) - batch_size = len(input_lengths) - pad_token_id = self.cfg["_pad_token_id"] + def resume_memory_occupation(self, tags: list[str] | None = None): + """Available tags for multi-stage resume: weights, kv_cache.""" + return self._make_request( + "resume_memory_occupation", + {"tags": tags}, + ) - # Verify inputs have correct padding - verify_right_padding(data, pad_value=pad_token_id) + def release_memory_weights(self): + from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS - # Original input length with padding - padded_input_length = input_ids.size(1) + return self.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - logger.debug( - f"[SGLang Worker] Rank {self.global_rank} batch_size: {batch_size}, padded_input_length: {padded_input_length}" + def release_memory_kv_cache_and_cuda_graph(self): + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, ) - if batch_size == 0: - raise ValueError("Empty batch received") + return self.release_memory_occupation( + tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH] + ) - context_length = self.sglang_cfg.get("context_length", None) + def resume_memory_weights(self): + from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS - # Create async tasks for all samples - tasks = [] - for i in range(batch_size): - input_len = input_lengths[i].item() + return self.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS]) - # Truncate input if it exceeds context_length - if context_length is not None and input_len >= context_length: - input_len = context_length - 1 + def resume_memory_kv_cache_and_cuda_graph(self): + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + ) - valid_input_ids = input_ids[i, :input_len].tolist() + return self.resume_memory_occupation( + tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH] + ) - # Build sampling params for this sample (with context_length adjustment) - sample_sampling_params = self._build_sampling_params( - greedy=greedy, - stop_strings=stop_strings, - max_new_tokens=None, - input_len=input_len, - context_length=context_length, - sample_index=i, - ) + def check_weights(self, action: str): + return self._make_request("weights_checker", {"action": action}) - tasks.append( - self._generate_single_sample( - input_ids=valid_input_ids, - sampling_params=sample_sampling_params, - stop_string=stop_strings[i], - ) - ) + 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, + }, + ) - # Execute all requests concurrently using the dedicated event loop thread + def destroy_weights_update_group(self, group_name): try: - all_results = self.async_loop_thread.run(self._generate_async(tasks)) - except Exception as e: - raise + return self._make_request( + "destroy_weights_update_group", + { + "group_name": group_name, + }, + ) + except requests.exceptions.RequestException: + # catch the case there the engine is just created and does not have the group. + pass - total_generated_tokens = sum(len(tokens) for tokens, _ in all_results) - avg_generation_length = ( - total_generated_tokens / batch_size if batch_size > 0 else 0 + 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, ) - # Process results - output_ids_list = [] - logprobs_list = [] - generation_lengths_list = [] - unpadded_sequence_lengths_list = [] - max_length = 0 - - # First pass: calculate max_length - for i, (new_tokens, new_logprobs) in enumerate(all_results): - input_len = input_lengths[i].item() - generation_length = len(new_tokens) - unpadded_length = input_len + generation_length - max_length = max(max_length, unpadded_length) - - total_length = max(max_length, padded_input_length) - - for i, (new_tokens, new_logprobs) in enumerate(all_results): - input_len = input_lengths[i].item() - generation_length = len(new_tokens) - unpadded_length = input_len + generation_length - - full_output = torch.full( - (total_length,), pad_token_id, dtype=input_ids.dtype - ) - full_output[:input_len] = input_ids[i][:input_len] - - # Add generated tokens after the original input - if new_tokens: - full_output[input_len : input_len + len(new_tokens)] = torch.tensor( - new_tokens, dtype=input_ids.dtype - ) - - # Construct logprobs: zeros for input tokens, actual logprobs for generated tokens - full_logprobs = torch.zeros(total_length, dtype=torch.float32) - if new_logprobs: - for idx, logprob in enumerate(new_logprobs): - position = input_len + idx - full_logprobs[position] = logprob - - output_ids_list.append(full_output) - logprobs_list.append(full_logprobs) - generation_lengths_list.append(generation_length) - unpadded_sequence_lengths_list.append(unpadded_length) - - # Stack into tensors - output_ids = torch.stack(output_ids_list) - logprobs = torch.stack(logprobs_list) - generation_lengths = torch.tensor(generation_lengths_list, dtype=torch.long) - unpadded_sequence_lengths = torch.tensor( - unpadded_sequence_lengths_list, dtype=torch.long - ) - logger.debug( - f"[SGLang Worker] Rank {self.global_rank} Generated {total_generated_tokens} tokens across {batch_size} samples (avg: {avg_generation_length:.1f} tokens/sample)" - ) - return BatchedDataDict[GenerationOutputSpec]( - { - "output_ids": output_ids, - "generation_lengths": generation_lengths, - "unpadded_sequence_lengths": unpadded_sequence_lengths, - "logprobs": logprobs, - } + 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 sleep(self): - # TODO - pass + def continue_generation(self): + response = requests.post(f"{self.server_base_url}/continue_generation", json={}) + response.raise_for_status() + return response - def wake_up(self, **kwargs): - # TODO - pass + def post_process_weights( + self, + restore_weights_before_load: bool = False, + post_process_quantization: bool = False, + ): + """Update model weights from tensor data. - def shutdown(self) -> bool: - """Shutdown the SGLang server process and cleanup async resources. + The HTTP server will only post meta data, and the real weights will be + copied directly from GPUs. - Returns: - bool: True if shutdown was successful, False otherwise + Note: The model should be on GPUs rather than CPU for this functionality + to work properly. If you encounter issues, ensure your model is loaded + on GPU devices rather than CPU. """ - if not self.is_model_owner: - if hasattr(self, "async_loop_thread"): - try: - self.async_loop_thread.shutdown() - logger.info( - f"[SGLang Worker] Rank {self.global_rank} Async loop thread shut down." - ) - except Exception as e: - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Error shutting down async loop thread: {e}" - ) - return True - - try: - # Only model owners started a server process; they require sglang for shutdown. - _, _, kill_process_tree = _require_sglang() - if hasattr(self, "session") and self.session is not None: - try: - - async def close_session(): - await self.session.close() - if self.connector is not None: - await self.connector.close() - - self.async_loop_thread.run(close_session()) - logger.info( - f"[SGLang Worker] Rank {self.global_rank} aiohttp session closed." - ) - except Exception as e: - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Error closing aiohttp session: {e}" - ) - - # Shutdown async loop thread after session cleanup - if hasattr(self, "async_loop_thread"): - try: - self.async_loop_thread.shutdown() - logger.info( - f"[SGLang Worker] Rank {self.global_rank} Async loop thread shut down." - ) - except Exception as e: - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Error shutting down async loop thread: {e}" - ) - - if not hasattr(self, "server_process") or self.server_process is None: - return True + return self._make_request( + "post_process_weights", + { + "restore_weights_before_load": restore_weights_before_load, + "post_process_quantization": post_process_quantization, + }, + ) - logger.info( - f"[SGLang Worker] Rank {self.global_rank} Shutting down server at {self.base_url}..." - ) + def start_profile( + self, + # The output directory + output_dir: str | None = None, + # If set, it profile as many as this number of steps. + # If it is set, profiling is automatically stopped after this step, and + # the caller doesn't need to run stop_profile. + start_step: int | None = None, + num_steps: int | None = None, + activities: list[str] | None = None, + profile_by_stage: bool = False, + with_stack: bool | None = None, + record_shapes: bool | None = None, + ): + response = requests.post( + f"{self.server_base_url}/start_profile", + json={ + "output_dir": output_dir, + "start_step": start_step, + "num_steps": num_steps, + "activities": activities, + "profile_by_stage": profile_by_stage, + "with_stack": with_stack, + "record_shapes": record_shapes, + }, + ) + response.raise_for_status() + return response - if self.server_process.is_alive(): - kill_process_tree(self.server_process.pid) + def stop_profile(self): + response = requests.post(f"{self.server_base_url}/stop_profile", json={}) + response.raise_for_status() + return response - # Wait for the process to terminate - self.server_process.join(timeout=5.0) + def _simulate_crash(self): + """Test-only: tear the engine down to simulate a crash. - if self.server_process.is_alive(): - return False - return True + 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() - except Exception as e: - logger.error( - f"[SGLang Worker] Rank {self.global_rank} Error during shutdown: {e}" - ) - return False + # --------------------------------------------------------------------------- + # Compatible with parent class or old interfaces + # --------------------------------------------------------------------------- + def get_base_url(self) -> str | None: + """Return the ``http://host:port`` base URL of this SGLang server. - def _make_request(self, endpoint: str, payload: Optional[dict] = None): - """Make a POST request to the specified endpoint with the given payload. + Only node-rank 0 owns the HTTP server; peer ranks return ``None`` so + callers can filter them out when collecting per-engine URLs. + """ + if self.node_rank != 0: + return None + return self.server_base_url - Args: - endpoint: The API endpoint to call - payload: The JSON payload to send (default: empty dict) + def invalidate_kv_cache(self) -> bool: + """Flush the cache of the server. Returns: - The JSON response from the server + ``True`` on a successful flush. Peer (non-node-0) ranks return + ``True`` since they do not own the HTTP server. + + Raises: + NewConnectionError: if the engine HTTP server is unreachable + (engine likely crashed); the caller cannot make progress + with stale KV state, so we surface the failure rather than + swallowing it. + TimeoutError: if the server keeps replying non-200 for the full + retry window — equivalent to a hang we shouldn't ignore. """ - # Use the stored base_url instead of constructing from server_args - url = f"{self.base_url}/{endpoint}" - headers = { - "Content-Type": "application/json; charset=utf-8", - } - response = requests.post(url, json=payload or {}, headers=headers, timeout=60) - response.raise_for_status() - return response.json() + if self.node_rank != 0: + return True + # flush cache will not return status_code 200 when there are pending requests + for _ in range(60): + try: + response = requests.get(f"{self.server_base_url}/flush_cache") + if response.status_code == 200: + return True + except NewConnectionError: + logger.exception("Connection error flushing cache") + raise + except Exception as e: + logger.info(f"Error flushing cache: {e}") + # Pace retries on both non-200 and exception paths; otherwise + # the 60 iterations fly by in milliseconds. + time.sleep(1) + raise TimeoutError("Timeout while flushing cache.") + + +# ---------------------------------------------------------------------------- +# Compute Server args +# ---------------------------------------------------------------------------- +def _compute_server_args( + gpus_per_node: int, + sglang_cfg, + rank, + dist_init_addr, + nccl_port, + host, + port, + base_gpu_id: int | None = None, + num_gpus_per_engine: int | None = None, +): + _gpus_per_engine = ( + num_gpus_per_engine or sglang_cfg["sglang_server"]["num_gpus_per_engine"] + ) + nnodes = max(1, _gpus_per_engine // gpus_per_node) + node_rank = rank % nnodes + base = ( + base_gpu_id + if base_gpu_id is not None + else get_base_gpu_id(gpus_per_node, sglang_cfg, rank) + ) + base = _to_local_gpu_id(base) + # ``_gpus_per_engine`` is the engine's total GPU count (TP × PP). When PP=1 + # (the historical default) this equals ``tp_size``; when ``pp_size > 1`` + # the engine spans ``tp_size * pp_size`` GPUs, so derive ``tp_size`` by + # dividing out PP. Falsy/0 ``pp_size`` is treated as 1. + _pp_size = sglang_cfg["sglang_cfg"].get("pp_size", 1) or 1 + kwargs = { + "model_path": sglang_cfg["sglang_cfg"]["model_path"], + "trust_remote_code": True, + "random_seed": sglang_cfg["sglang_cfg"]["random_seed"] + rank, + # memory + "enable_memory_saver": sglang_cfg["sglang_server"]["needs_offload"], + "enable_weights_cpu_backup": sglang_cfg["sglang_server"]["cpu_weight_backup"], + # distributed + "host": host, + "port": port, + "nccl_port": nccl_port, + "nnodes": nnodes, + "node_rank": node_rank, + "dist_init_addr": dist_init_addr, + "gpu_id_step": 1, + "base_gpu_id": base, + # parallel + "tp_size": _gpus_per_engine // _pp_size, + "dp_size": sglang_cfg["sglang_cfg"]["dp_size"], + "pp_size": _pp_size, + "ep_size": sglang_cfg["sglang_cfg"]["ep_size"], + # always skip warmup to prevent warmup timeout. + "skip_server_warmup": sglang_cfg["sglang_cfg"]["skip_server_warmup"], + # always enable draft weights cpu backup so that we run training without mtp weights. + "enable_draft_weights_cpu_backup": True, + } + + for key in [ + "dtype", + "kv_cache_dtype", + "context_length", + "max_running_requests", + "chunked_prefill_size", + "max_prefill_tokens", + "schedule_policy", + "schedule_conservativeness", + "cpu_offload_gb", + "log_level", + "mem_fraction_static", + "quantization", + "fp8_gemm_runner_backend", + "moe_runner_backend", + "allow_auto_truncate", + "disable_piecewise_cuda_graph", + "disable_cuda_graph", + # CUDA graph batch-size cap (Optional[int], default None). + "cuda_graph_max_bs", + # DP-attention switch (newer sglang forks): replicates attention along + # ``dp_size`` while keeping MoE/MLP under TP. + "enable_dp_attention", + # MoE all-to-all backend: "none" | "deepep" | "mooncake" | "mori" | + # "ascend_fuseep" | "flashinfer". Replaces the older + # ``enable_ep_moe`` boolean knob. + "moe_a2a_backend", + # DeepEP routing mode (used when ``moe_a2a_backend == "deepep"``): + # "auto" | "normal" | "low_latency". + "deepep_mode", + ]: + if key in sglang_cfg["sglang_cfg"]: + value = sglang_cfg["sglang_cfg"][key] + if key == "quantization" and isinstance(value, dict): + value = value.get("scheme") + kwargs[key] = value + + return kwargs diff --git a/nemo_rl/models/generation/sglang/utils.py b/nemo_rl/models/generation/sglang/utils.py deleted file mode 100644 index 7460302b5a..0000000000 --- a/nemo_rl/models/generation/sglang/utils.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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. - -import asyncio -import threading - - -class AsyncLoopThread: - """A background event loop thread for running async operations in Ray actors. - - This class creates a dedicated thread with its own event loop, allowing - synchronous Ray actor methods to execute async coroutines without blocking - the main actor thread. This is necessary because run_coroutine_threadsafe - requires the event loop to be in a different thread. - """ - - def __init__(self): - self.loop = asyncio.new_event_loop() - self._ready = threading.Event() - self._thread = threading.Thread(target=self._start_loop, daemon=True) - self._thread.start() - if not self._ready.wait(timeout=5.0): - raise RuntimeError("Event loop thread failed to start within 5 seconds") - - def _start_loop(self): - """Run the event loop in the background thread.""" - asyncio.set_event_loop(self.loop) - self._ready.set() - self.loop.run_forever() - - def run(self, coro): - """Schedule a coroutine onto the loop and block until it's done. - - Args: - coro: The coroutine to execute - - Returns: - The result of the coroutine - """ - if not self.loop.is_running(): - raise RuntimeError("Event loop is not running") - future = asyncio.run_coroutine_threadsafe(coro, self.loop) - result = future.result() - return result - - def shutdown(self): - """Shutdown the event loop and wait for the thread to finish.""" - if self.loop.is_running(): - self.loop.call_soon_threadsafe(self.loop.stop) - self._thread.join(timeout=2.0) - if not self.loop.is_closed(): - self.loop.close() diff --git a/nemo_rl/models/generation/sglang/utils/async_utils.py b/nemo_rl/models/generation/sglang/utils/async_utils.py new file mode 100644 index 0000000000..8d4c3b9d99 --- /dev/null +++ b/nemo_rl/models/generation/sglang/utils/async_utils.py @@ -0,0 +1,25 @@ +import asyncio +import threading + + +class AsyncLoopThread: + def __init__(self): + self.loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._start_loop, daemon=True) + self._thread.start() + + def _start_loop(self): + asyncio.set_event_loop(self.loop) + self.loop.run_forever() + + def run(self, coro): + # Schedule a coroutine onto the loop and block until it's done + return asyncio.run_coroutine_threadsafe(coro, self.loop).result() + + def close(self, timeout: float | None = 5.0) -> None: + """Stop the event loop, join the worker thread, and close the loop. Idempotent.""" + if self._thread.is_alive(): + self.loop.call_soon_threadsafe(self.loop.stop) + self._thread.join(timeout=timeout) + if not self.loop.is_closed(): + self.loop.close() diff --git a/nemo_rl/models/generation/sglang/utils/http_utils.py b/nemo_rl/models/generation/sglang/utils/http_utils.py new file mode 100644 index 0000000000..0a0f1b367d --- /dev/null +++ b/nemo_rl/models/generation/sglang/utils/http_utils.py @@ -0,0 +1,195 @@ +import asyncio +import json +import logging + +import httpx +import ray +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + +from nemo_rl.models.generation.sglang.config import SGLangConfig + +logger = logging.getLogger(__name__) + + +async def _post(client, url, payload, max_retries=10, action="post"): + retry_count = 0 + while retry_count < max_retries: + try: + if action in ("delete", "get"): + assert not payload + response = await getattr(client, action)(url) + else: + response = await getattr(client, action)(url, json=payload or {}) + response.raise_for_status() + try: + output = response.json() + except json.JSONDecodeError: + output = response.text + except Exception as e: + retry_count += 1 + + if isinstance(e, httpx.HTTPStatusError): + response_text = e.response.text + else: + response_text = None + + logger.info( + f"Error: {e}, retrying... (attempt {retry_count}/{max_retries}, url={url}, response={response_text})" + ) + if retry_count >= max_retries: + logger.info( + f"Max retries ({max_retries}) reached, failing... (url={url})" + ) + raise e + await asyncio.sleep(1) + continue + break + + return output + + +class HttpClient: + """HTTP client wrapper with optional Ray-based distributed POST dispatch.""" + + def __init__(self, args: SGLangConfig | None = None): + self._client: httpx.AsyncClient | None = None + self._client_concurrency: int = 0 + self._distributed_post_enabled: bool = False + self._post_actors: list[object] = [] + self._post_actor_idx: int = 0 + + if args is not None: + self.init(args) + + def init(self, args: SGLangConfig) -> None: + """Configure HTTP client limits and optional distributed POST actors.""" + server_cfg = args.get("sglang_server") or {} + if not server_cfg.get("num_gpus"): + return + + self._client_concurrency = ( + server_cfg["sglang_server_concurrency"] + * server_cfg["num_gpus"] + // server_cfg["num_gpus_per_engine"] + ) + + router_cfg = args.get("sglang_router") or {} + if router_cfg.get("use_distributed_post"): + self._init_ray_distributed_post(args) + self._distributed_post_enabled = True + + def _get_client(self) -> httpx.AsyncClient: + if self._client is None: + if self._client_concurrency > 0: + self._client = httpx.AsyncClient( + limits=httpx.Limits(max_connections=self._client_concurrency), + timeout=httpx.Timeout(None), + ) + else: + self._client = httpx.AsyncClient(timeout=httpx.Timeout(None)) + return self._client + + def _next_actor(self): + if not self._post_actors: + return None + actor = self._post_actors[self._post_actor_idx % len(self._post_actors)] + self._post_actor_idx = (self._post_actor_idx + 1) % len(self._post_actors) + return actor + + def _init_ray_distributed_post(self, args: SGLangConfig) -> None: + """Initialize one or more Ray async actors per node for HTTP POST.""" + if self._post_actors: + return # Already initialized + + # Discover alive, schedulable nodes. Filter out nodes with CPU=0 + # (e.g. an unschedulable head node) — placing an actor there would hang. + nodes = [ + n + for n in ray.nodes() + if n.get("Alive") and n.get("Resources", {}).get("CPU", 0) > 0 + ] + if not nodes: + raise RuntimeError("No alive Ray nodes to place HTTP POST actors.") + + @ray.remote + class _HttpPosterActor: + def __init__(self, concurrency: int): + # Lazy creation to this actor's event loop + self._client = httpx.AsyncClient( + limits=httpx.Limits(max_connections=max(1, concurrency)), + timeout=httpx.Timeout(None), + ) + + async def do_post(self, url, payload, max_retries=10, action="post"): + return await _post( + self._client, url, payload, max_retries, action=action + ) + + created = [] + per_actor_conc = max(1, (self._client_concurrency + len(nodes)) // len(nodes)) + + for node in nodes: + node_id = node["NodeID"] + scheduling = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) + for _ in range(args["sglang_server"]["num_gpus_per_engine"]): + actor = _HttpPosterActor.options( + name=None, + lifetime="detached", + scheduling_strategy=scheduling, + max_concurrency=per_actor_conc, + # Use tiny CPU to schedule + num_cpus=0.001, + ).remote(per_actor_conc) + created.append(actor) + + self._post_actors = created + + async def post(self, url, payload, max_retries=10, action="post"): + if self._distributed_post_enabled and self._post_actors: + try: + actor = self._next_actor() + if actor is not None: + # Use a thread to avoid blocking the event loop on ray.get + obj_ref = actor.do_post.remote( + url, payload, max_retries, action=action + ) + return await asyncio.to_thread(ray.get, obj_ref) + except Exception as e: + logger.info( + f"[http_utils] Distributed POST failed, falling back to local: {e} (url={url})" + ) + # fall through to local + + return await _post(self._get_client(), url, payload, max_retries, action=action) + + async def get(self, url): + response = await self._get_client().get(url) + response.raise_for_status() + output = response.json() + return output + + def shutdown(self) -> None: + """Kill HTTP POST actors created by this client.""" + if not self._post_actors: + return + + for actor in self._post_actors: + try: + ray.kill(actor) + except Exception: + pass + self._post_actors = [] + self._post_actor_idx = 0 + self._distributed_post_enabled = False + + async def aclose(self) -> None: + """Close local HTTP resources and kill distributed POST actors.""" + self.shutdown() + if self._client is not None: + await self._client.aclose() + self._client = None + + +def init_http_client(args: SGLangConfig) -> HttpClient: + """Create an HTTP client for SGLang requests.""" + return HttpClient(args) diff --git a/nemo_rl/models/generation/sglang/utils/ray_utils.py b/nemo_rl/models/generation/sglang/utils/ray_utils.py new file mode 100644 index 0000000000..180968b11b --- /dev/null +++ b/nemo_rl/models/generation/sglang/utils/ray_utils.py @@ -0,0 +1,168 @@ +import ipaddress +import os +import random +import socket + +import ray + +# 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 +# visibility themselves. +NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [ + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES", + "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES", + "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS", + "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR", +] + + +class RayActor: + """Base class for Ray actors providing node IP / free port helpers.""" + + @staticmethod + def _get_current_node_ip_and_free_port(start_port=10000, consecutive=1): + return get_current_node_ip(), get_free_port( + start_port=start_port, consecutive=consecutive + ) + + def get_master_addr_and_port(self): + return self.master_addr, self.master_port + + +@ray.remote +class Lock(RayActor): + def __init__(self): + self._locked = False # False: unlocked, True: locked + + def acquire(self): + """Try to acquire the lock. + + Returns True if acquired, False otherwise. Caller should retry until + it returns True. + """ + if not self._locked: + self._locked = True + return True + 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 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 + + +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) + return True + except OSError: + return False + except OverflowError: + return False + + +def get_host_info(): + hostname = socket.gethostname() + + def _is_loopback(ip): + return ip.startswith("127.") or ip == "::1" + + def _resolve_ip(family, test_target_ip): + """Attempt to get the local LAN IP for the specific family (IPv4/IPv6). + + Strategy: UDP Probe (Preferred) -> Hostname Resolution (Fallback) -> None. + """ + # Strategy 1: UDP Connect Probe (Most accurate, relies on routing table) + # Useful when the machine has a default gateway or internet access. + try: + with socket.socket(family, socket.SOCK_DGRAM) as s: + # The IP doesn't need to be reachable, but the routing table must exist. + s.connect((test_target_ip, 80)) + ip = s.getsockname()[0] + if not _is_loopback(ip): + return ip + except Exception: + pass # Route unreachable or network error, move to next strategy. + + # Strategy 2: Hostname Resolution (Fallback for offline clusters) + # Useful for offline environments where UDP connect fails but /etc/hosts is configured. + try: + # getaddrinfo allows specifying the family (AF_INET or AF_INET6) + # Result format: [(family, type, proto, canonname, sockaddr), ...] + infos = socket.getaddrinfo( + hostname, None, family=family, type=socket.SOCK_STREAM + ) + + for info in infos: + ip = info[4][0] # The first element of sockaddr is the IP + # Must filter out loopback addresses to avoid "127.0.0.1" issues + if not _is_loopback(ip): + return ip + except Exception: + pass + + return None + + prefer_ipv6 = os.getenv("PREFER_IPV6", "0").lower() in ("1", "true", "yes", "on") + local_ip = None + final_fallback = "127.0.0.1" + + if prefer_ipv6: + # [Strict Mode] IPv6 Only + # 1. Try UDP V6 Probe + # 2. Try Hostname Resolution (V6) + # If failed, fallback to V6 loopback. Never mix with V4. + local_ip = _resolve_ip(socket.AF_INET6, "2001:4860:4860::8888") + final_fallback = "::1" + else: + # [Strict Mode] IPv4 Only (Default) + # 1. Try UDP V4 Probe + # 2. Try Hostname Resolution (V4) + # If failed, fallback to V4 loopback. Never mix with V6. + local_ip = _resolve_ip(socket.AF_INET, "8.8.8.8") + final_fallback = "127.0.0.1" + + return hostname, local_ip or final_fallback + + +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=10000, 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 + + +def _wrap_ipv6(host): + """Wrap IPv6 address in [] if needed.""" + try: + ipaddress.IPv6Address(host.strip("[]")) + return f"[{host.strip('[]')}]" + except ipaddress.AddressValueError: + return host diff --git a/nemo_rl/models/generation/sglang/utils/router_utils.py b/nemo_rl/models/generation/sglang/utils/router_utils.py new file mode 100644 index 0000000000..701dc186ab --- /dev/null +++ b/nemo_rl/models/generation/sglang/utils/router_utils.py @@ -0,0 +1,37 @@ +import logging +import multiprocessing + +logger = logging.getLogger(__name__) + + +def run_router(args): + try: + from sglang_router.launch_router import launch_router + + router = launch_router(args) + if router is None: + return 1 + return 0 + except Exception: + # Runs inside a subprocess; surface the full traceback at ERROR level + # so it isn't filtered by INFO config, and re-raise so the subprocess + # exits non-zero (caller asserts on ``_process.is_alive()``). + logger.exception("sglang router failed to launch") + raise + + +def terminate_process(process: multiprocessing.Process, timeout: float = 1.0) -> None: + """Terminate a process gracefully, with forced kill as fallback. + + Args: + process: The process to terminate + timeout: Seconds to wait for graceful termination before forcing kill + """ + if not process.is_alive(): + return + + process.terminate() + process.join(timeout=timeout) + if process.is_alive(): + process.kill() + process.join() diff --git a/nemo_rl/models/generation/vllm/quantization/fp8.py b/nemo_rl/models/generation/vllm/quantization/fp8.py index b01e889139..ae496c45cd 100644 --- a/nemo_rl/models/generation/vllm/quantization/fp8.py +++ b/nemo_rl/models/generation/vllm/quantization/fp8.py @@ -644,8 +644,8 @@ def process_weights_after_loading_kv(self, layer) -> None: else: prob_scale = 1.0 - is_singleton_float = ( - lambda x: isinstance(x, float) + is_singleton_float = lambda x: ( + isinstance(x, float) or isinstance(x, torch.Tensor) and x.numel() == 1 and x.is_floating_point() diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 0faaad17a1..e9d3e3162a 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -781,6 +781,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 e4f324efed..ce94a02094 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -181,8 +181,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() @@ -202,11 +210,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"], @@ -214,10 +224,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 f6facfc748..921e72174a 100644 --- a/nemo_rl/models/policy/interfaces.py +++ b/nemo_rl/models/policy/interfaces.py @@ -189,12 +189,16 @@ def stream_weights_via_ipc_zmq( pass def stream_weights_via_http( - self, sglang_url_to_gpu_uuids: dict[str, list[str]] + self, + rollout_engine_urls: list[str], + num_gpus_per_engine: int, ) -> list[ray.ObjectRef]: - """Stream model weights to SGLang servers via HTTP API. + """Stream model weights to colocated SGLang engines via CUDA IPC over HTTP. Args: - sglang_url_to_gpu_uuids: Dict mapping SGLang server URL to list of GPU UUIDs it uses + rollout_engine_urls: ``http://host:port`` base URLs of each + engine's ``node_rank=0`` SGLang HTTP server. + num_gpus_per_engine: TP size per SGLang engine. """ raise NotImplementedError( "stream_weights_via_http is not implemented for this policy worker" diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index e21bd6dac6..23c7ecb335 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -883,16 +883,102 @@ def stream_weights_via_ipc_zmq( return futures def stream_weights_via_http( - self, sglang_url_to_gpu_uuids: dict[str, list[str]] + self, + rollout_engine_urls: list[str], + num_gpus_per_engine: int, + engine_gpu_counts: Optional[list[int]] = None, + engine_gpu_offsets: Optional[list[int]] = None, ) -> list[ray.ObjectRef]: - """Send the weights to SGLang servers via HTTP API. + """Send the weights to colocated SGLang engines via CUDA IPC over HTTP. Args: - sglang_url_to_gpu_uuids: Dict mapping SGLang server URL to list of GPU UUIDs it uses + 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. + num_gpus_per_engine: TP size per SGLang engine. Used as the + fallback dense layout when ``engine_gpu_counts`` is None. + engine_gpu_counts: Optional explicit per-engine GPU count. + engine_gpu_offsets: Optional explicit per-engine GPU offset. """ futures = self.worker_group.run_all_workers_single_data( "stream_weights_via_http", - sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids, + rollout_engine_urls=rollout_engine_urls, + num_gpus_per_engine=num_gpus_per_engine, + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + return futures + + 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. + + Megatron-only entry point. The FSDP path runs the same setup lazily + from inside ``stream_weights_via_http``. + """ + futures = self.worker_group.run_all_workers_single_data( + "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 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 diff --git a/nemo_rl/models/policy/torch_reductions_utils.py b/nemo_rl/models/policy/torch_reductions_utils.py new file mode 100644 index 0000000000..99a6757424 --- /dev/null +++ b/nemo_rl/models/policy/torch_reductions_utils.py @@ -0,0 +1,225 @@ +# Copyright 2023-2024 SGLang Team +# 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 io +from dataclasses import dataclass +from multiprocessing.reduction import ForkingPickler +from typing import Callable, List, Tuple, Union + +import pybase64 +import torch +from torch.multiprocessing import reductions + + +class MultiprocessingSerializer: # pragma: no cover + """Serialize/deserialize Python objects using ForkingPickler for IPC. + + This class enables serialization of objects (including CUDA tensors with IPC + handles) for transfer between processes via HTTP or other mechanisms. + + Original source (sglang v0.5.2): + https://github.com/sgl-project/sglang/blob/v0.5.2/python/sglang/srt/utils.py#L589-L623 + """ + + @staticmethod + def serialize(obj, output_str: bool = False): + """Serialize a Python object using ForkingPickler. + + Args: + obj: The object to serialize. + output_str (bool): If True, return a base64-encoded string instead of raw bytes. + + Returns: + bytes or str: The serialized object. + """ + buf = io.BytesIO() + ForkingPickler(buf).dump(obj) + buf.seek(0) + output = buf.read() + + if output_str: + # Convert bytes to base64-encoded string + output = pybase64.b64encode(output).decode("utf-8") + + return output + + @staticmethod + def deserialize(data): + """Deserialize a previously serialized object. + + Args: + data (bytes or str): The serialized data, optionally base64-encoded. + + Returns: + The deserialized Python object. + """ + if isinstance(data, str): + # Decode base64 string to bytes + data = pybase64.b64decode(data, validate=True) + + return ForkingPickler.loads(data) + + +def monkey_patch_torch_reductions(): + """Monkey patching before Torch https://github.com/pytorch/pytorch/pull/149248 is fixed.""" + if hasattr(reductions, "_reduce_tensor_original"): + return + reductions._reduce_tensor_original = reductions.reduce_tensor + reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor + + reductions.reduce_tensor = _reduce_tensor_modified + reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified + reductions.init_reductions() + + +# The signature has not been changed for years, and we will not need this when the next version is released, +# so it looks safe to use a constant. +_REDUCE_TENSOR_ARG_DEVICE_INDEX = 6 + + +def _reduce_tensor_modified(*args, **kwargs): + output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs) + output_args = _modify_tuple( + output_args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_to_uuid + ) + return output_fn, output_args + + +def _rebuild_cuda_tensor_modified(*args): + args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_from_maybe_uuid) + return reductions._rebuild_cuda_tensor_original(*args) + + +def _device_to_uuid(device: int) -> str: + return str(torch.cuda.get_device_properties(device).uuid) + + +def _device_from_maybe_uuid(device_maybe_uuid: Union[int, str]) -> int: + if isinstance(device_maybe_uuid, int): + return device_maybe_uuid + + if isinstance(device_maybe_uuid, str): + for device in range(torch.cuda.device_count()): + if str(torch.cuda.get_device_properties(device).uuid) == device_maybe_uuid: + return device + raise Exception("Invalid device_uuid=" + device_maybe_uuid) + + raise Exception(f"Unknown type: {device_maybe_uuid=}") + + +def _modify_tuple(t, index: int, modifier: Callable): + return *t[:index], modifier(t[index]), *t[index + 1 :] + + +@dataclass +class FlattenedTensorMetadata: + """Metadata for a tensor in a flattened bucket.""" + + name: str + shape: torch.Size + dtype: torch.dtype + start_idx: int + end_idx: int + numel: int + + +class FlattenedTensorBucket: + """A bucket that flattens multiple tensors into a single tensor. + + Provides efficient batched processing while preserving all metadata + needed for reconstruction. + """ + + # This field is solely for users of to check whether the class supports this feature + supports_multi_dtypes = True + + def __init__( + self, + named_tensors: List[Tuple[str, torch.Tensor]] = None, + flattened_tensor: torch.Tensor = None, + metadata: List[FlattenedTensorMetadata] = None, + ): + """Initialize a tensor bucket from a list of named tensors OR from pre-flattened data. + + Args: + named_tensors: List of (name, tensor) tuples (for creating new bucket) + flattened_tensor: Pre-flattened tensor (for reconstruction) + metadata: Pre-computed metadata (for reconstruction) + """ + if named_tensors is not None: + # Create bucket from named tensors + self.metadata: List[FlattenedTensorMetadata] = [None] * len(named_tensors) + self.flattened_tensor: torch.Tensor = None + + if not named_tensors: + raise ValueError("Cannot create empty tensor bucket") + + # Collect metadata and flatten tensors + current_idx = 0 + flattened_tensors: List[torch.Tensor] = [None] * len(named_tensors) + + for i, (name, tensor) in enumerate(named_tensors): + flattened = tensor.flatten().view(torch.uint8) + flattened_tensors[i] = flattened + + # Store metadata + + numel = flattened.numel() + metadata_obj = FlattenedTensorMetadata( + name=name, + shape=tensor.shape, + dtype=tensor.dtype, + start_idx=current_idx, + end_idx=current_idx + numel, + numel=numel, + ) + self.metadata[i] = metadata_obj + current_idx += numel + + # Concatenate all flattened tensors + self.flattened_tensor = torch.cat(flattened_tensors, dim=0) + else: + # Initialize from pre-flattened data + if flattened_tensor is None or metadata is None: + raise ValueError( + "Must provide either named_tensors or both flattened_tensor and metadata" + ) + self.flattened_tensor = flattened_tensor + self.metadata = metadata + + def get_flattened_tensor(self) -> torch.Tensor: + """Get the flattened tensor containing all bucket tensors.""" + return self.flattened_tensor + + def get_metadata(self) -> List[FlattenedTensorMetadata]: + """Get metadata for all tensors in the bucket.""" + return self.metadata + + def reconstruct_tensors(self) -> List[Tuple[str, torch.Tensor]]: + """Reconstruct original tensors from flattened tensor with optimized performance. + + Uses memory-efficient operations to minimize allocations and copies. + """ + # preallocate the result list + reconstructed = [None] * len(self.metadata) + + for i, meta in enumerate(self.metadata): + tensor = ( + self.flattened_tensor[meta.start_idx : meta.end_idx] + .view(meta.dtype) + .reshape(meta.shape) + ) + + reconstructed[i] = (meta.name, tensor) + + return reconstructed diff --git a/nemo_rl/models/policy/utils.py b/nemo_rl/models/policy/utils.py index bbd2e6d2f6..3bed6f4986 100644 --- a/nemo_rl/models/policy/utils.py +++ b/nemo_rl/models/policy/utils.py @@ -15,8 +15,9 @@ import gc import os import traceback +from datetime import timedelta from enum import Enum -from typing import Any, Dict, Optional, cast +from typing import Any, Dict, Iterable, Optional import requests import torch @@ -381,125 +382,272 @@ def rebuild_cuda_tensor_from_ipc( return func(*list_args) -def stream_weights_via_http_impl( - params_generator, - sglang_url_to_gpu_uuids: dict[str, list[str]], - rank: int, - worker_name: str, - current_device_uuid: str, -) -> None: - """Stream weights to SGLang servers via HTTP API (update_weights_from_tensor). - - Flow: Each rank creates IPC handler → gather handlers in rank order → send list → SGLang matches by tp_rank index +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 - Key points: - - Each rank creates handler on its own GPU - - Handlers are gathered in rank order: [rank0_handler, rank1_handler, ...] - - List index = rank = GPU ID - - SGLang automatically matches: handler = serialized_handlers[tp_rank] - Args: - params_generator: Generator yielding (name, tensor) pairs - sglang_url_to_gpu_uuids: Dict mapping SGLang server URL to list of GPU UUIDs it uses - rank: Worker rank for logging - worker_name: Name of the worker for logging - current_device_uuid: UUID of the current training worker's GPU +def connect_colocate_topology( + *, + engine_gpu_counts: list[int], + engine_gpu_offsets: Optional[list[int]] = None, + worker_state: dict, + monkey_patch_fn=None, +) -> 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. """ - from nemo_rl.models.generation.sglang.sglang_copied_utils import ( - MultiprocessingSerializer, - ) + 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 - print("[sglang refit details] entering stream_weights_via_http_impl") + if monkey_patch_fn is not None and not worker_state.get("_ipc_monkey_patched"): + monkey_patch_fn() + worker_state["_ipc_monkey_patched"] = True - target_urls = [ - url - for url, uuids in sglang_url_to_gpu_uuids.items() - if current_device_uuid in uuids + 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() + 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: + 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) + + +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 ] - if not target_urls: - raise RuntimeError( - f"{worker_name} (rank {rank}): No matching SGLang server found for GPU UUID {current_device_uuid}. " - f"Available servers: {list(sglang_url_to_gpu_uuids.keys())}" + 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) ) - if len(target_urls) > 1: - print( - f"[WARNING] {worker_name} (rank {rank}): GPU UUID {current_device_uuid} matches multiple SGLang servers: {target_urls}. " - f"Using the first one: {target_urls[0]}" + 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, + ) + + if my_rank != gather_src: + return + + num_dtypes = len(gathered[0]) + assert num_dtypes > 0 + for i in range(num_dtypes): + body = { + "serialized_named_tensors": [g[i] for g in gathered], + "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" ) - target_urls = [target_urls[0]] + if not success: + raise RuntimeError( + f"Weight sync failed on rollout engine: {error_msg}. " + f"Check SGLang version compatibility." + ) - base_url = target_urls[0] - url = f"{base_url}/update_weights_from_tensor" - sglang_gpu_uuids = sglang_url_to_gpu_uuids[base_url] - ipc_gather_group, ipc_gather_src, matching_ranks = _setup_ipc_gather_group( - rank, current_device_uuid, sglang_gpu_uuids, sglang_url_to_gpu_uuids - ) - print( - f"[sglang refit] {worker_name} (rank {rank}): ipc_gather_group={ipc_gather_group}, ipc_gather_src={ipc_gather_src}, matching_ranks={matching_ranks}" - ) - tensor_count = 0 +def stream_weights_via_http_impl( + 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, + worker_state: dict, + *, + engine_gpu_counts: Optional[list[int]] = None, + engine_gpu_offsets: Optional[list[int]] = None, +) -> None: + """Stream FSDP weights to colocated SGLang engines via CUDA IPC over HTTP. - try: - tensor_list = list(params_generator) - total_tensors = len(tensor_list) + 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. + """ + from nemo_rl.models.policy.torch_reductions_utils import ( + FlattenedTensorBucket, + MultiprocessingSerializer, + monkey_patch_torch_reductions, + ) - if rank == ipc_gather_src: - print( - f"[sglang refit details] {worker_name}: Starting weight update - " - f"Total parameters to update: {total_tensors}", - flush=True, - ) + rollout_engine_urls = list(rollout_engine_urls) - for idx, (name, tensor) in enumerate(tensor_list): - torch.cuda.current_stream().synchronize() - tensor = tensor.contiguous().cuda() + if engine_gpu_counts is None: + engine_gpu_counts = [num_gpus_per_engine] * len(rollout_engine_urls) + if engine_gpu_offsets is None: + engine_gpu_offsets = _derive_engine_gpu_offsets(engine_gpu_counts) - named_tensors = [(name, tensor)] - serialized_handler = MultiprocessingSerializer.serialize( - named_tensors, output_str=True - ) - # output_str=True ensures the return type is str - serialized_handler_str = cast(str, serialized_handler) - - gathered_handlers = _gather_ipc_handlers( - serialized_handler_str, - ipc_gather_group, - ipc_gather_src, - rank, - matching_ranks, - ) + connect_colocate_topology( + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + worker_state=worker_state, + monkey_patch_fn=monkey_patch_torch_reductions, + ) - if rank == ipc_gather_src and gathered_handlers is not None: - _send_tensor_to_sglang( - url, - name, - gathered_handlers, - tensor.shape, - str(tensor.dtype), - flush_cache=False, - ) - tensor_count += 1 + worker_state["weight_version"] = worker_state.get("weight_version", 0) + 1 + weight_version = worker_state["weight_version"] + gather_src = worker_state["_ipc_gather_src"] + gather_group = worker_state["_ipc_gather_group"] + engine_idx = worker_state["_ipc_engine_index"] + engine_url = rollout_engine_urls[engine_idx] if engine_idx is not None else None - del tensor, serialized_handler - if rank == ipc_gather_src: - del gathered_handlers - torch.cuda.empty_cache() + if gather_group is None: + # Placeholder rank not covered by any engine: drain quietly. + return - if rank == ipc_gather_src: - print( - f"[sglang refit details] {worker_name}: Weight update completed - " - f"Successfully updated {tensor_count}/{total_tensors} parameters to SGLang server: {base_url}", - flush=True, - ) - if tensor_count != total_tensors: - print( - f"[sglang refit details] {worker_name}: WARNING - Expected {total_tensors} tensors, " - f"but only sent {tensor_count}", - flush=True, + 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, + ) + + if dist.get_rank() == gather_src: + # Mirror SGLangGenerationWorker.flush_cache: the endpoint returns + # non-200 while requests are still pending, so retry up to 60s. + import time + + 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) + else: + raise TimeoutError(f"Timeout while flushing cache at {engine_url}.") except Exception as e: print( @@ -507,132 +655,443 @@ def stream_weights_via_http_impl( f"{traceback.format_exc()}" ) raise - finally: gc.collect() torch.cuda.empty_cache() -def _setup_ipc_gather_group( - rank: int, - current_device_uuid: str, - sglang_gpu_uuids: list[str], - sglang_url_to_gpu_uuids: dict[str, list[str]], -) -> tuple[Optional[dist.ProcessGroup], Optional[int], Optional[list[int]]]: - """Setup gather configuration for IPC handlers. +def _check_weight_sync_results(results: list) -> None: + from collections.abc import Mapping - Returns: - Tuple of (gather_group, gather_src_rank, matching_ranks) - - gather_group: None (use default FSDP group) - - gather_src_rank: The rank that will collect and send to SGLang server - - matching_ranks: List of ranks that belong to the same SGLang server + 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 + + if success is False: + raise RuntimeError( + f"SGLang weight sync failed on rollout engine: {error_msg}. " + "Check SGLang version compatibility." + ) + + +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: + """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. """ - if not dist.is_initialized(): - return None, None, None + import ray + + from nemo_rl.models.policy.torch_reductions_utils import ( + FlattenedTensorBucket, + MultiprocessingSerializer, + ) - world_size = dist.get_world_size() + gather_group = worker_state.get("_ipc_gather_group") + gather_src = worker_state.get("_ipc_gather_src") + engine_idx = worker_state.get("_ipc_engine_index") + + 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 + + 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() - all_ranks_uuids = [None] * world_size - dist.all_gather_object(all_ranks_uuids, current_device_uuid) + try: + for bucket in bucket_iterator: + if not bucket: + continue + + # No async-collective ``.wait()`` here — Megatron's AutoBridge + # yields plain ``torch.Tensor``, no DTensor wrapping. + + if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False): + by_dtype: dict = {"dtype": list(bucket)} + else: + 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) + ) - matching_ranks = [ - r for r, uuid in enumerate(all_ranks_uuids) if uuid in sglang_gpu_uuids - ] + 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, + ) - if len(matching_ranks) == 0: - return None, None, None + 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() - matching_ranks = sorted(matching_ranks) - gather_src = matching_ranks[0] + return None - return None, gather_src, matching_ranks +def find_free_port() -> int: + """Return a currently-free TCP port on the local node.""" + import socket -def _gather_ipc_handlers( - serialized_handler: str, - gather_group: Optional[dist.ProcessGroup], - gather_src: Optional[int], - rank: int, - matching_ranks: Optional[list[int]] = None, -) -> Optional[list[str]]: - """Gather IPC handlers from all ranks in the default FSDP group, then filter by server. + with socket.socket() as sock: + sock.bind(("", 0)) + return int(sock.getsockname()[1]) - Args: - serialized_handler: Serialized IPC handler from this rank - gather_group: Process group (None means use default FSDP group) - gather_src: Rank that will collect and filter handlers - rank: Current rank - matching_ranks: List of ranks that belong to the same SGLang server - Returns: - List of serialized handlers in rank order (only on gather_src rank), None otherwise - The list contains handlers from matching_ranks only, in rank order +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. """ - if gather_src is None: - return None + from torch.distributed.distributed_c10d import ( + Backend, + PrefixStore, + _new_process_group_helper, + _world, + default_pg_timeout, + rendezvous, + ) - if not dist.is_initialized(): - return None + assert (store is None) or (init_method is None), ( + "Cannot specify both init_method and store." + ) - world_size = dist.get_world_size() + 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" + pg, _ = _new_process_group_helper( + world_size, + rank, + [], + backend, + store, + group_name=group_name, + **{pg_options_kw: pg_options}, + timeout=timeout, + ) - all_handlers: list[Optional[str]] = [None for _ in range(world_size)] - dist.all_gather_object(all_handlers, serialized_handler) - all_handlers_str = cast(list[str], all_handlers) + # 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 - if rank == gather_src and matching_ranks is not None: - filtered_handlers: list[str] = [all_handlers_str[r] for r in matching_ranks] - return filtered_handlers - else: - return None +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. -def _send_tensor_to_sglang( - url: str, - tensor_name: str, - gathered_handlers: list[str], - shape: torch.Size, - dtype: str, - flush_cache: bool = False, + 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 + + master_address = ray._private.services.get_node_ip_address() + master_port = find_free_port() + 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: - """Send gathered IPC handlers to SGLang server via HTTP. + """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 - Key: gathered_handlers are in rank order [rank0, rank1, ...] - SGLang will automatically match: handler = serialized_handlers[tp_rank] - Args: - url: SGLang server URL - tensor_name: Name of the tensor - gathered_handlers: List of serialized IPC handlers in rank order - shape: Tensor shape - dtype: Tensor dtype - flush_cache: Whether to flush cache after this tensor (for last tensor) +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. """ - payload = { - "serialized_named_tensors": gathered_handlers, - "flush_cache": flush_cache, - } + return dict(policy_generation.sglang_cfg["sglang_cfg"].get("quantization") or {}) - try: - response = requests.post( - url, - json=payload, - headers={"Content-Type": "application/json"}, - timeout=120, + +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] + devices = [str(tensor.device) for _, tensor in bucket] + total_bytes = sum(t.numel() * t.element_size() for _, t in bucket) + print( + f"[BCAST bucket={bucket_idx}] n={len(bucket)} bytes={total_bytes} " + f"first={names[0]} last={names[-1]} devs={set(devices)} dtypes={set(dtypes)}", + flush=True, + ) + + print( + f"[BCAST bucket={bucket_idx}] acquiring rollout_engine_lock...", flush=True ) - response.raise_for_status() - except requests.exceptions.HTTPError as e: - error_msg = f"Failed to send tensor '{tensor_name}' to {url}: {e}" + while not ray.get(rollout_engine_lock.acquire.remote()): + _time.sleep(0.1) + print(f"[BCAST bucket={bucket_idx}] lock acquired", flush=True) try: - error_detail = response.text - error_msg += f"\nResponse status: {response.status_code}" - error_msg += f"\nResponse body: {error_detail[:500]}" - except: - pass - print(f"[sglang refit] {error_msg}", flush=True) - raise RuntimeError(error_msg) from e - except Exception as e: - raise RuntimeError( - f"Failed to send tensor '{tensor_name}' to {url}: {e}" - ) from e + print( + f"[BCAST bucket={bucket_idx}] kicking engine.update_weights_from_distributed.remote() RPCs...", + flush=True, + ) + 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 + ] + print( + f"[BCAST bucket={bucket_idx}] issuing {len(bucket)} dist.broadcast async_op calls...", + flush=True, + ) + handles = [] + for i, (_, tensor) in enumerate(bucket): + handles.append( + dist.broadcast( + tensor.data, 0, group=model_update_group, async_op=True + ) + ) + print( + f"[BCAST bucket={bucket_idx}] all {len(handles)} broadcasts launched; waiting...", + flush=True, + ) + for i, handle in enumerate(handles): + handle.wait() + if i == 0 or (i + 1) == len(handles): + print( + f"[BCAST bucket={bucket_idx}] handle.wait() done {i + 1}/{len(handles)}", + flush=True, + ) + print( + f"[BCAST bucket={bucket_idx}] all broadcasts complete; ray.get(refs)...", + flush=True, + ) + ray.get(refs) + print( + f"[BCAST bucket={bucket_idx}] engine RPCs returned (engine done loading)", + flush=True, + ) + finally: + ray.get(rollout_engine_lock.release.remote()) + print(f"[BCAST bucket={bucket_idx}] lock released", flush=True) 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 2fa8a8e604..cd20ec011e 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -911,12 +911,27 @@ def stream_weights_via_ipc_zmq( @wrap_with_nvtx_name("dtensor_policy_worker_v2/stream_weights_via_http") def stream_weights_via_http( self, - sglang_url_to_gpu_uuids: dict[str, list[str]], + rollout_engine_urls: list[str], + num_gpus_per_engine: int, + buffer_size_bytes: int = 512 * 1024 * 1024, + engine_gpu_counts: Optional[list[int]] = None, + engine_gpu_offsets: Optional[list[int]] = None, ) -> None: - """Stream model weights to SGLang servers via HTTP API. + """Stream FSDP weights to colocated SGLang engines via CUDA IPC over HTTP. Args: - sglang_url_to_gpu_uuids: Dict mapping SGLang server URL to list of GPU UUIDs it uses + 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. + num_gpus_per_engine: TP size per SGLang engine. Used as the + fallback when ``engine_gpu_counts`` is not provided (dense + layout: engine ``i`` owns ranks ``[i*K, (i+1)*K)``). + buffer_size_bytes: Max bucket size in bytes before flushing. + engine_gpu_counts: Optional explicit per-engine GPU count. + engine_gpu_offsets: Optional explicit per-engine GPU start + offset. Use together with ``engine_gpu_counts`` to express + placeholder gaps or heterogeneous TP sizes. """ # Manually move model to cuda for cpu offload case if self.cpu_offload: @@ -924,34 +939,20 @@ def stream_weights_via_http( from nemo_rl.models.policy.utils import stream_weights_via_http_impl - # Get current GPU UUID - current_device_uuid = self.report_device_id() + if not hasattr(self, "_ipc_worker_state"): + self._ipc_worker_state: dict = {} - def dtensor_params_generator(): - """Generator that yields (name, tensor) pairs, converting DTensors to local tensors.""" - state_dict_items = sorted( - self.model.state_dict().items(), key=lambda x: x[0] - ) - for name, tensor in state_dict_items: - if isinstance(tensor, DTensor): - # Convert DTensor to full tensor for streaming - full_tensor = tensor.full_tensor() - # Convert to target dtype - yield ( - name, - full_tensor.to(self.dtype, non_blocking=True).contiguous(), - ) - else: - # Convert to target dtype - yield name, tensor.to(self.dtype, non_blocking=True).contiguous() - - # Use the HTTP implementation stream_weights_via_http_impl( - params_generator=dtensor_params_generator(), - sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids, + params_generator=dtensor_params_generator(self.model, self.dtype), + rollout_engine_urls=rollout_engine_urls, + num_gpus_per_engine=num_gpus_per_engine, rank=self.rank, + world_size=torch.distributed.get_world_size(), worker_name=str(self), - current_device_uuid=current_device_uuid, + buffer_size_bytes=buffer_size_bytes, + worker_state=self._ipc_worker_state, + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, ) @torch.no_grad() @@ -1154,3 +1155,69 @@ 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, # noqa: ARG001 — accepted for dispatch parity +) -> bool: + """Refit colocated SGLang engines from the FSDP/DTensor policy. + + Reuses the existing ``stream_weights_via_http`` path, which embeds its + own pause + flush_cache lifecycle inside the SGLang HTTP gateway. The + ``buffer_size_bytes`` argument is accepted for parity with the Megatron + dispatch but is not consumed here — the FSDP path uses the worker-side + fixed buffer size. Includes the optional fault-tolerance recover prelude. + """ + 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: + # Topology refresh runs lazily inside ``stream_weights_via_http_impl`` + # the first time a covered rank enters it after the layout changes. + 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" + ) + + rollout_engine_urls = ray.get([e.get_base_url.remote() for e in rollout_engines]) + futures_train = policy.stream_weights_via_http( + rollout_engine_urls=rollout_engine_urls, + num_gpus_per_engine=policy_generation.num_gpus_per_engine, + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + ray.get(futures_train) + 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 cf8d74ee04..b6109e0c44 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 @@ ColocatablePolicyInterface, LogprobOutputSpec, ) -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.nsys import wrap_with_nvtx_name @@ -134,7 +141,7 @@ def __init__( self.rank = get_rank_safe() # Step 1: Setup distributed - setup_distributed() + setup_distributed(config) # Step 2: Validate and setup model paths hf_model_name, pretrained_path, pt_checkpoint_exists = validate_model_paths( @@ -232,6 +239,15 @@ 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 = {} + 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 + def enable_forward_pre_hook(self): assert isinstance(self.model, DistributedDataParallel) self.model.enable_forward_pre_hook() @@ -1069,6 +1085,188 @@ 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. + """ + from nemo_rl.models.policy.torch_reductions_utils import ( + monkey_patch_torch_reductions, + ) + + 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, + monkey_patch_fn=monkey_patch_torch_reductions, + ) + + @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( @@ -1607,3 +1805,162 @@ 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 flush_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, + ) + + print("[REFIT-DIST 0] entering refit_sglang_distributed", flush=True) + sglang_quant = get_sglang_quantization_cfg(policy_generation) + target_precision = sglang_quant.get("scheme", "bf16") + + print("[REFIT-DIST 1] fetch_updatable_engines_with_recover...", flush=True) + ( + rollout_engines, + rollout_engine_lock, + num_new_engines, + engine_gpu_counts, + _engine_gpu_offsets, + ) = fetch_updatable_engines_with_recover(policy_generation) + print( + f"[REFIT-DIST 1] fetched: num_new_engines={num_new_engines} " + f"engine_gpu_counts={engine_gpu_counts}", + flush=True, + ) + + if num_new_engines > 0: + print( + "[REFIT-DIST 2] connect_sglang_rollout_engines_distributed (NCCL group setup)...", + flush=True, + ) + policy.connect_sglang_rollout_engines_distributed( + rollout_engines=rollout_engines, + engine_gpu_counts=engine_gpu_counts, + ) + print( + "[REFIT-DIST 2] connect_sglang_rollout_engines_distributed done", flush=True + ) + 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 flush_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 + print(f"[REFIT-DIST 3] pause_generation(mode={pause_mode!r})...", flush=True) + policy_generation.pause_generation(mode=pause_mode) + print("[REFIT-DIST 3] pause_generation done", flush=True) + if pause_mode != "in_place": + print("[REFIT-DIST 4] invalidate_kv_cache...", flush=True) + policy_generation.invalidate_kv_cache() + print("[REFIT-DIST 4] invalidate_kv_cache done", flush=True) + try: + print( + "[REFIT-DIST 5] update_weights_to_sglang_distributed (kicks engine RPC + NCCL broadcast)...", + flush=True, + ) + 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, + ) + print( + f"[REFIT-DIST 5] update_weights_to_sglang_distributed returned futures (n={len(futures) if futures is not None else 0}); ray.get...", + flush=True, + ) + ray.get(futures) + print("[REFIT-DIST 5] ray.get(futures) done", flush=True) + print("[REFIT-DIST 6] post_process_weights...", flush=True) + policy_generation.post_process_weights() + print("[REFIT-DIST 6] post_process_weights done", flush=True) + finally: + print("[REFIT-DIST 7] continue_generation (finally)...", flush=True) + policy_generation.continue_generation() + print("[REFIT-DIST 7] continue_generation done", flush=True) + print("[REFIT-DIST 8] refit_sglang_distributed complete", flush=True) + 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..5bece2a543 --- /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 ( + build_dynamic_skip_substrings, + quantize_mxfp8, + should_quantize, + SOURCE_FP8_SCALE_KEY_SUFFIX, + 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/pyproject.toml b/pyproject.toml index 17830c391c..bc4bd543e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,17 +13,23 @@ readme = { file = "README.md", content-type = "text/markdown" } name = "nemo-rl" dynamic = ["version", "readme"] description = "NeMo RL: A Scalable and Efficient Post-Training Library for Models Ranging from 1 GPU to 1000s, and from Tiny to >100B Parameters" -requires-python = ">=3.12" +requires-python = ">=3.13.9,<3.14" license = { text = "Apache 2.0" } +maintainers = [ + { name = "Yi-Fu Wu", email = "yifuw@nvidia.com" }, + { name = "Terry Kong", email = "terryk@nvidia.com" }, + { name = "Yuki Huang", email = "yukih@nvidia.com" }, + { name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }, +] dependencies = [ "setuptools", "pip", # Required for frozen environments; uv venv --seed may not reliably install pip "ninja", # for flash-attn parallel build - "torch==2.10.0", + "torch==2.11.0", "triton; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", "colored==2.2.3", "ray[default]==2.54.0", - "wandb>=0.25.0", + "wandb>=0.25.1", "numpy", "datasets>=4.0.0", "rich", @@ -41,36 +47,82 @@ dependencies = [ "matplotlib", "plotly", "sympy>=1.14.0", - "pillow>=11.3.0", - "torchvision==0.25.0", - "transformers==5.3.0", + "pillow>=12.1.1", + "torchvision==0.26.0", + "transformers==5.6.0", "num2words>=0.5.14", # for SmolVLM - "mlflow>=3.5.0,<3.6.0", - "nvidia-nvshmem-cu12; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # for deep_ep build + "mlflow>=3.12.0", + "nvidia-nvshmem-cu13; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # for deep_ep build "swanlab", "pyzmq", "decord2", - - "nccl4py", # for non-colocated refit - "cuda-bindings", # for non-colocated refit - "pybase64", # for sglang refit - "nvidia-cudnn-cu12==9.19.0.56", # for transformer-engine no build isolation + "soundfile>=0.13.1", + "nccl4py; sys_platform != 'darwin'", # for non-colocated refit + "cuda-bindings; sys_platform != 'darwin'", # for non-colocated refit + "pybase64", # for sglang refit + "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", # for transformer-engine no build isolation + # tilelang — replacement Triton kernel mamba-ssm requires when + # Triton >= 3.4.0 on Hopper, see github.com/state-spaces/mamba#640. + # Without this, qwen3.5 / nano-v3 / moonlight megatron recipes + # crash at first gated-chunk backward with a RuntimeError pointing + # at this exact pip install. Linux x86_64 only — mamba-ssm itself + # is gated to that pair. + "tilelang; sys_platform == 'linux' and platform_machine == 'x86_64'", + # Data-plane stack — promoted to base so worker venvs (built by + # nemo_rl.utils.venvs.create_local_venv via bare `uv sync`, no extras) + # automatically include them. Removes the need for a `[data-plane]` + # extra and the corresponding plumbing in the per-worker venv builder. + "tensordict", + # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore + # refactor: `clear` switched from unanchored `remove_by_regex` to + # exact-key `batch_remove`, which fixes a collateral-key-deletion bug + # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released. + "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39", + # Backs data_plane.backend="mooncake_cpu". Default backend is "simple" + # (in-process), but the mooncake_cpu path needs the `mooncake_master` + # binary that ships in this wheel at /mooncake/. Bundled + # with TQ rather than gated behind an extra so worker venvs (built + # without extras) can be flipped to mooncake_cpu via config alone. + # PyPI's base `mooncake-transfer-engine` is cu12-only (links + # libcudart.so.12), which breaks on cu13 containers. Upstream now also + # publishes a cu13 variant as a separate distribution name + # `mooncake-transfer-engine-cuda13` (same `mooncake/` import namespace, + # store.so linked against libcudart.so.13). Resolve from PyPI rather + # than the GitHub release URL — the wheel is byte-identical (verified + # sha256), and PyPI's CDN is far more reliable than github releases + # from compute nodes. + # Upstream publishes both x86_64 and aarch64 wheels (see uv.lock). CI's + # build-container runner is aarch64 (uv reports aarch64-unknown-linux-gnu), + # so the marker must include aarch64 — otherwise mooncake is silently + # excluded from the resolution during the Docker build. + "mooncake-transfer-engine-cuda13==0.3.10.post2 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", ] [project.optional-dependencies] -fsdp = ["flash-attn==2.8.1", "mamba-ssm", "causal-conv1d"] +fsdp = [ + # +cu13 wheels from GitHub match torch cu130; PyPI often resolves to +cu12 (libcudart.so.12). + # https://github.com/Dao-AILab/flash-attention/releases/tag/v2.8.1 + "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'", + "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'", + "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')", + "mamba-ssm", + "causal-conv1d", +] automodel = [ "nemo-automodel[moe]", # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108 # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 - "flash-attn==2.8.1", + "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'", + "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'", + "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')", "transformers>=5.3.0", "mamba-ssm", "causal-conv1d", "nv-grouped-gemm", - "transformer-engine[pytorch]>=2.9.0a0,<2.12.0", - "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'", ] vllm = [ "cuda-python", @@ -78,15 +130,25 @@ vllm = [ # deep_ep also needs libibverbs-dev # sudo apt-get update # sudo apt-get install libibverbs-dev - "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", - "vllm==0.17.0", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'", + # Default wheels on GitHub are cu130. See v0.20.0 release assets: + # https://github.com/vllm-project/vllm/releases/tag/v0.20.0 + "vllm @ https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'", + "vllm @ https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'", + "vllm==0.20.0 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')", "num2words>=0.5.14", - "flashinfer-python==0.6.4", + "flashinfer-python==0.6.8.post1", + "flashinfer-cubin==0.6.8.post1", "nvidia-cutlass-dsl>=4.4.0.dev1", ] sglang = [ "sglang", - "sgl-kernel", # Must be a direct dep so [tool.uv.sources] VCS override applies (transitive deps don't use sources) + "flashinfer-python==0.6.11.post1", + "flashinfer-cubin==0.6.11.post1", + "kernels>=0.12.0,<0.13", + "sglang-kernel==0.4.2.post2", # Direct dep for explicit version control + "sglang-router", # Used by nemo_rl.models.generation.sglang.sglang_router for multi-engine routing ] mcore = [ # also need cudnn (https://developer.nvidia.com/cudnn-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=20.04&target_type=deb_network) @@ -98,16 +160,27 @@ mcore = [ # This dependency also needs to be compatible with the spec in Megatron-Bridge/pyproject.toml. # It is specified here since we don't directly use Megatron-Bridge/pyproject.toml, but a proxy setup.py+pyproject.toml combo # outside to allow "optionally" installing the megatron path. It's simpler to deal with transformer-engine here in the NeMo RL pyproject.toml - "transformer-engine[pytorch]==2.12.0", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", "megatron-core", "megatron-bridge", + "mamba-ssm", + "causal-conv1d", + "nvidia-modelopt[torch]; sys_platform != 'darwin'", + "onnxscript", # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108 # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 - "flash-attn==2.8.1", - "emerging-optimizers==0.1.0", - "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", + "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'", + "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'", + "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')", + "emerging-optimizers==0.2.0", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'", ] +modelopt = ["nvidia-modelopt"] +nvrx = [ + "nvidia-resiliency-ext", +] # for ft_launcher (fault-tolerant training launcher) nemo_gym = ["nemo_gym"] [dependency-groups] @@ -115,7 +188,7 @@ nemo_gym = ["nemo_gym"] # This is a default group so that we install these even with bare `uv sync` build = [ # Build requirement for TE - "torch==2.10.0", + "torch==2.11.0", # Build requirement for TE "setuptools", "packaging", @@ -148,42 +221,47 @@ dev = [ "pyrefly==0.24.2", ] test = [ - "pytest>=7.0.0", + "pytest>=8.4.2", "pytest-timeout", "pytest-cov", "pytest-asyncio", "pytest-testmon", + "pytest-shard", ] [tool.uv.sources] -megatron-core = { workspace = true } +megatron-core = { path = "3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM", editable = true } nemo-automodel = { path = "3rdparty/Automodel-workspace/Automodel", editable = true } megatron-bridge = { path = "3rdparty/Megatron-Bridge-workspace", editable = true } nemo_gym = { workspace = true } nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2c71bb1186e97ccbf96d6844c" } -# torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels +# torch/torchaudio/torchvision/triton all come from the torch index in order to pick up aarch64 wheels torch = [ - { index = "pytorch-cu129", marker = "sys_platform != 'darwin'" }, + { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, { index = "pypi", marker = "sys_platform == 'darwin'" }, ] torchvision = [ - { index = "pytorch-cu129", marker = "sys_platform != 'darwin'" }, + { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, + { index = "pypi", marker = "sys_platform == 'darwin'" }, +] +torchaudio = [ + { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, { index = "pypi", marker = "sys_platform == 'darwin'" }, ] triton = [ - { index = "pytorch-cu129", marker = "sys_platform != 'darwin'" }, + { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" }, { index = "pypi", marker = "sys_platform == 'darwin'" }, ] causal-conv1d = { git = "https://github.com/Dao-AILab/causal-conv1d", rev = "67e0a9dfe1518fc0036444e9ab5fe06ab78299e0" } mamba-ssm = { git = "https://github.com/state-spaces/mamba.git", rev = "d68d16ed7d5d5164eb5a57c0285f3b7eb8394ec1" } nv-grouped-gemm = { git = "https://github.com/fanshiqing/grouped_gemm", tag = "v1.1.4.post7" } -# From JustinTong0323/sglang branch update-transformers-v5 (sgl-project/sglang#17784) -sglang = { git = "https://github.com/JustinTong0323/sglang.git", rev = "70aa688742dd2b75bf9e8e980249303f39295b0d", subdirectory = "python" } -sgl-kernel = { git = "https://github.com/JustinTong0323/sglang.git", rev = "70aa688742dd2b75bf9e8e980249303f39295b0d", subdirectory = "sgl-kernel" } +emerging-optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" } +nvidia-modelopt = { git = "https://github.com/NVIDIA/Model-Optimizer", rev = "905018803414702e414a86716484ed4115b37ba6" } +nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "15a851565a4ce846c04431ecb0cf09903ab4837e" } +sglang = { git = "https://github.com/sgl-project/sglang.git", branch = "sglang-miles", subdirectory = "python" } [tool.uv.workspace] members = [ - "3rdparty/Megatron-LM-workspace", "3rdparty/Gym-workspace/Gym", # Research projects are also added here in order for them to share the global root level uv.lock. # If we don't do this, the research projects do not see the global uv.lock, and may mistakenly @@ -199,13 +277,12 @@ url = "https://pypi.org/simple" explicit = true [[tool.uv.index]] -name = "pytorch-cu129" -url = "https://download.pytorch.org/whl/cu129" +name = "pytorch-cu130" +url = "https://download.pytorch.org/whl/cu130" explicit = true [tool.uv] preview = true # Enable preview features like extra-build-dependencies -extra-build-variables = { sgl-kernel = { CMAKE_BUILD_PARALLEL_LEVEL = "24", FLASHINFER_CUDA_ARCH_LIST = "9.0a 10.0a", CMAKE_ARGS = "-DCMAKE_POLICY_VERSION_MINIMUM=3.5" } } no-build-isolation-package = [ "transformer-engine-torch", "transformer-engine", @@ -215,7 +292,6 @@ no-build-isolation-package = [ "deep_gemm", "deep_ep", "nv-grouped-gemm", # from mlm (added here to make sure it's built no isolation since mlm workspace uses setup.py) - "sgl-kernel", ] # Always apply the build group since dependencies like TE/mcore/nemo-run require build dependencies # and this lets us assume they are implicitly installed with a simply `uv sync`. Ideally, we'd @@ -231,45 +307,64 @@ link-mode = "copy" # The timm override is needed because current automodel pins to 1.0.16. This can be removed once we move ToT automodel # The nvidia-modelopt override is needed because mcore is still on 0.33 override-dependencies = [ - "transformer-engine[pytorch]==2.12.0", - "nvidia-cudnn-cu12==9.19.0.56", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", + "nvidia-cublas==13.3.0.5; sys_platform != 'darwin'", + "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", "opencv-python-headless>=4.11.0", "timm<=1.0.22", "nvidia-modelopt[torch]>=0.39.0", - "torch==2.10.0", - "torchaudio==2.10.0", + "torch==2.11.0", + "torchaudio==2.11.0", # sglang has conflicting llguidance versions than vllm, so enforcing vllm's version since it's newer "llguidance>=1.3.0,<1.4.0", # Override setuptools range in other dependencies to address CVE GHSA-58pv-8j8x-9vj2 "setuptools>=80.10.2", - "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", - # Pin flashinfer globally — flashinfer-python must match flashinfer-cubin at runtime, and - # they're resolved independently by uv so ranges risk version mismatch. - # When changing this version, check what each backend expects: - # vllm extra (this file, [project.optional-dependencies].vllm): flashinfer-python==0.6.4 - # sglang dependency-metadata (this file, [[tool.uv.dependency-metadata]] name="sglang"): flashinfer_python==0.6.4, flashinfer_cubin==0.6.4 - # megatron-core (3rdparty/Megatron-LM-workspace/Megatron-LM/pyproject.toml): flashinfer-python~=0.5.0 - "flashinfer-python==0.6.4", - "flashinfer-cubin==0.6.4", - # sglang pins nvidia-cutlass-dsl==4.2.1, conflicting with flashinfer 0.6.4 (>=4.3.4) and vllm (>=4.4.0.dev1). - # Override to >=4.2.1 so uv can resolve to a version satisfying all three. - "nvidia-cutlass-dsl>=4.2.1", + "emerging-optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0", + "nvidia-resiliency-ext @ git+https://github.com/NVIDIA/nvidia-resiliency-ext.git@15a851565a4ce846c04431ecb0cf09903ab4837e", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'", + "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'", + # Keep FlashInfer aligned with sglang-miles' python/pyproject.toml. + "flashinfer-python==0.6.11.post1", + "flashinfer-cubin==0.6.11.post1", + # SGLang leaves kernels unbounded, but Transformers 5.6.0 expects the 0.12 API. + "kernels>=0.12.0,<0.13", + # Override to >=4.4.1 so uv can resolve to a version satisfying both vllm and sglang. + "nvidia-cutlass-dsl>=4.4.1", # Relax megatron-core workspace member's opentelemetry-api ceiling (<1.34) for protobuf 6.x compat with ray "opentelemetry-api>=1.33.1", - # vLLM 0.17.0 code is compatible with transformers v5 but the PyPI metadata still declares <5. - # Override until vllm officially relaxes the constraint (https://github.com/vllm-project/vllm/issues/30466). - "transformers==5.3.0", + # Keep Transformers aligned with sglang-miles' python/pyproject.toml. + "transformers==5.6.0", + # Override sglang's xgrammar==0.1.32 to address CVE GHSA-7rgv-gqhr-fxg3 + "xgrammar==0.1.33", + # Override dependencies to address CVEs + "mlflow>=3.12.0", + # Override outlines for Python 3.13 support + "outlines>=0.2.0", + # Upgrade pytest to 9.0.3 + "pytest>=9.0.3", + # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs + # numpy>=2.1.0 via onnx → ml-dtypes. Override globally so the data-plane + # extra composes with mcore/automodel without version-mirroring TQ's + # requirements.txt. Forward-compatible across TQ minor bumps. + "numpy>=2.1.0", ] + # CVE fixes constraint-dependencies = [ "brotli>=1.2.0", # Address CVE GHSA-2qfp-q593-8484 "starlette>=0.49.1", # Address CVE GHSA-7f5h-v6xp-fcq8 - "urllib3>=2.6.3", # Address CVE GHSA-38jv-5279-wg99 + "urllib3>=2.7.0", # Address CVE GHSA-38jv-5279-wg99 "aiohttp>=3.13.3", # Address CVE GHSA-mqqc-3gqh-h2x8 - "pyasn1>=0.6.2", # Address CVE GHSA-63vm-454h-vhhq + "pyasn1>=0.6.3", # Address CVE GHSA-jr27-m4p2-rc6r "wheel>=0.46.2", # Address CVE GHSA-8rrh-rw8j-w5fx "protobuf>=6.33.5", # Address CVE GHSA-7gcm-g887-7qv7 "python-multipart>=0.0.22", # Address CVE GHSA-wp53-j4wj-2cfg + "pygments>=2.20.0", # Address CVE GHSA-5239-wwwm-4pmq + "cbor2>=5.9.0", # Address CVE GHSA-3c37-wwvx-h642 + "onnx>=1.21.0rc4", # Address CVE GHSA-hqmj-h5c6-369m + "cryptography>=46.0.6", # Address CVE GHSA-6w46-j5rx-g56g + "orjson>=3.11.6", # Address CVE GHSA-hx9q-6w63-j58v + "pyjwt>=2.12.0", # Address CVE GHSA-752w-5fwx-jx9f ] conflicts = [ @@ -319,7 +414,6 @@ transformer-engine-torch = [{ requirement = "torch", match-runtime = true }] mamba-ssm = [{ requirement = "torch", match-runtime = true }] causal-conv1d = [{ requirement = "torch", match-runtime = true }] nv-grouped-gemm = [{ requirement = "torch", match-runtime = true }] -sgl-kernel = [{ requirement = "torch", match-runtime = true }] # Needed when building from source [[tool.uv.dependency-metadata]] @@ -351,88 +445,20 @@ version = "v2.0.0+7b6b556" requires-dist = ["torch", "packaging", "ninja"] [[tool.uv.dependency-metadata]] -name = "nv-grouped-gemm" -# This version has to match the version in the commit/rev/tag used -version = "v1.1.4.post7" -requires-dist = ["setuptools", "wheel", "torch", "numpy"] - +name = "transformer-engine" +version = "2.14.1+366798e" +requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[tool.uv.dependency-metadata]] -name = "sgl-kernel" -# This version has to match the version in the commit/rev/tag used -version = "0.3.21" -requires-dist = ["torch", "scikit-build-core", "wheel"] +name = "transformer-engine-torch" +version = "2.14.1+366798e" +requires-dist = ["torch", "transformer-engine"] [[tool.uv.dependency-metadata]] -name = "sglang" -# VCS install from JustinTong0323/sglang@update-transformers-v5 -# Version is dynamic (setuptools-scm), so uv cannot resolve deps from the VCS source automatically. -# This requires-dist list must be kept in sync with the fork's python/pyproject.toml [project].dependencies. -# Source: https://github.com/JustinTong0323/sglang/blob/70aa688742dd2b75bf9e8e980249303f39295b0d/python/pyproject.toml -version = "0.5.7.dev0" -requires-dist = [ - "IPython", - "aiohttp", - "apache-tvm-ffi>=0.1.5,<0.2", - "anthropic>=0.20.0", - "blobfile==3.0.0", - "build", - "compressed-tensors", - "cuda-python==12.9", - "decord2", - "datasets", - "einops", - "fastapi", - "flashinfer_python==0.6.4", - "flashinfer_cubin==0.6.4", - "gguf", - "interegular", - "llguidance>=0.7.11,<0.8.0", - "modelscope", - "msgspec", - "ninja", - "numpy", - "nvidia-cutlass-dsl>=4.3.4", - "nvidia-ml-py", - "openai-harmony==0.0.4", - "openai==2.6.1", - "orjson", - "outlines==0.1.11", - "packaging", - "partial_json_parser", - "pillow", - "prometheus-client>=0.20.0", - "psutil", - "py-spy", - "pybase64", - "pydantic", - "python-multipart", - "pyzmq>=25.1.2", - "quack-kernels==0.2.4", - "requests", - "scipy", - "sentencepiece", - "setproctitle", - "sgl-kernel==0.3.21", - "soundfile==0.13.1", - "tiktoken", - "timm==1.0.16", - "torch_memory_saver==0.0.9", - "torch==2.9.1", - "torchao==0.9.0", - "torchaudio==2.9.1", - "torchcodec==0.8.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", - "torchvision", - "tqdm", - "transformers==5.3.0", - "uvicorn", - "uvloop", - "xgrammar==0.1.27", - "smg-grpc-proto>=0.3.3", - "grpcio>=1.78.0", - "grpcio-reflection>=1.78.0", - "grpcio-health-checking>=1.78.0", -] +name = "nv-grouped-gemm" +# This version has to match the version in the commit/rev/tag used +version = "v1.1.4.post7" +requires-dist = ["setuptools", "wheel", "torch", "numpy"] [[tool.uv.dependency-metadata]] name = "megatron-bridge" @@ -464,18 +490,43 @@ requires-dist = [ # Non-workspace path deps cannot depend on workspace members (uv name-shadowing restriction). "qwen-vl-utils", # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13 - "transformer-engine[pytorch,core_cu12]", + "transformer-engine[pytorch,core_cu13]", "mamba-ssm", - "nvidia-resiliency-ext~=0.5.0", + "nvidia-resiliency-ext", "causal-conv1d", "flash-linear-attention", "timm", "open-clip-torch>=3.2.0", - "mlflow>=3.5.0", + "mlflow>=3.9.0", "comet-ml>=3.50.0", "torch>=2.6.0", ] +# Override logsage metadata to remove numpy<=2.0.2 and pandas<=2.3.3 upper bounds +# (numpy cap conflicts with onnx>=1.21.0rc4 CVE fix via nvidia-resiliency-ext) +# Tracking: https://github.com/NVIDIA/nvidia-resiliency-ext/issues/301 +[[tool.uv.dependency-metadata]] +name = "logsage" +version = "0.1.5" +requires-dist = [ + "drain3>=0.9.11,<0.10.0", + "langchain>=0.3.27,<0.4.0", + "langchain-core>=0.3.0,<1.0.0", + "langchain-nvidia-ai-endpoints>=0.3.18,<0.4.0", + "nh3>=0.3.1,<0.4.0", + "numpy", + "pandas", + "pydantic-settings>=2.11.0,<3.0.0", + "requests>=2.32.5,<3.0.0", +] + +# Override drain3 to relax cachetools==4.2.1 pin (conflicts with mlflow's cachetools>=5.0.0) +# Tracking: https://github.com/logpai/Drain3/issues/119 +[[tool.uv.dependency-metadata]] +name = "drain3" +version = "0.9.11" +requires-dist = ["jsonpickle", "cachetools>=4.2.1"] + [tool.black] line-length = 120 include = '\.pyi?$' @@ -488,8 +539,7 @@ exclude = ''' ''' [tool.pytest.ini_options] -addopts = "--testmon --durations=100 -s -rA -x" -#addopts = "--durations=100 -s -rA -x" +addopts = "--durations=100 -s -rA -x" testpaths = ["tests"] python_files = "test_*.py" markers = [ @@ -499,6 +549,7 @@ markers = [ "automodel: marks tests that require the automodel extra", "vllm: marks tests that require the vllm extra", "sglang: marks tests that require the sglang extra", + "nemo_gym: marks tests that require the nemo_gym extra", ] [tool.pyrefly] diff --git a/pyrefly.toml b/pyrefly.toml index 98b2161906..fe4c86c26f 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -110,7 +110,6 @@ project-includes = [ "nemo_rl/models/dtensor/parallelize.py", "nemo_rl/models/generation/__init__.py", "nemo_rl/models/generation/interfaces.py", - "nemo_rl/models/generation/sglang/__init__.py", "nemo_rl/models/generation/sglang/config.py", "nemo_rl/models/generation/vllm/__init__.py", "nemo_rl/models/generation/vllm/config.py", diff --git a/tests/functional/grpo_sglang.sh b/tests/functional/grpo_sglang.sh index fd268d0250..41a6d2db7f 100755 --- a/tests/functional/grpo_sglang.sh +++ b/tests/functional/grpo_sglang.sh @@ -20,7 +20,7 @@ mkdir -p $EXP_DIR $LOG_DIR cd $PROJECT_ROOT uv run --group test coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo.py \ - --config $PROJECT_ROOT/examples/configs/grpo_math_1B_sglang.yaml \ + --config $PROJECT_ROOT/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.yaml \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ grpo.num_generations_per_prompt=4 \ diff --git a/tests/test_suites/llm/grpo-qwen3-0.6b-1n8g-sglang.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.sh similarity index 80% rename from tests/test_suites/llm/grpo-qwen3-0.6b-1n8g-sglang.sh rename to tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.sh index 10ff34699c..b92891150e 100755 --- a/tests/test_suites/llm/grpo-qwen3-0.6b-1n8g-sglang.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.sh @@ -2,12 +2,14 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env +export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-4,5,6,7} + # ===== BEGIN CONFIG ===== NUM_NODES=1 -STEPS_PER_RUN=500 -MAX_STEPS=500 +STEPS_PER_RUN=450 +MAX_STEPS=450 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=180 # bumped from 120: ~18.5s/step without piecewise CUDA graphs +NUM_MINUTES=150 # ~13.7s/step without piecewise CUDA graphs # ===== END CONFIG ===== exit_if_max_steps_reached @@ -32,12 +34,12 @@ uv run examples/run_grpo.py \ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # Only run metrics if the target step is reached +# Same thresholds as the 1n8g fsdp2tp1-sglang recipe for alignment verification if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ - 'mean(data["timing/train/total_step_time"], 2) < 30' + 'mean(data["timing/train/total_step_time"], 2) < 25' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" fi - diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm.sh old mode 100755 new mode 100644 similarity index 89% rename from tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.sh rename to tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm.sh index 77a7896f73..d68fc68d3d --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-vllm.sh @@ -2,12 +2,14 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env +export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-6,7} + # ===== BEGIN CONFIG ===== NUM_NODES=1 STEPS_PER_RUN=450 MAX_STEPS=450 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=150 # bumped from 120: ~13.7s/step without piecewise CUDA graphs +NUM_MINUTES=150 # ===== END CONFIG ===== exit_if_max_steps_reached @@ -32,7 +34,6 @@ uv run examples/run_grpo.py \ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # Only run metrics if the target step is reached -# Using the same metrics thresholds as the vllm version to verify alignment if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ @@ -41,5 +42,3 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" fi - - diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index d0b9a76c09..5502678638 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -11,8 +11,7 @@ tests/test_suites/llm/grpo-gemma3-1b-it-1n8g-fsdp2tp1.sh tests/test_suites/llm/prorlv2-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1.sh # SGLang backend -tests/test_suites/llm/grpo-qwen3-0.6b-1n8g-sglang.sh -tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-fsdp2tp1-sglang.sh +tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n2g-fsdp2tp1-sglang.sh # Dtensor (Qwen/Qwen2.5-7B-Instruct) tests/test_suites/llm/grpo-qwen2.5-7b-instruct-4n8g-fsdp2tp4.v3.sh diff --git a/tests/unit/L0_Unit_Tests_Other.sh b/tests/unit/L0_Unit_Tests_Other.sh index 1a9f574ea2..5713d75c6c 100644 --- a/tests/unit/L0_Unit_Tests_Other.sh +++ b/tests/unit/L0_Unit_Tests_Other.sh @@ -57,14 +57,6 @@ else uv run --extra vllm bash -x ./tests/run_unit.sh "${TEST_PATHS[@]}" "${IGNORE[@]}" "${EXCLUDED_UNIT_TESTS[@]}" --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --vllm-only fi -# Check and run sglang tests -exit_code=$(cd ${PROJECT_ROOT}/tests && uv run --extra sglang pytest "${TEST_PATHS[@]}" "${IGNORE[@]}" "${EXCLUDED_UNIT_TESTS[@]}" --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?) -if [[ $exit_code -eq 5 ]]; then - echo "No sglang tests to run" -else - uv run --extra sglang bash -x ./tests/run_unit.sh "${TEST_PATHS[@]}" "${IGNORE[@]}" "${EXCLUDED_UNIT_TESTS[@]}" --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only -fi - # Skip research tests in fast mode if [[ "${FAST:-0}" != "1" ]]; then for i in research/*/tests/unit; do diff --git a/tests/unit/L0_Unit_Tests_Policy.sh b/tests/unit/L0_Unit_Tests_Policy.sh index f19691c421..382ffada5a 100644 --- a/tests/unit/L0_Unit_Tests_Policy.sh +++ b/tests/unit/L0_Unit_Tests_Policy.sh @@ -55,12 +55,4 @@ if [[ $exit_code -eq 5 ]]; then echo "No vllm tests to run" else uv run --extra vllm bash -x ./tests/run_unit.sh "${TEST_PATHS[@]}" "${IGNORE[@]}" "${EXCLUDED_UNIT_TESTS[@]}" --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --vllm-only -fi - -# Check and run sglang tests -exit_code=$(cd ${PROJECT_ROOT}/tests && uv run --extra sglang pytest "${TEST_PATHS[@]}" "${IGNORE[@]}" "${EXCLUDED_UNIT_TESTS[@]}" --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?) -if [[ $exit_code -eq 5 ]]; then - echo "No sglang tests to run" -else - uv run --extra sglang bash -x ./tests/run_unit.sh "${TEST_PATHS[@]}" "${IGNORE[@]}" "${EXCLUDED_UNIT_TESTS[@]}" --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only -fi +fi \ No newline at end of file diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 2ddbf001c9..feedc79bf1 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -874,338 +874,6 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_multi_node(): setup(master_config, tokenizer, dataset, None) -@pytest.mark.parametrize( - "colocated_inference, expected_parallel", - [(True, 0.0), (False, True)], -) -def test_setup_sglang_sets_model_path_and_parallel_flag( - monkeypatch, colocated_inference, expected_parallel -): - from nemo_rl.algorithms import grpo as grpo_mod - - logged = {} - - class DummyLogger: - def log_hyperparams(self, *_args, **_kwargs): - pass - - def log_metrics(self, metrics, *_args, **_kwargs): - logged["metrics"] = metrics - - class DummyCheckpointer: - def get_latest_checkpoint_path(self): - return None - - def load_training_info(self, _path): - return None - - def get_resume_paths(self, _path): - return None, None - - class DummyLoader: - def __init__(self, *_args, **_kwargs): - pass - - def __len__(self): - return 1 - - def load_state_dict(self, _state): - pass - - class DummyCluster: - def __init__(self, *_args, **_kwargs): - pass - - def world_size(self): - return 1 - - def get_master_address_and_port(self): - return "127.0.0.1", 1234 - - class DummyPolicy: - def print_node_ip_and_gpu_id(self): - pass - - def init_collective(self, *_args, **_kwargs): - return [] - - def prepare_refit_info(self): - return {} - - class DummySGLangGeneration: - def finish_generation(self): - pass - - def prepare_refit_info(self, _state): - pass - - def init_collective(self, *_args, **_kwargs): - return [] - - monkeypatch.setattr(grpo_mod, "Logger", lambda *_args, **_kwargs: DummyLogger()) - monkeypatch.setattr( - grpo_mod, "CheckpointManager", lambda *_args, **_kwargs: DummyCheckpointer() - ) - monkeypatch.setattr( - grpo_mod, "ClippedPGLossFn", lambda *_args, **_kwargs: MagicMock() - ) - monkeypatch.setattr(grpo_mod, "StatefulDataLoader", DummyLoader) - monkeypatch.setattr(grpo_mod, "RayVirtualCluster", DummyCluster) - monkeypatch.setattr(grpo_mod, "Policy", lambda *_args, **_kwargs: DummyPolicy()) - monkeypatch.setattr( - grpo_mod, - "SGLangGeneration", - lambda *_args, **_kwargs: DummySGLangGeneration(), - ) - monkeypatch.setattr(grpo_mod.ray, "get", lambda x: x) - - generation_resources = { - "gpus_per_node": 1, - "num_nodes": 1, - } - if colocated_inference: - generation_resources = {"gpus_per_node": None, "num_nodes": None} - - master_config = { - "policy": { - "model_name": "fake-model", - "train_global_batch_size": 1, - "train_micro_batch_size": 1, - "max_total_sequence_length": 8, - "make_sequence_length_divisible_by": 1, - "dtensor_cfg": {"enabled": False}, - "megatron_cfg": {"enabled": False, "pipeline_model_parallel_size": 1}, - "generation": { - "temperature": 1.0, - "top_p": 1.0, - "top_k": None, - "backend": "sglang", - "colocated": { - "enabled": colocated_inference, - "resources": generation_resources, - }, - "sglang_cfg": { - "gpus_per_server": 1, - "dp_size": 1, - "pp_size": 1, - "ep_size": 1, - }, - }, - }, - "loss_fn": { - "force_on_policy_ratio": False, - "use_importance_sampling_correction": False, - }, - "env": {}, - "grpo": { - "seed": 1, - "num_prompts_per_step": 1, - "num_generations_per_prompt": 1, - "max_num_steps": 1, - "max_num_epochs": 1, - "val_period": 0, - "val_batch_size": 1, - "val_at_start": False, - "val_at_end": False, - "max_val_samples": 1, - "use_dynamic_sampling": False, - "batch_multiplier": 1, - "normalize_rewards": False, - "use_leave_one_out_baseline": False, - "reward_scaling": {"enabled": False}, - "reward_shaping": {"enabled": False}, - "overlong_filtering": False, - }, - "data": { - "shuffle": False, - "num_workers": 0, - "env_name": None, - "use_multiple_dataloader": False, - }, - "logger": {"num_val_samples_to_print": 0}, - "checkpointing": {"enabled": False}, - "cluster": {"num_nodes": 1, "gpus_per_node": 4}, - } - - tokenizer = MagicMock() - dataset = MagicMock() - dataset.__len__ = MagicMock(return_value=1) - - grpo_mod.setup(master_config, tokenizer, dataset, None) - - assert ( - master_config["policy"]["generation"]["sglang_cfg"]["model_path"] - == master_config["policy"]["model_name"] - ) - assert logged["metrics"]["parallel_init_enabled"] == expected_parallel - - -def test_refit_policy_generation_sglang_colocated_http(monkeypatch): - from nemo_rl.algorithms import grpo as grpo_mod - - calls = { - "prepare_for_generation_tags": [], - "invalidate_kv_cache": 0, - "stream_weights_via_http": [], - "offload_before_refit": 0, - "offload_after_refit": 0, - } - - class DummySGLangGeneration: - def prepare_for_generation(self, tags=None): - calls["prepare_for_generation_tags"].append(tags) - - def get_sglang_url_to_gpu_uuids(self): - return {"http://localhost:12345": ["gpu-uuid-0"]} - - def invalidate_kv_cache(self): - calls["invalidate_kv_cache"] += 1 - return True - - class DummyPolicy: - def offload_before_refit(self): - calls["offload_before_refit"] += 1 - - def offload_after_refit(self): - calls["offload_after_refit"] += 1 - - def get_free_memory_bytes(self): - return 1024 * 1024 * 1024 - - def stream_weights_via_http(self, sglang_url_to_gpu_uuids): - calls["stream_weights_via_http"].append(sglang_url_to_gpu_uuids) - return ["ok"] - - monkeypatch.setattr(grpo_mod, "SGLangGeneration", DummySGLangGeneration) - monkeypatch.setattr(grpo_mod.ray, "get", lambda x: x) - - grpo_mod.refit_policy_generation( - policy=DummyPolicy(), - policy_generation=DummySGLangGeneration(), - colocated_inference=True, - ) - - assert calls["offload_before_refit"] == 1 - assert calls["offload_after_refit"] == 1 - assert calls["invalidate_kv_cache"] == 1 - assert calls["stream_weights_via_http"] == [ - {"http://localhost:12345": ["gpu-uuid-0"]} - ] - assert calls["prepare_for_generation_tags"] == [["weights"], ["kv_cache"]] - - -def test_refit_policy_generation_sglang_non_colocated_raises(monkeypatch): - from nemo_rl.algorithms import grpo as grpo_mod - - class DummySGLangGeneration: - pass - - monkeypatch.setattr(grpo_mod, "SGLangGeneration", DummySGLangGeneration) - - with pytest.raises(NotImplementedError): - grpo_mod.refit_policy_generation( - policy=object(), - policy_generation=DummySGLangGeneration(), - colocated_inference=False, - ) - - -def test_grpo_train_collects_generation_logger_metrics( - monkeypatch, mock_grpo_components -): - from nemo_rl.algorithms import grpo as grpo_mod - - policy_generation = MagicMock() - policy_generation.clear_logger_metrics = MagicMock() - policy_generation.get_logger_metrics = MagicMock( - return_value={"pending_requests": 1} - ) - policy_generation.prepare_for_generation = MagicMock() - policy_generation.finish_generation = MagicMock() - - mock_batch = next(iter(mock_grpo_components["train_dataloader"])) - mock_rollout_metrics = {"gen_kl_error": 0.0, "mean_gen_tokens_per_sample": 2.0} - - def fake_batched_message_log_to_flat_message(*_args, **_kwargs): - flat = BatchedDataDict( - { - "token_ids": torch.tensor([[1, 2]]), - "advantages": torch.tensor([[0.5, 0.5]]), - "generation_logprobs": torch.tensor([[0.0, 0.0]]), - "token_loss_mask": torch.tensor([[1, 1]]), - "content": ["ok"], - } - ) - return flat, torch.tensor([2]) - - monkeypatch.setattr( - grpo_mod, - "batched_message_log_to_flat_message", - fake_batched_message_log_to_flat_message, - ) - monkeypatch.setattr( - grpo_mod, "_should_use_async_rollouts", lambda *_args, **_kwargs: True - ) - monkeypatch.setattr( - grpo_mod, - "run_async_multi_turn_rollout", - lambda *_args, **_kwargs: (mock_batch, mock_rollout_metrics), - ) - monkeypatch.setattr( - grpo_mod, - "run_multi_turn_rollout", - lambda *_args, **_kwargs: (mock_batch, mock_rollout_metrics), - ) - monkeypatch.setattr( - grpo_mod, - "calculate_baseline_and_std_per_prompt", - lambda *_args, **_kwargs: (torch.tensor([0.1]), torch.tensor([1.0])), - ) - monkeypatch.setattr( - grpo_mod, "refit_policy_generation", lambda *_args, **_kwargs: None - ) - monkeypatch.setattr( - grpo_mod, "print_performance_metrics", lambda *_args, **_kwargs: {} - ) - monkeypatch.setattr( - grpo_mod, "maybe_gpu_profile_step", lambda *_args, **_kwargs: None - ) - monkeypatch.setattr( - grpo_mod, - "compute_and_apply_seq_logprob_error_masking", - lambda *_args, **_kwargs: (0.0, 0, 0.0), - ) - - master_config = mock_grpo_components["master_config"] - master_config["grpo"]["max_num_steps"] = 1 - master_config["grpo"]["max_num_epochs"] = 1 - master_config["grpo"]["val_period"] = 0 - master_config["grpo"]["val_at_start"] = False - master_config["grpo"]["use_dynamic_sampling"] = False - - grpo_mod.grpo_train( - mock_grpo_components["policy"], - policy_generation, - mock_grpo_components["train_dataloader"], - mock_grpo_components["val_dataloader"], - mock_grpo_components["tokenizer"], - mock_grpo_components["loss_fn"], - mock_grpo_components["task_to_env"], - mock_grpo_components["val_task_to_env"], - mock_grpo_components["logger"], - mock_grpo_components["checkpointer"], - _default_grpo_save_state(), - master_config, - ) - - assert policy_generation.clear_logger_metrics.called - assert policy_generation.get_logger_metrics.called - assert any( - "generation_logger_metrics" in call.args[0] - for call in mock_grpo_components["logger"].log_metrics.call_args_list - ) - - @pytest.fixture def mock_grpo_components(): # Create mock components diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index ebc5569f86..fc56320592 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -499,8 +499,8 @@ def mock_2gpu_distributed_env(): # Create the 2D mesh that acts like a dictionary (this is so we can test DTensorPolicyWorker with TP > 1) mesh_2d = unittest.mock.MagicMock() - mesh_2d.__getitem__.side_effect = ( - lambda key: dp_mesh if key == "dp" else tp_mesh if key == "tp" else None + mesh_2d.__getitem__.side_effect = lambda key: ( + dp_mesh if key == "dp" else tp_mesh if key == "tp" else None ) mesh_2d.device_type = "cuda" diff --git a/tests/unit/excluded_unit_tests.sh b/tests/unit/excluded_unit_tests.sh index 38a806f33c..ce54e7bef8 100644 --- a/tests/unit/excluded_unit_tests.sh +++ b/tests/unit/excluded_unit_tests.sh @@ -123,19 +123,6 @@ EXCLUDED_UNIT_TESTS=( --deselect=tests/unit/models/generation/test_vllm_generation.py::test_vllm_generation_with_hf_training_colocated[False-False-bfloat16-True] --deselect=tests/unit/models/generation/test_vllm_generation.py::test_vllm_generation_with_hf_training_colocated[True-False-bfloat16-True] - # test_sglang_generation.py — keep 2 key tests - # Kept: test_sglang_policy_generation (basic generation correctness), - # test_sglang_generation_with_hf_training_colocated (generation + HF training A+B) - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_missing_required_config_key - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_top_p_top_k_validation - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_worker_seed_behavior - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_policy_tensor_parallel - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_generate_text - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_http_server - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_non_divisible_batch_handling - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_generation_with_hf_training_non_colocated - --deselect=tests/unit/models/generation/test_sglang_generation.py::test_sglang_weight_update_and_prefix_cache_reset - # test_vllm_utils.py — exclude only the @vllm-marked test (rest are cheap) --deselect=tests/unit/models/generation/test_vllm_utils.py::test_vllm_speculative_decoding_patch_still_needed diff --git a/tests/unit/models/generation/sglang/__init__.py b/tests/unit/models/generation/sglang/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/models/generation/sglang/_megatron_helpers.py b/tests/unit/models/generation/sglang/_megatron_helpers.py new file mode 100644 index 0000000000..9a90783d6a --- /dev/null +++ b/tests/unit/models/generation/sglang/_megatron_helpers.py @@ -0,0 +1,500 @@ +# 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 config builders for Megatron + SGLang weight-update / generation tests. + +Exposes: + +- ``MEGATRON_CFGS`` / ``SGLANG_CFGS`` — the parametrization matrix the user + asked for ("ep2 pp2 / tp2 pp2 / tp2 ep2 pp2" Megatron × "tp4 ep4 dp4 + --enable-dp-attention / tp4 ep2 dp4 --enable-dp-attention / tp2 ep2 pp2" + SGLang). +- ``make_policy_config(...)`` — produces a ``PolicyConfig`` dict suitable for + ``nemo_rl.models.policy.lm_policy.Policy`` against the + Qwen3-30B-A3B-Instruct-2507 model (MoE, supports EP > 1). +- ``make_sglang_cfg(...)`` — produces an SGLang generation config compatible + with ``SGLangGeneration``. +- ``required_world_size(...)`` / ``min_dp_for_megatron(...)`` — helpers used + by the test fixtures to size ``RayVirtualCluster`` and to skip cleanly when + the host doesn't have enough GPUs. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +# Qwen3 tokenizer ids (shared across the Qwen3 family — verified against +# ``Qwen/Qwen3-30B-A3B-Instruct-2507``). Hard-coding keeps these fixtures +# importable without spinning up a tokenizer. +PAD_TOKEN_ID = 151643 # Qwen3 ``<|endoftext|>`` +EOS_TOKEN_ID = 151645 # Qwen3 ``<|im_end|>`` +WEIGHT_UPDATE_PRECISION_ENV = "NEMO_RL_SGLANG_WEIGHT_UPDATE_PRECISION" +SGLANG_MOE_RUNNER_BACKEND_ENV = "NEMO_RL_SGLANG_MOE_RUNNER_BACKEND" +SGLANG_FP8_GEMM_RUNNER_BACKEND_ENV = "NEMO_RL_SGLANG_FP8_GEMM_RUNNER_BACKEND" + + +def weight_update_precision() -> str: + precision = os.environ.get(WEIGHT_UPDATE_PRECISION_ENV, "bf16").lower() + if precision not in {"bf16", "mxfp8"}: + raise ValueError( + f"{WEIGHT_UPDATE_PRECISION_ENV} must be 'bf16' or 'mxfp8', " + f"got {precision!r}" + ) + return precision + + +def _mxfp8_quantization_cfg() -> dict[str, Any]: + from nemo_rl.models.generation.sglang.mxfp8_quantization_core import ( + MXFP8_QUANTIZATION_CONFIG, + ) + + return { + "scheme": "mxfp8", + **MXFP8_QUANTIZATION_CONFIG, + } + + +# --------------------------------------------------------------------------- +# Parametrization matrix +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class MegatronShape: + """One Megatron parallelism shape under test. + + ``ep`` is the expert-model-parallel size; it must divide the data-parallel + size, so the minimum DP equals ``max(ep, 1)``. World size is therefore + ``tp * pp * cp * dp``. + """ + + id: str + tp: int = 1 + pp: int = 1 + ep: int = 1 + cp: int = 1 + # Optional dp_size override; if None, we use ``max(ep, 1)``. + min_dp: int | None = None + + +@dataclass(frozen=True) +class SGLangShape: + """One SGLang engine shape under test. + + Mirrors SGLang ``ServerArgs`` knobs: ``tp_size`` controls the actual GPU + count of one engine; ``dp_size`` is meaningful only with + ``--enable-dp-attention`` and refers to attention DP within those + ``tp_size`` GPUs (i.e. it does not change the GPU count). ``pp_size`` + multiplies the per-engine GPU count. + """ + + id: str + tp_size: int + ep_size: int = 1 + dp_size: int = 1 + pp_size: int = 1 + enable_dp_attention: bool = False + # Some shapes (PP > 1, dp-attention) require piecewise CUDA graph off in + # SGLang upstream; we just always disable both graph paths in tests. + disable_cuda_graph: bool = True + + @property + def num_gpus_per_engine(self) -> int: + return self.tp_size * self.pp_size + + +MEGATRON_CFGS: tuple[MegatronShape, ...] = ( + MegatronShape(id="mcore_ep2_pp2", tp=1, pp=2, ep=2), + MegatronShape(id="mcore_tp2_pp2", tp=2, pp=2, ep=1), + MegatronShape(id="mcore_tp2_ep2_pp2", tp=2, pp=2, ep=2), +) + +SGLANG_CFGS: tuple[SGLangShape, ...] = ( + SGLangShape( + id="sgl_tp4_ep4_dp4_dpattn", + tp_size=4, + ep_size=4, + dp_size=4, + enable_dp_attention=True, + ), + SGLangShape( + id="sgl_tp4_ep2_dp4_dpattn", + tp_size=4, + ep_size=2, + dp_size=4, + enable_dp_attention=True, + ), + # ``dp_size=2`` partitions the 4-GPU TP world into two attention DP + # sub-groups of 2 GPUs each (instead of dp=4 = one rank per group). + # Both ``ep=2`` and ``ep=4`` fanouts are exercised. + SGLangShape( + id="sgl_tp4_ep2_dp2_dpattn", + tp_size=4, + ep_size=2, + dp_size=2, + enable_dp_attention=True, + ), + SGLangShape( + id="sgl_tp4_ep4_dp2_dpattn", + tp_size=4, + ep_size=4, + dp_size=2, + enable_dp_attention=True, + ), + SGLangShape( + id="sgl_tp2_ep2_pp2", + tp_size=2, + ep_size=2, + dp_size=1, + pp_size=2, + enable_dp_attention=False, + ), + # ``tp=2 pp=2 ep=2 dp=2 + --enable-dp-attention``: same physical 4 GPU + # layout as ``sgl_tp2_ep2_pp2`` but with DP attention turned on, which + # routes through a different MoE token dispatcher and avoids the + # standard-dispatcher illegal-mem-access path that crashes the + # no-dpattn variant on Qwen3-MoE. + SGLangShape( + id="sgl_tp2_ep2_dp2_pp2_dpattn", + tp_size=2, + ep_size=2, + dp_size=2, + pp_size=2, + enable_dp_attention=True, + ), +) + +# Single-GPU shapes used for the smallest end-to-end variant: one Megatron +# trainer rank (DP=1, TP=PP=EP=1) feeding one single-GPU SGLang engine +# (TP=1). Kept out of ``MEGATRON_CFGS`` / ``SGLANG_CFGS`` so they don't +# multiply through the full cartesian product — they're paired explicitly +# in the test parametrization. +MEGATRON_DP1: MegatronShape = MegatronShape(id="mcore_dp1", tp=1, pp=1, ep=1) +SGLANG_TP1: SGLangShape = SGLangShape( + id="sgl_tp1", + tp_size=1, + ep_size=1, + dp_size=1, + pp_size=1, + enable_dp_attention=False, +) + + +# --------------------------------------------------------------------------- +# Sizing helpers +# --------------------------------------------------------------------------- +def min_dp_for_megatron(shape: MegatronShape) -> int: + """Smallest data-parallel size that satisfies ``EP | DP``.""" + if shape.min_dp is not None: + return shape.min_dp + return max(shape.ep, 1) + + +def megatron_world_size(shape: MegatronShape) -> int: + return shape.tp * shape.pp * shape.cp * min_dp_for_megatron(shape) + + +def required_world_size( + *, megatron: MegatronShape, sglang: SGLangShape, colocated: bool +) -> int: + """Total GPUs needed for one ``(megatron, sglang, mode)`` triple. + + Colocate: trainer and SGLang share the same physical GPUs, so the answer + is ``max(megatron, sglang)``. Disaggregate: trainer + inference are on + disjoint placement groups, so the answer is ``megatron + sglang``. + """ + m = megatron_world_size(megatron) + s = sglang.num_gpus_per_engine # one engine; multi-engine clusters scale this + return max(m, s) if colocated else m + s + + +# --------------------------------------------------------------------------- +# Policy config builder (Megatron training side) +# --------------------------------------------------------------------------- +def make_policy_config( + *, + model_path: str, + megatron: MegatronShape, + colocated: bool, + max_seq_len: int = 1024, + train_micro_batch_size: int = 1, +) -> dict[str, Any]: + """Build a ``PolicyConfig`` dict for ``lm_policy.Policy`` (Megatron backend). + + The returned dict mirrors the keys ``Policy.__init__`` and the megatron + worker's ``validate_and_set_config`` read. The ``generation`` block is + populated only as a stub — the real generator is constructed separately + via ``SGLangGeneration``; ``Policy`` only needs ``generation.colocated`` + set so the megatron setup honours the colocate/disaggregate mode. + """ + # Megatron requires ``global_batch_size`` to be divisible by + # ``micro_batch_size * data_parallel_size``. Our DP comes from EP (since + # EP must divide DP) — see ``min_dp_for_megatron``. The roundtrip test + # never runs an actual training step so the smallest legal global batch + # size will do. + dp_size = min_dp_for_megatron(megatron) + train_global_batch_size = train_micro_batch_size * dp_size + use_mxfp8 = weight_update_precision() == "mxfp8" + return { + "model_name": model_path, + "tokenizer": {"name": model_path}, + "train_global_batch_size": train_global_batch_size, + "train_micro_batch_size": train_micro_batch_size, + "logprob_batch_size": train_micro_batch_size, + "precision": "bfloat16", + "max_total_sequence_length": max_seq_len, + "make_sequence_length_divisible_by": megatron.tp if megatron.tp > 1 else 1, + "max_grad_norm": 1.0, + "offload_optimizer_for_logprob": False, + "refit_buffer_size_gb": 1, + # No DTensor — pure Megatron path. + "dtensor_cfg": {"enabled": False}, + "megatron_cfg": { + "enabled": True, + # ``train_iters`` is required by ``_validate_training_config``. + # The roundtrip test never actually steps the optimizer (we only + # exercise the refit path) so any positive integer works; pick a + # small one and keep the lr_decay_iters below in sync. + "train_iters": 10, + "empty_unused_memory_level": 1, + "activation_checkpointing": False, + "converter_type": "Qwen3MoeForCausalLM", + "tensor_model_parallel_size": megatron.tp, + "expert_tensor_parallel_size": 1, + "expert_model_parallel_size": megatron.ep, + "pipeline_model_parallel_size": megatron.pp, + "num_layers_in_first_pipeline_stage": None, + "num_layers_in_last_pipeline_stage": None, + "context_parallel_size": megatron.cp, + "pipeline_dtype": "bfloat16", + "sequence_parallel": megatron.tp > 1, + "freeze_moe_router": True, + "moe_router_dtype": "fp64", + "moe_router_load_balancing_type": "none", + "moe_router_bias_update_rate": 0.0, + "moe_permute_fusion": False, + "apply_rope_fusion": True, + # Megatron-Core only fuses bias for gelu/swiglu/quick_geglu. + # Leave the bias-activation fusion off so any swap of activation + # function (e.g. silu) doesn't trip the fused kernel's check. + "bias_activation_fusion": False, + "defer_fp32_logits": False, + "moe_per_layer_logging": False, + "moe_enable_deepep": False, + "moe_token_dispatcher_type": "alltoall", + "moe_shared_expert_overlap": False, + "peft": {"enabled": False}, + "optimizer": { + "optimizer": "adam", + "lr": 5.0e-6, + "min_lr": 5.0e-7, + "weight_decay": 0.0, + "bf16": True, + "fp16": False, + "params_dtype": "float32", + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_eps": 1e-8, + "sgd_momentum": 0.9, + "use_distributed_optimizer": True, + "use_precision_aware_optimizer": True, + "clip_grad": 1.0, + "optimizer_cpu_offload": False, + "optimizer_offload_fraction": 0.0, + }, + "scheduler": { + "start_weight_decay": 0.0, + "end_weight_decay": 0.0, + "weight_decay_incr_style": "constant", + "lr_decay_style": "constant", + "lr_decay_iters": 10, + "lr_warmup_iters": 0, + "lr_warmup_init": 5.0e-7, + }, + "distributed_data_parallel_config": { + "grad_reduce_in_fp32": False, + "overlap_grad_reduce": False, + "overlap_param_gather": False, + "use_custom_fsdp": False, + "data_parallel_sharding_strategy": "optim_grads_params", + }, + "fp8_cfg": { + "enabled": use_mxfp8, + "fp8": "e4m3", + "fp8_recipe": "mxfp8" if use_mxfp8 else "blockwise", + "fp8_param": False, + }, + "env_vars": ( + {"NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1"} + if use_mxfp8 + else None + ), + }, + "dynamic_batching": {"enabled": False}, + "sequence_packing": { + "enabled": True, + "train_mb_tokens": max_seq_len * train_micro_batch_size, + "logprob_mb_tokens": max_seq_len * train_micro_batch_size, + "algorithm": "modified_first_fit_decreasing", + "sequence_length_round": 64, + }, + # Stub generation block; ``Policy`` only reads ``colocated.enabled``. + "generation": { + "backend": "sglang", + "max_new_tokens": 16, + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "stop_token_ids": None, + "stop_strings": None, + "colocated": { + "enabled": colocated, + "resources": {"gpus_per_node": None, "num_nodes": None}, + }, + }, + } + + +# --------------------------------------------------------------------------- +# SGLang generation config builder +# --------------------------------------------------------------------------- +def make_sglang_cfg( + *, + model_path: str, + sglang: SGLangShape, + colocated: bool, + max_seq_len: int = 1024, + pad_token_id: int = PAD_TOKEN_ID, + eos_token_id: int = EOS_TOKEN_ID, +) -> dict[str, Any]: + """Build the SGLang generation config consumed by ``SGLangGeneration``. + + Field names track ``nemo_rl.models.generation.sglang.config.SGLangConfig`` + and ``SglangSpecificArgs``; SGLang-side flags (``enable_dp_attention``, + ``enable_ep_moe``, ``ep_size``, ``dp_size``, ``pp_size``) match upstream + ``ServerArgs`` (see ``sglang/python/sglang/srt/server_args.py``). + """ + runtime_model_path = model_path + quantization_cfg: dict[str, Any] | None = None + if weight_update_precision() == "mxfp8": + from nemo_rl.models.generation.sglang.mxfp8_setup import ( + ensure_mxfp8_checkpoint, + ) + + quantization_cfg = _mxfp8_quantization_cfg() + runtime_model_path = ensure_mxfp8_checkpoint( + model_path=model_path, + quantization_cfg=quantization_cfg, + ) + + sglang_cfg: dict[str, Any] = { + "model_path": runtime_model_path, + "dtype": "bfloat16", + "random_seed": 42, + "context_length": max_seq_len, + "log_level": "warning", + "skip_server_warmup": True, + "dp_size": sglang.dp_size, + "pp_size": sglang.pp_size, + "ep_size": sglang.ep_size, + "disable_piecewise_cuda_graph": True, + "disable_cuda_graph": sglang.disable_cuda_graph, + # Keep the static pool small enough for MXFP8 weight_checker, which + # may need a large temporary dequant buffer for expert weights. + "mem_fraction_static": 0.3, + } + if sglang.enable_dp_attention: + sglang_cfg["enable_dp_attention"] = True + if moe_runner_backend := os.environ.get(SGLANG_MOE_RUNNER_BACKEND_ENV): + sglang_cfg["moe_runner_backend"] = moe_runner_backend + if fp8_gemm_runner_backend := os.environ.get(SGLANG_FP8_GEMM_RUNNER_BACKEND_ENV): + sglang_cfg["fp8_gemm_runner_backend"] = fp8_gemm_runner_backend + if quantization_cfg is not None: + sglang_cfg["quantization"] = quantization_cfg + # NOTE: ``enable_ep_moe`` was removed in newer sglang versions; EP MoE is + # now activated implicitly when ``ep_size > 1`` (or explicitly via + # ``moe_a2a_backend``). We rely on the implicit path here. + + weight_transfer_mode = "ipc" if colocated else "broadcast" + + return { + "backend": "sglang", + "model_name": runtime_model_path, + "model_path": runtime_model_path, + "tokenizer": {"name": runtime_model_path}, + "dtype": "bfloat16", + "max_new_tokens": 16, + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "stop_token_ids": [eos_token_id], + "stop_strings": None, + "_pad_token_id": pad_token_id, + "sglang_cfg": sglang_cfg, + "sglang_server": { + # Total inference-side GPUs in this SGLang group. We always launch + # exactly one engine per parametrize variant, so num_gpus equals + # one engine's GPU count. + "num_gpus": sglang.num_gpus_per_engine, + "num_gpus_per_engine": sglang.num_gpus_per_engine, + # Offload is only meaningful in colocate mode, where the trainer + # and engine share the same GPU and the engine has to release + # its weights / kv / cuda_graph to free room for Megatron. + # In disaggregate the engine owns its GPU outright; turning + # ``torch_memory_saver`` off here also avoids the side effect + # where it forces NCCL to fall back to the ``P2P/IPC`` + # transport (which fails on hosts without inter-GPU P2P). + "needs_offload": colocated, + "cpu_weight_backup": False, + "sglang_server_concurrency": 64, + "pause_generation_mode": "retract", + "weight_transfer_mode": weight_transfer_mode, + }, + "sglang_router": { + "sglang_router_ip": None, + "sglang_router_port": None, + }, + "sglang_kwargs": {}, + } + + +# --------------------------------------------------------------------------- +# Misc +# --------------------------------------------------------------------------- +@dataclass +class TestTriple: + """A single (megatron, sglang, mode) parametrize variant.""" + + megatron: MegatronShape + sglang: SGLangShape + colocated: bool + extra_marks: list[Any] = field(default_factory=list) + + @property + def id(self) -> str: + mode = "colo" if self.colocated else "disag" + return f"{mode}-{self.megatron.id}-{self.sglang.id}" + + +def all_triples() -> list[TestTriple]: + """Full Cartesian product the user asked for: 2 modes × 3 mcore × 3 sgl.""" + out: list[TestTriple] = [] + for colocated in (True, False): + for m in MEGATRON_CFGS: + for s in SGLANG_CFGS: + out.append(TestTriple(megatron=m, sglang=s, colocated=colocated)) + return out diff --git a/tests/unit/models/generation/sglang/_nemotron_slicer.py b/tests/unit/models/generation/sglang/_nemotron_slicer.py new file mode 100644 index 0000000000..a4a75ec9bd --- /dev/null +++ b/tests/unit/models/generation/sglang/_nemotron_slicer.py @@ -0,0 +1,232 @@ +# 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. + +"""Slice Nemotron-3-Nano-30B-A3B-BF16 down to its first ``MEMEM*`` block. + +The full model is far too large for a unit test (~60GB BF16). The Nemotron-H +hybrid layer pattern is:: + + MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME + +The first segment ``MEMEM*`` covers six layers (Mamba / MoE alternating with +one attention layer ``*``). Keeping just those six layers gives a model that +still exercises every Nemotron module class (Mamba, MoE, attention) under +both colocate and disaggregate refit, while small enough to fit on the test +hosts we use (~8 GPU H100). + +This module is invoked from a session-scoped fixture in ``conftest.py``; the +sliced checkpoint is cached at ``${HF_HOME or ~/.cache/huggingface}/ +nemotron-3-nano-sliced-MEMEM`` so the heavy work only runs once per host. +""" + +from __future__ import annotations + +import json +import os +import re +import shutil +from pathlib import Path + +# Source HF repo to slice from. +SOURCE_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + +# Layer pattern of the slice we keep. ``MEMEM*EMEMEM*E`` = 14 layers (the +# first 14 characters of the Nemotron-3-Nano-30B-A3B hybrid pattern). +# +# Why 14 specifically? +# 1. Sglang's NemotronH model implementation has a corner case where, with +# PP > 1, **every** PP rank must get at least one attention layer +# (``*``). Otherwise the weight checker's GPU→CPU copy of the rank's +# ``embed_tokens.weight`` fails with ``CUDA error: invalid argument``. +# 2. Megatron requires ``num_layers % pp_size == 0`` (or an explicit +# ``|`` separator in the hybrid layer pattern). 14 splits cleanly 7+7 +# for PP=2 — the only PP value used by our megatron configs. +# 3. The first 14 chars of the upstream pattern are ``MEMEM*EMEMEM*E``, +# which under a 7+7 PP split puts exactly one ``*`` in each half +# (rank 0: ``MEMEM*E``; rank 1: ``MEMEM*E`` — same shape). +# +# A 6-layer slice (``MEMEM*``) trips (1); a 13-layer slice (``MEMEM*EMEMEM*``) +# satisfies (1) but trips (2). 14 is the smallest layer count satisfying +# both. +SLICED_PATTERN = "MEMEM*EMEMEM*E" +SLICED_NUM_LAYERS = len(SLICED_PATTERN) + +# Cache directory inside the user's HF cache root. The trailing suffix is +# included in the directory name so a future change to ``SLICED_PATTERN`` +# auto-creates a fresh cache instead of silently reusing a stale slice. +SLICED_DIR_NAME = f"nemotron-3-nano-sliced-{SLICED_PATTERN.replace('*', 'A')}" + + +def _hf_cache_root() -> Path: + """Return the same root that HuggingFace would use, honoring HF_HOME.""" + root = os.environ.get("HF_HOME") or str(Path.home() / ".cache" / "huggingface") + return Path(root) + + +def sliced_model_path() -> Path: + """Absolute path to the sliced checkpoint (whether or not it exists yet). + + Lives under ``${HF_HOME or ~/.cache/huggingface}/hub/`` so it sits next to + the other HF Hub-cached models (``models----/...``) instead of + polluting the parent ``huggingface/`` directory. + """ + return _hf_cache_root() / "hub" / SLICED_DIR_NAME + + +# --------------------------------------------------------------------------- +# Layer-key matching +# --------------------------------------------------------------------------- +# Match weight keys that live under a transformer layer indexed by an integer. +# Examples: +# model.layers.0.self_attn.q_proj.weight (layer 0) +# model.layers.42.mlp.experts.7.w1.weight (layer 42) +# backbone.layers.5.mixer.A_log (Mamba style) +# +# Anything that does not match (embeddings, lm_head, final norm, rotary cache) +# is a non-layer tensor and is preserved unchanged. +_LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)\.") + + +def _layer_index_for_key(key: str) -> int | None: + m = _LAYER_RE.search(key) + return int(m.group(1)) if m else None + + +def _should_keep_key(key: str, kept_layers: int) -> bool: + idx = _layer_index_for_key(key) + return idx is None or idx < kept_layers + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- +def ensure_sliced_model(force: bool = False) -> Path: + """Materialize the sliced checkpoint and return its path. + + Idempotent: if the target directory already contains a populated + ``config.json`` and weight files we assume a previous run produced it. + Set ``force=True`` to nuke and rebuild. + """ + out_dir = sliced_model_path() + if not force and out_dir.is_dir() and (out_dir / "config.json").is_file(): + # Cheap sanity check: at least one safetensors shard. + if any(out_dir.glob("*.safetensors")): + return out_dir + + # Heavy imports kept local so importing this module costs nothing. + from huggingface_hub import snapshot_download + from safetensors import safe_open + from safetensors.torch import save_file + + if force and out_dir.exists(): + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # 1. Pull every file we need from HF (config, tokenizer, safetensors, + # safetensors index). This populates the standard HF cache; we then + # re-emit a slimmed copy at ``out_dir``. + src = Path( + snapshot_download( + repo_id=SOURCE_MODEL_ID, + allow_patterns=[ + "*.json", + "*.txt", + "*.model", + "*.safetensors", + "*.py", # NemotronH ships custom modeling/configuration code + "tokenizer.*", + "special_tokens_map.json", + "generation_config.json", + ], + ) + ) + + # 2. Patch config.json: shrink to ``SLICED_NUM_LAYERS`` and rewrite the + # hybrid layer pattern. + config_path = src / "config.json" + with open(config_path) as f: + cfg = json.load(f) + + # Both keys exist on Nemotron-H configs in the wild (older checkpoints + # used ``hybrid_override_pattern``, newer ones ``hybrid_layer_pattern``); + # rewrite whichever is set so we don't depend on which one HF picked. + for pattern_key in ("hybrid_override_pattern", "hybrid_layer_pattern"): + if pattern_key in cfg: + cfg[pattern_key] = SLICED_PATTERN + cfg["num_hidden_layers"] = SLICED_NUM_LAYERS + # Some converters look at this instead. + if "num_layers" in cfg: + cfg["num_layers"] = SLICED_NUM_LAYERS + + with open(out_dir / "config.json", "w") as f: + json.dump(cfg, f, indent=2) + + # 3. Copy non-weight files verbatim (tokenizer, generation_config, ...). + for item in src.iterdir(): + if item.suffix == ".safetensors" or item.name == "config.json": + continue + if item.name.endswith(".safetensors.index.json"): + continue + if item.is_file(): + shutil.copy2(item, out_dir / item.name) + + # 4. Walk every safetensors shard and write a pruned copy. We keep tensors + # that either don't reference a layer index, or reference layer < N. + src_index_path = src / "model.safetensors.index.json" + if src_index_path.is_file(): + with open(src_index_path) as f: + src_index = json.load(f) + weight_map: dict[str, str] = src_index.get("weight_map", {}) + # Group by source shard to amortize the open() cost. + per_shard: dict[str, list[str]] = {} + for k, shard in weight_map.items(): + per_shard.setdefault(shard, []).append(k) + new_weight_map: dict[str, str] = {} + new_total_size = 0 + for shard, keys in per_shard.items(): + keep = [k for k in keys if _should_keep_key(k, SLICED_NUM_LAYERS)] + if not keep: + continue + with safe_open(src / shard, framework="pt") as reader: + tensors = {k: reader.get_tensor(k) for k in keep} + metadata = reader.metadata() or {} + save_file(tensors, str(out_dir / shard), metadata=metadata) + for k in keep: + new_weight_map[k] = shard + new_total_size += tensors[k].numel() * tensors[k].element_size() + del tensors + with open(out_dir / "model.safetensors.index.json", "w") as f: + json.dump( + { + "metadata": {"total_size": new_total_size}, + "weight_map": new_weight_map, + }, + f, + indent=2, + ) + else: + # Single-file checkpoint path. + single = src / "model.safetensors" + with safe_open(single, framework="pt") as reader: + keep = [k for k in reader.keys() if _should_keep_key(k, SLICED_NUM_LAYERS)] + tensors = {k: reader.get_tensor(k) for k in keep} + metadata = reader.metadata() or {} + save_file(tensors, str(out_dir / "model.safetensors"), metadata=metadata) + + return out_dir + + +if __name__ == "__main__": # pragma: no cover — manual one-shot use + path = ensure_sliced_model(force="--force" in os.sys.argv) + print(path) diff --git a/tests/unit/models/generation/sglang/_qwen3_slicer.py b/tests/unit/models/generation/sglang/_qwen3_slicer.py new file mode 100644 index 0000000000..5a27402d81 --- /dev/null +++ b/tests/unit/models/generation/sglang/_qwen3_slicer.py @@ -0,0 +1,180 @@ +# 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. + +"""Slice ``Qwen/Qwen3-30B-A3B-Instruct-2507`` down to the first N layers. + +The full model has 48 transformer layers and 30B parameters (~60 GiB bf16), +which doesn't fit on a single H200 with the Megatron-side full +params+grads+optimizer footprint (~240 GiB total) needed by the test +fixtures' single-GPU variants. + +Slicing to ``SLICED_NUM_LAYERS`` keeps the same architecture (Qwen3MoE, +128 experts per layer, 8 active per token) but drops all but the first N +transformer blocks. This: + +* Preserves every module class the refit path exercises (attention, MoE + router, MoE experts). +* Brings the parameter count to roughly ``base + N * per_layer ≈ + embeddings + small`` — small enough for the ``mcore_dp1``/``sgl_tp1`` + parametrization to fit on one GPU. +* Leaves the tokenizer, generation config, and all non-weight files + untouched. + +Idempotent: the sliced checkpoint is cached at +``${HF_HOME or ~/.cache/huggingface}/qwen3-30b-a3b-sliced-N``. Set +``force=True`` to rebuild. +""" + +from __future__ import annotations + +import json +import os +import re +import shutil +from pathlib import Path + +# Source HF repo to slice from. +SOURCE_MODEL_ID = "Qwen/Qwen3-30B-A3B-Instruct-2507" + +# Number of transformer layers to keep (full model has 48). +SLICED_NUM_LAYERS = 4 + +# Tensors with names like ``model.layers..<...>`` carry per-layer +# weights. Anything that doesn't match (embeddings, norm, lm_head, +# rotary buffers) is kept unconditionally. +_LAYER_RE = re.compile(r"\bmodel\.layers\.(\d+)\.") + + +def sliced_model_path() -> Path: + """Where the sliced checkpoint lives on disk.""" + hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface") + return Path(hf_home) / "hub" / f"qwen3-30b-a3b-sliced-{SLICED_NUM_LAYERS}" + + +def _layer_index_for_key(key: str) -> int | None: + m = _LAYER_RE.search(key) + return int(m.group(1)) if m else None + + +def _should_keep_key(key: str, kept_layers: int) -> bool: + idx = _layer_index_for_key(key) + return idx is None or idx < kept_layers + + +def ensure_sliced_model(force: bool = False) -> Path: + """Materialize the sliced Qwen3-30B-A3B checkpoint and return its path. + + Idempotent: if the target directory already contains a populated + ``config.json`` and at least one safetensors shard we assume a + previous run produced it. Set ``force=True`` to nuke and rebuild. + """ + out_dir = sliced_model_path() + if not force and out_dir.is_dir() and (out_dir / "config.json").is_file(): + if any(out_dir.glob("*.safetensors")): + return out_dir + + # Heavy imports kept local so importing this module costs nothing. + from huggingface_hub import snapshot_download + from safetensors import safe_open + from safetensors.torch import save_file + + if force and out_dir.exists(): + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # 1. Pull every file we need from HF (config, tokenizer, safetensors, + # index). This populates the standard HF cache; we then re-emit a + # slimmed copy at ``out_dir``. + src = Path( + snapshot_download( + repo_id=SOURCE_MODEL_ID, + allow_patterns=[ + "*.json", + "*.txt", + "*.model", + "*.safetensors", + "tokenizer.*", + "special_tokens_map.json", + "generation_config.json", + ], + ) + ) + + # 2. Patch config.json: shrink ``num_hidden_layers``. + config_path = src / "config.json" + with open(config_path) as f: + cfg = json.load(f) + cfg["num_hidden_layers"] = SLICED_NUM_LAYERS + if "num_layers" in cfg: + cfg["num_layers"] = SLICED_NUM_LAYERS + with open(out_dir / "config.json", "w") as f: + json.dump(cfg, f, indent=2) + + # 3. Copy non-weight files verbatim (tokenizer, generation_config, ...). + for item in src.iterdir(): + if item.suffix == ".safetensors" or item.name == "config.json": + continue + if item.name.endswith(".safetensors.index.json"): + continue + if item.is_file(): + shutil.copy2(item, out_dir / item.name) + + # 4. Walk every safetensors shard and write a pruned copy. We keep tensors + # that either don't reference a layer index, or reference layer < N. + src_index_path = src / "model.safetensors.index.json" + if src_index_path.is_file(): + with open(src_index_path) as f: + src_index = json.load(f) + weight_map: dict[str, str] = src_index.get("weight_map", {}) + per_shard: dict[str, list[str]] = {} + for k, shard in weight_map.items(): + per_shard.setdefault(shard, []).append(k) + new_weight_map: dict[str, str] = {} + new_total_size = 0 + for shard, keys in per_shard.items(): + keep = [k for k in keys if _should_keep_key(k, SLICED_NUM_LAYERS)] + if not keep: + continue + with safe_open(src / shard, framework="pt") as reader: + tensors = {k: reader.get_tensor(k) for k in keep} + metadata = reader.metadata() or {} + save_file(tensors, str(out_dir / shard), metadata=metadata) + for k in keep: + new_weight_map[k] = shard + new_total_size += tensors[k].numel() * tensors[k].element_size() + del tensors + with open(out_dir / "model.safetensors.index.json", "w") as f: + json.dump( + { + "metadata": {"total_size": new_total_size}, + "weight_map": new_weight_map, + }, + f, + indent=2, + ) + else: + # Single-file checkpoint path. + single = next(src.glob("*.safetensors"), None) + if single is None: + raise RuntimeError( + f"no safetensors files found under {src}; cannot slice" + ) + with safe_open(single, framework="pt") as reader: + keys = list(reader.keys()) + keep = [k for k in keys if _should_keep_key(k, SLICED_NUM_LAYERS)] + tensors = {k: reader.get_tensor(k) for k in keep} + metadata = reader.metadata() or {} + save_file(tensors, str(out_dir / single.name), metadata=metadata) + + return out_dir diff --git a/tests/unit/models/generation/sglang/conftest.py b/tests/unit/models/generation/sglang/conftest.py new file mode 100644 index 0000000000..b3fea83088 --- /dev/null +++ b/tests/unit/models/generation/sglang/conftest.py @@ -0,0 +1,115 @@ +# 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. + +"""Conftest for sglang tests — real Ray, real SGLang. + +Tests in this directory exercise the sglang generation modules using real Ray +actors and real SGLang servers. The conftest stubs non-sglang heavy +dependencies but lets sglang imports resolve naturally against the installed +package. +""" + +import importlib.machinery +import os +import sys +from unittest.mock import MagicMock + +# Set default GPU devices before any CUDA/Ray initialisation. +# The remote cluster reserves GPUs 4-7 for this work. +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "4,5,6,7") + +# Use system Python for all Ray actors (uv not configured in container). +os.environ.setdefault("NEMO_RL_PY_EXECUTABLES_SYSTEM", "1") + +# Disable sglang's per-GPU memory imbalance check — when running tests on a +# shared host other processes may already hold memory on some of our GPUs. +os.environ.setdefault("SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK", "false") + +# Skip the runtime ``sglang-kernel`` distribution version check. The container +# ships the pre-built ``sgl-kernel`` (the older distribution name) bound +# against torch 2.10; the renamed ``sglang-kernel`` dist isn't installed so +# the assert would unconditionally raise. The kernel itself is fine — only +# the metadata lookup fails. +os.environ.setdefault("SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK", "1") + +# Ensure the test directory is on sys.path so helpers.py is importable. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# --------------------------------------------------------------------------- +# Stub heavy modules NOT installed in the sglang test environment. +# sglang is NOT stubbed — we test against a real server. +# --------------------------------------------------------------------------- +_STUB_MODULES = [ + "decord", + "vllm", + "vllm.sampling_params", + "vllm.lora", + "vllm.lora.request", + "wandb", +] + +# Transformer Engine is conditionally stubbed: the sglang-only test image +# historically didn't ship a working TE, but the e2e image (where we also +# exercise the Megatron policy worker) does. Probe with importlib so we only +# replace it with a MagicMock when real TE genuinely cannot be imported — +# stubbing real TE turns ``transformer_engine.pytorch`` into a MagicMock which +# makes downstream ``from transformer_engine.pytorch.tensor import ...`` fail +# with "X is not a package" inside megatron.core. +import importlib.util as _importlib_util + +if _importlib_util.find_spec("transformer_engine") is None: + _STUB_MODULES += [ + "transformer_engine", + "transformer_engine.common", + "transformer_engine.pytorch", + ] + +for _mod in _STUB_MODULES: + if _mod in sys.modules: + continue + stub = MagicMock() + # importlib.util.find_spec requires __spec__ to be a real ModuleSpec. + stub.__spec__ = importlib.machinery.ModuleSpec(_mod, loader=None) + stub.__name__ = _mod + sys.modules[_mod] = stub + +import pytest +import ray + +from nemo_rl.models.generation.sglang.sglang_router import RouterActor + + +# --------------------------------------------------------------------------- +# Session-scoped fixtures +# --------------------------------------------------------------------------- +@pytest.fixture(scope="session") +def ray_cluster(): + """Initialise Ray once for the entire test session.""" + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + yield + ray.shutdown() + + +@pytest.fixture(scope="session") +def router(ray_cluster): + """Start a real sglang router that lives for the session.""" + actor = RouterActor.remote() + ip, port = ray.get(actor.start.remote({})) + yield {"actor": actor, "ip": ip, "port": port} + try: + ray.get(actor.stop.remote()) + except Exception: + pass + ray.kill(actor) diff --git a/tests/unit/models/generation/sglang/helpers.py b/tests/unit/models/generation/sglang/helpers.py new file mode 100644 index 0000000000..bec0df332d --- /dev/null +++ b/tests/unit/models/generation/sglang/helpers.py @@ -0,0 +1,175 @@ +# 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 helpers for sglang tests. + +Kept in a regular module (not conftest.py) so test files can import it +directly. conftest.py also imports from here for fixture definitions. +""" + +import os + +import ray + +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, +) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +MODEL_PATH = "Qwen/Qwen3-0.6B" + +# Qwen3-0.6B model dimensions (verified against HuggingFace config) +HIDDEN_SIZE = 1024 +INTERMEDIATE_SIZE = 3072 +NUM_ATTENTION_HEADS = 16 +NUM_KV_HEADS = 8 +HEAD_DIM = 128 +QKV_OUTPUT_DIM = NUM_ATTENTION_HEADS * HEAD_DIM + 2 * NUM_KV_HEADS * HEAD_DIM # 4096 + + +# --------------------------------------------------------------------------- +# Config helpers +# --------------------------------------------------------------------------- +DEFAULT_GPUS_PER_NODE = 4 + + +def make_sglang_cfg( + model_path=MODEL_PATH, + tp_size=1, + num_gpus=4, + router_ip=None, + router_port=None, +): + return { + "sglang_cfg": { + "model_path": model_path, + "random_seed": 42, + "dp_size": 1, + "pp_size": 1, + "ep_size": 1, + "skip_server_warmup": True, + "dtype": "bfloat16", + "context_length": 1024, + "log_level": "warning", + "disable_piecewise_cuda_graph": True, + "disable_cuda_graph": True, + "mem_fraction_static": 0.3, + }, + "sglang_server": { + "num_gpus": num_gpus, + "num_gpus_per_engine": tp_size, + "needs_offload": True, + "cpu_weight_backup": False, + "sglang_server_concurrency": 64, + "pause_generation_mode": "retract", + }, + "sglang_router": { + "sglang_router_ip": router_ip, + "sglang_router_port": router_port, + }, + } + + +def make_actor_env_vars(): + """Build env-vars dict for SGLang worker actors.""" + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd: + env_vars["CUDA_VISIBLE_DEVICES"] = cvd + return env_vars + + +def create_worker(router_info, base_gpu_id=0, tp_size=1, rank=0): + """Create and initialise a real SGLangGenerationWorker Ray actor. + + Returns the actor handle after ``init`` completes. + """ + gpus_per_node = DEFAULT_GPUS_PER_NODE + sglang_cfg = make_sglang_cfg( + tp_size=tp_size, + router_ip=router_info["ip"], + router_port=router_info["port"], + ) + + worker = SGLangGenerationWorker.options( + num_cpus=0.2, + num_gpus=0.2, + runtime_env={"env_vars": make_actor_env_vars()}, + ).remote( + gpus_per_node, + sglang_cfg, + rank=rank, + base_gpu_id=base_gpu_id, + num_gpus_per_engine=tp_size, + ) + + 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) + + ray.get( + worker.init.remote( + dist_init_addr=f"{host_ip}:{dist_init_port}", + port=port, + nccl_port=nccl_port, + router_ip=router_info["ip"], + router_port=router_info["port"], + ) + ) + return worker + + +def make_generation_sampling_params( + max_new_tokens=16, + temperature=0.0, + top_p=1.0, + stop=None, +): + """Build sampling_params dict for generate_one_sample / router /generate.""" + params = { + "temperature": temperature, + "max_new_tokens": max_new_tokens, + "top_p": top_p, + "no_stop_trim": True, + "spaces_between_special_tokens": False, + } + if stop is not None: + params["stop"] = stop + return params + + +# --------------------------------------------------------------------------- +# HTTP helpers for tests that want an explicit status-code check +# --------------------------------------------------------------------------- +def post_and_assert_200(base_url, endpoint, payload=None): + """POST ``payload`` to ``{base_url}/{endpoint}`` and assert HTTP 200. + + Tests that exercise ``release_memory_occupation`` / ``resume_memory_occupation`` + use this instead of ``_make_request`` so the 200 check is visible in the + test body (``_make_request`` consumes the status code inside + ``raise_for_status()`` and returns only the parsed JSON). + """ + import requests + + resp = requests.post(f"{base_url}/{endpoint}", json=payload or {}) + assert resp.status_code == 200, ( + f"POST {endpoint} expected 200, got {resp.status_code}: {resp.text}" + ) + return resp.json() diff --git a/tests/unit/models/generation/sglang/pytest.ini b/tests/unit/models/generation/sglang/pytest.ini new file mode 100644 index 0000000000..adf05beca2 --- /dev/null +++ b/tests/unit/models/generation/sglang/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +markers = + sglang: tests requiring sglang and GPU resources diff --git a/tests/unit/models/generation/sglang/repro_pp_colocate_weight_update.py b/tests/unit/models/generation/sglang/repro_pp_colocate_weight_update.py new file mode 100644 index 0000000000..0283787e19 --- /dev/null +++ b/tests/unit/models/generation/sglang/repro_pp_colocate_weight_update.py @@ -0,0 +1,691 @@ +#!/usr/bin/env python3 +# 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 +"""Standalone colocate/CUDA-IPC repro for SGLang weight update. + +This is the colocate companion to ``repro_pp_weight_update.py``. It avoids +Megatron and the nemo-rl policy worker, but it keeps the production colocate +transport: + +1. A Ray actor hosts an ``sglang.srt.entrypoints.engine.Engine``. +2. The main process is a mock trainer on the same visible GPU. +3. SGLang snapshots its initial weights and randomizes them. Optionally, the + ``test offload/onload`` branch offloads weights/KV and onloads weights before + the refit. +4. The trainer loads the sliced Qwen checkpoint, builds HF weight buckets, and + sends them through ``send_hf_buckets_via_ipc_actor_impl``. +5. SGLang receives the IPC refit, post-processes weights, and compares back to + the snapshot. + +Usage inside the nemo-rl container: + + CUDA_VISIBLE_DEVICES=0 python \ + tests/unit/models/generation/sglang/repro_pp_colocate_weight_update.py \ + --pp 1 --tp 1 --dp 1 --dtype bfloat16 --test-offload-onload + + CUDA_VISIBLE_DEVICES=0,1,2,3 python \ + tests/unit/models/generation/sglang/repro_pp_colocate_weight_update.py \ + --pp 2 --tp 2 --ep 2 --dp 2 --enable-dp-attention \ + --dtype bfloat16 --moe-runner-backend triton --test-offload-onload +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +import traceback +from collections.abc import Iterable, Mapping +from pathlib import Path +from typing import Any + +import ray +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from safetensors import safe_open + +_DTYPE_FROM_STR = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + "float64": torch.float64, + "int8": torch.int8, + "uint8": torch.uint8, + "int32": torch.int32, + "int64": torch.int64, + "bool": torch.bool, +} + +_SAFETENSORS_TO_TORCH_NAME = { + "F64": "float64", + "F32": "float32", + "F16": "float16", + "BF16": "bfloat16", + "I64": "int64", + "I32": "int32", + "I16": "int16", + "I8": "int8", + "U8": "uint8", + "BOOL": "bool", +} + +_MOE_RUNNER_BACKEND_CHOICES = ( + "auto", + "deep_gemm", + "triton", + "triton_kernel", + "flashinfer_trtllm", + "flashinfer_trtllm_routed", + "flashinfer_cutlass", + "flashinfer_mxfp4", + "flashinfer_cutedsl", + "cutlass", + "marlin", +) + + +def _ckpt_path(model_path: str | None = None) -> Path: + if model_path is not None: + path = Path(model_path).expanduser() + if not path.is_dir(): + sys.exit(f"model path not found: {path}") + return path + + hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface") + path = Path(hf_home) / "hub" / "qwen3-30b-a3b-sliced-4" + if not path.is_dir(): + sys.exit( + f"sliced ckpt not found at {path}; produce it first via " + "tests/unit/models/generation/sglang/_qwen3_slicer.py" + ) + return path + + +def _collect_specs(ckpt: Path) -> tuple[list[tuple[str, str, list[int]]], dict[str, str]]: + """Return (ordered list of (name, dtype_str, shape), name -> shard file).""" + index_path = ckpt / "model.safetensors.index.json" + if index_path.is_file(): + with open(index_path) as f: + weight_map: dict[str, str] = json.load(f)["weight_map"] + else: + single = next(ckpt.glob("*.safetensors")).name + with safe_open(ckpt / single, framework="pt") as reader: + weight_map = {name: single for name in reader.keys()} + + per_shard: dict[str, list[str]] = {} + for name, shard in weight_map.items(): + per_shard.setdefault(shard, []).append(name) + + specs: list[tuple[str, str, list[int]]] = [] + for shard in sorted(per_shard): + with safe_open(ckpt / shard, framework="pt") as reader: + for name in sorted(per_shard[shard]): + tensor_slice = reader.get_slice(name) + st_dtype = tensor_slice.get_dtype() + if st_dtype not in _SAFETENSORS_TO_TORCH_NAME: + raise RuntimeError(f"unmapped safetensors dtype {st_dtype!r} for {name}") + specs.append( + ( + name, + _SAFETENSORS_TO_TORCH_NAME[st_dtype], + list(tensor_slice.get_shape()), + ) + ) + + specs.sort(key=lambda x: x[0]) + return specs, weight_map + + +def _result_failed(result: Any) -> tuple[bool, str]: + if isinstance(result, tuple) and len(result) >= 2: + return result[0] is False, str(result[1]) + if isinstance(result, Mapping): + success = result.get("success", True) + message = ( + result.get("error_message") + or result.get("error") + or result.get("message") + or "unknown error" + ) + return success is False, str(message) + if hasattr(result, "success"): + message = getattr(result, "error_message", None) or getattr( + result, "message", "unknown error" + ) + return result.success is False, str(message) + return False, "" + + +def _raise_if_failed(result: Any, action: str) -> Any: + failed, message = _result_failed(result) + if failed: + raise RuntimeError(f"{action} failed: {message}") + return result + + +def _text(output: Any) -> str: + if isinstance(output, list): + output = output[0] + if isinstance(output, Mapping): + return str(output.get("text", repr(output))) + return repr(output) + + +@ray.remote(num_gpus=1) +class _SGLangEngineActor: + """Ray actor facade matching SGLangGenerationWorker's weight-update API.""" + + def __init__( + self, + *, + ckpt: str, + pp: int, + tp: int, + ep: int, + dp: int, + enable_dp_attention: bool, + dtype: str, + mem_fraction_static: float, + moe_runner_backend: str, + ): + os.environ["SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK"] = "false" + + try: + from nemo_rl.models.generation.sglang.sglang_worker import ( + _apply_sglang_compat_patches, + ) + + _apply_sglang_compat_patches() + except Exception as exc: # noqa: BLE001 + print(f"[actor] WARN: failed to apply nemo-rl SGLang patches: {exc}", flush=True) + + from sglang.srt.entrypoints.engine import Engine + + kwargs: dict[str, Any] = { + "model_path": ckpt, + "tp_size": tp, + "pp_size": pp, + "ep_size": ep, + "dp_size": dp, + "dtype": dtype, + "mem_fraction_static": mem_fraction_static, + "log_level": "info", + "random_seed": 42, + "disable_cuda_graph": True, + "enable_memory_saver": True, + "enable_weights_cpu_backup": False, + "moe_runner_backend": moe_runner_backend, + } + if enable_dp_attention: + kwargs["enable_dp_attention"] = True + + print( + f"[actor] starting Engine pp={pp} tp={tp} ep={ep} dp={dp} " + f"enable_dp_attention={enable_dp_attention} dtype={dtype} " + f"moe_runner_backend={moe_runner_backend} " + f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}", + flush=True, + ) + self.engine = Engine(**kwargs) + print("[actor] engine up", flush=True) + try: + server_args = self.engine.tokenizer_manager.server_args + print( + f"[actor] resolved server_args dtype={getattr(server_args, 'dtype', '')!r} " + f"moe_runner_backend={getattr(server_args, 'moe_runner_backend', '')!r}", + flush=True, + ) + except Exception as exc: # noqa: BLE001 + print(f"[actor] WARN: could not introspect server_args: {exc!r}", flush=True) + + def generate(self, prompt: str, sampling_params: dict[str, Any]) -> Any: + return self.engine.generate(prompt, sampling_params=sampling_params) + + def check_weights(self, action: str) -> Any: + from sglang.srt.managers.io_struct import CheckWeightsReqInput + + result = self.engine.loop.run_until_complete( + self.engine.tokenizer_manager.check_weights( + CheckWeightsReqInput(action=action) + ) + ) + return _raise_if_failed(result, f"check_weights({action!r})") + + def offload_weights(self) -> Any: + from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS + + self.engine.flush_cache() + result = self.engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + return _raise_if_failed(result, "offload_weights") + + def offload_kv(self) -> Any: + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + ) + + self.engine.flush_cache() + result = self.engine.release_memory_occupation( + tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH] + ) + return _raise_if_failed(result, "offload_kv") + + def onload_weights(self) -> Any: + from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS + + result = self.engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + return _raise_if_failed(result, "onload_weights") + + def onload_kv(self) -> Any: + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + ) + + result = self.engine.resume_memory_occupation( + tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH] + ) + return _raise_if_failed(result, "onload_kv") + + def post_process_weights( + self, + *, + restore_weights_before_load: bool = False, + post_process_quantization: bool = True, + ) -> Any: + from sglang.srt.managers.io_struct import PostProcessWeightsReqInput + + result = self.engine.loop.run_until_complete( + self.engine.tokenizer_manager.post_process_weights( + PostProcessWeightsReqInput( + restore_weights_before_load=restore_weights_before_load, + post_process_quantization=post_process_quantization, + ) + ) + ) + return _raise_if_failed(result, "post_process_weights") + + def update_weights_from_tensor( + self, + serialized_named_tensors: list[str], + load_format: str | None = None, + flush_cache: bool = False, + weight_version: str | None = None, + ) -> Any: + del weight_version + result = self.engine.update_weights_from_tensor( + named_tensors=serialized_named_tensors, + load_format=load_format, + flush_cache=flush_cache, + ) + return _raise_if_failed(result, "update_weights_from_tensor") + + def shutdown(self) -> None: + if self.engine is not None: + self.engine.shutdown() + self.engine = None + + +class _MockTrainer: + """Load HF checkpoint tensors on one GPU and yield byte-bounded buckets.""" + + def __init__( + self, + *, + ckpt: Path, + device: str, + target_dtype: torch.dtype, + specs: list[tuple[str, str, list[int]]], + weight_map: dict[str, str], + ): + self._tensors: list[tuple[str, torch.Tensor]] = [] + for idx, (name, dtype_str, _shape) in enumerate(specs): + with safe_open(ckpt / weight_map[name], framework="pt") as reader: + tensor = reader.get_tensor(name) + dtype = target_dtype if tensor.is_floating_point() else _DTYPE_FROM_STR[dtype_str] + self._tensors.append((name, tensor.to(device=device, dtype=dtype).contiguous())) + if idx % 25 == 0 or idx == len(specs) - 1: + print(f"[trainer] loaded {idx + 1}/{len(specs)} {name}", flush=True) + + def iter_buckets(self, buffer_size_bytes: int) -> Iterable[list[tuple[str, torch.Tensor]]]: + bucket: list[tuple[str, torch.Tensor]] = [] + bucket_size = 0 + for name, tensor in self._tensors: + 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 _mock_trainer_rank_main( + *, + rank: int, + world_size: int, + master_addr: str, + master_port: int, + ckpt_str: str, + specs: list[tuple[str, str, list[int]]], + weight_map: dict[str, str], + target_dtype_name: str, + buffer_size_bytes: int, + ray_address: str, + ray_namespace: str, + engine_name: str, + engine_gpu_counts: list[int], + engine_gpu_offsets: list[int], +) -> None: + """Mock one Megatron rank and use the production colocated IPC topology.""" + try: + torch.cuda.set_device(rank % torch.cuda.device_count()) + + from nemo_rl.models.policy.torch_reductions_utils import ( + monkey_patch_torch_reductions, + ) + from nemo_rl.models.policy.utils import ( + connect_colocate_topology, + send_hf_buckets_via_ipc_actor_impl, + ) + + ray.init( + address=ray_address, + namespace=ray_namespace, + ignore_reinit_error=True, + log_to_driver=True, + ) + engine = ray.get_actor(engine_name, namespace=ray_namespace) + + dist.init_process_group( + backend="gloo", + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=rank, + ) + + worker_state: dict[str, Any] = {} + connect_colocate_topology( + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + worker_state=worker_state, + monkey_patch_fn=monkey_patch_torch_reductions, + ) + + device = f"cuda:{rank % torch.cuda.device_count()}" + print( + f"[trainer rank={rank}] connected colocate topology " + f"engine_gpu_counts={engine_gpu_counts} offsets={engine_gpu_offsets} " + f"device={device}", + flush=True, + ) + trainer = _MockTrainer( + ckpt=Path(ckpt_str), + device=device, + target_dtype=_DTYPE_FROM_STR[target_dtype_name], + specs=specs, + weight_map=weight_map, + ) + + print(f"[trainer rank={rank}] send_hf_buckets_via_ipc_actor_impl...", flush=True) + t0 = time.time() + send_hf_buckets_via_ipc_actor_impl( + bucket_iterator=trainer.iter_buckets(buffer_size_bytes), + rollout_engines=[engine], + worker_state=worker_state, + weight_version=1, + ) + print( + f"[trainer rank={rank}] send_hf_buckets done in {time.time() - t0:.1f}s", + flush=True, + ) + except Exception: # noqa: BLE001 + traceback.print_exc() + os._exit(1) + finally: + if dist.is_initialized(): + try: + group = locals().get("worker_state", {}).get("_ipc_gather_group") + if group is not None: + dist.destroy_process_group(group) + except Exception: + pass + try: + dist.destroy_process_group() + except Exception: + pass + if ray.is_initialized(): + ray.shutdown() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--ep", type=int, default=1) + parser.add_argument("--dp", type=int, default=1) + parser.add_argument("--enable-dp-attention", action="store_true") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=("bfloat16", "float16", "float32")) + parser.add_argument("--buffer-size-bytes", type=int, default=512 * 1024 * 1024) + parser.add_argument("--master-port", type=int, default=29565) + parser.add_argument("--mem-fraction-static", type=float, default=0.3) + parser.add_argument( + "--model-path", + type=str, + default=None, + help="HF model/checkpoint directory. Defaults to the sliced Qwen repro checkpoint.", + ) + parser.add_argument( + "--test-offload-onload", + action="store_true", + help="Enable the 'test offload/onload' branch around the IPC weight update.", + ) + parser.add_argument( + "--moe-runner-backend", + "--sglang-moe-runner-backend", + dest="moe_runner_backend", + type=str, + default="triton", + choices=_MOE_RUNNER_BACKEND_CHOICES, + help="SGLang MoE runner backend passed as Engine(moe_runner_backend=...).", + ) + args = parser.parse_args() + + n_engine = args.pp * args.tp if args.enable_dp_attention else args.pp * args.tp * args.dp + n_trainer = n_engine + n_avail = torch.cuda.device_count() + if n_avail < n_engine: + sys.exit(f"need at least {n_engine} visible GPU(s), have {n_avail}") + + torch.cuda.set_device(0) + ckpt = _ckpt_path(args.model_path) + specs, weight_map = _collect_specs(ckpt) + + print(f"[main] ckpt={ckpt}", flush=True) + print( + f"[main] pp={args.pp} tp={args.tp} ep={args.ep} dp={args.dp} " + f"enable_dp_attention={args.enable_dp_attention} dtype={args.dtype} " + f"n_engine={n_engine} n_trainer={n_trainer} " + f"moe_runner_backend={args.moe_runner_backend}", + flush=True, + ) + print( + f"[main] branch={'test offload/onload' if args.test_offload_onload else 'baseline'}", + flush=True, + ) + print(f"[main] CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}", flush=True) + print(f"[main] {len(specs)} weight tensors to round-trip", flush=True) + + engine = None + trainer_procs: list[mp.Process] = [] + ray_namespace = f"repro_pp_colocate_{os.getpid()}" + engine_name = f"sglang_engine_{os.getpid()}" + try: + ray_ctx = ray.init( + num_gpus=n_avail, + namespace=ray_namespace, + include_dashboard=False, + ignore_reinit_error=True, + log_to_driver=True, + ) + ray_address = ( + ray_ctx.address_info.get("gcs_address") + or ray_ctx.address_info.get("address") + or "auto" + ) + from nemo_rl.models.generation.sglang.utils.ray_utils import ( + NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, + ) + + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | { + "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true", + "NCCL_CUMEM_ENABLE": "0", + } + if os.environ.get("CUDA_VISIBLE_DEVICES"): + env_vars["CUDA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] + + engine = _SGLangEngineActor.options( + num_gpus=0.2, + name=engine_name, + lifetime="detached", + runtime_env={"env_vars": env_vars}, + ).remote( + ckpt=str(ckpt), + pp=args.pp, + tp=args.tp, + ep=args.ep, + dp=args.dp, + enable_dp_attention=args.enable_dp_attention, + dtype=args.dtype, + mem_fraction_static=args.mem_fraction_static, + moe_runner_backend=args.moe_runner_backend, + ) + + prompt = "The capital of France is" + sampling = {"max_new_tokens": 16, "temperature": 0.0} + + out_pre = ray.get(engine.generate.remote(prompt, sampling)) + print(f"[main] pre-update text={_text(out_pre)!r}", flush=True) + + print("[main] snapshot...", flush=True) + ray.get(engine.check_weights.remote("snapshot")) + + print("[main] reset_tensors (randomize sglang weights)...", flush=True) + ray.get(engine.check_weights.remote("reset_tensors")) + out_random = ray.get(engine.generate.remote(prompt, sampling)) + print(f"[main] post-reset text={_text(out_random)!r}", flush=True) + if _text(out_random) == _text(out_pre): + print( + "[main] WARN: post-reset output matches pre; proceeding with weight compare", + flush=True, + ) + + if args.test_offload_onload: + print("[main] offload weights...", flush=True) + ray.get(engine.offload_weights.remote()) + print("[main] offload kv+cuda_graph...", flush=True) + ray.get(engine.offload_kv.remote()) + + if args.test_offload_onload: + print("[main] onload weights...", flush=True) + ray.get(engine.onload_weights.remote()) + + print("[main] starting mock trainer ranks...", flush=True) + t0 = time.time() + ctx = mp.get_context("spawn") + for rank in range(n_trainer): + proc = ctx.Process( + target=_mock_trainer_rank_main, + kwargs={ + "rank": rank, + "world_size": n_trainer, + "master_addr": "127.0.0.1", + "master_port": args.master_port, + "ckpt_str": str(ckpt), + "specs": specs, + "weight_map": weight_map, + "target_dtype_name": args.dtype, + "buffer_size_bytes": args.buffer_size_bytes, + "ray_address": ray_address, + "ray_namespace": ray_namespace, + "engine_name": engine_name, + "engine_gpu_counts": [n_engine], + "engine_gpu_offsets": [0], + }, + ) + proc.start() + trainer_procs.append(proc) + + failed = False + for proc in trainer_procs: + proc.join() + if proc.exitcode != 0: + failed = True + print( + f"[main] trainer process pid={proc.pid} failed exitcode={proc.exitcode}", + flush=True, + ) + if failed: + sys.exit(1) + print( + f"[main] all mock trainer ranks sent buckets in {time.time() - t0:.1f}s", + flush=True, + ) + + print("[main] post_process_weights...", flush=True) + ray.get( + engine.post_process_weights.remote( + restore_weights_before_load=False, + post_process_quantization=True, + ) + ) + + print("[main] compare against snapshot...", flush=True) + ray.get(engine.check_weights.remote("compare")) + print("[main] compare passed.", flush=True) + + if args.test_offload_onload: + print("[main] onload kv+cuda_graph...", flush=True) + ray.get(engine.onload_kv.remote()) + + out_post = ray.get(engine.generate.remote(prompt, sampling)) + print(f"[main] post-update text={_text(out_post)!r}", flush=True) + + if _text(out_pre) != _text(out_post): + print("[main] FAIL: pre/post generation outputs differ", flush=True) + print(f" pre : {_text(out_pre)!r}", flush=True) + print(f" post: {_text(out_post)!r}", flush=True) + sys.exit(1) + + print("[main] PASS: snapshot+reset+refit+compare colocate roundtrip", flush=True) + finally: + for proc in trainer_procs: + if proc.is_alive(): + proc.terminate() + proc.join(timeout=10) + if engine is not None: + try: + ray.get(engine.shutdown.remote(), timeout=30) + except Exception as exc: # noqa: BLE001 + print(f"[main] WARN: actor shutdown failed: {exc}", flush=True) + try: + ray.kill(engine, no_restart=True) + except Exception: + pass + if ray.is_initialized(): + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/unit/models/generation/sglang/repro_pp_weight_update.py b/tests/unit/models/generation/sglang/repro_pp_weight_update.py new file mode 100644 index 0000000000..5af36588bd --- /dev/null +++ b/tests/unit/models/generation/sglang/repro_pp_weight_update.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python3 +# 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 +"""Standalone repro for the sglang weight-update failure with ``pp_size > 1``. + +Hypothesis under test +--------------------- +``model_runner.init_weights_update_group`` (sglang) computes +``rank = rank_offset + self.tp_rank``, ignoring ``self.pp_rank`` / +``self.dp_rank``. When sglang runs with ``pp_size > 1``, multiple engine +workers pick the same NCCL rank and the rendezvous either deadlocks or +the subsequent broadcast lands on the wrong device. + +Test plan +--------- +Stage 1 (sanity) : ``--pp 1 --tp 1 --dp 1`` — 1 trainer + 1 engine rank. + Pre/post ``generate`` must produce identical text. +Stage 2 (bug) : ``--pp 2 --tp 1 --dp 1`` — 1 trainer + 2 engine ranks. + Expected to hang at ``init_weights_update_group`` or + fail mid-broadcast with CUDA invalid-argument. + +Usage (inside the sglang-nemorl-e2e-zhw container):: + + python tests/unit/models/generation/sglang/repro_pp_weight_update.py --pp 1 + python tests/unit/models/generation/sglang/repro_pp_weight_update.py --pp 2 + +We deliberately do NOT use Megatron / nemo-rl plumbing — only the bare +sglang Engine API and ``torch.distributed`` from a hand-rolled trainer +process. Weights come from the cached sliced Qwen3-30B-A3B checkpoint +(produced by ``_qwen3_slicer.ensure_sliced_model``). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +import traceback +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from safetensors import safe_open + +_DTYPE_FROM_STR = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + "float64": torch.float64, + "int8": torch.int8, + "uint8": torch.uint8, + "int32": torch.int32, + "int64": torch.int64, + "bool": torch.bool, +} + +# safetensors uses its own dtype tag strings; map to torch dtype names. +_SAFETENSORS_TO_TORCH_NAME = { + "F64": "float64", + "F32": "float32", + "F16": "float16", + "BF16": "bfloat16", + "I64": "int64", + "I32": "int32", + "I16": "int16", + "I8": "int8", + "U8": "uint8", + "BOOL": "bool", +} + + +def _ckpt_path() -> Path: + hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface") + p = Path(hf_home) / "hub" / "qwen3-30b-a3b-sliced-4" + if not p.is_dir(): + sys.exit( + f"sliced ckpt not found at {p}; produce it first via " + "tests/unit/models/generation/sglang/_qwen3_slicer.py" + ) + return p + + +def _collect_specs(ckpt: Path) -> tuple[list[tuple[str, str, list[int]]], dict[str, str]]: + """Return (ordered list of (name, dtype_str, shape), name -> shard filename). + + Order is deterministic (sorted by name) so the trainer side and the + engine side can iterate in lockstep without any explicit handshake. + """ + index_path = ckpt / "model.safetensors.index.json" + if index_path.is_file(): + with open(index_path) as f: + weight_map: dict[str, str] = json.load(f)["weight_map"] + else: + single = next(ckpt.glob("*.safetensors")).name + with safe_open(ckpt / single, framework="pt") as r: + weight_map = {k: single for k in r.keys()} + + # Read header (shape/dtype) for each tensor without materialising the + # storage. ``safe_open`` exposes ``get_slice`` which is metadata-only. + per_shard: dict[str, list[str]] = {} + for name, shard in weight_map.items(): + per_shard.setdefault(shard, []).append(name) + + specs: list[tuple[str, str, list[int]]] = [] + for shard in sorted(per_shard.keys()): + with safe_open(ckpt / shard, framework="pt") as r: + for name in sorted(per_shard[shard]): + sl = r.get_slice(name) + st_dtype = sl.get_dtype() # safetensors tag, e.g. "BF16" + if st_dtype not in _SAFETENSORS_TO_TORCH_NAME: + raise RuntimeError(f"unmapped safetensors dtype {st_dtype!r} for {name}") + dtype_str = _SAFETENSORS_TO_TORCH_NAME[st_dtype] + shape = list(sl.get_shape()) + specs.append((name, dtype_str, shape)) + specs.sort(key=lambda x: x[0]) + return specs, weight_map + + +def _trainer_proc( + *, + trainer_gpu: int, + world_size: int, + master_addr: str, + master_port: int, + ckpt_str: str, + specs: list, + weight_map: dict, + group_name: str, + rank: int = 0, +): + """Run as a separate process; broadcasts each tensor to the engine.""" + try: + # Trainer is launched by mp.spawn with CUDA_VISIBLE_DEVICES already + # restricted by the parent — ``trainer_gpu`` is a *visible* index. + torch.cuda.set_device(trainer_gpu) + + from nemo_rl.models.policy.utils import init_process_group + + print( + f"[trainer rank={rank}] init_process_group " + f"world={world_size} master={master_addr}:{master_port} group={group_name}", + flush=True, + ) + pg = init_process_group( + backend="nccl", + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=rank, + group_name=group_name, + ) + print(f"[trainer rank={rank}] joined group; about to broadcast {len(specs)} tensors", flush=True) + + ckpt = Path(ckpt_str) + device = f"cuda:{trainer_gpu}" + for i, (name, dtype_str, shape) in enumerate(specs): + shard = ckpt / weight_map[name] + with safe_open(shard, framework="pt") as r: + t = r.get_tensor(name) + t = t.to(device=device, dtype=_DTYPE_FROM_STR[dtype_str]).contiguous() + dist.broadcast(t, src=0, group=pg) + del t + if i % 25 == 0 or i == len(specs) - 1: + print( + f"[trainer rank={rank}] broadcast {i+1}/{len(specs)} {name} shape={shape}", + flush=True, + ) + print(f"[trainer rank={rank}] all broadcasts complete", flush=True) + dist.destroy_process_group(pg) + except Exception: # noqa: BLE001 + # Crash visibly — parent watches exit code. + traceback.print_exc() + os._exit(1) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--dp", type=int, default=1) + parser.add_argument("--master-port", type=int, default=29555) + parser.add_argument( + "--max-tensors", + type=int, + default=0, + help="If > 0, only round-trip the first N tensors (faster smoke test).", + ) + parser.add_argument( + "--moe-runner-backend", + type=str, + default=None, + help="Pass through to sglang Engine (e.g. flashinfer_trtllm_routed, triton, auto).", + ) + args = parser.parse_args() + + n_engine = args.pp * args.tp * args.dp + n_trainer = 1 + world_size = n_engine + n_trainer + rank_offset = n_trainer + group_name = "weight_update_group" + + n_avail = torch.cuda.device_count() + if n_avail < n_engine + 1: + sys.exit(f"need at least {n_engine + 1} GPUs visible, have {n_avail}") + + ckpt = _ckpt_path() + print(f"[main] ckpt={ckpt}", flush=True) + print( + f"[main] pp={args.pp} tp={args.tp} dp={args.dp} " + f"n_engine={n_engine} n_trainer={n_trainer} world={world_size}", + flush=True, + ) + + specs, weight_map = _collect_specs(ckpt) + if args.max_tensors > 0: + specs = specs[: args.max_tensors] + print(f"[main] {len(specs)} weight tensors to round-trip", flush=True) + + # GPU layout: caller controls absolute device ids via CUDA_VISIBLE_DEVICES. + # Within that view, engine takes visible [0, n_engine) and trainer takes n_engine. + print(f"[main] CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}", flush=True) + + from sglang.srt.entrypoints.engine import Engine + + engine_kwargs = {} + if args.moe_runner_backend is not None: + engine_kwargs["moe_runner_backend"] = args.moe_runner_backend + + engine = Engine( + model_path=str(ckpt), + tp_size=args.tp, + pp_size=args.pp, + dp_size=args.dp, + mem_fraction_static=0.5, + log_level="info", + random_seed=42, + disable_cuda_graph=True, + **engine_kwargs, + ) + print("[main] engine up", flush=True) + + # Pre-update generation (greedy so we can compare deterministically). + sampling = {"max_new_tokens": 16, "temperature": 0.0} + prompt = "The capital of France is" + + def _text(x): + if isinstance(x, list): + x = x[0] + if isinstance(x, dict): + return x.get("text", repr(x)) + return repr(x) + + out_pre = engine.generate(prompt, sampling_params=sampling) + print(f"[main] pre-update text={_text(out_pre)!r}", flush=True) + + # Snapshot the freshly-loaded weights, then randomize them so the + # subsequent refit MUST overwrite them. Without this step the + # post-update equality check trivially passes regardless of whether + # the broadcast loop actually delivered anything. + from sglang.srt.managers.io_struct import ( + CheckWeightsReqInput, + PostProcessWeightsReqInput, + ) + + def _check_weights(action: str): + res = engine.loop.run_until_complete( + engine.tokenizer_manager.check_weights(CheckWeightsReqInput(action=action)) + ) + # tokenizer_manager.check_weights returns _Communicator.merge_results, + # which is a ``(all_success: bool, joined_message: str)`` tuple — NOT + # a CheckWeightsReqOutput. The HTTP layer would convert + # ``success=False`` to a 400; this in-process path does not, so we + # have to raise ourselves to avoid a failed compare masquerading as + # success. + success, message = res + if not success: + raise RuntimeError(f"check_weights({action!r}) failed: {message}") + return res + + def _post_process_weights(): + """Re-run sglang's ``process_weights_after_loading`` hook on every module + after the broadcast/load loop. Without this, flashinfer trtllm MoE + layers keep canonical-shape weights from the broadcast and never get + re-packed into the block layout that the kernel expects — so post- + update inference crashes / produces garbage. nemo-rl's production + refit dispatch path always calls this (see + ``megatron_policy_worker.py:1870/1958``); the bare sglang + ``update_weights_from_distributed`` API does NOT do it implicitly. + """ + res = engine.loop.run_until_complete( + engine.tokenizer_manager.post_process_weights( + PostProcessWeightsReqInput( + restore_weights_before_load=False, + post_process_quantization=True, + ) + ) + ) + success, message = res + if not success: + raise RuntimeError(f"post_process_weights failed: {message}") + return res + + print("[main] snapshot...", flush=True) + _check_weights("snapshot") + print("[main] reset_tensors (randomize sglang weights)...", flush=True) + _check_weights("reset_tensors") + out_random = engine.generate(prompt, sampling_params=sampling) + print(f"[main] post-reset text={_text(out_random)!r}", flush=True) + if _text(out_random) == _text(out_pre): + print( + "[main] WARN: post-reset output matches pre — reset_tensors may " + "have been a no-op for this prompt; proceeding anyway", + flush=True, + ) + + # Spawn the trainer FIRST so it is waiting on the rendezvous when the + # engine workers join. Trainer is on the GPU just past the engine block. + trainer_gpu = n_engine + ctx = mp.get_context("spawn") + p = ctx.Process( + target=_trainer_proc, + kwargs=dict( + trainer_gpu=trainer_gpu, + world_size=world_size, + master_addr="127.0.0.1", + master_port=args.master_port, + ckpt_str=str(ckpt), + specs=specs, + weight_map=weight_map, + group_name=group_name, + rank=0, + ), + ) + p.start() + + print( + f"[main] engine.init_weights_update_group(rank_offset={rank_offset}, world={world_size}) ...", + flush=True, + ) + t0 = time.time() + success, msg = engine.init_weights_update_group( + master_address="127.0.0.1", + master_port=args.master_port, + rank_offset=rank_offset, + world_size=world_size, + group_name=group_name, + backend="nccl", + ) + print( + f"[main] init_weights_update_group => success={success} msg={msg} ({time.time()-t0:.1f}s)", + flush=True, + ) + if not success: + if p.is_alive(): + p.terminate() + sys.exit(f"init_weights_update_group failed: {msg}") + + print("[main] starting weight update loop", flush=True) + t0 = time.time() + for i, (name, dtype_str, shape) in enumerate(specs): + engine.update_weights_from_distributed( + names=[name], + dtypes=[dtype_str], + shapes=[shape], + group_name=group_name, + flush_cache=False, + ) + if i % 25 == 0 or i == len(specs) - 1: + print( + f"[main] update {i+1}/{len(specs)} {name} (elapsed {time.time()-t0:.1f}s)", + flush=True, + ) + print(f"[main] weight updates complete in {time.time()-t0:.1f}s", flush=True) + + # Tear down the engine side of the weight-update NCCL group so the + # trainer's dist.destroy_process_group(pg) can pair with it. Without + # this the trainer's destroy hangs (engine never matches it) and the + # script aborts before reaching the snapshot compare below. + print("[main] destroy_weights_update_group...", flush=True) + engine.destroy_weights_update_group(group_name=group_name) + + # Re-run ``process_weights_after_loading`` on every module so the + # flashinfer trtllm MoE layers re-pack the just-broadcast canonical + # weights into the 4-D block layout the kernel expects. nemo-rl's + # production refit dispatch does this via + # ``policy_generation.post_process_weights()``; bare sglang's + # ``update_weights_from_distributed`` does NOT do it implicitly. + print("[main] post_process_weights...", flush=True) + _post_process_weights() + print("[main] post_process_weights done.", flush=True) + + p.join(timeout=120) + if p.is_alive(): + print("[main] trainer still alive after timeout; terminating", flush=True) + p.terminate() + p.join() + if p.exitcode not in (0, None): + sys.exit(f"trainer exited with code {p.exitcode}") + + # Compare current weights against the snapshot — strongest correctness + # signal: raises if any tensor differs from the freshly-loaded original. + print("[main] compare against snapshot...", flush=True) + _check_weights("compare") + print("[main] compare passed.", flush=True) + + out_post = engine.generate(prompt, sampling_params=sampling) + print(f"[main] post-update text={_text(out_post)!r}", flush=True) + + pre_text = _text(out_pre) + post_text = _text(out_post) + if pre_text == post_text: + print("[main] PASS: snapshot+reset+broadcast+compare roundtrip", flush=True) + else: + print("[main] FAIL: pre/post generation outputs differ", flush=True) + print(f" pre : {pre_text!r}") + print(f" post: {post_text!r}") + sys.exit(1) + + +if __name__ == "__main__": + mp.set_start_method("spawn", force=True) + main() diff --git a/tests/unit/models/generation/sglang/smoke_routed_engine.py b/tests/unit/models/generation/sglang/smoke_routed_engine.py new file mode 100644 index 0000000000..730468b4b2 --- /dev/null +++ b/tests/unit/models/generation/sglang/smoke_routed_engine.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Smoke test: bring up sglang Engine on the sliced 4-layer Qwen3-30B-A3B +checkpoint with a configurable ``moe_runner_backend`` and run one greedy +generate. No weight update at all — just verifies the engine can stand up +under this backend. +""" + +import argparse +import os +import sys +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--moe-runner-backend", + type=str, + default="flashinfer_trtllm_routed", + help="sglang moe_runner_backend (e.g. flashinfer_trtllm, flashinfer_trtllm_routed, triton, auto).", + ) + args = parser.parse_args() + ckpt = ( + Path(os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")) + / "hub" + / "qwen3-30b-a3b-sliced-4" + ) + if not ckpt.is_dir(): + sys.exit(f"sliced ckpt not found at {ckpt}") + print(f"[smoke] ckpt={ckpt}", flush=True) + print( + f"[smoke] CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}", + flush=True, + ) + print(f"[smoke] moe_runner_backend={args.moe_runner_backend}", flush=True) + + from sglang.srt.entrypoints.engine import Engine + + engine = Engine( + model_path=str(ckpt), + tp_size=1, + pp_size=1, + dp_size=1, + mem_fraction_static=0.5, + log_level="info", + random_seed=42, + disable_cuda_graph=True, + dtype="bfloat16", + moe_runner_backend=args.moe_runner_backend, + ) + print("[smoke] engine up", flush=True) + + # Confirm the actually-resolved dtype + MoE backend the engine picked. + try: + sa = engine.tokenizer_manager.server_args + print( + f"[smoke] resolved dtype={getattr(sa, 'dtype', '')!r} " + f"moe_runner_backend={getattr(sa, 'moe_runner_backend', '')!r}", + flush=True, + ) + except Exception as e: + print(f"[smoke] could not introspect server_args: {e!r}", flush=True) + + prompt = "The capital of France is" + out = engine.generate( + prompt, sampling_params={"max_new_tokens": 16, "temperature": 0.0} + ) + print(f"[smoke] generate output: {out!r}", flush=True) + print("[smoke] PASS — engine started and produced output", flush=True) + + engine.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/unit/models/generation/sglang/test_megatron_sglang_generation.py b/tests/unit/models/generation/sglang/test_megatron_sglang_generation.py new file mode 100644 index 0000000000..c8c5311b40 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_megatron_sglang_generation.py @@ -0,0 +1,529 @@ +# 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. + +"""Generation tests for the Megatron + SGLang stack. + +Mirrors ``tests/unit/models/generation/sglang/test_sglang_generation.py`` but: + + • Uses the sliced ``nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`` checkpoint + (only the first ``MEMEM*`` block; see ``_nemotron_slicer.py``). + • Runs every ``SGLangGeneration`` API check across the same three SGLang + shapes used in the weight-update test: + - tp=4 ep=4 dp=4 --enable-dp-attention + - tp=4 ep=2 dp=4 --enable-dp-attention + - tp=2 ep=2 pp=2 + • Brings up a real Megatron ``Policy`` alongside SGLang for each of the + three Megatron parallelism shapes (``ep2 pp2`` / ``tp2 pp2`` / + ``tp2 ep2 pp2``), in both ``colocate`` and ``disaggregate`` modes — that + way each generation test runs against a non-trivial trainer next to it + and the same ``PolicyConfig`` schema as the weight-update test. + +Tests follow the structure of ``test_sglang_generation.py``: + + • ``generate()`` — output keys, shape, determinism, truncation, logprobs, + max_new_tokens cap, batched prompts, empty input, stop strings. + • ``generate_async()`` — single-sample yield, agreement with sync. + • ``generate_one_sample()`` — return tuple shape and types. + • Memory cycle via worker API and via direct HTTP 200 + top-level API. + • ``invalidate_kv_cache()`` aggregator + after-generate flush_cache pacing. +""" + +from __future__ import annotations + +import asyncio +import gc + +import pytest +import ray +import torch +from _megatron_helpers import ( + EOS_TOKEN_ID, + MEGATRON_CFGS, + PAD_TOKEN_ID, + SGLANG_CFGS, + TestTriple, + make_policy_config, + make_sglang_cfg, + megatron_world_size, + required_world_size, +) +from _nemotron_slicer import ensure_sliced_model +from helpers import make_generation_sampling_params, post_and_assert_200 + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.sglang.sglang_generation import ( + SGLangGeneration, + generate_one_sample, +) + +pytestmark = pytest.mark.sglang + +TOTAL_AVAILABLE_GPUS = 8 + + +# --------------------------------------------------------------------------- +# Cartesian product → pytest params (same matrix as the weight-update test) +# --------------------------------------------------------------------------- +def _build_params() -> list[pytest.param]: + out: list[pytest.param] = [] + for colocated in (True, False): + for m in MEGATRON_CFGS: + for s in SGLANG_CFGS: + triple = TestTriple(megatron=m, sglang=s, colocated=colocated) + marks: list = [] + need = required_world_size( + megatron=m, sglang=s, colocated=colocated + ) + if need > TOTAL_AVAILABLE_GPUS: + marks.append( + pytest.mark.skip( + reason=( + f"{triple.id} needs {need} GPUs but only " + f"{TOTAL_AVAILABLE_GPUS} are available" + ) + ) + ) + out.append(pytest.param(triple, id=triple.id, marks=marks)) + return out + + +# --------------------------------------------------------------------------- +# Sliced-model + tokenizer fixtures +# --------------------------------------------------------------------------- +@pytest.fixture(scope="session") +def sliced_model_path() -> str: + return str(ensure_sliced_model()) + + +@pytest.fixture(scope="session") +def tokenizer(sliced_model_path): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(sliced_model_path, trust_remote_code=True) + + +# --------------------------------------------------------------------------- +# Cluster + Policy + SGLang fixture (parametrised). Same shape as the +# weight-update test's ``env`` fixture, but exposed as ``sglang_gen`` so the +# (already long) generation test bodies stay readable. +# --------------------------------------------------------------------------- +@pytest.fixture(params=_build_params()) +def sglang_gen(request, ray_cluster, sliced_model_path): + triple: TestTriple = request.param + m, s, colocated = triple.megatron, triple.sglang, triple.colocated + + train_world = megatron_world_size(m) + sglang_world = s.num_gpus_per_engine + + if colocated: + bundle_count = max(train_world, sglang_world) + train_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[bundle_count], + use_gpus=True, + max_colocated_worker_groups=2, + num_gpus_per_node=bundle_count, + name=f"gen-colo-{triple.id}", + ) + sglang_cluster = train_cluster + else: + train_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[train_world], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=train_world, + name=f"gen-disag-train-{triple.id}", + ) + sglang_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[sglang_world], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=sglang_world, + name=f"gen-disag-infer-{triple.id}", + ) + + sglang_cfg = make_sglang_cfg( + model_path=sliced_model_path, + sglang=s, + colocated=colocated, + ) + gen = SGLangGeneration(sglang_cluster, sglang_cfg) + gen.finish_generation() + + # Build the Megatron policy. Even though the generation tests don't refit, + # bringing the trainer up exercises the same setup path as the weight- + # update test and ensures generation correctness with a colocated/ + # disaggregated trainer present. + from transformers import AutoTokenizer + + from nemo_rl.models.policy.lm_policy import Policy + + tok = AutoTokenizer.from_pretrained(sliced_model_path, trust_remote_code=True) + policy_cfg = make_policy_config( + model_path=sliced_model_path, + megatron=m, + colocated=colocated, + ) + policy = Policy( + cluster=train_cluster, + config=policy_cfg, + tokenizer=tok, + init_optimizer=False, + init_reference_model=False, + ) + + state_dict_info = policy.prepare_refit_info() + gen.prepare_refit_info(state_dict_info) + + # Bring SGLang back up so generate() works. + gen.prepare_for_generation() + + yield gen + + try: + policy.shutdown() + except Exception: + pass + try: + gen.shutdown() + except Exception: + pass + try: + train_cluster.shutdown() + except Exception: + pass + if not colocated: + try: + sglang_cluster.shutdown() + except Exception: + pass + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_input(tokenizer, prompt, pad_length=None): + token_ids = tokenizer.encode(prompt) + input_length = len(token_ids) + if pad_length and pad_length > input_length: + token_ids = token_ids + [tokenizer.pad_token_id] * (pad_length - input_length) + return BatchedDataDict( + { + "input_ids": torch.tensor([token_ids], dtype=torch.long), + "input_lengths": torch.tensor([input_length], dtype=torch.long), + } + ) + + +def _make_batch(tokenizer, prompts, pad_length=None): + all_ids = [] + all_lengths = [] + max_len = 0 + for p in prompts: + ids = tokenizer.encode(p) + all_ids.append(ids) + all_lengths.append(len(ids)) + max_len = max(max_len, len(ids)) + if pad_length: + max_len = max(max_len, pad_length) + padded = [ids + [tokenizer.pad_token_id] * (max_len - len(ids)) for ids in all_ids] + return BatchedDataDict( + { + "input_ids": torch.tensor(padded, dtype=torch.long), + "input_lengths": torch.tensor(all_lengths, dtype=torch.long), + } + ) + + +# =================================================================== +# Tests: SGLangGeneration.generate() +# =================================================================== +def test_generate_returns_batched_data_dict(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Hello") + result = sglang_gen.generate(data, greedy=True) + for key in ( + "output_ids", + "logprobs", + "generation_lengths", + "unpadded_sequence_lengths", + "truncated", + ): + assert key in result, f"Missing key: {key}" + + +def test_generate_output_ids_shape(sglang_gen, tokenizer): + data = _make_input(tokenizer, "The capital of France is") + result = sglang_gen.generate(data, greedy=True) + assert result["output_ids"].dim() == 2 + assert result["output_ids"].shape[0] == 1 + gen_len = result["generation_lengths"][0].item() + input_len = data["input_lengths"][0].item() + assert result["unpadded_sequence_lengths"][0].item() == input_len + gen_len + + +def test_generate_greedy_determinism(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Once upon a time") + r1 = sglang_gen.generate(data, greedy=True) + r2 = sglang_gen.generate(data, greedy=True) + assert torch.equal(r1["output_ids"], r2["output_ids"]), ( + "Greedy generation is not deterministic" + ) + + +def test_generate_truncation_flag(sglang_gen, tokenizer): + orig = sglang_gen.sglang_cfg["max_new_tokens"] + sglang_gen.sglang_cfg["max_new_tokens"] = 1 + try: + data = _make_input(tokenizer, "Tell me a very long story about dragons and") + result = sglang_gen.generate(data, greedy=True) + gen_len = result["generation_lengths"][0].item() + assert gen_len == 1, f"Expected 1 token, got {gen_len}" + assert result["truncated"][0].item() is True, "Expected truncated=True" + finally: + sglang_gen.sglang_cfg["max_new_tokens"] = orig + + +def test_generate_logprobs_valid(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Hello world") + result = sglang_gen.generate(data, greedy=True) + gen_len = result["generation_lengths"][0].item() + input_len = data["input_lengths"][0].item() + lps = result["logprobs"][0, input_len : input_len + gen_len] + assert torch.isfinite(lps).all(), "Logprobs contain NaN or Inf" + assert (lps <= 0.0).all(), "Logprobs should be non-positive" + + +def test_generate_respects_max_new_tokens(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Count from 1 to 100:") + result = sglang_gen.generate(data, greedy=True) + max_new = sglang_gen.sglang_cfg["max_new_tokens"] + gen_len = result["generation_lengths"][0].item() + assert gen_len <= max_new, f"gen_len={gen_len} > max_new_tokens={max_new}" + + +def test_generate_batch_multiple_samples(sglang_gen, tokenizer): + prompts = [ + "Hello, my name is", + "The capital of France is", + "What is 2 plus 2?", + ] + data = _make_batch(tokenizer, prompts) + result = sglang_gen.generate(data, greedy=True) + assert result["output_ids"].shape[0] == 3 + assert result["generation_lengths"].shape[0] == 3 + for i in range(3): + gen_len = result["generation_lengths"][i].item() + assert gen_len > 0, f"Sample {i} generated 0 tokens" + + +def test_generate_empty_input(sglang_gen): + data = BatchedDataDict( + { + "input_ids": torch.zeros((0, 0), dtype=torch.long), + "input_lengths": torch.zeros(0, dtype=torch.long), + } + ) + result = sglang_gen.generate(data, greedy=True) + assert result["output_ids"].shape[0] == 0 + + +def test_generate_with_stop_strings(sglang_gen, tokenizer): + orig_stop = sglang_gen.sglang_cfg.get("stop_strings") + sglang_gen.sglang_cfg["stop_strings"] = ["\n"] + try: + data = _make_input(tokenizer, "List:\n1. Apple\n2.") + result = sglang_gen.generate(data, greedy=True) + gen_len = result["generation_lengths"][0].item() + max_new = sglang_gen.sglang_cfg["max_new_tokens"] + # Soft check — model could emit \n on the very first token. + assert gen_len <= max_new + finally: + sglang_gen.sglang_cfg["stop_strings"] = orig_stop + + +# =================================================================== +# Tests: SGLangGeneration.generate_async() +# =================================================================== +def test_generate_async_yields_single_sample(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Hello") + + async def _run(): + results = [] + async for idx, batch in sglang_gen.generate_async(data, greedy=True): + results.append((idx, batch)) + return results + + results = asyncio.run(_run()) + assert len(results) == 1 + idx, batch = results[0] + assert idx == 0 + assert "output_ids" in batch + assert batch["generation_lengths"][0].item() > 0 + + +def test_generate_async_output_matches_generate(sglang_gen, tokenizer): + data = _make_input(tokenizer, "The answer is") + sync_result = sglang_gen.generate(data, greedy=True) + + async def _run(): + async for _, batch in sglang_gen.generate_async(data, greedy=True): + return batch + return None + + async_result = asyncio.run(_run()) + assert async_result is not None + + sync_len = sync_result["generation_lengths"][0].item() + async_len = async_result["generation_lengths"][0].item() + assert sync_len == async_len, f"sync={sync_len} vs async={async_len}" + + input_len = data["input_lengths"][0].item() + sync_tokens = sync_result["output_ids"][0, input_len : input_len + sync_len] + async_tokens = async_result["output_ids"][0, input_len : input_len + async_len] + assert torch.equal(sync_tokens, async_tokens), ( + "generate() and generate_async() produced different tokens" + ) + + +# =================================================================== +# Tests: generate_one_sample() — the underlying async function +# =================================================================== +def test_generate_one_sample_returns_correct_tuple(sglang_gen, tokenizer): + sp = make_generation_sampling_params(max_new_tokens=5, temperature=0.0) + input_ids = tokenizer.encode("The capital of France is") + + result = asyncio.run( + generate_one_sample( + sglang_gen.router_ip, sglang_gen.router_port, sp, input_ids, index=42 + ) + ) + assert len(result) == 4 + idx, tokens, logprobs, truncated = result + assert idx == 42 + assert isinstance(tokens, list) and len(tokens) > 0 + assert isinstance(logprobs, list) and len(logprobs) == len(tokens) + assert isinstance(truncated, bool) + assert all(isinstance(t, int) for t in tokens) + assert all(isinstance(lp, float) for lp in logprobs) + + +# =================================================================== +# Tests: memory cycle (engine API → HTTP 200 → top-level API) +# =================================================================== +def test_generate_after_memory_cycle(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Two plus two equals") + r_before = sglang_gen.generate(data, greedy=True) + + for engine in sglang_gen.engines: + ray.get(engine.release_memory_weights.remote()) + ray.get(engine.release_memory_kv_cache_and_cuda_graph.remote()) + ray.get(engine.resume_memory_weights.remote()) + ray.get(engine.resume_memory_kv_cache_and_cuda_graph.remote()) + + r_after = sglang_gen.generate(data, greedy=True) + assert torch.equal(r_before["output_ids"], r_after["output_ids"]), ( + "Generation output changed after memory cycle" + ) + + +def test_generate_after_memory_cycle_via_http_200(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Two plus two equals") + r_before = sglang_gen.generate(data, greedy=True) + + for engine in sglang_gen.engines: + base_url = ray.get(engine.get_base_url.remote()) + assert base_url is not None + + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + base_url, "release_memory_occupation", {"tags": ["weights"]} + ) + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + base_url, + "release_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + post_and_assert_200( + base_url, "resume_memory_occupation", {"tags": ["weights"]} + ) + post_and_assert_200( + base_url, + "resume_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + + r_after = sglang_gen.generate(data, greedy=True) + assert torch.equal(r_before["output_ids"], r_after["output_ids"]), ( + "Generation output changed after HTTP-driven memory cycle" + ) + + +def test_generate_after_memory_cycle_top_level_api(sglang_gen, tokenizer): + data = _make_input(tokenizer, "Two plus two equals") + + r_before = sglang_gen.generate(data, greedy=True) + input_len = data["input_lengths"][0].item() + gen_len_before = r_before["generation_lengths"][0].item() + assert gen_len_before > 0, "generate() before memory cycle produced 0 tokens" + tokens_before = r_before["output_ids"][0, input_len : input_len + gen_len_before] + assert (tokens_before != PAD_TOKEN_ID).all(), "before: generated tokens contain pad" + + sglang_gen.offload_weights() + sglang_gen.offload_kv() + sglang_gen.onload_weights() + sglang_gen.onload_kv() + + r_after = sglang_gen.generate(data, greedy=True) + gen_len_after = r_after["generation_lengths"][0].item() + assert gen_len_after > 0, "generate() after memory cycle produced 0 tokens" + tokens_after = r_after["output_ids"][0, input_len : input_len + gen_len_after] + assert (tokens_after != PAD_TOKEN_ID).all(), "after: generated tokens contain pad" + + assert gen_len_before == gen_len_after, ( + f"Different generation_lengths before vs. after: " + f"before={gen_len_before}, after={gen_len_after}" + ) + assert torch.equal(r_before["output_ids"], r_after["output_ids"]), ( + "Generation output changed after top-level offload/onload cycle" + ) + + +# =================================================================== +# Tests: invalidate_kv_cache aggregator +# =================================================================== +def test_invalidate_kv_cache_aggregator(sglang_gen): + """``invalidate_kv_cache`` fans out to every engine and reduces with + ``all(results)``. Verifies True on a healthy cluster across the full + parametrize matrix (single-engine TP/PP/EP variants and dp-attention + variants alike).""" + assert sglang_gen.invalidate_kv_cache() is True + + +def test_invalidate_kv_cache_after_generate(sglang_gen, tokenizer): + """Most likely path to surface the flush_cache pacing bug — sglang's + ``/flush_cache`` may transiently return non-200 while draining the + just-completed generation's queue, so the worker's retry loop must + actually wait between attempts.""" + data = _make_input(tokenizer, "Two plus two equals") + sglang_gen.generate(data, greedy=True) + assert sglang_gen.invalidate_kv_cache() is True + + +__all__ = [ + "PAD_TOKEN_ID", + "EOS_TOKEN_ID", +] diff --git a/tests/unit/models/generation/sglang/test_megatron_sglang_weight_update.py b/tests/unit/models/generation/sglang/test_megatron_sglang_weight_update.py new file mode 100644 index 0000000000..344064adf5 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_megatron_sglang_weight_update.py @@ -0,0 +1,627 @@ +# 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. + +"""End-to-end weight update tests: real Megatron Policy → real SGLangGeneration. + +Mirrors ``tests/unit/models/generation/sglang/test_weight_update_real.py`` but +replaces the ``MockFSDPWorker`` (FSDP / DTensor) trainer with a *real* Megatron +``Policy`` using ``nemo_rl.models.policy.workers.megatron_policy_worker``. + +Cross-product the user asked for: + + • Mode: ``colocate`` (IPC) and ``disaggregate`` (NCCL broadcast) + • Megatron parallelism: + - ep2 pp2 (TP=1, PP=2, EP=2, DP=2 → 4 GPUs) + - tp2 pp2 (TP=2, PP=2, EP=1, DP=1 → 4 GPUs) + - tp2 ep2 pp2 (TP=2, PP=2, EP=2, DP=2 → 8 GPUs) + • SGLang shape: + - tp=4 ep=4 dp=4 --enable-dp-attention (4 engine GPUs) + - tp=4 ep=2 dp=4 --enable-dp-attention (4 engine GPUs) + - tp=2 ep=2 pp=2 (4 engine GPUs) + +Model: ``Qwen/Qwen3-30B-A3B-Instruct-2507``. The full HF checkpoint is +resolved against ``HF_HOME``. This is an MoE model (128 experts, 8 +active per token, 48 layers) so the EP > 1 ``MEGATRON_CFGS`` variants +(``mcore_ep2_pp2`` / ``mcore_tp2_ep2_pp2``) are valid here. Megatron-Bridge +auto-detects the architecture via the HF ``architectures`` field +(``Qwen3MoeForCausalLM``), so no custom converter wiring is required. + +Tests: + + * ``test_weight_update_roundtrip`` — snapshot → reset → offload → onload + weights → refit (Megatron streams to SGLang) → compare → onload kv. + * ``test_weight_update_roundtrip_with_router_generation`` — same flow, + but bracketed by router-driven greedy ``generate()`` calls; every + HTTP call is asserted 200 by reaching into the per-worker endpoints. +""" + +from __future__ import annotations + +import gc +import os + +import pytest +import ray +import torch +from _megatron_helpers import ( + EOS_TOKEN_ID, + MEGATRON_CFGS, + MEGATRON_DP1, + PAD_TOKEN_ID, + SGLANG_CFGS, + SGLANG_TP1, + TestTriple, + make_policy_config, + make_sglang_cfg, + megatron_world_size, + required_world_size, +) +from _qwen3_slicer import ensure_sliced_model +from helpers import post_and_assert_200 + +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration + +pytestmark = pytest.mark.sglang + +# --------------------------------------------------------------------------- +# Per-host GPU budget. The host's Ray cluster decides what is actually +# available; ``cluster.world_size()`` checks happen inside +# ``RayVirtualCluster``. We use this constant only to pre-skip variants whose +# parallelism layout is impossible on this host (avoid ResourceInsufficient). +# --------------------------------------------------------------------------- +TOTAL_AVAILABLE_GPUS = 8 + + +# --------------------------------------------------------------------------- +# Cartesian product → pytest params +# --------------------------------------------------------------------------- +def _build_params() -> list[pytest.param]: + out: list[pytest.param] = [] + # Cartesian product over the multi-GPU shapes. + pairs: list[tuple] = [(m, s) for m in MEGATRON_CFGS for s in SGLANG_CFGS] + # Plus one explicit single-GPU pairing: Megatron DP=1 (TP=PP=EP=1) → + # SGLang TP=1. Adding it as a standalone pair (rather than expanding + # MEGATRON_CFGS / SGLANG_CFGS) keeps the matrix from blowing up with + # combinations no one asked for. + pairs.append((MEGATRON_DP1, SGLANG_TP1)) + for colocated in (True, False): + for m, s in pairs: + triple = TestTriple(megatron=m, sglang=s, colocated=colocated) + marks: list = [] + need = required_world_size( + megatron=m, sglang=s, colocated=colocated + ) + if need > TOTAL_AVAILABLE_GPUS: + marks.append( + pytest.mark.skip( + reason=( + f"{triple.id} needs {need} GPUs but only " + f"{TOTAL_AVAILABLE_GPUS} are available" + ) + ) + ) + out.append(pytest.param(triple, id=triple.id, marks=marks)) + return out + + +# --------------------------------------------------------------------------- +# Model-path fixture +# --------------------------------------------------------------------------- +# Default to the upstream ``Qwen/Qwen3-4B`` HF id (resolved by transformers / +# AutoBridge against ``HF_HOME``). Override via ``QWEN3_TEST_MODEL_PATH`` for +# offline / mirror scenarios where the suite needs to point at a local +# checkpoint instead. +_QWEN3_TEST_MODEL_PATH_ENV = "QWEN3_TEST_MODEL_PATH" + + +@pytest.fixture(scope="session") +def sliced_model_path() -> str: + """Resolve the sliced Qwen3-30B-A3B model path once per session. + + Slices ``Qwen/Qwen3-30B-A3B-Instruct-2507`` down to the first + ``_qwen3_slicer.SLICED_NUM_LAYERS`` transformer blocks (default 4) + so that the single-GPU ``mcore_dp1`` / ``sgl_tp1`` parametrizations + fit in one H200's memory budget while still exercising the full + Qwen3MoE module set (attention, MoE router, experts). + + Honours ``QWEN3_TEST_MODEL_PATH`` as an override for offline or + mirrored snapshots; the slicer is then bypassed entirely. + """ + override = os.environ.get(_QWEN3_TEST_MODEL_PATH_ENV) + if override: + return override + return str(ensure_sliced_model()) + + +# --------------------------------------------------------------------------- +# Cluster + SGLang fixture (parametrised) +# +# Note: Megatron ``Policy`` is **not** built in the fixture — see the test +# bodies. The required sequence per the GRPO colocate refit is: +# +# snapshot → reset → offload sglang weights+kv+cuda_graph +# → create megatron Policy → onload sglang weights → refit +# → compare → onload sglang kv+cuda_graph +# +# Disaggregate is similar but skips the offload/onload dance because trainer +# and inference live on different GPUs: +# +# snapshot → reset → create megatron Policy → refit → compare +# +# Putting the Policy creation in the test body lets us observe each step in +# the order GRPO actually exercises in production. +# --------------------------------------------------------------------------- +@pytest.fixture(params=_build_params()) +def env(request, ray_cluster, sliced_model_path): + """Materialize ``(triple, sglang_gen, train_cluster, sliced_model_path)``. + + Colocate: a single ``RayVirtualCluster`` with ``max_colocated_worker_groups + = 2`` so trainer + SGLang share the same placement-group bundles. + + Disaggregate: two clusters — one for the Megatron trainer, one for the + SGLang engines. They live in separate placement groups but in the same + Ray runtime. + """ + triple: TestTriple = request.param + m, s, colocated = triple.megatron, triple.sglang, triple.colocated + + train_world = megatron_world_size(m) + sglang_world = s.num_gpus_per_engine + + # --- build clusters ---------------------------------------------------- + if colocated: + # Single shared cluster. Bundles equal max(train, sglang) GPUs because + # both worker groups need to fit on the same placement group. + bundle_count = max(train_world, sglang_world) + train_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[bundle_count], + use_gpus=True, + max_colocated_worker_groups=2, + num_gpus_per_node=bundle_count, + name=f"colo-{triple.id}", + ) + sglang_cluster = train_cluster + else: + train_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[train_world], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=train_world, + name=f"disag-train-{triple.id}", + ) + sglang_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[sglang_world], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=sglang_world, + name=f"disag-infer-{triple.id}", + ) + + # --- build SGLangGeneration; weights stay live on GPU until the test + # explicitly offloads them. The fixture deliberately does not + # pre-offload — some sglang weight_checker paths assert that the + # backing storage is still resident at snapshot time. + sglang_cfg = make_sglang_cfg( + model_path=sliced_model_path, + sglang=s, + colocated=colocated, + ) + sglang_gen = SGLangGeneration(sglang_cluster, sglang_cfg) + + yield triple, sglang_gen, train_cluster, sliced_model_path + + # --- teardown ---------------------------------------------------------- + try: + sglang_gen.shutdown() + except Exception: + pass + try: + train_cluster.shutdown() + except Exception: + pass + if not colocated: + try: + sglang_cluster.shutdown() + except Exception: + pass + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _refit_buffer_bytes() -> int: + """Per-bucket buffer size for the streaming refit. 256 MiB is plenty for + the sliced 14-layer model and matches the order-of-magnitude used in the + DTensor mock test.""" + return 256 * 1024 * 1024 + + +def _refit_megatron_to_sglang(*, policy, policy_generation, colocated: bool) -> None: + """Drive one full Megatron → SGLang refit using the production helpers. + + Mirrors the dispatch logic in + ``nemo_rl.algorithms.grpo._refit_sglang_dispatch`` so we exercise exactly + the code path GRPO uses, parametrised on ``colocated``. + """ + from nemo_rl.models.policy.workers import megatron_policy_worker as _backend + + helper = ( + _backend.refit_sglang_colocated + if colocated + else _backend.refit_sglang_distributed + ) + helper( + policy=policy, + policy_generation=policy_generation, + buffer_size_bytes=_refit_buffer_bytes(), + ) + + +def _build_policy(*, train_cluster, megatron_shape, model_path, colocated: bool): + """Construct a real Megatron ``Policy`` for the test body. + + Kept out of the fixture because the colocate flow requires sglang's + weights to be offloaded *first* (so megatron can claim GPU memory + without OOM); that ordering is most readable inside the test body. + """ + from transformers import AutoTokenizer + + from nemo_rl.models.policy.lm_policy import Policy + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + policy_cfg = make_policy_config( + model_path=model_path, + megatron=megatron_shape, + colocated=colocated, + ) + return Policy( + cluster=train_cluster, + config=policy_cfg, + tokenizer=tokenizer, + init_optimizer=False, + init_reference_model=False, + ) + + +def _shutdown_policy(policy) -> None: + try: + policy.shutdown() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +def test_weight_update_roundtrip(env): + """Full Megatron→SGLang refit roundtrip with snapshot/reset/compare. + + Sequencing per the GRPO colocate refit contract: + + Colocate (sglang and megatron share GPUs): + 1. snapshot — capture sglang's freshly-loaded weights + 2. reset_tensors — overwrite weights with random data + 3. offload sglang weights + kv + cuda_graph to CPU + 4. create megatron Policy ← claims the now-free GPU mem + 5. onload sglang weights ← weights are random/invalid here + 6. refit (megatron → sglang) ← restores them + 7. compare against snapshot ← passes ⇔ refit landed correctly + 8. onload sglang kv + cuda_graph + + Disaggregate (sglang and megatron on disjoint GPUs): + 1. snapshot + 2. reset_tensors + 3. create megatron Policy (no offload needed; different GPUs) + 4. refit (megatron → sglang via NCCL broadcast) + 5. compare + + Uses the production refit helpers + (``megatron_policy_worker.refit_sglang_{colocated,distributed}``) — no + naive reimplementation. + """ + triple, sglang_gen, train_cluster, model_path = env + colocated = triple.colocated + + # 1. Snapshot sglang's freshly-loaded weights. + print("[STEP 1] Snapshotting original sglang weights...", flush=True) + sglang_gen.check_weights("snapshot") + print("[STEP 1] Snapshot complete.", flush=True) + + # 2. Randomize sglang's weights — refit MUST overwrite them. Without + # this step ``compare`` in step 7 trivially passes regardless of + # whether refit actually copied anything. + print("[STEP 2] Randomizing (reset_tensors) sglang weights...", flush=True) + sglang_gen.check_weights("reset_tensors") + print("[STEP 2] Reset complete.", flush=True) + + if colocated: + # 3. Offload sglang weights + kv + cuda_graph to CPU so megatron + # can claim GPU memory at construction time. + print("[STEP 3] Offloading sglang weights+kv+cuda_graph...", flush=True) + sglang_gen.offload_weights() + sglang_gen.offload_kv() + print("[STEP 3] Offload complete.", flush=True) + + # 4. Create the megatron trainer. + print( + f"[STEP 4] Creating Megatron Policy ({triple.megatron.id})...", + flush=True, + ) + policy = _build_policy( + train_cluster=train_cluster, + megatron_shape=triple.megatron, + model_path=model_path, + colocated=colocated, + ) + try: + # The refit drivers expect both sides to have exchanged the + # state-dict shape via ``prepare_refit_info`` once at startup. + state_dict_info = policy.prepare_refit_info() + sglang_gen.prepare_refit_info(state_dict_info) + print("[STEP 4] Megatron Policy ready.", flush=True) + + if colocated: + # 5. Onload sglang weights so refit (CUDA IPC) can target them. + print("[STEP 5] Onloading sglang weight buffers...", flush=True) + sglang_gen.onload_weights() + print("[STEP 5] Onload weights complete.", flush=True) + + # 6. Refit weights through the production helper. + print( + f"[STEP 6] Refitting Megatron → SGLang via " + f"{'IPC (colocated)' if colocated else 'NCCL broadcast (disaggregate)'}...", + flush=True, + ) + _refit_megatron_to_sglang( + policy=policy, policy_generation=sglang_gen, colocated=colocated + ) + print("[STEP 6] Refit complete.", flush=True) + + # 7. Compare current sglang weights against the step-1 snapshot. + # Passes ⇔ megatron streamed the correct values back. + print("[STEP 7] Comparing current weights against snapshot...", flush=True) + sglang_gen.check_weights("compare") + print("[STEP 7] Compare passed.", flush=True) + + if colocated: + # 8. Onload sglang kv + cuda_graph for next-time inference. + print("[STEP 8] Onloading sglang kv+cuda_graph...", flush=True) + sglang_gen.onload_kv() + print("[STEP 8] Roundtrip complete.", flush=True) + finally: + _shutdown_policy(policy) + + +def test_weight_update_roundtrip_with_router_generation(env): + """Full refit roundtrip with router-driven generation and per-worker 200 checks. + + Same colocate / disaggregate sequencing as ``test_weight_update_roundtrip`` + (snapshot → reset → [offload] → create trainer → [onload weights] → refit + → compare → [onload kv]), but additionally: + + 1. *Generation through the router.* Both the pre-snapshot and + post-onload_kv generations go through ``sglang_gen.generate(..., + greedy=True)`` which routes via the SGLang router. With + ``greedy=True`` the two token sequences must be identical token-by- + token across the roundtrip. + 2. *Per-worker HTTP 200 checks for the refit cycle.* Instead of calling + ``sglang_gen.check_weights / offload_* / onload_*`` (which use Ray + actor methods that consume the status code), this test iterates + ``sglang_gen.engines`` and drives the equivalent HTTP endpoints on + every worker directly via ``post_and_assert_200``. + """ + from transformers import AutoTokenizer + + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + triple, sglang_gen, train_cluster, model_path = env + colocated = triple.colocated + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + + # Sanity: router must be set so ``generate`` actually routes. + assert sglang_gen.router_ip is not None and sglang_gen.router_port is not None, ( + "router_ip/router_port not set on sglang_gen — generate() would not route" + ) + print( + f"[setup] {triple.id} router=http://{sglang_gen.router_ip}:{sglang_gen.router_port}", + flush=True, + ) + + engines = [e for e in sglang_gen.engines if e is not None] + assert len(engines) >= 1, "sglang_gen has no engines" + base_urls = ray.get([e.get_base_url.remote() for e in engines]) + assert all(u is not None for u in base_urls), f"missing base_url in {base_urls}" + print(f"[setup] {len(engines)} worker(s); base_urls={base_urls}", flush=True) + + # --- Per-worker HTTP helpers ----------------------------------------------- + def _http_check_weights_all(action: str) -> None: + for url in base_urls: + post_and_assert_200(url, "weights_checker", {"action": action}) + + def _http_release_weights_all() -> None: + for engine, url in zip(engines, base_urls): + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + url, "release_memory_occupation", {"tags": ["weights"]} + ) + + def _http_release_kv_all() -> None: + for engine, url in zip(engines, base_urls): + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + url, + "release_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + + def _http_resume_weights_all() -> None: + for url in base_urls: + post_and_assert_200(url, "resume_memory_occupation", {"tags": ["weights"]}) + + def _http_resume_kv_all() -> None: + for url in base_urls: + post_and_assert_200( + url, + "resume_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + + # --- Router-based greedy generation --------------------------------------- + test_prompt = "The capital of France is" + input_ids = tokenizer.encode(test_prompt, add_special_tokens=True) + input_len = len(input_ids) + data = BatchedDataDict( + { + "input_ids": torch.tensor([input_ids], dtype=torch.long), + "input_lengths": torch.tensor([input_len], dtype=torch.long), + } + ) + + def _generate(tag: str) -> list[int]: + result = sglang_gen.generate(data, greedy=True) + for key in ( + "output_ids", + "generation_lengths", + "unpadded_sequence_lengths", + "logprobs", + ): + assert key in result, f"[{tag}] generate() output missing key: {key}" + gen_len = int(result["generation_lengths"][0].item()) + assert gen_len > 0, ( + f"[{tag}] generate() returned 0 tokens (no new tokens generated)" + ) + tokens = result["output_ids"][0, input_len : input_len + gen_len].tolist() + assert all(isinstance(t, int) for t in tokens), ( + f"[{tag}] output tokens should be ints, got {tokens!r}" + ) + text = tokenizer.decode(tokens, skip_special_tokens=True) + assert len(text) > 0, f"[{tag}] decoded generated text is empty" + print(f"[{tag}] gen_len={gen_len} tokens={tokens} text={text!r}", flush=True) + return tokens + + # --- Generation BEFORE snapshot (via router) ------------------------------- + print("[PRE] Router greedy generate() before snapshot...", flush=True) + tokens_before = _generate("PRE") + + # 1. snapshot via per-worker HTTP weights_checker + print("[STEP 1] Snapshotting weights (HTTP weights_checker×workers)...", flush=True) + _http_check_weights_all("snapshot") + print("[STEP 1] Snapshot complete.", flush=True) + + # 2. reset_tensors via per-worker HTTP weights_checker — refit MUST + # overwrite these random values for ``compare`` and the post-greedy + # token sequence to match. + print("[STEP 2] Randomizing weights (HTTP reset_tensors×workers)...", flush=True) + _http_check_weights_all("reset_tensors") + print("[STEP 2] Reset complete.", flush=True) + + if colocated: + # 3. offload sglang weights + kv + cuda_graph + print( + "[STEP 3] Offloading weights+kv+cuda_graph (HTTP release_memory_occupation×workers)...", + flush=True, + ) + _http_release_weights_all() + _http_release_kv_all() + print("[STEP 3] Offload complete.", flush=True) + + # 4. create the megatron trainer + print(f"[STEP 4] Creating Megatron Policy ({triple.megatron.id})...", flush=True) + policy = _build_policy( + train_cluster=train_cluster, + megatron_shape=triple.megatron, + model_path=model_path, + colocated=colocated, + ) + try: + state_dict_info = policy.prepare_refit_info() + sglang_gen.prepare_refit_info(state_dict_info) + print("[STEP 4] Megatron Policy ready.", flush=True) + + if colocated: + # 5. onload sglang weight buffers + print( + "[STEP 5] Onloading weights (HTTP resume_memory_occupation×workers)...", + flush=True, + ) + _http_resume_weights_all() + print("[STEP 5] Onload weights complete.", flush=True) + + # 6. refit through production helper + print( + f"[STEP 6] Refitting Megatron → SGLang via " + f"{'IPC (colocated)' if colocated else 'NCCL broadcast (disaggregate)'}...", + flush=True, + ) + _refit_megatron_to_sglang( + policy=policy, policy_generation=sglang_gen, colocated=colocated + ) + print("[STEP 6] Refit complete.", flush=True) + + # 7. compare vs snapshot + print( + "[STEP 7] Compare vs snapshot (HTTP weights_checker×workers)...", + flush=True, + ) + _http_check_weights_all("compare") + print("[STEP 7] Compare passed.", flush=True) + + if colocated: + # 8. onload sglang kv + cuda_graph + print( + "[STEP 8] Onloading kv+cuda_graph (HTTP resume_memory_occupation×workers)...", + flush=True, + ) + _http_resume_kv_all() + print("[STEP 8] Onload kv complete.", flush=True) + finally: + _shutdown_policy(policy) + + # --- Generation AFTER onload_kv (via router) ------------------------------- + print("[POST] Router greedy generate() after onload_kv...", flush=True) + tokens_after = _generate("POST") + + # --- Sanity & strict equality --------------------------------------------- + assert len(tokens_before) > 0, "generate() returned no tokens before roundtrip" + assert len(tokens_after) > 0, "generate() returned no tokens after roundtrip" + assert len(tokens_before) == len(tokens_after), ( + f"Different number of generated tokens before vs. after: " + f"before={len(tokens_before)}, after={len(tokens_after)}" + ) + assert tokens_before == tokens_after, ( + "Greedy tokens changed across the refit roundtrip:\n" + f" before (pre-snapshot): {tokens_before}\n" + f" after (post-onload_kv): {tokens_after}" + ) + print( + f"[ASSERT] Greedy tokens match before vs. after roundtrip " + f"(n={len(tokens_before)} tokens, both non-empty, both via router).", + flush=True, + ) + + +# Surface the constants so callers (or rerun tooling) can sanity-check that +# ``PAD_TOKEN_ID`` / ``EOS_TOKEN_ID`` line up with what the sliced tokenizer +# reports. Kept at module scope rather than inside the test body so it's +# visible in ``pytest --collect-only``. +__all__ = [ + "PAD_TOKEN_ID", + "EOS_TOKEN_ID", + "test_weight_update_roundtrip", + "test_weight_update_roundtrip_with_router_generation", +] diff --git a/tests/unit/models/generation/sglang/test_pure_nccl_broadcast_smoke.py b/tests/unit/models/generation/sglang/test_pure_nccl_broadcast_smoke.py new file mode 100644 index 0000000000..99e6693dc0 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_pure_nccl_broadcast_smoke.py @@ -0,0 +1,213 @@ +# 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. + +"""Minimal cross-process NCCL broadcast smoke test. + +Two Ray actors, one GPU each, no SGLang, no Megatron, no AutoBridge. +The only nemo-rl code in the loop is ``init_process_group`` from +``nemo_rl.models.policy.utils``. + +Pass / fail mapping: + +* PASS → the host supports cross-process NCCL between two GPUs + (P2P/IPC, SHM, or NET, whichever NCCL picks). Any failure of the + full disag refit is in our higher-level code, not the transport. +* FAIL with ``Cuda failure 'invalid argument'`` → cross-GPU + ``cudaIpcOpenMemHandle`` is blocked at the kernel layer (typically + IOMMU). The full disag refit is therefore not a code bug; the host + needs an admin-level fix or NCCL has to be forced onto a different + transport (``NCCL_P2P_DISABLE=1`` etc.). +""" + +from __future__ import annotations + +import os +import socket + +import pytest +import ray +import torch + +pytestmark = pytest.mark.sglang + + +# --------------------------------------------------------------------------- +# NcclWorker — minimal Ray actor that holds one GPU and a custom NCCL group +# --------------------------------------------------------------------------- +@ray.remote(num_gpus=1) +class NcclWorker: + """Single-GPU actor that participates in one cross-process NCCL group.""" + + def __init__(self) -> None: + self._pg = None + + def get_node_ip(self) -> str: + import ray as _ray + + return _ray.util.get_node_ip_address() + + def find_free_port(self) -> int: + with socket.socket() as s: + s.bind(("", 0)) + return int(s.getsockname()[1]) + + def device_info(self) -> dict: + import torch as _torch + + return { + "current_device": _torch.cuda.current_device(), + "device_count": _torch.cuda.device_count(), + "device_name": _torch.cuda.get_device_name(0), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), + } + + def setup_default_pg(self) -> None: + """Init a 1-rank gloo default PG to mimic real-world conditions where + the trainer or engine has already initialized the world before + building the cross-process NCCL group on top.""" + import torch.distributed as dist + + if dist.is_initialized(): + torch.cuda.set_device(0) + return + with socket.socket() as s: + s.bind(("", 0)) + port = s.getsockname()[1] + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + world_size=1, + rank=0, + ) + torch.cuda.set_device(0) + + def init_nccl_group( + self, + master_addr: str, + master_port: int, + rank: int, + world_size: int, + group_name: str, + ) -> None: + """Build the cross-process NCCL group via the in-tree helper.""" + from nemo_rl.models.policy.utils import init_process_group + + self._pg = init_process_group( + backend="nccl", + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=rank, + group_name=group_name, + ) + + def broadcast_send(self, n: int = 8, fill: float = 1.5) -> list: + """Rank 0: broadcast a known tiny bf16 tensor.""" + import torch.distributed as dist + + tensor = torch.full((n,), fill, dtype=torch.bfloat16, device="cuda:0") + dist.broadcast(tensor, src=0, group=self._pg) + torch.cuda.synchronize() + return tensor.cpu().tolist() + + def broadcast_recv(self, n: int = 8) -> list: + """Rank 1: receive into a fresh buffer and return it.""" + import torch.distributed as dist + + tensor = torch.empty((n,), dtype=torch.bfloat16, device="cuda:0") + dist.broadcast(tensor, src=0, group=self._pg) + torch.cuda.synchronize() + return tensor.cpu().tolist() + + def shutdown(self) -> None: + import torch.distributed as dist + + if self._pg is not None: + try: + dist.destroy_process_group(self._pg) + except Exception: + pass + self._pg = None + if dist.is_initialized(): + try: + dist.destroy_process_group() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# The test +# --------------------------------------------------------------------------- +def test_pure_nccl_broadcast(ray_cluster): + """Two GPUs, two Ray actors, one bf16 tensor. Pass ⇔ host transport works. + + No SGLang, no Megatron, no AutoBridge — the only nemo-rl symbol used + is ``init_process_group``. If this fails the same way the real disag + refit does (``Cuda failure 'invalid argument'``), the failure is at + NCCL transport, not in our higher-level code. + """ + if torch.cuda.device_count() < 2: + pytest.skip("test requires 2 GPUs visible to Ray") + + rank0 = NcclWorker.remote() + rank1 = NcclWorker.remote() + + try: + # 1. 1-rank default PG on each side (mimics Megatron / SGLang startup). + ray.get( + [rank0.setup_default_pg.remote(), rank1.setup_default_pg.remote()] + ) + info0, info1 = ray.get( + [rank0.device_info.remote(), rank1.device_info.remote()] + ) + print(f"[pure-nccl] rank0={info0}") + print(f"[pure-nccl] rank1={info1}") + + # 2. Pick a master address+port on rank 0. + master_addr, master_port = ray.get( + [rank0.get_node_ip.remote(), rank0.find_free_port.remote()] + ) + print(f"[pure-nccl] master={master_addr}:{master_port}") + + # 3. Bring up the cross-process NCCL group on both ranks in parallel. + ray.get( + [ + rank0.init_nccl_group.remote( + master_addr, master_port, 0, 2, "smoke-pure" + ), + rank1.init_nccl_group.remote( + master_addr, master_port, 1, 2, "smoke-pure" + ), + ] + ) + print("[pure-nccl] cross-process NCCL group up on both ranks") + + # 4. Broadcast (rank 0 → rank 1) and verify. + sent_fut = rank0.broadcast_send.remote() + recv_fut = rank1.broadcast_recv.remote() + sent, received = ray.get([sent_fut, recv_fut]) + print(f"[pure-nccl] sent={sent}") + print(f"[pure-nccl] received={received}") + assert sent == received, ( + f"broadcast bytes diverged:\n sent={sent}\n received={received}" + ) + finally: + try: + ray.get([rank0.shutdown.remote(), rank1.shutdown.remote()]) + except Exception: + pass + for actor in (rank0, rank1): + try: + ray.kill(actor) + except Exception: + pass diff --git a/tests/unit/models/generation/sglang/test_ray_second_nccl_group_broadcast.py b/tests/unit/models/generation/sglang/test_ray_second_nccl_group_broadcast.py new file mode 100644 index 0000000000..c9038cd90e --- /dev/null +++ b/tests/unit/models/generation/sglang/test_ray_second_nccl_group_broadcast.py @@ -0,0 +1,311 @@ +# 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. + +"""Two Ray actors, two NCCL groups, one bf16 broadcast. + +This is a narrow transport demo for the Megatron -> SGLang weight-update +failure mode: + +1. Start two Ray actors, each reserving one GPU. +2. Initialize a normal/default NCCL process group across the two actors. +3. Initialize a second side-by-side NCCL group with the nemo-rl + ``init_process_group`` helper, the same helper used for distributed + SGLang weight updates after Megatron has already initialized torch.distributed. +4. Broadcast a tiny bf16 tensor through the second group. + +If this fails only when NCCL P2P/NVSHM is enabled, the problem is already +reproducible without SGLang or Megatron. If it passes, the real failure needs +additional topology from SGLang/Megatron. +""" + +from __future__ import annotations + +import argparse +import os +import socket +from datetime import timedelta +from typing import Any + +import pytest +import ray +import torch + +pytestmark = pytest.mark.sglang + +os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + + +def _find_free_port() -> int: + with socket.socket() as sock: + sock.bind(("", 0)) + return int(sock.getsockname()[1]) + + +@ray.remote(num_gpus=1) +class TwoGroupNcclActor: + """Single-GPU actor holding one default PG and one custom weight-update PG.""" + + def __init__(self) -> None: + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + self._rank: int | None = None + self._weight_update_pg = None + + def node_ip(self) -> str: + import ray as _ray + + return _ray.util.get_node_ip_address() + + def device_info(self) -> dict[str, Any]: + return { + "rank": self._rank, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), + "device_count": torch.cuda.device_count(), + "current_device": torch.cuda.current_device() + if torch.cuda.is_available() + else None, + "device_name": torch.cuda.get_device_name(0) + if torch.cuda.is_available() + else None, + "nccl_p2p_disable": os.environ.get("NCCL_P2P_DISABLE", ""), + "nccl_shm_disable": os.environ.get("NCCL_SHM_DISABLE", ""), + "nccl_cumem_enable": os.environ.get("NCCL_CUMEM_ENABLE", ""), + "nccl_debug": os.environ.get("NCCL_DEBUG", ""), + } + + def init_default_nccl_pg( + self, + *, + master_addr: str, + master_port: int, + rank: int, + world_size: int, + ) -> dict[str, Any]: + import torch.distributed as dist + + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + torch.cuda.set_device(0) + self._rank = rank + if not dist.is_initialized(): + dist.init_process_group( + backend="nccl", + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=rank, + timeout=timedelta(seconds=120), + ) + dist.barrier(device_ids=[0]) + return self.device_info() + + def default_group_all_reduce(self) -> float: + import torch.distributed as dist + + tensor = torch.tensor( + [float((self._rank or 0) + 1)], dtype=torch.float32, device="cuda:0" + ) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + torch.cuda.synchronize() + return float(tensor.item()) + + def init_weight_update_nccl_pg( + self, + *, + master_addr: str, + master_port: int, + rank: int, + world_size: int, + group_name: str, + ) -> None: + from nemo_rl.models.policy.utils import init_process_group + + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + torch.cuda.set_device(0) + self._weight_update_pg = init_process_group( + backend="nccl", + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=rank, + group_name=group_name, + timeout=timedelta(seconds=120), + ) + + def broadcast_weight_update_tensor( + self, + *, + n: int, + src: int = 0, + async_op: bool = True, + ) -> dict[str, Any]: + import torch.distributed as dist + + if self._weight_update_pg is None: + raise RuntimeError("weight-update NCCL group is not initialized") + + rank = self._rank + if rank is None: + raise RuntimeError("default NCCL group is not initialized") + + if rank == src: + tensor = (torch.arange(n, device="cuda:0", dtype=torch.float32) % 97).to( + torch.bfloat16 + ) + else: + tensor = torch.empty((n,), dtype=torch.bfloat16, device="cuda:0") + + work = dist.broadcast( + tensor, + src=src, + group=self._weight_update_pg, + async_op=async_op, + ) + if async_op: + work.wait() + torch.cuda.synchronize() + + return { + "rank": rank, + "numel": tensor.numel(), + "dtype": str(tensor.dtype), + "device": str(tensor.device), + "first16": tensor[:16].float().cpu().tolist(), + "last16": tensor[-16:].float().cpu().tolist(), + "checksum": float(tensor.float().sum().item()), + } + + def shutdown(self) -> None: + import torch.distributed as dist + + if self._weight_update_pg is not None: + try: + dist.destroy_process_group(self._weight_update_pg) + except Exception: + pass + self._weight_update_pg = None + if dist.is_initialized(): + try: + dist.destroy_process_group() + except Exception: + pass + + +def run_two_actor_second_group_demo(*, tensor_numel: int = 2048) -> tuple[dict, dict]: + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + + if torch.cuda.device_count() < 2: + raise RuntimeError("demo requires at least 2 GPUs visible to Ray") + + rank0 = TwoGroupNcclActor.remote() + rank1 = TwoGroupNcclActor.remote() + try: + master_addr = ray.get(rank0.node_ip.remote()) + default_port = _find_free_port() + weight_update_port = _find_free_port() + + print( + f"[two-group-demo] default_pg=tcp://{master_addr}:{default_port} " + f"weight_update_pg=tcp://{master_addr}:{weight_update_port}", + flush=True, + ) + + infos = ray.get( + [ + rank0.init_default_nccl_pg.remote( + master_addr=master_addr, + master_port=default_port, + rank=0, + world_size=2, + ), + rank1.init_default_nccl_pg.remote( + master_addr=master_addr, + master_port=default_port, + rank=1, + world_size=2, + ), + ] + ) + print(f"[two-group-demo] default_pg_infos={infos}", flush=True) + + reduced = ray.get( + [ + rank0.default_group_all_reduce.remote(), + rank1.default_group_all_reduce.remote(), + ] + ) + print(f"[two-group-demo] default_pg_all_reduce={reduced}", flush=True) + assert reduced == [3.0, 3.0] + + group_name = "weight-update-two-ray-actors" + ray.get( + [ + rank0.init_weight_update_nccl_pg.remote( + master_addr=master_addr, + master_port=weight_update_port, + rank=0, + world_size=2, + group_name=group_name, + ), + rank1.init_weight_update_nccl_pg.remote( + master_addr=master_addr, + master_port=weight_update_port, + rank=1, + world_size=2, + group_name=group_name, + ), + ] + ) + print("[two-group-demo] weight-update NCCL group ready", flush=True) + + sent, received = ray.get( + [ + rank0.broadcast_weight_update_tensor.remote(n=tensor_numel), + rank1.broadcast_weight_update_tensor.remote(n=tensor_numel), + ] + ) + print(f"[two-group-demo] sent={sent}", flush=True) + print(f"[two-group-demo] received={received}", flush=True) + comparable_keys = ("numel", "dtype", "first16", "last16", "checksum") + assert {k: sent[k] for k in comparable_keys} == { + k: received[k] for k in comparable_keys + } + return sent, received + finally: + try: + ray.get([rank0.shutdown.remote(), rank1.shutdown.remote()]) + except Exception: + pass + for actor in (rank0, rank1): + try: + ray.kill(actor) + except Exception: + pass + + +def test_two_ray_actors_second_nccl_group_broadcast(ray_cluster): + run_two_actor_second_group_demo() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tensor-numel", type=int, default=2048) + args = parser.parse_args() + try: + run_two_actor_second_group_demo(tensor_numel=args.tensor_numel) + finally: + if ray.is_initialized(): + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/unit/models/generation/sglang/test_ray_wrapped_repro_weight_update.py b/tests/unit/models/generation/sglang/test_ray_wrapped_repro_weight_update.py new file mode 100644 index 0000000000..ca8b21a3ba --- /dev/null +++ b/tests/unit/models/generation/sglang/test_ray_wrapped_repro_weight_update.py @@ -0,0 +1,428 @@ +# 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. + +"""Ray-wrapped version of ``repro_pp_weight_update.py``. + +The standalone repro is real SGLang plus a mock Megatron sender: + +* receiver: a real bare ``sglang.srt.entrypoints.engine.Engine`` +* sender: a hand-written trainer that reads real HF safetensors and broadcasts + them with ``nemo_rl.models.policy.utils.init_process_group`` + +This test keeps those two roles, but wraps both sides in Ray actors. It answers +one narrow question: does adding Ray actor boundaries to the standalone repro +trigger the same NCCL failure as the real Megatron -> SGLangGeneration UT? +""" + +from __future__ import annotations + +import argparse +import os +import socket +from pathlib import Path +from typing import Any + +import pytest +import ray +import torch + +pytestmark = pytest.mark.sglang + +os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + +_DTYPE_FROM_STR = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, + "float64": torch.float64, + "int8": torch.int8, + "uint8": torch.uint8, + "int32": torch.int32, + "int64": torch.int64, + "bool": torch.bool, +} + + +def _find_free_port() -> int: + with socket.socket() as sock: + sock.bind(("", 0)) + return int(sock.getsockname()[1]) + + +def _select_specs( + specs: list[tuple[str, str, list[int]]], + *, + param_name: str | None, + max_tensors: int, +) -> list[tuple[str, str, list[int]]]: + if param_name: + matches = [spec for spec in specs if spec[0] == param_name] + if not matches: + available = ", ".join(name for name, _, _ in specs[:20]) + raise RuntimeError( + f"param {param_name!r} not found; first available params: {available}" + ) + return matches + if max_tensors > 0: + return specs[:max_tensors] + return specs + + +def _ray_get_best_effort(ref, *, label: str, timeout_s: float = 5.0): + try: + return ray.get(ref, timeout=timeout_s) + except Exception as exc: + print(f"[ray-wrapped-repro] cleanup {label} did not finish: {exc!r}", flush=True) + return None + + +def _describe_ref(label: str, ref, ready_refs: set) -> str: + if ref not in ready_refs: + try: + ray.cancel(ref, force=True) + except Exception: + pass + return f"{label}=pending" + try: + return f"{label}={ray.get(ref)!r}" + except Exception as exc: + return f"{label}=raised {exc!r}" + + +@ray.remote(num_gpus=1) +class RaySGLangEngine: + """Real bare SGLang Engine running inside one Ray actor.""" + + def __init__(self) -> None: + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + self._engine = None + + def start(self, *, model_path: str, tp: int, pp: int, dp: int) -> dict[str, Any]: + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + from sglang.srt.entrypoints.engine import Engine + + self._engine = Engine( + model_path=model_path, + tp_size=tp, + pp_size=pp, + dp_size=dp, + mem_fraction_static=0.5, + log_level="info", + random_seed=42, + disable_cuda_graph=True, + ) + return { + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), + "device_count": torch.cuda.device_count(), + "device_name": torch.cuda.get_device_name(0), + "nccl_cumem_enable": os.environ.get("NCCL_CUMEM_ENABLE", ""), + } + + def init_weight_update_group( + self, + *, + master_addr: str, + master_port: int, + rank_offset: int, + world_size: int, + group_name: str, + ) -> tuple[bool, str]: + if self._engine is None: + raise RuntimeError("engine is not started") + return self._engine.init_weights_update_group( + master_address=master_addr, + master_port=master_port, + rank_offset=rank_offset, + world_size=world_size, + group_name=group_name, + backend="nccl", + ) + + def update_one( + self, + *, + name: str, + dtype: str, + shape: list[int], + group_name: str, + ): + if self._engine is None: + raise RuntimeError("engine is not started") + return self._engine.update_weights_from_distributed( + names=[name], + dtypes=[dtype], + shapes=[shape], + group_name=group_name, + flush_cache=False, + ) + + def destroy_weight_update_group(self, *, group_name: str): + if self._engine is not None: + return self._engine.destroy_weights_update_group(group_name=group_name) + return False, "engine is not started" + + def shutdown(self) -> None: + if self._engine is not None: + try: + self._engine.shutdown() + finally: + self._engine = None + + +@ray.remote(num_gpus=1) +class RayMockMegatronSender: + """Mock Megatron rank 0 running inside one Ray actor.""" + + def __init__(self) -> None: + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + self._pg = None + + def node_ip(self) -> str: + import ray as _ray + + return _ray.util.get_node_ip_address() + + def device_info(self) -> dict[str, Any]: + return { + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", ""), + "device_count": torch.cuda.device_count(), + "current_device": torch.cuda.current_device() + if torch.cuda.is_available() + else None, + "device_name": torch.cuda.get_device_name(0), + "nccl_cumem_enable": os.environ.get("NCCL_CUMEM_ENABLE", ""), + } + + def init_weight_update_group( + self, + *, + master_addr: str, + master_port: int, + world_size: int, + group_name: str, + ) -> dict[str, Any]: + from nemo_rl.models.policy.utils import init_process_group + + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + torch.cuda.set_device(0) + self._pg = init_process_group( + backend="nccl", + init_method=f"tcp://{master_addr}:{master_port}", + world_size=world_size, + rank=0, + group_name=group_name, + ) + return self.device_info() + + def broadcast_one( + self, + *, + ckpt: str, + weight_map: dict[str, str], + name: str, + dtype: str, + group_name: str, + ) -> dict[str, Any]: + import torch.distributed as dist + from safetensors import safe_open + + if self._pg is None: + raise RuntimeError("weight-update process group is not initialized") + + shard = Path(ckpt) / weight_map[name] + with safe_open(shard, framework="pt") as reader: + tensor = reader.get_tensor(name) + tensor = tensor.to(device="cuda:0", dtype=_DTYPE_FROM_STR[dtype]).contiguous() + dist.broadcast(tensor, src=0, group=self._pg) + torch.cuda.synchronize() + return { + "name": name, + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "checksum": float(tensor.float().sum().item()), + } + + def destroy_weight_update_group(self) -> None: + import torch.distributed as dist + + if self._pg is not None: + try: + dist.destroy_process_group(self._pg) + finally: + self._pg = None + + +def run_ray_wrapped_repro( + *, + tp: int = 1, + pp: int = 1, + dp: int = 1, + param_name: str | None = "model.norm.weight", + max_tensors: int = 1, +) -> list[dict[str, Any]]: + os.environ.setdefault("NCCL_CUMEM_ENABLE", "0") + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + + if torch.cuda.device_count() < 2: + raise RuntimeError("ray-wrapped repro requires at least 2 visible GPUs") + + from repro_pp_weight_update import _ckpt_path, _collect_specs + + ckpt = _ckpt_path() + specs, weight_map = _collect_specs(ckpt) + selected_specs = _select_specs( + specs, + param_name=param_name, + max_tensors=max_tensors, + ) + print( + f"[ray-wrapped-repro] ckpt={ckpt} selected={len(selected_specs)} " + f"first={selected_specs[0][0]} last={selected_specs[-1][0]}", + flush=True, + ) + + group_name = "ray_wrapped_repro_weight_update" + world_size = 2 + rank_offset = 1 + + engine = RaySGLangEngine.remote() + sender = RayMockMegatronSender.remote() + try: + engine_info = ray.get( + engine.start.remote(model_path=str(ckpt), tp=tp, pp=pp, dp=dp) + ) + sender_info = ray.get(sender.device_info.remote()) + print(f"[ray-wrapped-repro] engine={engine_info}", flush=True) + print(f"[ray-wrapped-repro] sender={sender_info}", flush=True) + + master_addr = ray.get(sender.node_ip.remote()) + master_port = _find_free_port() + print( + f"[ray-wrapped-repro] init group tcp://{master_addr}:{master_port}", + flush=True, + ) + sender_ref = sender.init_weight_update_group.remote( + master_addr=master_addr, + master_port=master_port, + world_size=world_size, + group_name=group_name, + ) + engine_ref = engine.init_weight_update_group.remote( + master_addr=master_addr, + master_port=master_port, + rank_offset=rank_offset, + world_size=world_size, + group_name=group_name, + ) + sender_group_info, engine_group_result = ray.get([sender_ref, engine_ref]) + print( + f"[ray-wrapped-repro] sender_group={sender_group_info} " + f"engine_group={engine_group_result}", + flush=True, + ) + success, message = engine_group_result + if not success: + raise RuntimeError(f"engine init_weights_update_group failed: {message}") + + results: list[dict[str, Any]] = [] + for i, (name, dtype, shape) in enumerate(selected_specs, start=1): + recv_ref = engine.update_one.remote( + name=name, + dtype=dtype, + shape=shape, + group_name=group_name, + ) + send_ref = sender.broadcast_one.remote( + ckpt=str(ckpt), + weight_map=weight_map, + name=name, + dtype=dtype, + group_name=group_name, + ) + ready, pending = ray.wait([send_ref, recv_ref], num_returns=2, timeout=60) + if pending: + ready_refs = set(ready) + details = ", ".join( + [ + _describe_ref("send", send_ref, ready_refs), + _describe_ref("recv", recv_ref, ready_refs), + ] + ) + raise RuntimeError( + f"timed out waiting for weight update {name!r}; {details}" + ) + send_result, recv_result = ray.get([send_ref, recv_ref]) + print( + f"[ray-wrapped-repro] {i}/{len(selected_specs)} {name} " + f"send={send_result} recv={recv_result}", + flush=True, + ) + success, message = recv_result + if not success: + raise RuntimeError(f"engine update_one failed for {name!r}: {message}") + results.append(send_result) + return results + finally: + _ray_get_best_effort( + engine.destroy_weight_update_group.remote(group_name=group_name), + label="engine.destroy_weight_update_group", + ) + _ray_get_best_effort( + sender.destroy_weight_update_group.remote(), + label="sender.destroy_weight_update_group", + ) + _ray_get_best_effort(engine.shutdown.remote(), label="engine.shutdown") + for actor in (engine, sender): + try: + ray.kill(actor) + except Exception: + pass + + +def test_ray_wrapped_sglang_engine_mock_megatron_weight_update(ray_cluster): + results = run_ray_wrapped_repro() + assert results + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--dp", type=int, default=1) + parser.add_argument("--param-name", type=str, default="model.norm.weight") + parser.add_argument( + "--max-tensors", + type=int, + default=1, + help="Used only when --param-name is empty.", + ) + args = parser.parse_args() + + param_name = args.param_name or None + try: + run_ray_wrapped_repro( + tp=args.tp, + pp=args.pp, + dp=args.dp, + param_name=param_name, + max_tensors=args.max_tensors, + ) + finally: + if ray.is_initialized(): + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/unit/models/generation/sglang/test_sglang_distributed_broadcast_smoke.py b/tests/unit/models/generation/sglang/test_sglang_distributed_broadcast_smoke.py new file mode 100644 index 0000000000..ddc0553299 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_distributed_broadcast_smoke.py @@ -0,0 +1,284 @@ +# 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. + +"""Smoke test: cross-process NCCL broadcast trainer → SGLang without Megatron. + +Replaces the heavy Megatron ``Policy`` with a single-GPU Ray actor +(``MockTrainer``) that mimics only the trainer-side half of +``refit_sglang_distributed``: + + 1. Initialize a 1-rank ``gloo`` default torch process group, so that the + cross-process ``init_process_group`` (the helper we're testing) runs + *after* a default PG is already up — exactly the condition that holds + during a real Megatron refit. + 2. Call ``connect_rollout_engines_from_distributed`` to bring up the + trainer ↔ engine NCCL group. + 3. Run ``broadcast_hf_buckets_via_distributed_impl`` for one tiny + bucket containing one fake tensor. + +A pass here isolates the trainer ↔ engine NCCL transport from any +Megatron-specific machinery (AutoBridge export, refit-buffer sizing, +mcore TP/PP collectives, etc.). A failure points at either the in-tree +``init_process_group`` helper or the cross-process NCCL channel +establishment (P2P/IPC vs SHM transport). +""" + +from __future__ import annotations + +import gc +import os +import socket + +import pytest +import ray +import torch +from _megatron_helpers import SGLANG_TP1, make_sglang_cfg +from _nemotron_slicer import ensure_sliced_model + +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration + +pytestmark = pytest.mark.sglang + + +# --------------------------------------------------------------------------- +# Sliced-model fixture (mirror of test_megatron_sglang_weight_update.py) +# --------------------------------------------------------------------------- +_NEMOTRON_TEST_MODEL_PATH_ENV = "NEMOTRON_TEST_MODEL_PATH" + + +@pytest.fixture(scope="session") +def sliced_model_path() -> str: + override = os.environ.get(_NEMOTRON_TEST_MODEL_PATH_ENV) + if override: + return override + return str(ensure_sliced_model()) + + +# --------------------------------------------------------------------------- +# MockTrainer — one Ray actor on one GPU, mirrors the trainer-side ops +# --------------------------------------------------------------------------- +@ray.remote(num_gpus=1) +class MockTrainer: + """Single-GPU mock of a Megatron rank-0 trainer. + + Holds the ``model_update_group`` returned by ``connect_rollout_engines_ + from_distributed`` in the actor's process so it does not need to cross + Ray's serialization boundary. + """ + + def __init__(self) -> None: + self._model_update_group = None + + def setup_default_pg(self) -> dict: + """Initialize a 1-rank gloo default PG (mimics Megatron's startup).""" + import torch.distributed as dist + + if not dist.is_initialized(): + with socket.socket() as sock: + sock.bind(("", 0)) + port = sock.getsockname()[1] + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", str(port)) + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + world_size=1, + rank=0, + ) + torch.cuda.set_device(0) + return { + "rank": dist.get_rank(), + "world_size": dist.get_world_size(), + "device": str(torch.cuda.current_device()), + "device_name": torch.cuda.get_device_name(0), + } + + def connect_to_engine( + self, + rollout_engines: list, + engine_gpu_counts: list, + group_name: str, + ) -> None: + """Bring up the trainer ↔ engine NCCL group via the production helper.""" + from nemo_rl.models.policy.utils import ( + connect_rollout_engines_from_distributed, + ) + + self._model_update_group = connect_rollout_engines_from_distributed( + group_name=group_name, + rollout_engines=rollout_engines, + engine_gpu_counts=engine_gpu_counts, + ) + + def broadcast_one_bucket( + self, + rollout_engines: list, + rollout_engine_lock, + group_name: str, + weight_version: int, + param_name: str, + shape: tuple, + dtype_str: str, + ) -> None: + """Drive the production broadcast helper for a single fake tensor.""" + from nemo_rl.models.policy.utils import ( + broadcast_hf_buckets_via_distributed_impl, + ) + + target_dtype = getattr(torch, dtype_str) + tensor = torch.empty(shape, dtype=target_dtype, device="cuda:0") + tensor.fill_(1.0) + + bucket_iter = iter([[(param_name, tensor)]]) + broadcast_hf_buckets_via_distributed_impl( + bucket_iterator=bucket_iter, + rollout_engines=rollout_engines, + rollout_engine_lock=rollout_engine_lock, + group_name=group_name, + model_update_group=self._model_update_group, + weight_version=weight_version, + ) + + def shutdown(self) -> None: + import torch.distributed as dist + + if self._model_update_group is not None: + try: + dist.destroy_process_group(self._model_update_group) + except Exception: + pass + self._model_update_group = None + if dist.is_initialized(): + try: + dist.destroy_process_group() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Fixture: real SGLangGeneration on its own RayVirtualCluster +# --------------------------------------------------------------------------- +@pytest.fixture +def smoke_env(ray_cluster, sliced_model_path): + """Build a single-engine ``SGLangGeneration`` (sgl_tp1) on its own cluster.""" + sglang_cluster = RayVirtualCluster( + bundle_ct_per_node_list=[1], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=1, + name="smoke-sglang", + ) + sglang_cfg = make_sglang_cfg( + model_path=sliced_model_path, + sglang=SGLANG_TP1, + colocated=False, + ) + sglang_gen = SGLangGeneration(sglang_cluster, sglang_cfg) + yield sglang_gen, sglang_cluster + try: + sglang_gen.shutdown() + except Exception: + pass + try: + sglang_cluster.shutdown() + except Exception: + pass + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# The test +# --------------------------------------------------------------------------- +def test_disag_broadcast_smoke(smoke_env): + """Single-bucket NCCL broadcast: MockTrainer → real sglang engine. + + Sequencing: + 1. MockTrainer.setup_default_pg — 1-rank gloo default PG + 2. fetch_updatable_engines_with_recover (engine + lock + gpu counts) + 3. MockTrainer.connect_to_engine — cross-process NCCL group + 4. MockTrainer.broadcast_one_bucket — single-tensor bucket via prod helper + + The fake tensor uses an unrecognized parameter name (``_smoke_test_param``) + so SGLang's ``model.load_weights`` will return ``(False, "...")`` *after* + the broadcast itself completes. We assert only that the broadcast succeeds + (no NCCL transport error) — verifying the load-weights body is the real + test's job. + """ + from nemo_rl.models.policy.utils import fetch_updatable_engines_with_recover + + sglang_gen, _sglang_cluster = smoke_env + + # 1. Pull engine + lock + per-engine GPU count from the same path the + # production refit uses, so any future changes to the engine-discovery + # API stay in sync. + ( + rollout_engines, + rollout_engine_lock, + _num_new, + engine_gpu_counts, + _engine_gpu_offsets, + ) = fetch_updatable_engines_with_recover(sglang_gen) + rollout_engines = [e for e in rollout_engines if e is not None] + assert len(rollout_engines) == 1, ( + f"smoke test expects one engine, got {len(rollout_engines)}" + ) + assert rollout_engine_lock is not None, "rollout_engine_lock not set" + print(f"[smoke] engines={len(rollout_engines)} gpu_counts={engine_gpu_counts}") + + # 2. Spawn MockTrainer on a separate GPU. Ray's pool has the GPUs the + # sglang engine didn't claim. + mock_trainer = MockTrainer.remote() + try: + info = ray.get(mock_trainer.setup_default_pg.remote()) + print(f"[smoke] MockTrainer default PG up: {info}") + + # 3. Stand up the trainer ↔ engine NCCL group. + ray.get( + mock_trainer.connect_to_engine.remote( + rollout_engines=rollout_engines, + engine_gpu_counts=engine_gpu_counts, + group_name="smoke-group", + ) + ) + print("[smoke] connect_rollout_engines_from_distributed succeeded") + + # 4. Drive the production broadcast helper for one tiny bucket. + # Failure here is the same NCCL ``invalid argument`` we'd see in the + # full Megatron→SGLang refit if the cross-process NCCL transport is + # broken (e.g. P2P/IPC channel mapping). + ray.get( + mock_trainer.broadcast_one_bucket.remote( + rollout_engines=rollout_engines, + rollout_engine_lock=rollout_engine_lock, + group_name="smoke-group", + weight_version=1, + param_name="_smoke_test_param.weight", + shape=(8,), + dtype_str="bfloat16", + ) + ) + print("[smoke] broadcast_hf_buckets_via_distributed_impl succeeded") + finally: + try: + ray.get(mock_trainer.shutdown.remote()) + except Exception: + pass + try: + ray.kill(mock_trainer) + except Exception: + pass diff --git a/tests/unit/models/generation/sglang/test_sglang_generation.py b/tests/unit/models/generation/sglang/test_sglang_generation.py new file mode 100644 index 0000000000..d2ea7c462b --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_generation.py @@ -0,0 +1,504 @@ +# 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. + +"""Generation tests using a real SGLangGeneration instance. + +Spins up a real RayVirtualCluster + SGLangGeneration (router + workers) +and tests ``generate()``, ``generate_async()``, and the underlying +``generate_one_sample()`` function against a live Qwen3-0.6B model. + +Parametrised over two configurations (both use 4 GPUs total): + • tp4_1server — 1 server × TP=4 + • tp2_2servers — 2 servers × TP=2 + +Model: Qwen/Qwen3-0.6B +""" + +import asyncio +import gc + +import pytest +import ray +import torch +from helpers import ( + make_generation_sampling_params, + post_and_assert_200, +) + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.sglang.sglang_generation import ( + SGLangGeneration, + generate_one_sample, +) + +MODEL_PATH = "Qwen/Qwen3-4B" + +pytestmark = pytest.mark.sglang + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +PAD_TOKEN_ID = 151643 +EOS_TOKEN_ID = 151645 + + +# --------------------------------------------------------------------------- +# SGLang config for SGLangGeneration (mirrors existing test pattern) +# --------------------------------------------------------------------------- +def _make_sglang_generation_cfg(pad_token_id=PAD_TOKEN_ID, tp_size=1): + return { + "backend": "sglang", + "model_name": MODEL_PATH, + "model_path": MODEL_PATH, + "tokenizer": {"name": MODEL_PATH}, + "dtype": "bfloat16", + "max_new_tokens": 16, + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "stop_token_ids": [EOS_TOKEN_ID], + "stop_strings": None, + "_pad_token_id": pad_token_id, + "sglang_cfg": { + "model_path": MODEL_PATH, + "dtype": "bfloat16", + "random_seed": 42, + "context_length": 1024, + "log_level": "info", + "skip_server_warmup": True, + "dp_size": 1, + "pp_size": 1, + "ep_size": 1, + "disable_piecewise_cuda_graph": True, + "disable_cuda_graph": False, + "mem_fraction_static": 0.3, + }, + "sglang_server": { + "num_gpus": 4, + "num_gpus_per_engine": tp_size, + "needs_offload": True, + "cpu_weight_backup": True, + "sglang_server_concurrency": 64, + "pause_generation_mode": "retract", + }, + "sglang_router": { + "sglang_router_ip": None, + "sglang_router_port": None, + }, + "sglang_kwargs": {}, + } + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- +@pytest.fixture(scope="module") +def tokenizer(): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + + +@pytest.fixture( + scope="module", + params=[ + pytest.param({"tp_size": 4, "num_servers": 1}, id="tp4_1server"), + pytest.param({"tp_size": 2, "num_servers": 2}, id="tp2_2servers"), + ], +) +def sglang_gen(request, ray_cluster, tokenizer): + """Real SGLangGeneration: RayVirtualCluster → router → engines. + + Parametrised over tp4_1server (1 server × TP=4) and tp2_2servers + (2 servers × TP=2). All variants use 4 GPUs. + """ + tp_size = request.param["tp_size"] + cluster = RayVirtualCluster( + bundle_ct_per_node_list=[4], + use_gpus=True, + max_colocated_worker_groups=1, + num_gpus_per_node=4, + name=f"gen-test-{request.param['num_servers']}srv-tp{tp_size}", + ) + sglang_cfg = _make_sglang_generation_cfg( + pad_token_id=tokenizer.pad_token_id, + tp_size=tp_size, + ) + + gen = SGLangGeneration(cluster, sglang_cfg) + yield gen + try: + gen.shutdown() + except Exception: + pass + try: + cluster.shutdown() + except Exception: + pass + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_input(tokenizer, prompt, pad_length=None): + """Tokenize a prompt → BatchedDataDict for generate().""" + token_ids = tokenizer.encode(prompt) + input_length = len(token_ids) + if pad_length and pad_length > input_length: + token_ids = token_ids + [tokenizer.pad_token_id] * (pad_length - input_length) + return BatchedDataDict( + { + "input_ids": torch.tensor([token_ids], dtype=torch.long), + "input_lengths": torch.tensor([input_length], dtype=torch.long), + } + ) + + +def _make_batch(tokenizer, prompts, pad_length=None): + """Tokenize multiple prompts → single BatchedDataDict.""" + all_ids = [] + all_lengths = [] + max_len = 0 + for p in prompts: + ids = tokenizer.encode(p) + all_ids.append(ids) + all_lengths.append(len(ids)) + max_len = max(max_len, len(ids)) + + if pad_length: + max_len = max(max_len, pad_length) + + padded = [] + for ids in all_ids: + padded.append(ids + [tokenizer.pad_token_id] * (max_len - len(ids))) + + return BatchedDataDict( + { + "input_ids": torch.tensor(padded, dtype=torch.long), + "input_lengths": torch.tensor(all_lengths, dtype=torch.long), + } + ) + + +# =================================================================== +# Tests: SGLangGeneration.generate() +# =================================================================== + + +def test_generate_returns_batched_data_dict(sglang_gen, tokenizer): + """generate() returns BatchedDataDict with all required output keys.""" + data = _make_input(tokenizer, "Hello") + result = sglang_gen.generate(data, greedy=True) + + for key in [ + "output_ids", + "logprobs", + "generation_lengths", + "unpadded_sequence_lengths", + "truncated", + ]: + assert key in result, f"Missing key: {key}" + + +def test_generate_output_ids_shape(sglang_gen, tokenizer): + """output_ids has shape (batch_size, total_length) with correct padding.""" + data = _make_input(tokenizer, "The capital of France is") + result = sglang_gen.generate(data, greedy=True) + + assert result["output_ids"].dim() == 2 + assert result["output_ids"].shape[0] == 1 # batch_size + gen_len = result["generation_lengths"][0].item() + input_len = data["input_lengths"][0].item() + assert result["unpadded_sequence_lengths"][0].item() == input_len + gen_len + + +def test_generate_greedy_determinism(sglang_gen, tokenizer): + """Same prompt + greedy=True → identical output_ids across two calls.""" + data = _make_input(tokenizer, "Once upon a time") + r1 = sglang_gen.generate(data, greedy=True) + r2 = sglang_gen.generate(data, greedy=True) + + assert torch.equal(r1["output_ids"], r2["output_ids"]), ( + "Greedy generation is not deterministic" + ) + + +def test_generate_truncation_flag(sglang_gen, tokenizer): + """When max_new_tokens is small, truncated=True.""" + # Temporarily reduce max_new_tokens + orig = sglang_gen.sglang_cfg["max_new_tokens"] + sglang_gen.sglang_cfg["max_new_tokens"] = 1 + try: + data = _make_input(tokenizer, "Tell me a very long story about dragons and") + result = sglang_gen.generate(data, greedy=True) + gen_len = result["generation_lengths"][0].item() + assert gen_len == 1, f"Expected 1 token, got {gen_len}" + assert result["truncated"][0].item() is True, "Expected truncated=True" + finally: + sglang_gen.sglang_cfg["max_new_tokens"] = orig + + +def test_generate_logprobs_valid(sglang_gen, tokenizer): + """Logprobs are finite, non-positive at generated positions.""" + data = _make_input(tokenizer, "Hello world") + result = sglang_gen.generate(data, greedy=True) + + gen_len = result["generation_lengths"][0].item() + input_len = data["input_lengths"][0].item() + lps = result["logprobs"][0, input_len : input_len + gen_len] + + assert torch.isfinite(lps).all(), "Logprobs contain NaN or Inf" + assert (lps <= 0.0).all(), "Logprobs should be non-positive" + + +def test_generate_respects_max_new_tokens(sglang_gen, tokenizer): + """generation_lengths ≤ max_new_tokens for all samples.""" + data = _make_input(tokenizer, "Count from 1 to 100:") + result = sglang_gen.generate(data, greedy=True) + + max_new = sglang_gen.sglang_cfg["max_new_tokens"] + gen_len = result["generation_lengths"][0].item() + assert gen_len <= max_new, f"gen_len={gen_len} > max_new_tokens={max_new}" + + +def test_generate_batch_multiple_samples(sglang_gen, tokenizer): + """Batch of 3 prompts: all produce valid output.""" + prompts = [ + "Hello, my name is", + "The capital of France is", + "What is 2 plus 2?", + ] + data = _make_batch(tokenizer, prompts) + result = sglang_gen.generate(data, greedy=True) + + assert result["output_ids"].shape[0] == 3 + assert result["generation_lengths"].shape[0] == 3 + for i in range(3): + gen_len = result["generation_lengths"][i].item() + assert gen_len > 0, f"Sample {i} generated 0 tokens" + + +def test_generate_empty_input(sglang_gen): + """Empty batch → empty BatchedDataDict with zero-size tensors.""" + data = BatchedDataDict( + { + "input_ids": torch.zeros((0, 0), dtype=torch.long), + "input_lengths": torch.zeros(0, dtype=torch.long), + } + ) + result = sglang_gen.generate(data, greedy=True) + assert result["output_ids"].shape[0] == 0 + + +def test_generate_with_stop_strings(sglang_gen, tokenizer): + """Stop string causes early termination.""" + orig_stop = sglang_gen.sglang_cfg.get("stop_strings") + sglang_gen.sglang_cfg["stop_strings"] = ["\n"] + try: + data = _make_input(tokenizer, "List:\n1. Apple\n2.") + result = sglang_gen.generate(data, greedy=True) + gen_len = result["generation_lengths"][0].item() + max_new = sglang_gen.sglang_cfg["max_new_tokens"] + # If stop string triggered, generation should be shorter than max + # (this is a soft check — the model might produce \n on first token) + assert gen_len <= max_new + finally: + sglang_gen.sglang_cfg["stop_strings"] = orig_stop + + +# =================================================================== +# Tests: SGLangGeneration.generate_async() +# =================================================================== + + +def test_generate_async_yields_single_sample(sglang_gen, tokenizer): + """generate_async() with batch_size=1 yields (0, BatchedDataDict).""" + data = _make_input(tokenizer, "Hello") + + async def _run(): + results = [] + async for idx, batch in sglang_gen.generate_async(data, greedy=True): + results.append((idx, batch)) + return results + + results = asyncio.run(_run()) + assert len(results) == 1 + idx, batch = results[0] + assert idx == 0 + assert "output_ids" in batch + assert batch["generation_lengths"][0].item() > 0 + + +def test_generate_async_output_matches_generate(sglang_gen, tokenizer): + """Same prompt, greedy: generate() and generate_async() produce same tokens.""" + data = _make_input(tokenizer, "The answer is") + sync_result = sglang_gen.generate(data, greedy=True) + + async def _run(): + results = [] + async for _, batch in sglang_gen.generate_async(data, greedy=True): + results.append(batch) + return results[0] + + async_result = asyncio.run(_run()) + + sync_len = sync_result["generation_lengths"][0].item() + async_len = async_result["generation_lengths"][0].item() + assert sync_len == async_len, f"sync={sync_len} vs async={async_len}" + + input_len = data["input_lengths"][0].item() + sync_tokens = sync_result["output_ids"][0, input_len : input_len + sync_len] + async_tokens = async_result["output_ids"][0, input_len : input_len + async_len] + assert torch.equal(sync_tokens, async_tokens), ( + "generate() and generate_async() produced different tokens" + ) + + +# =================================================================== +# Tests: generate_one_sample() — the underlying async function +# =================================================================== + + +def test_generate_one_sample_returns_correct_tuple(sglang_gen, tokenizer): + """generate_one_sample() returns (index, tokens, logprobs, truncated).""" + sp = make_generation_sampling_params(max_new_tokens=5, temperature=0.0) + input_ids = tokenizer.encode("The capital of France is") + + result = asyncio.run( + generate_one_sample( + sglang_gen.router_ip, sglang_gen.router_port, sp, input_ids, index=42 + ) + ) + + assert len(result) == 4 + idx, tokens, logprobs, truncated = result + assert idx == 42 + assert isinstance(tokens, list) and len(tokens) > 0 + assert isinstance(logprobs, list) and len(logprobs) == len(tokens) + assert isinstance(truncated, bool) + assert all(isinstance(t, int) for t in tokens) + assert all(isinstance(lp, float) for lp in logprobs) + + +def test_generate_after_memory_cycle(sglang_gen, tokenizer): + """Generate → offload/onload → generate → same greedy output.""" + data = _make_input(tokenizer, "Two plus two equals") + r_before = sglang_gen.generate(data, greedy=True) + + # Offload and onload weights + KV on all engines + for engine in sglang_gen.engines: + ray.get(engine.release_memory_weights.remote()) + ray.get(engine.release_memory_kv_cache_and_cuda_graph.remote()) + ray.get(engine.resume_memory_weights.remote()) + ray.get(engine.resume_memory_kv_cache_and_cuda_graph.remote()) + + r_after = sglang_gen.generate(data, greedy=True) + + assert torch.equal(r_before["output_ids"], r_after["output_ids"]), ( + "Generation output changed after memory cycle" + ) + + +def test_generate_after_memory_cycle_via_http_200(sglang_gen, tokenizer): + """Generate → offload/onload via direct HTTP (asserting 200) → generate → same greedy output.""" + data = _make_input(tokenizer, "Two plus two equals") + r_before = sglang_gen.generate(data, greedy=True) + + for engine in sglang_gen.engines: + base_url = ray.get(engine.get_base_url.remote()) + assert base_url is not None + + # Release weights (flush_cache first, mirroring release_memory_occupation) + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + base_url, "release_memory_occupation", {"tags": ["weights"]} + ) + # Release KV cache + CUDA graphs + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + base_url, + "release_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + # Resume weights + post_and_assert_200(base_url, "resume_memory_occupation", {"tags": ["weights"]}) + # Resume KV cache + CUDA graphs + post_and_assert_200( + base_url, + "resume_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + + r_after = sglang_gen.generate(data, greedy=True) + + assert torch.equal(r_before["output_ids"], r_after["output_ids"]), ( + "Generation output changed after HTTP-driven memory cycle" + ) + + +def test_generate_after_memory_cycle_top_level_api(sglang_gen, tokenizer): + """Generate -> top-level offload/onload -> generate -> same greedy output.""" + data = _make_input(tokenizer, "Two plus two equals") + + r_before = sglang_gen.generate(data, greedy=True) + input_len = data["input_lengths"][0].item() + gen_len_before = r_before["generation_lengths"][0].item() + assert gen_len_before > 0, "generate() before memory cycle produced 0 tokens" + tokens_before = r_before["output_ids"][0, input_len : input_len + gen_len_before] + assert (tokens_before != PAD_TOKEN_ID).all(), "before: generated tokens contain pad" + + # Full offload + onload cycle using the top-level SGLangGeneration API. + sglang_gen.offload_weights() + sglang_gen.offload_kv() + sglang_gen.onload_weights() + sglang_gen.onload_kv() + + r_after = sglang_gen.generate(data, greedy=True) + gen_len_after = r_after["generation_lengths"][0].item() + assert gen_len_after > 0, "generate() after memory cycle produced 0 tokens" + tokens_after = r_after["output_ids"][0, input_len : input_len + gen_len_after] + assert (tokens_after != PAD_TOKEN_ID).all(), "after: generated tokens contain pad" + + assert gen_len_before == gen_len_after, ( + f"Different generation_lengths before vs. after: " + f"before={gen_len_before}, after={gen_len_after}" + ) + assert torch.equal(r_before["output_ids"], r_after["output_ids"]), ( + "Generation output changed after top-level offload/onload cycle" + ) + + +def test_invalidate_kv_cache_aggregator(sglang_gen): + """SGLangGeneration.invalidate_kv_cache() fans out to every engine and + reduces with ``all(results)``. Verifies True on a healthy cluster + (single-engine TP=4 and two-engine TP=2 variants both covered via the + fixture parametrization). + """ + assert sglang_gen.invalidate_kv_cache() is True + + +def test_invalidate_kv_cache_after_generate(sglang_gen, tokenizer): + """invalidate_kv_cache after a generate() call still returns True. + + Exercises the path most likely to surface the flush_cache pacing bug: + sglang's /flush_cache endpoint may return non-200 transiently while + draining the just-completed generation's queue, so the worker's retry + loop must actually wait between attempts. + """ + data = _make_input(tokenizer, "Two plus two equals") + sglang_gen.generate(data, greedy=True) + assert sglang_gen.invalidate_kv_cache() is True diff --git a/tests/unit/models/generation/sglang/test_sglang_launch.py b/tests/unit/models/generation/sglang/test_sglang_launch.py new file mode 100644 index 0000000000..d58549368d --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_launch.py @@ -0,0 +1,96 @@ +# 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. + +"""Tests for the SGLangGeneration init chain — multi-worker orchestration, +router integration, and the Lock actor. + +Instead of instantiating the full SGLangGeneration class (which needs +RayVirtualCluster), we test each component that __init__ wires together: + • multiple SGLangGenerationWorker actors started in parallel + • all workers registered with the real router + • the Lock Ray actor used for rollout_engine_lock +""" + +import pytest +import ray +import requests +from helpers import create_worker + +from nemo_rl.models.generation.sglang.utils.ray_utils import Lock + +pytestmark = pytest.mark.sglang + + +# ------------------------------------------------------------------ +# Multi-worker orchestration +# ------------------------------------------------------------------ +@pytest.fixture(scope="module") +def two_workers(ray_cluster, router): + """Start two TP=1 workers on GPUs 0 and 1.""" + w0 = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + w1 = create_worker(router, base_gpu_id=1, tp_size=1, rank=1) + yield [w0, w1] + for w in [w0, w1]: + try: + ray.get(w.shutdown.remote()) + except Exception: + pass + + +def test_multiple_workers_init(two_workers): + """Two workers start successfully on separate GPUs.""" + for w in two_workers: + assert ray.get(w.health_generate.remote()) is True + + +def test_workers_register_with_router(two_workers, router): + """Both workers appear in the router's /workers list.""" + resp = requests.get(f"http://{router['ip']}:{router['port']}/workers", timeout=10) + assert resp.status_code == 200 + workers_list = resp.json().get("workers", []) + assert len(workers_list) >= 2 + + +def test_workers_have_distinct_urls(two_workers): + """Each worker reports a unique base URL.""" + urls = [ray.get(w.get_base_url.remote()) for w in two_workers] + assert len(set(urls)) == 2 + for url in urls: + assert url.startswith("http://") + + +# ------------------------------------------------------------------ +# Lock actor +# ------------------------------------------------------------------ +def test_lock_actor_acquire_release(ray_cluster): + """Lock.acquire / release round-trip works.""" + lock = Lock.options(num_cpus=0.1, num_gpus=0).remote() + try: + assert ray.get(lock.acquire.remote()) is True + ray.get(lock.release.remote()) + finally: + ray.kill(lock) + + +def test_lock_actor_mutual_exclusion(ray_cluster): + """A second acquire fails while the lock is held.""" + lock = Lock.options(num_cpus=0.1, num_gpus=0).remote() + try: + assert ray.get(lock.acquire.remote()) is True + assert ray.get(lock.acquire.remote()) is False # already held + ray.get(lock.release.remote()) + assert ray.get(lock.acquire.remote()) is True # free again + ray.get(lock.release.remote()) + finally: + ray.kill(lock) diff --git a/tests/unit/models/generation/sglang/test_sglang_router.py b/tests/unit/models/generation/sglang/test_sglang_router.py new file mode 100644 index 0000000000..79dc3ed59b --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_router.py @@ -0,0 +1,95 @@ +# 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. + +"""Tests for RouterActor lifecycle — start, port allocation, stop. + +All tests use a real Ray cluster and a real sglang_router subprocess. +Each test creates its own RouterActor to avoid cross-test interference. +""" + +import pytest +import ray +import requests + +from nemo_rl.models.generation.sglang.sglang_router import RouterActor +from nemo_rl.models.generation.sglang.utils.ray_utils import find_available_port + +pytestmark = pytest.mark.sglang + + +def _start_and_cleanup(actor, router_cfg): + """Start a router, return (ip, port), register cleanup on failure.""" + ip, port = ray.get(actor.start.remote(router_cfg)) + return ip, port + + +def _stop_router(actor): + try: + ray.get(actor.stop.remote()) + except Exception: + pass + ray.kill(actor) + + +def test_start_returns_ip_and_port(ray_cluster): + """RouterActor.start returns a (str, int) tuple.""" + actor = RouterActor.remote() + try: + ip, port = _start_and_cleanup(actor, {}) + assert isinstance(ip, str) and len(ip) > 0 + assert isinstance(port, int) and port > 0 + finally: + _stop_router(actor) + + +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) + actor = RouterActor.remote() + try: + ip, port = _start_and_cleanup(actor, {"sglang_router_port": configured_port}) + assert port == configured_port + finally: + _stop_router(actor) + + +def test_start_finds_port_when_not_configured(ray_cluster): + """When sglang_router_port is None, the router picks one automatically.""" + actor = RouterActor.remote() + try: + ip, port = _start_and_cleanup(actor, {}) + assert isinstance(port, int) and port > 0 + finally: + _stop_router(actor) + + +def test_stop_terminates_process(ray_cluster): + """stop() completes without error after a successful start.""" + actor = RouterActor.remote() + ray.get(actor.start.remote({})) + ray.get(actor.stop.remote()) # should not raise + ray.kill(actor) + + +def test_router_serves_workers_endpoint(ray_cluster): + """A started router exposes the /workers HTTP endpoint.""" + actor = RouterActor.remote() + try: + ip, port = _start_and_cleanup(actor, {}) + resp = requests.get(f"http://{ip}:{port}/workers", timeout=10) + assert resp.status_code == 200 + data = resp.json() + assert "workers" in data + finally: + _stop_router(actor) diff --git a/tests/unit/models/generation/sglang/test_sglang_shutdown_and_recover.py b/tests/unit/models/generation/sglang/test_sglang_shutdown_and_recover.py new file mode 100644 index 0000000000..d9f0888550 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_shutdown_and_recover.py @@ -0,0 +1,89 @@ +# 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. + +"""Tests for worker shutdown, router un-registration, and recovery. + +Each test creates a **fresh** worker so that shutdown / crash is +non-destructive to the rest of the session. +""" + +import time + +import pytest +import ray +import requests +from helpers import create_worker + +pytestmark = pytest.mark.sglang + + +def _get_worker_count(router): + """Get the number of workers registered with the router.""" + resp = requests.get(f"http://{router['ip']}:{router['port']}/workers", timeout=10) + return len(resp.json().get("workers", [])) + + +def _wait_for_worker_count(router, expected, timeout=15): + """Poll until the router reports the expected worker count.""" + deadline = time.time() + timeout + while time.time() < deadline: + if _get_worker_count(router) == expected: + return True + time.sleep(1) + return False + + +# ------------------------------------------------------------------ +# shutdown +# ------------------------------------------------------------------ +def test_shutdown_worker(ray_cluster, router): + """Worker shutdown completes without error.""" + worker = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + ray.get(worker.shutdown.remote()) # should not raise + + +def test_shutdown_unregisters_from_router(ray_cluster, router): + """After shutdown the worker is no longer in the router's list.""" + count_before_create = _get_worker_count(router) + worker = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + + # Wait for the worker to appear in the router + assert _wait_for_worker_count(router, count_before_create + 1), ( + f"Worker never appeared in router (expected {count_before_create + 1}, " + f"got {_get_worker_count(router)})" + ) + + ray.get(worker.shutdown.remote()) + + # Wait for the worker to disappear + assert _wait_for_worker_count(router, count_before_create), ( + f"Worker still in router after shutdown (expected {count_before_create}, " + f"got {_get_worker_count(router)})" + ) + + +def test_new_worker_after_shutdown(ray_cluster, router): + """A new worker can be created on the same GPU after shutdown.""" + w1 = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + ray.get(w1.shutdown.remote()) + + w2 = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + assert ray.get(w2.health_generate.remote()) is True + ray.get(w2.shutdown.remote()) + + +def test_simulate_crash(ray_cluster, router): + """_simulate_crash (which calls shutdown) does not raise.""" + worker = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + ray.get(worker._simulate_crash.remote()) diff --git a/tests/unit/models/generation/sglang/test_sglang_worker_init.py b/tests/unit/models/generation/sglang/test_sglang_worker_init.py new file mode 100644 index 0000000000..b234b1177b --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_worker_init.py @@ -0,0 +1,67 @@ +# 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. + +"""Tests for SGLangGenerationWorker.init — server launch and router registration. + +Uses a real Ray cluster, a real sglang router, and a real SGLang server +(Qwen3-0.6B, TP=1). A module-scoped worker is shared across all tests +in this file. +""" + +import pytest +import ray +import requests +from helpers import create_worker + +pytestmark = pytest.mark.sglang + + +@pytest.fixture(scope="module") +def worker(ray_cluster, router): + """Create a single TP=1 worker for this module's tests.""" + w = create_worker(router, base_gpu_id=0, tp_size=1, rank=0) + yield w + try: + ray.get(w.shutdown.remote()) + except Exception: + pass + + +# ------------------------------------------------------------------ +def test_init_server_healthy(worker): + """After init, the underlying SGLang server is healthy.""" + result = ray.get(worker.health_generate.remote()) + assert result is True + + +def test_init_sets_base_url(worker): + """get_base_url returns a valid http:// URL after init.""" + url = ray.get(worker.get_base_url.remote()) + assert url is not None + assert url.startswith("http://") + + +def test_init_registers_with_router(worker, router): + """The worker registers itself with the session router on init.""" + resp = requests.get(f"http://{router['ip']}:{router['port']}/workers", timeout=10) + assert resp.status_code == 200 + workers_list = resp.json().get("workers", []) + # At least one worker should be registered + assert len(workers_list) >= 1 + + +def test_health_generate_returns_true(worker): + """health_generate succeeds multiple times (idempotent).""" + for _ in range(3): + assert ray.get(worker.health_generate.remote()) is True diff --git a/tests/unit/models/generation/sglang/test_sglang_worker_memory.py b/tests/unit/models/generation/sglang/test_sglang_worker_memory.py new file mode 100644 index 0000000000..33590b03e3 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_worker_memory.py @@ -0,0 +1,212 @@ +# 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. + +"""Tests for SGLangGenerationWorker memory management: +flush_cache, release_memory_occupation, resume_memory_occupation. + +Uses a real SGLang server (Qwen3-0.6B), parametrised over two +configurations so the same tests exercise both a single-worker TP=4 +setup and a two-worker TP=2 setup: + + • tp4 — 1 worker × TP=4 + • tp2_2workers — 2 workers × TP=2 (the memory tests target worker 0, + but both workers share the router) + +Each test is fully self-contained — it leaves the server in the same +state it found it. +""" + +import pytest +import ray +from helpers import create_worker, post_and_assert_200 + +pytestmark = pytest.mark.sglang + + +@pytest.fixture( + scope="module", + params=[ + pytest.param({"tp_size": 4, "num_workers": 1}, id="tp4"), + pytest.param({"tp_size": 2, "num_workers": 2}, id="tp2_2workers"), + ], +) +def worker(request, ray_cluster, router): + """Worker(s) dedicated to memory tests. + + For ``tp4`` a single TP=4 worker is created. For ``tp2_2workers`` + two TP=2 workers share the same router (mirroring the 2-servers + configuration exercised elsewhere); memory tests run against the + first worker but the second is kept alive so the router has the + multi-worker topology in place. + """ + tp_size = request.param["tp_size"] + num_workers = request.param["num_workers"] + + workers = [] + for rank in range(num_workers): + workers.append( + create_worker( + router, + base_gpu_id=rank * tp_size, + tp_size=tp_size, + rank=rank, + ) + ) + + yield workers[0] + + for w in workers: + try: + ray.get(w.shutdown.remote()) + except Exception: + pass + + +# ------------------------------------------------------------------ +# flush_cache +# ------------------------------------------------------------------ +def test_flush_cache_success(worker): + """flush_cache returns without error on a healthy server.""" + ray.get(worker.flush_cache.remote()) + + +# ------------------------------------------------------------------ +# invalidate_kv_cache (worker-level) +# ------------------------------------------------------------------ +def test_invalidate_kv_cache_success(worker): + """invalidate_kv_cache returns True on a healthy server.""" + assert ray.get(worker.invalidate_kv_cache.remote()) is True + + +def test_invalidate_kv_cache_after_resume(worker): + """invalidate_kv_cache succeeds after a release → resume round-trip. + + Exercises the same retry path that flush_cache hits — sglang's + /flush_cache endpoint can return non-200 transiently while the queue + drains, so the loop must pace its retries. + """ + ray.get(worker.release_memory_weights.remote()) + ray.get(worker.resume_memory_weights.remote()) + assert ray.get(worker.invalidate_kv_cache.remote()) is True + + +def test_invalidate_kv_cache_repeated(worker): + """Back-to-back invalidate_kv_cache calls all return True (no state leak).""" + for _ in range(3): + assert ray.get(worker.invalidate_kv_cache.remote()) is True + + +# ------------------------------------------------------------------ +# release / resume — weights (self-contained) +# ------------------------------------------------------------------ +def test_release_and_resume_memory_weights(worker): + """release_memory_weights followed by resume succeeds.""" + ray.get(worker.release_memory_weights.remote()) + ray.get(worker.resume_memory_weights.remote()) + + +# ------------------------------------------------------------------ +# release / resume — KV cache + CUDA graphs (self-contained) +# ------------------------------------------------------------------ +def test_release_and_resume_memory_kv_cache_and_cuda_graph(worker): + """release then resume KV cache + CUDA graphs succeeds.""" + ray.get(worker.release_memory_kv_cache_and_cuda_graph.remote()) + ray.get(worker.resume_memory_kv_cache_and_cuda_graph.remote()) + + +# ------------------------------------------------------------------ +# full offload / onload cycle +# ------------------------------------------------------------------ +def test_full_offload_onload_cycle(worker): + """Full offload (weights then KV) then onload (weights then KV) works.""" + ray.get(worker.release_memory_weights.remote()) + ray.get(worker.release_memory_kv_cache_and_cuda_graph.remote()) + ray.get(worker.resume_memory_weights.remote()) + ray.get(worker.resume_memory_kv_cache_and_cuda_graph.remote()) + + +def test_health_after_memory_cycle(worker): + """health_generate passes after a full offload / onload cycle.""" + ray.get(worker.release_memory_weights.remote()) + ray.get(worker.release_memory_kv_cache_and_cuda_graph.remote()) + ray.get(worker.resume_memory_weights.remote()) + ray.get(worker.resume_memory_kv_cache_and_cuda_graph.remote()) + assert ray.get(worker.health_generate.remote()) is True + + +def test_flush_cache_after_resume(worker): + """flush_cache succeeds after a release → resume round-trip.""" + ray.get(worker.release_memory_weights.remote()) + ray.get(worker.resume_memory_weights.remote()) + ray.get(worker.flush_cache.remote()) + + +# ------------------------------------------------------------------ +# Equivalent tests using _make_request directly — verify HTTP 200 +# ------------------------------------------------------------------ +def test_offload_onload_via_http_200(worker): + """Full offload/onload cycle driven by direct HTTP POST, asserting 200. + + Uses ``post_and_assert_200`` (which checks ``resp.status_code == 200`` + explicitly) rather than ``_make_request`` — ``_make_request`` throws the + status code away inside ``raise_for_status()`` so callers cannot inspect it. + """ + base_url = ray.get(worker.get_base_url.remote()) + assert base_url is not None + + # Release weights (flush_cache first, mirroring release_memory_occupation) + ray.get(worker.flush_cache.remote()) + post_and_assert_200(base_url, "release_memory_occupation", {"tags": ["weights"]}) + # Release KV cache + CUDA graphs + ray.get(worker.flush_cache.remote()) + post_and_assert_200( + base_url, "release_memory_occupation", {"tags": ["kv_cache", "cuda_graph"]} + ) + # Resume weights + post_and_assert_200(base_url, "resume_memory_occupation", {"tags": ["weights"]}) + # Resume KV cache + CUDA graphs + post_and_assert_200( + base_url, "resume_memory_occupation", {"tags": ["kv_cache", "cuda_graph"]} + ) + assert ray.get(worker.health_generate.remote()) is True + + +def test_double_offload_onload_cycle(worker): + """Two back-to-back full offload/onload cycles via direct HTTP, asserting 200 on every call. + + Exercises the same endpoints twice to catch state leaks across cycles. + Uses ``post_and_assert_200`` so each of the eight POSTs explicitly + verifies ``resp.status_code == 200``. + """ + base_url = ray.get(worker.get_base_url.remote()) + assert base_url is not None + + for _ in range(2): + ray.get(worker.flush_cache.remote()) + post_and_assert_200( + base_url, "release_memory_occupation", {"tags": ["weights"]} + ) + ray.get(worker.flush_cache.remote()) + post_and_assert_200( + base_url, + "release_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + post_and_assert_200(base_url, "resume_memory_occupation", {"tags": ["weights"]}) + post_and_assert_200( + base_url, + "resume_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + assert ray.get(worker.health_generate.remote()) is True diff --git a/tests/unit/models/generation/sglang/test_utils_smoke.py b/tests/unit/models/generation/sglang/test_utils_smoke.py new file mode 100644 index 0000000000..e3ff386733 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_utils_smoke.py @@ -0,0 +1,90 @@ +# 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. + +"""Smoke tests for utility modules (ray_utils, misc, async_utils). + +These tests verify basic functionality of helper utilities and do NOT +require a running SGLang server or GPU. +""" + +import multiprocessing + +from nemo_rl.models.generation.sglang.utils.ray_utils import ( + _wrap_ipv6, + find_available_port, + get_host_info, + is_port_available, +) +from nemo_rl.models.generation.sglang.utils.router_utils import ( + terminate_process, +) +from nemo_rl.models.policy.torch_reductions_utils import ( + MultiprocessingSerializer, +) + + +# --------------------------------------------------------------------------- +# 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" + assert _wrap_ipv6("10.0.0.1") == "10.0.0.1" + assert _wrap_ipv6("127.0.0.1") == "127.0.0.1" + + +def test_wrap_ipv6_brackets_ipv6(): + """IPv6 addresses are wrapped in [] by _wrap_ipv6, idempotently.""" + # Bare IPv6 → wrapped. + assert _wrap_ipv6("::1") == "[::1]" + assert _wrap_ipv6("2001:db8::1") == "[2001:db8::1]" + # Already-bracketed input stays a single pair of brackets. + assert _wrap_ipv6("[::1]") == "[::1]" + assert _wrap_ipv6("[2001:db8::1]") == "[2001:db8::1]" + + +def test_get_host_info_returns_tuple(): + """get_host_info returns (hostname, ip_address) strings.""" + hostname, ip = get_host_info() + assert isinstance(hostname, str) and len(hostname) > 0 + assert isinstance(ip, str) and len(ip) > 0 + + +# --------------------------------------------------------------------------- +# misc +# --------------------------------------------------------------------------- +def test_serializer_roundtrip(): + """serialize → deserialize returns the original object.""" + obj = {"key": "value", "numbers": [1, 2, 3], "nested": {"a": True}} + serialized = MultiprocessingSerializer.serialize(obj, output_str=True) + assert isinstance(serialized, str) and len(serialized) > 0 + deserialized = MultiprocessingSerializer.deserialize(serialized) + assert deserialized == obj + + +def test_terminate_process_already_dead(): + """terminate_process does not raise when the process is already dead.""" + p = multiprocessing.Process(target=lambda: None) + p.start() + p.join() + # Process has already exited — should be a harmless no-op + terminate_process(p) diff --git a/tests/unit/models/generation/sglang/test_weight_update_real.py b/tests/unit/models/generation/sglang/test_weight_update_real.py new file mode 100644 index 0000000000..b340c1406c --- /dev/null +++ b/tests/unit/models/generation/sglang/test_weight_update_real.py @@ -0,0 +1,537 @@ +# 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. + +""" +End-to-end weight update tests using SGLangGeneration + mock FSDP trainer. + +Verifies the full weight-streaming path: + 1. SGLangGeneration.check_weights("snapshot") — save original weights + 2. SGLangGeneration.check_weights("reset_tensors") — randomize weights + 3. Mock FSDP trainer streams Qwen3-1.7B weights via stream_weights_via_http_impl + 4. SGLangGeneration.check_weights("compare") — verify restored weights + +Parametrised over two configurations (both require 4 GPUs): + • 1 server × TP=4 — single-server high-TP + • 2 servers × TP=2 — multi-server routing + +Model: Qwen/Qwen3-1.7B +""" + +import gc +import os + +import pytest +import ray +import torch +import torch.distributed as dist +from helpers import make_actor_env_vars, post_and_assert_200 +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, +) + +pytestmark = pytest.mark.sglang + +MODEL_PATH = "Qwen/Qwen3-4B" +PAD_TOKEN_ID = 151643 +EOS_TOKEN_ID = 151645 + + +# --------------------------------------------------------------------------- +# SGLang config builder +# --------------------------------------------------------------------------- +def _make_sglang_cfg(tp_size, pad_token_id=PAD_TOKEN_ID): + return { + "backend": "sglang", + "model_name": MODEL_PATH, + "model_path": MODEL_PATH, + "tokenizer": {"name": MODEL_PATH}, + "dtype": "bfloat16", + "max_new_tokens": 16, + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "stop_token_ids": [EOS_TOKEN_ID], + "stop_strings": None, + "_pad_token_id": pad_token_id, + "sglang_cfg": { + "model_path": MODEL_PATH, + "dtype": "bfloat16", + "random_seed": 42, + "context_length": 1024, + "log_level": "warning", + "skip_server_warmup": True, + "dp_size": 1, + "pp_size": 1, + "ep_size": 1, + "disable_piecewise_cuda_graph": True, + "disable_cuda_graph": True, + "mem_fraction_static": 0.3, + }, + "sglang_server": { + "num_gpus": 4, + "num_gpus_per_engine": tp_size, + "needs_offload": True, + "cpu_weight_backup": False, + "sglang_server_concurrency": 64, + "pause_generation_mode": "retract", + }, + "sglang_router": { + "sglang_router_ip": None, + "sglang_router_port": None, + }, + "sglang_kwargs": {}, + } + + +# --------------------------------------------------------------------------- +# Mock FSDP trainer worker +# --------------------------------------------------------------------------- +@ray.remote(num_cpus=0.1) +class MockFSDPWorker: + """Simulates one FSDP rank for weight streaming. + + Loads the full model on a single GPU and calls the real + ``stream_weights_via_http_impl`` to send weights to SGLang servers. + """ + + def init(self, rank, world_size, master_addr, master_port, model_path, gpu_index): + os.environ["MASTER_ADDR"] = master_addr + os.environ["MASTER_PORT"] = str(master_port) + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(gpu_index) + os.environ["WORLD_SIZE"] = str(world_size) + + self.rank = rank + self.gpu_index = gpu_index + self.dtype = torch.bfloat16 + + torch.cuda.set_device(gpu_index) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + from transformers import AutoModelForCausalLM + + device = torch.device(f"cuda:{gpu_index}") + self.model = AutoModelForCausalLM.from_pretrained( + model_path, + torch_dtype=self.dtype, + trust_remote_code=True, + ).to(device) + + from nemo_rl.utils.nvml import get_device_uuid + + self.device_uuid = get_device_uuid(gpu_index) + + def get_device_uuid(self): + return self.device_uuid + + def stream_weights(self, rollout_engines, num_gpus_per_engine): + from nemo_rl.models.policy.utils import stream_weights_via_http_impl + + if not hasattr(self, "_ipc_worker_state"): + self._ipc_worker_state = {} + + from nemo_rl.models.policy.workers.dtensor_policy_worker_v2 import ( + dtensor_params_generator, + ) + + rollout_engine_urls = ray.get( + [e.get_base_url.remote() for e in rollout_engines] + ) + + stream_weights_via_http_impl( + params_generator=dtensor_params_generator(self.model, self.dtype), + rollout_engine_urls=rollout_engine_urls, + num_gpus_per_engine=num_gpus_per_engine, + rank=self.rank, + world_size=dist.get_world_size(), + worker_name=f"MockFSDPWorker-{self.rank}", + buffer_size_bytes=512 * 1024 * 1024, + worker_state=self._ipc_worker_state, + ) + + def shutdown(self): + if dist.is_initialized(): + dist.destroy_process_group() + if hasattr(self, "model"): + del self.model + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- +@pytest.fixture( + params=[ + pytest.param({"tp_size": 4, "num_servers": 1}, id="tp4_1server"), + pytest.param({"tp_size": 2, "num_servers": 2}, id="tp2_2servers"), + ] +) +def sglang_gen(request, ray_cluster): + """Real SGLangGeneration: RayVirtualCluster → router → engines.""" + cfg = request.param + tp_size = cfg["tp_size"] + + cluster = RayVirtualCluster( + bundle_ct_per_node_list=[4], + use_gpus=True, + max_colocated_worker_groups=2, + num_gpus_per_node=4, + name="weight-update-test", + ) + sglang_cfg = _make_sglang_cfg(tp_size) + + gen = SGLangGeneration(cluster, sglang_cfg) + yield gen + + try: + gen.shutdown() + except Exception: + pass + try: + cluster.shutdown() + except Exception: + pass + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture +def mock_trainer(ray_cluster, sglang_gen): + """4 MockFSDPWorker actors with torch.distributed (gloo), each loading Qwen3-1.7B. + + Actors are launched into the SGLang cluster's placement group using + PlacementGroupSchedulingStrategy with fractional GPU (num_gpus=0.2), so + they co-reside with the SGLang worker (which also takes 0.2) on the same + bundles. This matches the nemo_rl colocated-mode pattern; the + PG's bundles have ``CPU: max_colocated_worker_groups`` capacity (=2) to + fit both worker groups. + """ + host_ip = get_host_info()[1] + master_port = find_available_port(29500) + env_vars = make_actor_env_vars() + + pg = sglang_gen.cluster.get_placement_groups()[0] + + workers = [] + for rank in range(4): + w = MockFSDPWorker.options( + num_cpus=0.2, + num_gpus=0.2, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=rank, + ), + runtime_env={"env_vars": env_vars}, + ).remote() + workers.append(w) + + # All workers must init simultaneously (gloo rendezvous). + ray.get( + [ + w.init.remote( + rank=rank, + world_size=4, + master_addr=host_ip, + master_port=master_port, + model_path=MODEL_PATH, + gpu_index=rank, + ) + for rank, w in enumerate(workers) + ] + ) + + yield workers + + for w in workers: + try: + ray.get(w.shutdown.remote()) + except Exception: + pass + ray.kill(w) + gc.collect() + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +def test_weight_update_roundtrip(sglang_gen, mock_trainer): + """Snapshot -> reset -> offload -> update -> compare -> onload_kv. + + Exercises the full colocated-refit memory dance: + snapshot -> reset_tensors -> offload_weights -> offload_kv -> + onload_weights -> update_weights (stream) -> check compare -> + onload_kv. + """ + # 1. Snapshot original Qwen3-1.7B weights. + print("[STEP 1/7] Snapshotting original weights...", flush=True) + sglang_gen.check_weights("snapshot") + print("[STEP 1/7] Snapshot complete.", flush=True) + + # 2. Randomize all model weights on the SGLang servers. + print("[STEP 2/7] Randomizing (reset_tensors) model weights...", flush=True) + sglang_gen.check_weights("reset_tensors") + print("[STEP 2/7] Reset complete.", flush=True) + + # 3. Offload weights and KV cache to CPU (refit prelude). + print("[STEP 3/7] Offloading weights and KV cache to CPU...", flush=True) + sglang_gen.offload_weights() + sglang_gen.offload_kv() + print("[STEP 3/7] Offload complete.", flush=True) + + # 4. Onload weight buffers back to GPU so IPC handles can target them. + print("[STEP 4/7] Onloading weight buffers back to GPU...", flush=True) + sglang_gen.onload_weights() + print("[STEP 4/7] Onload weights complete.", flush=True) + + # 5. All 4 mock FSDP workers stream weights simultaneously via CUDA IPC over HTTP. + print( + "[STEP 5/7] Streaming weights from mock FSDP workers via CUDA IPC...", + flush=True, + ) + rollout_engines = sglang_gen.rollout_engines + num_gpus_per_engine = sglang_gen.num_gpus_per_engine + ray.get( + [ + w.stream_weights.remote(rollout_engines, num_gpus_per_engine) + for w in mock_trainer + ] + ) + print("[STEP 5/7] Weight streaming complete.", flush=True) + + # 6. Compare current weights against snapshot - raises on mismatch. + print("[STEP 6/7] Comparing current weights against snapshot...", flush=True) + sglang_gen.check_weights("compare") + print("[STEP 6/7] Compare passed.", flush=True) + + # 7. Onload KV cache to finish the refit cycle. + print("[STEP 7/7] Onloading KV cache to finish refit cycle...", flush=True) + sglang_gen.onload_kv() + print("[STEP 7/7] Roundtrip complete.", flush=True) + + +# --------------------------------------------------------------------------- +# Test: roundtrip + router-based generate() + greedy before/after comparison +# --------------------------------------------------------------------------- +def test_weight_update_roundtrip_with_router_generation(sglang_gen, mock_trainer): + """Full refit roundtrip with generation via router and per-worker HTTP 200 checks. + + Differs from ``test_weight_update_roundtrip`` in two ways: + + 1. *Generation through the router.* Both the pre-snapshot and post-onload_kv + generations go through ``sglang_gen.generate(..., greedy=True)`` which + calls ``generate_one_sample(router_ip, router_port, ...)`` — i.e. an + HTTP POST to ``http://{router_ip}:{router_port}/generate``, not to + any individual server. (``sglang_gen.generate`` internally calls + ``resp.raise_for_status()`` so a successful return implies HTTP 200.) + Parametrised over ``tp4_1server`` and ``tp2_2servers``; both configs + share the same router, so the same generation path is exercised in + both. + 2. *Per-worker HTTP 200 checks for the refit cycle.* Instead of calling + ``sglang_gen.check_weights(...)`` / ``offload_weights`` / ``offload_kv`` / + ``onload_weights`` / ``onload_kv``, this test iterates + ``sglang_gen.engines`` and drives the equivalent HTTP endpoints on + **every worker** directly via ``post_and_assert_200`` (same pattern as + ``tests/unit/models/generation/sglang/test_sglang_worker_memory.py``). + That way every single memory/weights transition is verified to return + ``resp.status_code == 200`` — ``_make_request`` would hide the status + behind ``raise_for_status()``. + + Strict outer check: with ``temperature=0.0`` the mock FSDP trainer streams + the original Qwen3-1.7B weights back, so pre- and post-roundtrip greedy + token sequences must match exactly. + """ + from transformers import AutoTokenizer + + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + + # Sanity: router endpoint is set so sglang_gen.generate actually routes. + assert sglang_gen.router_ip is not None and sglang_gen.router_port is not None, ( + "router_ip/router_port not set on sglang_gen — generate() would not route" + ) + print( + f"[setup] Router endpoint: http://{sglang_gen.router_ip}:{sglang_gen.router_port}", + flush=True, + ) + + # All logical-engine node-0 actors (one per SGLang server). + engines = [e for e in sglang_gen.engines if e is not None] + assert len(engines) >= 1, "sglang_gen has no engines" + base_urls = ray.get([e.get_base_url.remote() for e in engines]) + assert all(u is not None for u in base_urls), f"missing base_url in {base_urls}" + print(f"[setup] {len(engines)} worker(s); base_urls={base_urls}", flush=True) + + # --- Per-worker HTTP helpers ----------------------------------------------- + def _http_check_weights_all(action: str): + """POST /weights_checker on every worker, asserting 200 each time.""" + for url in base_urls: + post_and_assert_200(url, "weights_checker", {"action": action}) + + def _http_release_weights_all(): + """Flush cache + POST /release_memory_occupation(tags=[weights]) per worker.""" + for engine, url in zip(engines, base_urls): + ray.get(engine.flush_cache.remote()) + post_and_assert_200(url, "release_memory_occupation", {"tags": ["weights"]}) + + def _http_release_kv_all(): + """Flush cache + POST /release_memory_occupation(tags=[kv_cache, cuda_graph]).""" + for engine, url in zip(engines, base_urls): + ray.get(engine.flush_cache.remote()) + post_and_assert_200( + url, + "release_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + + def _http_resume_weights_all(): + """POST /resume_memory_occupation(tags=[weights]) per worker.""" + for url in base_urls: + post_and_assert_200(url, "resume_memory_occupation", {"tags": ["weights"]}) + + def _http_resume_kv_all(): + """POST /resume_memory_occupation(tags=[kv_cache, cuda_graph]) per worker.""" + for url in base_urls: + post_and_assert_200( + url, + "resume_memory_occupation", + {"tags": ["kv_cache", "cuda_graph"]}, + ) + + # --- Router-based greedy generation --------------------------------------- + test_prompt = "The capital of France is" + input_ids = tokenizer.encode(test_prompt, add_special_tokens=True) + input_len = len(input_ids) + + data = BatchedDataDict( + { + "input_ids": torch.tensor([input_ids], dtype=torch.long), + "input_lengths": torch.tensor([input_len], dtype=torch.long), + } + ) + + def _generate(tag): + result = sglang_gen.generate(data, greedy=True) + for key in ( + "output_ids", + "generation_lengths", + "unpadded_sequence_lengths", + "logprobs", + ): + assert key in result, f"[{tag}] generate() output missing key: {key}" + gen_len = int(result["generation_lengths"][0].item()) + assert gen_len > 0, ( + f"[{tag}] generate() returned 0 tokens (no new tokens generated)" + ) + tokens = result["output_ids"][0, input_len : input_len + gen_len].tolist() + assert all(isinstance(t, int) for t in tokens), ( + f"[{tag}] output tokens should be ints, got {tokens!r}" + ) + text = tokenizer.decode(tokens, skip_special_tokens=True) + assert len(text) > 0, f"[{tag}] decoded generated text is empty" + print(f"[{tag}] gen_len={gen_len} tokens={tokens} text={text!r}", flush=True) + return tokens + + # --- Generation BEFORE snapshot (via router) ------------------------------- + print("[PRE] Router greedy generate() before snapshot...", flush=True) + tokens_before = _generate("PRE") + + # --- Steps 1-7, every HTTP call asserted 200 ------------------------------ + print( + "[STEP 1/7] Snapshotting original weights (HTTP weights_checker×workers)...", + flush=True, + ) + _http_check_weights_all("snapshot") + print("[STEP 1/7] Snapshot complete.", flush=True) + + print( + "[STEP 2/7] Randomizing weights (HTTP weights_checker reset_tensors×workers)...", + flush=True, + ) + _http_check_weights_all("reset_tensors") + print("[STEP 2/7] Reset complete.", flush=True) + + print( + "[STEP 3/7] Offloading weights + KV (HTTP release_memory_occupation×workers)...", + flush=True, + ) + _http_release_weights_all() + _http_release_kv_all() + print("[STEP 3/7] Offload complete.", flush=True) + + print( + "[STEP 4/7] Onloading weights (HTTP resume_memory_occupation weights×workers)...", + flush=True, + ) + _http_resume_weights_all() + print("[STEP 4/7] Onload weights complete.", flush=True) + + print( + "[STEP 5/7] Streaming weights from mock FSDP workers via CUDA IPC...", + flush=True, + ) + rollout_engines = sglang_gen.rollout_engines + num_gpus_per_engine = sglang_gen.num_gpus_per_engine + ray.get( + [ + w.stream_weights.remote(rollout_engines, num_gpus_per_engine) + for w in mock_trainer + ] + ) + print("[STEP 5/7] Weight streaming complete.", flush=True) + + print( + "[STEP 6/7] Compare vs snapshot (HTTP weights_checker compare×workers)...", + flush=True, + ) + _http_check_weights_all("compare") + print("[STEP 6/7] Compare passed.", flush=True) + + print( + "[STEP 7/7] Onloading KV (HTTP resume_memory_occupation kv×workers)...", + flush=True, + ) + _http_resume_kv_all() + print("[STEP 7/7] Roundtrip complete.", flush=True) + + # --- Generation AFTER onload_kv (via router) ------------------------------- + print("[POST] Router greedy generate() after onload_kv...", flush=True) + tokens_after = _generate("POST") + + # --- Sanity: generate() actually produced new tokens on BOTH runs ---------- + assert len(tokens_before) > 0, "generate() returned no tokens before roundtrip" + assert len(tokens_after) > 0, "generate() returned no tokens after roundtrip" + assert len(tokens_before) == len(tokens_after), ( + f"Different number of generated tokens before vs. after: " + f"before={len(tokens_before)}, after={len(tokens_after)}" + ) + + # --- Strict equality (greedy => deterministic across roundtrip) ------------ + assert tokens_before == tokens_after, ( + "Greedy tokens changed across the refit roundtrip:\n" + f" before (pre-snapshot): {tokens_before}\n" + f" after (post-onload_kv): {tokens_after}" + ) + print( + f"[ASSERT] Greedy tokens match before vs. after roundtrip " + f"(n={len(tokens_before)} tokens, both non-empty, both via router).", + flush=True, + ) diff --git a/tests/unit/models/generation/test_sglang_generation.py b/tests/unit/models/generation/test_sglang_generation.py deleted file mode 100644 index 508bd214d1..0000000000 --- a/tests/unit/models/generation/test_sglang_generation.py +++ /dev/null @@ -1,928 +0,0 @@ -# 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. - -"""Unit tests for SGLang generation backend. - -These tests verify that the SGLang generation backend produces sane outputs. -While not true unit tests, they validate the generation quality in unit test runs. -""" - -import gc -from copy import deepcopy - -import pytest -import ray -import torch - -from nemo_rl.algorithms.utils import get_tokenizer -from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.virtual_cluster import RayVirtualCluster -from nemo_rl.models.generation.sglang import SGLangConfig, SGLangGeneration - -model_name = "Qwen/Qwen3-0.6B" - -# Define basic SGLang test config -basic_sglang_test_config: SGLangConfig = { - "backend": "sglang", - "model_name": model_name, - "model_path": model_name, - "tokenizer": { - "name": model_name, - }, - "dtype": "bfloat16", - "max_new_tokens": 5, # Small number of tokens for testing - "temperature": 1.0, - "top_p": 1.0, - "top_k": None, - "stop_token_ids": None, - "stop_strings": None, - "sglang_cfg": { - "model_path": model_name, - "gpus_per_server": 2, - "dtype": "bfloat16", - "context_length": 1024, - "log_level": "warning", - "skip_server_warmup": True, - "enable_memory_saver": False, - "dp_size": 1, - "pp_size": 1, - "ep_size": 1, - "mem_fraction_static": 0.7, - "disable_piecewise_cuda_graph": True, - }, - "colocated": { - "enabled": True, - "resources": { - "gpus_per_node": None, - "num_nodes": None, - }, - }, - "sglang_kwargs": {}, -} - -# Basic DTensor test config for Policy tests -basic_dtensor_test_config = { - "model_name": model_name, - "tokenizer": { - "name": model_name, - }, - "train_global_batch_size": 1, - "train_micro_batch_size": 1, - "learning_rate": 5e-6, - "logprob_batch_size": 1, - "max_new_tokens": 16, - "do_sample": False, - "precision": "float32", - "offload_optimizer_for_logprob": False, - "optimizer": { - "name": "torch.optim.AdamW", - "kwargs": { - "lr": 5e-6, - "weight_decay": 0.01, - "betas": [0.9, 0.999], - "eps": 1e-8, - }, - }, - "dtensor_cfg": { - "_v2": True, # Use DTensorPolicyWorkerV2 for stream_weights_via_http - "enabled": True, - "cpu_offload": False, - "sequence_parallel": False, - "activation_checkpointing": False, - "tensor_parallel_size": 2, - "context_parallel_size": 1, - "custom_parallel_plan": None, - }, - "dynamic_batching": { - "enabled": True, - "train_mb_tokens": 40, - "logprob_mb_tokens": 40, - "sequence_length_round": 4, - }, - "sequence_packing": { - "enabled": False, - }, - "max_grad_norm": 1.0, - "make_sequence_length_divisible_by": 1, - "generation": deepcopy(basic_sglang_test_config), -} - - -def configure_sglang_config( - config: SGLangConfig, tokenizer, is_eval=True -) -> SGLangConfig: - """Apply specific configurations to SGLang config.""" - config = deepcopy(config) - config["_pad_token_id"] = tokenizer.pad_token_id - if config["stop_token_ids"] is None: - config["stop_token_ids"] = [tokenizer.eos_token_id] - return config - - -@pytest.fixture(scope="function") -def cluster(): - """Create a virtual cluster for testing with 2 GPUs.""" - virtual_cluster = RayVirtualCluster( - bundle_ct_per_node_list=[2], - use_gpus=True, - max_colocated_worker_groups=2, - num_gpus_per_node=2, - name="sglang-test-cluster", - ) - yield virtual_cluster - virtual_cluster.shutdown() - - -@pytest.fixture(scope="function") -def tokenizer(): - """Initialize tokenizer for the test model.""" - tokenizer = get_tokenizer(basic_sglang_test_config["tokenizer"]) - return tokenizer - - -@pytest.fixture(scope="function") -def policy(cluster, tokenizer): - """Initialize the SGLang policy.""" - sglang_config = deepcopy(basic_sglang_test_config) - sglang_config = configure_sglang_config(sglang_config, tokenizer) - p = SGLangGeneration(cluster, sglang_config) - yield p - try: - p.shutdown() - gc.collect() - torch.cuda.empty_cache() - except Exception as e: - print(f"Error during policy cleanup: {e}") - - -@pytest.fixture(scope="function") -def test_input_data(tokenizer): - """Create test input data for inference.""" - test_prompts = [ - "Hello, my name is", - "The capital of France is", - ] - - # Tokenize prompts - encodings = tokenizer( - test_prompts, - padding="max_length", - max_length=20, - truncation=True, - return_tensors="pt", - padding_side="right", - ) - - # Calculate input lengths from attention mask - input_lengths = encodings["attention_mask"].sum(dim=1).to(torch.int32) - - # Create input data dictionary - return BatchedDataDict( - { - "input_ids": encodings["input_ids"], - "input_lengths": input_lengths, - } - ) - - -@pytest.fixture(scope="function") -def policy_cluster_separate(): - """Create a virtual cluster for the Policy, using 2 GPUs.""" - cluster = RayVirtualCluster( - bundle_ct_per_node_list=[2], - use_gpus=True, - max_colocated_worker_groups=2, - num_gpus_per_node=2, - name="sglang-test-policy-cluster-separate", - ) - yield cluster - try: - cluster.shutdown() - except Exception as e: - print(f"Error during policy_cluster_separate shutdown: {e}") - - -def get_generation_cluster_separate(num_gpus_per_node: int = 2) -> RayVirtualCluster: - """Create a virtual cluster for the SGLangGeneration policy.""" - return RayVirtualCluster( - bundle_ct_per_node_list=[num_gpus_per_node], - use_gpus=True, - max_colocated_worker_groups=1, - num_gpus_per_node=num_gpus_per_node, - name="sglang-test-generation-cluster-separate", - ) - - -# ============================================================================= -# Basic Configuration Tests -# ============================================================================= - - -@pytest.mark.sglang -@pytest.mark.timeout(120) -def test_sglang_missing_required_config_key(cluster, tokenizer): - """Test that an error is raised when a required config key is missing.""" - # SGLang requires sglang_cfg to be present - incomplete_config = deepcopy(basic_sglang_test_config) - incomplete_config = configure_sglang_config(incomplete_config, tokenizer) - del incomplete_config["sglang_cfg"] - - with pytest.raises((KeyError, ValueError, AssertionError, TypeError)): - SGLangGeneration(cluster, incomplete_config) - - -@pytest.mark.sglang -def test_sglang_top_p_top_k_validation(cluster, tokenizer): - """Test that top_p and top_k values are accepted by SGLang. - - Note: SGLang may have different validation thresholds than vLLM. - This test verifies that reasonable sampling parameters are accepted. - """ - # Test that reasonable top_p and top_k values are accepted - config = deepcopy(basic_sglang_test_config) - config["top_p"] = 0.95 - config["top_k"] = 50 - config = configure_sglang_config(config, tokenizer) - - policy = None - try: - policy = SGLangGeneration(cluster, config) - print("Successfully initialized with top_p=0.95 and top_k=50") - except Exception as e: - pytest.fail(f"Should not raise error with reasonable sampling params: {e}") - finally: - if policy: - policy.shutdown() - gc.collect() - torch.cuda.empty_cache() - - -# ============================================================================= -# Basic Generation Tests -# ============================================================================= - - -@pytest.mark.sglang -@pytest.mark.timeout(180) -def test_sglang_policy_generation(policy, test_input_data, tokenizer): - """Test SGLang policy generation capabilities.""" - print("Testing SGLang generation...") - outputs = policy.generate(test_input_data) - - # Validate outputs format - assert "output_ids" in outputs, "output_ids not found in generation output" - assert "logprobs" in outputs, "logprobs not found in generation output" - assert "generation_lengths" in outputs, ( - "generation_lengths not found in generation output" - ) - assert "unpadded_sequence_lengths" in outputs, ( - "unpadded_sequence_lengths not found in generation output" - ) - - # Validate outputs shape and content - assert outputs["output_ids"].shape[0] == len(test_input_data["input_ids"]), ( - "Wrong batch size in output" - ) - assert outputs["generation_lengths"].shape[0] == len( - test_input_data["input_ids"] - ), "Wrong batch size in generation_lengths" - - # Decode and check outputs - generated_sequences = outputs["output_ids"] - generated_texts = tokenizer.batch_decode( - generated_sequences, skip_special_tokens=True - ) - - print(f"Generated texts: {generated_texts}") - - # All texts should have a non-zero length - assert all(len(text) > 0 for text in generated_texts), ( - "Some generated texts are empty" - ) - - -@pytest.mark.sglang -def test_sglang_worker_seed_behavior(cluster, tokenizer): - """ - Test that different workers generate different outputs for identical prompts due to different seeds. - This ensures proper randomization across distributed workers for diverse exploration in RLHF. - - Key: Use gpus_per_server=1 to create 2 independent SGLang servers (each with its own seed), - rather than 1 server with TP=2. - """ - from nemo_rl.algorithms.grpo import refit_policy_generation - from nemo_rl.models.policy.lm_policy import Policy - - unique_prompts = [ - "Hello, my name is", - "The capital of France is", - ] - - # Create a batch where each prompt appears twice - # When sharded, different workers will get the same prompt - duplicated_prompts = unique_prompts + unique_prompts - - # Tokenize prompts - encodings = tokenizer( - duplicated_prompts, - padding="max_length", - max_length=20, - truncation=True, - return_tensors="pt", - padding_side="right", - ) - - input_lengths = encodings["attention_mask"].sum(dim=1).to(torch.int32) - - # Create input data dictionary - duplicated_batch = BatchedDataDict( - { - "input_ids": encodings["input_ids"], - "input_lengths": input_lengths, - } - ) - - # Test with gpus_per_server=1 to create 2 independent servers with different seeds - print("Creating SGLang policy with gpus_per_server=1 (2 independent servers)...") - sglang_config = deepcopy(basic_sglang_test_config) - # Use gpus_per_server=1 to create 2 independent SGLang servers - sglang_config["sglang_cfg"]["gpus_per_server"] = 1 - sglang_config = configure_sglang_config(sglang_config, tokenizer) - - policy = SGLangGeneration(cluster, sglang_config) - policy.finish_generation() - - dtensor_config = deepcopy(basic_dtensor_test_config) - dtensor_config["dtensor_cfg"]["tensor_parallel_size"] = 1 # Match gpus_per_server - lm_policy = Policy(cluster, dtensor_config, tokenizer) - - state_dict_info = lm_policy.prepare_refit_info() - policy.prepare_refit_info(state_dict_info) - - print("Refitting SGLang policy...") - refit_policy_generation(lm_policy, policy, sglang_config["colocated"]["enabled"]) - - try: - # Generate with duplicated prompts - print("Running generation with duplicated prompts...") - outputs = policy.generate(duplicated_batch, greedy=False) - - # Decode the generated sequences - gen_texts = tokenizer.batch_decode( - outputs["output_ids"], skip_special_tokens=True - ) - - print(f"Generated texts with duplicated prompts: {gen_texts}") - - # Check if the duplicated prompts generated different texts - # The first half and second half should be different due to different worker seeds - first_half = gen_texts[: len(unique_prompts)] - second_half = gen_texts[len(unique_prompts) :] - - print(f"First worker outputs: {first_half}") - print(f"Second worker outputs: {second_half}") - - # At least one of the pairs should be different due to different seeds - assert first_half != second_half, ( - "Different workers should generate different outputs for identical prompts due to different seeds" - ) - - finally: - # Clean up resources - if "policy" in locals() and hasattr(policy, "shutdown"): - policy.shutdown() - if "lm_policy" in locals() and hasattr(lm_policy, "shutdown"): - lm_policy.shutdown() - - # Force garbage collection - gc.collect() - torch.cuda.empty_cache() - - -@pytest.mark.sglang -def test_sglang_policy_tensor_parallel(cluster, tokenizer): - """Test SGLang policy with tensor parallelism > 1 (gpus_per_server=2).""" - # Configure with gpus_per_server=2 for tensor parallelism - tp_config = deepcopy(basic_sglang_test_config) - tp_config = configure_sglang_config(tp_config, tokenizer) - tp_config["sglang_cfg"]["gpus_per_server"] = 2 # TP=2 - - sglang_policy = None - try: - sglang_policy = SGLangGeneration(cluster, tp_config) - - # Create simple test input - test_prompts = ["Hello, my name is", "The capital of France is"] - encodings = tokenizer( - test_prompts, - padding="max_length", - max_length=10, - truncation=True, - return_tensors="pt", - padding_side="right", - ) - - test_input_data = BatchedDataDict( - { - "input_ids": encodings["input_ids"], - "input_lengths": encodings["attention_mask"].sum(dim=1).to(torch.int32), - } - ) - - # Test generation with tensor parallelism - outputs = sglang_policy.generate(test_input_data) - - sglang_policy.finish_generation() - sglang_policy.prepare_for_generation() - - # Test generation again after cache reset - outputs = sglang_policy.generate(test_input_data) - - assert "output_ids" in outputs, "output_ids not found in generation output" - assert outputs["output_ids"].shape[0] == 2, "Wrong batch size in output" - - # Decode and check output - generated_text = tokenizer.decode( - outputs["output_ids"][0], skip_special_tokens=True - ) - print(f"Generated text with TP=2: {generated_text}") - assert len(generated_text) > 0, "Generated text is empty" - - finally: - # Clean up resources - if sglang_policy: - sglang_policy.shutdown() - gc.collect() - torch.cuda.empty_cache() - - -@pytest.mark.sglang -def test_sglang_generate_text(cluster, tokenizer): - """Test that SGLang can generate coherent text. - - Note: SGLang doesn't have a generate_text method like vLLM, - so we use generate + tokenizer decode to verify text generation. - """ - # Prepare test data - test_prompts = [ - "Hello, my name is", - "The capital of France is", - ] - - encodings = tokenizer( - test_prompts, - padding="max_length", - max_length=10, - truncation=True, - return_tensors="pt", - padding_side="right", - ) - - test_input_data = BatchedDataDict( - { - "input_ids": encodings["input_ids"], - "input_lengths": encodings["attention_mask"].sum(dim=1).to(torch.int32), - } - ) - - # Create SGLang config with gpus_per_server=2 (using tensor parallelism) - sglang_config = deepcopy(basic_sglang_test_config) - sglang_config["sglang_cfg"]["gpus_per_server"] = 2 - sglang_config = configure_sglang_config(sglang_config, tokenizer, is_eval=True) - - # Ensure correct model - assert sglang_config["model_name"] == "Qwen/Qwen3-0.6B", ( - "Model name should be Qwen/Qwen3-0.6B to get expected output" - ) - - sglang_generation = None - try: - # Create SGLang generation - sglang_generation = SGLangGeneration(cluster, sglang_config) - - # Generate with greedy decoding for deterministic output - output = sglang_generation.generate(test_input_data, greedy=True) - - # Decode generated text - generated_texts = tokenizer.batch_decode( - output["output_ids"], skip_special_tokens=True - ) - - print(f"Generated texts: {generated_texts}") - - # Verify we got non-empty text for each prompt - for i, text in enumerate(generated_texts): - assert len(text) > len(test_prompts[i]), ( - f"Generated text should be longer than input prompt: {text}" - ) - # Verify the generated text starts with or contains the prompt - print(f"Prompt: {test_prompts[i]} -> Generated: {text}") - - finally: - # Clean up - if sglang_generation: - sglang_generation.shutdown() - gc.collect() - torch.cuda.empty_cache() - - -def _wait_for_sglang_http_server_spinup(base_url: str): - """Wait for the SGLang HTTP server to be ready.""" - import time - - import requests - - max_wait = 60 # 60 seconds max wait - start = time.time() - while time.time() - start < max_wait: - try: - response = requests.get(f"{base_url}/health_generate", timeout=5) - if response.status_code == 200: - return - except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): - pass - time.sleep(1) - raise TimeoutError(f"SGLang server at {base_url} did not start within {max_wait}s") - - -@pytest.mark.sglang -def test_sglang_http_server(cluster, tokenizer): - """Test that SGLang HTTP server works with direct API calls. - - SGLang exposes a /generate endpoint that accepts input_ids and sampling_params. - This test verifies we can make direct HTTP requests to the SGLang server. - """ - import requests - - # Create SGLang config - sglang_config = deepcopy(basic_sglang_test_config) - sglang_config = configure_sglang_config(sglang_config, tokenizer, is_eval=True) - - # Ensure correct model for reproducible output - assert sglang_config["model_name"] == "Qwen/Qwen3-0.6B", ( - "Model name should be Qwen/Qwen3-0.6B to get expected output" - ) - - sglang_generation = None - try: - # Create SGLang generation (this starts the servers) - sglang_generation = SGLangGeneration(cluster, sglang_config) - - # Get server URLs - base_urls = sglang_generation.get_sglang_server_urls() - print(f"SGLang server URLs: {base_urls}") - assert len(base_urls) >= 1, "Should have at least one SGLang server" - - # Wait for server to be ready - _wait_for_sglang_http_server_spinup(base_urls[0]) - - # Prepare input - tokenize "count to 5" - test_prompt = "count to 5" - input_ids = tokenizer.encode(test_prompt, add_special_tokens=True) - - # Build request payload for SGLang /generate endpoint - payload = { - "input_ids": input_ids, - "sampling_params": { - "temperature": 0.0, # Greedy for determinism - "top_p": 1.0, - "max_new_tokens": 5, - }, - "return_logprob": True, - } - - # Make request to SGLang server - response = requests.post( - url=f"{base_urls[0]}/generate", - json=payload, - headers={"Content-Type": "application/json"}, - timeout=30, - ) - actual_result = response.json() - print(f"SGLang response: {actual_result}") - - # Verify response structure - assert response.status_code == 200, f"Expected 200, got {response.status_code}" - assert "meta_info" in actual_result, "Response should contain meta_info" - - meta_info = actual_result["meta_info"] - assert "output_token_logprobs" in meta_info, ( - "meta_info should contain output_token_logprobs" - ) - - # Verify we got some generated tokens - output_token_logprobs = meta_info["output_token_logprobs"] - assert len(output_token_logprobs) > 0, ( - "Should have generated at least one token" - ) - - # Each entry should be [logprob, token_id] - first_token_info = output_token_logprobs[0] - assert len(first_token_info) >= 2, ( - "Each token info should have logprob and token_id" - ) - - logprob = first_token_info[0] - token_id = first_token_info[1] - assert isinstance(logprob, float), "Logprob should be a float" - assert isinstance(token_id, int), "Token ID should be an int" - - print(f"First generated token: id={token_id}, logprob={logprob}") - - # Decode the generated tokens to verify text output - generated_token_ids = [item[1] for item in output_token_logprobs] - generated_text = tokenizer.decode(generated_token_ids, skip_special_tokens=True) - print(f"Generated text: {generated_text}") - - finally: - # Clean up - if sglang_generation: - sglang_generation.shutdown() - gc.collect() - torch.cuda.empty_cache() - - -@pytest.mark.sglang -@pytest.mark.timeout(180) -def test_sglang_non_divisible_batch_handling(policy): - """Test that SGLang generation handles non divisible input batches correctly.""" - empty_batch = BatchedDataDict( - { - "input_ids": torch.zeros((1, 1), dtype=torch.long), - "input_lengths": torch.ones(1, dtype=torch.long), - } - ) - - outputs = policy.generate(empty_batch) - - required_keys = [ - "output_ids", - "logprobs", - "generation_lengths", - "unpadded_sequence_lengths", - ] - assert all(key in outputs for key in required_keys), ( - "Missing required output fields" - ) - assert all(outputs[key].shape[0] == 1 for key in required_keys), ( - "Output tensors should have batch dimension of 1" - ) - - -# ============================================================================= -# Policy Integration Tests -# ============================================================================= - - -@pytest.mark.sglang -@pytest.mark.timeout(300) -def test_sglang_generation_with_hf_training_colocated(cluster, tokenizer): - """Test that DTensor policy can work together with colocated SGLang policy.""" - from nemo_rl.algorithms.grpo import refit_policy_generation - from nemo_rl.models.policy.lm_policy import Policy - - sglang_config = deepcopy(basic_sglang_test_config) - sglang_config = configure_sglang_config(sglang_config, tokenizer) - - dtensor_config = deepcopy(basic_dtensor_test_config) - dtensor_config["train_global_batch_size"] = 4 - dtensor_config["dtensor_cfg"]["_v2"] = ( - True # Use DTensorPolicyWorkerV2 for stream_weights_via_http - ) - - sglang_policy = None - lm_policy = None - - try: - print("Creating SGLang policy...") - sglang_policy = SGLangGeneration(cluster, sglang_config) - sglang_policy.finish_generation() - - print("Creating DTensor policy...") - lm_policy = Policy(cluster, dtensor_config, tokenizer) - - print("Preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() - sglang_policy.prepare_refit_info(state_dict_info) - - print("Refitting SGLang policy...") - refit_policy_generation( - lm_policy, sglang_policy, sglang_config["colocated"]["enabled"] - ) - - # Test generation - test_prompts = ["Hello, my name is", "The capital of France is"] - encodings = tokenizer( - test_prompts, - padding="max_length", - max_length=20, - truncation=True, - return_tensors="pt", - padding_side="right", - ) - test_input_data = BatchedDataDict( - { - "input_ids": encodings["input_ids"], - "input_lengths": encodings["attention_mask"].sum(dim=1).to(torch.int32), - } - ) - - outputs = sglang_policy.generate(test_input_data, greedy=True) - assert "output_ids" in outputs, "output_ids not found in generation output" - - generated_texts = tokenizer.batch_decode( - outputs["output_ids"], skip_special_tokens=True - ) - print(f"Generated texts: {generated_texts}") - - finally: - if sglang_policy: - sglang_policy.shutdown() - if lm_policy and hasattr(lm_policy, "shutdown"): - lm_policy.shutdown() - - -@pytest.mark.skip(reason="Non-colocated mode not implemented for SGLang") -@pytest.mark.timeout(300) -@pytest.mark.sglang -def test_sglang_generation_with_hf_training_non_colocated( - policy_cluster_separate, tokenizer -): - """Test that DTensor policy can work together with non-colocated SGLang policy.""" - from nemo_rl.algorithms.grpo import refit_policy_generation - from nemo_rl.models.policy.lm_policy import Policy - - generation_cluster_separate = get_generation_cluster_separate(2) - - sglang_config = deepcopy(basic_sglang_test_config) - sglang_config = configure_sglang_config(sglang_config, tokenizer) - sglang_config["colocated"]["enabled"] = False - - dtensor_config = deepcopy(basic_dtensor_test_config) - dtensor_config["generation"]["colocated"]["enabled"] = False - dtensor_config["train_global_batch_size"] = 4 - dtensor_config["dtensor_cfg"]["_v2"] = ( - True # Use DTensorPolicyWorkerV2 for stream_weights_via_http - ) - - sglang_policy = None - lm_policy = None - - try: - print("Creating SGLang policy...") - sglang_policy = SGLangGeneration(generation_cluster_separate, sglang_config) - sglang_policy.finish_generation() - - print("Creating DTensor policy...") - lm_policy = Policy(policy_cluster_separate, dtensor_config, tokenizer) - - # Initialize collective communication - ip, port = policy_cluster_separate.get_master_address_and_port() - train_world_size = policy_cluster_separate.world_size() - inference_world_size = generation_cluster_separate.world_size() - world_size = train_world_size + inference_world_size - - futures_train = lm_policy.init_collective( - ip, port, world_size=world_size, train_world_size=train_world_size - ) - futures_inference = sglang_policy.init_collective( - ip, port, world_size=world_size, train_world_size=train_world_size - ) - ray.get(futures_train + futures_inference) - - # Prepare refit info - state_dict_info = lm_policy.prepare_refit_info() - sglang_policy.prepare_refit_info(state_dict_info) - - print("Refitting SGLang policy...") - refit_policy_generation(lm_policy, sglang_policy, False) - - # Test generation - test_prompts = ["Hello, my name is", "The capital of France is"] - encodings = tokenizer( - test_prompts, - padding="max_length", - max_length=20, - truncation=True, - return_tensors="pt", - padding_side="right", - ) - test_input_data = BatchedDataDict( - { - "input_ids": encodings["input_ids"], - "input_lengths": encodings["attention_mask"].sum(dim=1).to(torch.int32), - } - ) - - outputs = sglang_policy.generate(test_input_data, greedy=True) - assert "output_ids" in outputs, "output_ids not found in generation output" - - finally: - if sglang_policy: - sglang_policy.shutdown() - if lm_policy and hasattr(lm_policy, "shutdown"): - lm_policy.shutdown() - try: - generation_cluster_separate.shutdown() - except Exception as e: - print(f"Error during generation_cluster_separate shutdown: {e}") - - -@pytest.mark.sglang -@pytest.mark.timeout(180) -def test_sglang_weight_update_and_prefix_cache_reset(cluster, tokenizer): - """Test that the SGLang prefix cache is correctly reset when weights change.""" - from nemo_rl.models.policy.lm_policy import Policy - - sglang_config = deepcopy(basic_sglang_test_config) - sglang_config = configure_sglang_config(sglang_config, tokenizer, is_eval=True) - - dtensor_config = basic_dtensor_test_config - - sglang_policy = None - lm_policy = None - - try: - print("Creating DTensor policy...") - lm_policy = Policy(cluster, dtensor_config, tokenizer) - - print("Creating SGLang policy...") - sglang_policy = SGLangGeneration(cluster, sglang_config) - - print("Preparing refit info...") - state_dict_info = lm_policy.prepare_refit_info() - sglang_policy.prepare_refit_info(state_dict_info) - - # Prepare input data - text = "Answer the question. What is 2+2?" - test_prompt = [text, text] - encodings = tokenizer( - test_prompt, - padding=True, - return_tensors="pt", - padding_side="right", - ) - input_ids = encodings["input_ids"] - input_lengths = encodings["attention_mask"].sum(dim=1).to(torch.int32) - test_input_data = BatchedDataDict( - {"input_ids": input_ids, "input_lengths": input_lengths} - ) - - print("Running Generation 1 (Initial)...") - sglang_policy.prepare_for_generation() - outputs1 = sglang_policy.generate(test_input_data, greedy=True) - logprob1 = outputs1["logprobs"][0, input_lengths[0]].item() - print(f"Logprob of first generated token (Run 1): {logprob1}") - - print("Adding noise to weights in HF policy...") - ray.get( - [ - worker._add_noise_to_weights.remote() - for worker in lm_policy.worker_group.workers - ] - ) - - print("Updating SGLang weights from DTensor policy via HTTP...") - # Get SGLang server URL to GPU UUID mapping - sglang_url_to_gpu_uuids = sglang_policy.get_sglang_url_to_gpu_uuids() - print(f"SGLang URL to GPU UUIDs: {sglang_url_to_gpu_uuids}") - - # Stream weights via HTTP (CUDA IPC) - ray.get(lm_policy.stream_weights_via_http(sglang_url_to_gpu_uuids)) - - print("Running Generation 2 (Weights Updated)...") - outputs2 = sglang_policy.generate(test_input_data, greedy=True) - logprob2 = outputs2["logprobs"][0, input_lengths[0]].item() - print(f"Logprob of first generated token (Run 2): {logprob2}") - assert logprob2 != logprob1, "Logprobs should be different after weight update." - - print("Resetting SGLang prefix cache...") - sglang_policy.finish_generation() - sglang_policy.prepare_for_generation() - - print("Running Generation 3 (Cache Reset)...") - outputs3 = sglang_policy.generate(test_input_data, greedy=True) - logprob3 = outputs3["logprobs"][0, input_lengths[0]].item() - print(f"Logprob of first generated token (Run 3): {logprob3}") - - print("Prefix cache reset verified successfully.") - - finally: - print("Cleaning up resources...") - if sglang_policy: - sglang_policy.shutdown() - if lm_policy: - lm_policy.shutdown() - gc.collect() - torch.cuda.empty_cache() diff --git a/tests/unit/models/generation/test_sglang_utils.py b/tests/unit/models/generation/test_sglang_utils.py deleted file mode 100644 index 396530edd7..0000000000 --- a/tests/unit/models/generation/test_sglang_utils.py +++ /dev/null @@ -1,49 +0,0 @@ -# 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. - -"""Unit tests for SGLang utilities. - -These tests verify that the SGLang utilities work as expected. -""" - -import pytest - -from nemo_rl.models.generation.sglang.utils import AsyncLoopThread - - -def test_async_loop_thread_run_returns_result(): - loop_thread = AsyncLoopThread() - - async def sample(): - return 42 - - try: - assert loop_thread.run(sample()) == 42 - finally: - loop_thread.shutdown() - - -def test_async_loop_thread_run_when_stopped_raises(): - loop_thread = AsyncLoopThread() - loop_thread.shutdown() - - async def sample(): - return 1 - - with pytest.raises(RuntimeError, match="Event loop is not running"): - coro = sample() - try: - loop_thread.run(coro) - finally: - coro.close() diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 1166c5ceea..956739e827 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -875,9 +875,11 @@ def test_topk_post_processor_no_packing(self, mock_topk, mock_tp_rank, mock_tp_g mock_data_dict = MagicMock() mock_data_dict.__getitem__ = MagicMock( - side_effect=lambda key: torch.tensor([[1, 2, 3, 4, 5]]) - if key == "input_ids" - else torch.tensor([5]) + side_effect=lambda key: ( + torch.tensor([[1, 2, 3, 4, 5]]) + if key == "input_ids" + else torch.tensor([5]) + ) ) mock_topk_vals = torch.randn(1, 5, k) @@ -919,9 +921,11 @@ def test_topk_post_processor_with_packing( mock_data_dict = MagicMock() mock_data_dict.__getitem__ = MagicMock( - side_effect=lambda key: torch.tensor([[1, 2, 3, 4, 5, 0, 0, 0]]) - if key == "input_ids" - else torch.tensor([5]) + side_effect=lambda key: ( + torch.tensor([[1, 2, 3, 4, 5, 0, 0, 0]]) + if key == "input_ids" + else torch.tensor([5]) + ) ) mock_topk_vals = torch.randn(1, 8, k) @@ -967,9 +971,9 @@ def test_topk_cp_without_packing_raises( mock_data_dict = MagicMock() mock_data_dict.__getitem__ = MagicMock( - side_effect=lambda key: torch.tensor([[1, 2, 3]]) - if key == "input_ids" - else torch.tensor([3]) + side_effect=lambda key: ( + torch.tensor([[1, 2, 3]]) if key == "input_ids" else torch.tensor([3]) + ) ) mock_topk.return_value = ( @@ -1015,9 +1019,11 @@ def test_topk_cp_with_packing_single_sequence( mock_data_dict = MagicMock() mock_data_dict.__getitem__ = MagicMock( - side_effect=lambda key: torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) - if key == "input_ids" - else torch.tensor([8]) + side_effect=lambda key: ( + torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + if key == "input_ids" + else torch.tensor([8]) + ) ) # distributed_vocab_topk returns local (CP-sharded) results @@ -1081,9 +1087,11 @@ def test_topk_cp_with_packing_multiple_sequences( mock_data_dict = MagicMock() mock_data_dict.__getitem__ = MagicMock( - side_effect=lambda key: torch.zeros(2, unpacked_seqlen, dtype=torch.long) - if key == "input_ids" - else torch.tensor([seq1_len, seq2_len]) + side_effect=lambda key: ( + torch.zeros(2, unpacked_seqlen, dtype=torch.long) + if key == "input_ids" + else torch.tensor([seq1_len, seq2_len]) + ) ) # distributed_vocab_topk returns local (CP-sharded) results diff --git a/tests/unit/models/policy/test_policy_utils.py b/tests/unit/models/policy/test_policy_utils.py deleted file mode 100644 index 5fbcf8e86e..0000000000 --- a/tests/unit/models/policy/test_policy_utils.py +++ /dev/null @@ -1,224 +0,0 @@ -# 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. - -import sys -import types -from unittest.mock import MagicMock - -import pytest -import requests -import torch - -from nemo_rl.models.policy import utils as policy_utils - -pytestmark = pytest.mark.sglang - - -def test_setup_ipc_gather_group_returns_none_when_dist_uninit(monkeypatch): - monkeypatch.setattr(policy_utils.dist, "is_initialized", lambda: False) - - group, src, ranks = policy_utils._setup_ipc_gather_group( - rank=0, - current_device_uuid="uuid0", - sglang_gpu_uuids=["uuid0"], - sglang_url_to_gpu_uuids={"http://sglang": ["uuid0"]}, - ) - - assert group is None - assert src is None - assert ranks is None - - -def test_setup_ipc_gather_group_selects_matching_ranks(monkeypatch): - all_ranks = ["uuid0", "uuid1", "uuid2", "uuid3"] - - monkeypatch.setattr(policy_utils.dist, "is_initialized", lambda: True) - monkeypatch.setattr(policy_utils.dist, "get_world_size", lambda: 4) - monkeypatch.setattr(policy_utils.dist, "get_rank", lambda: 1) - - def fake_all_gather_object(output_list, _value): - for idx, item in enumerate(all_ranks): - output_list[idx] = item - - monkeypatch.setattr(policy_utils.dist, "all_gather_object", fake_all_gather_object) - - group, src, ranks = policy_utils._setup_ipc_gather_group( - rank=1, - current_device_uuid="uuid1", - sglang_gpu_uuids=["uuid1", "uuid3"], - sglang_url_to_gpu_uuids={"http://sglang": ["uuid1", "uuid3"]}, - ) - - assert group is None - assert src == 1 - assert ranks == [1, 3] - - -def test_gather_ipc_handlers_returns_filtered_on_src(monkeypatch): - handlers = ["h0", "h1", "h2", "h3"] - monkeypatch.setattr(policy_utils.dist, "is_initialized", lambda: True) - monkeypatch.setattr(policy_utils.dist, "get_world_size", lambda: 4) - - def fake_all_gather_object(output_list, _value): - for idx, item in enumerate(handlers): - output_list[idx] = item - - monkeypatch.setattr(policy_utils.dist, "all_gather_object", fake_all_gather_object) - - gathered = policy_utils._gather_ipc_handlers( - serialized_handler="h1", - gather_group=None, - gather_src=0, - rank=0, - matching_ranks=[0, 2], - ) - - assert gathered == ["h0", "h2"] - - -def test_gather_ipc_handlers_non_src_returns_none(monkeypatch): - monkeypatch.setattr(policy_utils.dist, "is_initialized", lambda: True) - monkeypatch.setattr(policy_utils.dist, "get_world_size", lambda: 2) - monkeypatch.setattr(policy_utils.dist, "all_gather_object", lambda *_args: None) - - gathered = policy_utils._gather_ipc_handlers( - serialized_handler="h1", - gather_group=None, - gather_src=0, - rank=1, - matching_ranks=[0, 1], - ) - - assert gathered is None - - -def test_send_tensor_to_sglang_http_error(monkeypatch): - response = MagicMock() - response.raise_for_status.side_effect = requests.exceptions.HTTPError("boom") - response.status_code = 500 - response.text = "error" - monkeypatch.setattr( - policy_utils.requests, "post", lambda *_args, **_kwargs: response - ) - - with pytest.raises(RuntimeError, match="Failed to send tensor 'w'"): - policy_utils._send_tensor_to_sglang( - url="http://sglang/update", - tensor_name="w", - gathered_handlers=["h0"], - shape=torch.Size([1]), - dtype="torch.float32", - ) - - -def test_send_tensor_to_sglang_generic_error(monkeypatch): - def raise_error(*_args, **_kwargs): - raise RuntimeError("network down") - - monkeypatch.setattr(policy_utils.requests, "post", raise_error) - - with pytest.raises(RuntimeError, match="Failed to send tensor 'w'"): - policy_utils._send_tensor_to_sglang( - url="http://sglang/update", - tensor_name="w", - gathered_handlers=["h0"], - shape=torch.Size([1]), - dtype="torch.float32", - ) - - -def test_stream_weights_via_http_impl_no_matching_url(monkeypatch): - monkeypatch.setattr(policy_utils.torch.cuda, "empty_cache", lambda: None) - - with pytest.raises(RuntimeError, match="No matching SGLang server"): - policy_utils.stream_weights_via_http_impl( - params_generator=iter([]), - sglang_url_to_gpu_uuids={"http://sglang": ["uuid0"]}, - rank=0, - worker_name="worker", - current_device_uuid="uuid1", - ) - - -def test_stream_weights_via_http_impl_sends_tensors(monkeypatch): - def params_generator(): - yield "w1", torch.tensor([1.0]) - yield "w2", torch.tensor([2.0]) - - dummy_module = types.ModuleType( - "nemo_rl.models.generation.sglang.sglang_copied_utils" - ) - - class DummySerializer: - @staticmethod - def serialize(*_args, **_kwargs): - return "handler" - - dummy_module.MultiprocessingSerializer = DummySerializer - monkeypatch.setitem( - sys.modules, - "nemo_rl.models.generation.sglang.sglang_copied_utils", - dummy_module, - ) - monkeypatch.setattr(policy_utils.torch.cuda, "empty_cache", lambda: None) - monkeypatch.setattr( - policy_utils.torch.cuda, - "current_stream", - lambda: types.SimpleNamespace(synchronize=lambda: None), - ) - monkeypatch.setattr( - policy_utils.torch.Tensor, "cuda", lambda self: self, raising=False - ) - - send_calls = [] - - def fake_send_tensor_to_sglang( - url, name, gathered_handlers, shape, dtype, flush_cache=False - ): - send_calls.append( - { - "url": url, - "name": name, - "handlers": gathered_handlers, - "shape": shape, - "dtype": dtype, - "flush_cache": flush_cache, - } - ) - - monkeypatch.setattr( - policy_utils, - "_setup_ipc_gather_group", - lambda *_args, **_kwargs: (None, 0, [0]), - ) - monkeypatch.setattr( - policy_utils, "_gather_ipc_handlers", lambda *_args, **_kwargs: ["handler"] - ) - monkeypatch.setattr( - policy_utils, "_send_tensor_to_sglang", fake_send_tensor_to_sglang - ) - - policy_utils.stream_weights_via_http_impl( - params_generator=params_generator(), - sglang_url_to_gpu_uuids={ - "http://sglang-a": ["uuid0"], - "http://sglang-b": ["uuid0"], - }, - rank=0, - worker_name="worker", - current_device_uuid="uuid0", - ) - - assert [call["name"] for call in send_calls] == ["w1", "w2"] - assert all(call["handlers"] == ["handler"] for call in send_calls)