diff --git a/examples/configs/grpo_math_1B_sglang.yaml b/examples/configs/grpo_math_1B_sglang.yaml index 66ee4f35bcb..55849666671 100644 --- a/examples/configs/grpo_math_1B_sglang.yaml +++ b/examples/configs/grpo_math_1B_sglang.yaml @@ -31,10 +31,14 @@ policy: rollout_health_check_timeout: 60 rollout_health_check_first_wait: 60 # Weight precision for rollout/refit. scheme=bf16 (default) sends BF16 - # HF tensors; scheme=mxfp8 boots SGLang from an MXFP8 HF checkpoint and - # quantizes refit tensors online (see SglangQuantizationConfig). + # HF tensors; scheme=mxfp8/nvfp4 boots SGLang from a matching quantized + # HF checkpoint and quantizes refit tensors online. quantization: scheme: bf16 + extra_high_precision_layers_hf: [] + num_layers_at_start_in_bf16: 0 + num_layers_at_end_in_bf16: 0 + modules_to_not_convert: [] sglang_server_config: needs_offload: true cpu_weight_backup: true diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index b0ff19e58e8..194b7c29425 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1266,22 +1266,33 @@ def init_vllm_then_policy(): if "model_path" not in generation_config["sglang_cfg"]: generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"] - # If MXFP8 is requested, ensure SGLang boots from an MXFP8 HF - # checkpoint. This must happen before ``init_sglang`` so the engine - # loads quantized weights. + # Quantized online refit requires SGLang to boot from a checkpoint + # with the exact same tensor layout and high-precision exclusions. + # Resolve it before ``init_sglang`` starts the engines. 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, - ) + from nemo_rl.models.generation.sglang.quantization_utils import ( + ensure_sglang_quantized_checkpoint, + get_sglang_quantization_scheme, + validate_sglang_quantized_refit_backend, + ) - mxfp8_path = ensure_mxfp8_checkpoint( + sglang_quantization_scheme = get_sglang_quantization_scheme( + sglang_quantization_cfg + ) + validate_sglang_quantized_refit_backend( + scheme=sglang_quantization_scheme, + use_megatron=bool( + policy_config.get("megatron_cfg", {}).get("enabled", False) + ), + ) + generation_config["sglang_cfg"]["model_path"] = ( + ensure_sglang_quantized_checkpoint( model_path=generation_config["sglang_cfg"]["model_path"], - quantization_cfg=sglang_quantization_cfg, + quantization_config=sglang_quantization_cfg, ) - generation_config["sglang_cfg"]["model_path"] = mxfp8_path + ) policy_generation, policy = initialize_generation_with_policy( init_generation_fn=init_sglang, diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index f47a8472d2f..1d17f25cf3e 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -634,6 +634,31 @@ def initialize_generation_with_policy( if "model_path" not in generation_config["sglang_cfg"]: generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"] + sglang_quantization_cfg = ( + generation_config["sglang_cfg"].get("quantization") or {} + ) + from nemo_rl.models.generation.sglang.quantization_utils import ( + ensure_sglang_quantized_checkpoint, + get_sglang_quantization_scheme, + validate_sglang_quantized_refit_backend, + ) + + sglang_quantization_scheme = get_sglang_quantization_scheme( + sglang_quantization_cfg + ) + validate_sglang_quantized_refit_backend( + scheme=sglang_quantization_scheme, + use_megatron=bool( + policy_config.get("megatron_cfg", {}).get("enabled", False) + ), + ) + generation_config["sglang_cfg"]["model_path"] = ( + ensure_sglang_quantized_checkpoint( + model_path=generation_config["sglang_cfg"]["model_path"], + quantization_config=sglang_quantization_cfg, + ) + ) + policy_generation, policy, value_model = initialize_generation_with_policy( init_generation_fn=init_sglang, generation_name="SGLang", diff --git a/nemo_rl/models/generation/sglang/config.py b/nemo_rl/models/generation/sglang/config.py index d08f5385bd3..a5f4ecf49ae 100644 --- a/nemo_rl/models/generation/sglang/config.py +++ b/nemo_rl/models/generation/sglang/config.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, NotRequired, TypedDict +from typing import Any, Literal, NotRequired, TypedDict from nemo_rl.models.generation.interfaces import GenerationConfig @@ -21,15 +21,18 @@ 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="mxfp8"`` or ``scheme="nvfp4"`` to boot SGLang from the + corresponding quantized HF checkpoint and quantize HF tensors during + online refit. High-precision exclusions are shared by offline conversion + and online refit. """ - scheme: str # "bf16" | "mxfp8" - weight_block_size: list[int] - scale_fmt: str + scheme: Literal["bf16", "mxfp8", "nvfp4"] + # HF module-name substrings that the checkpoint loader and refit both skip. modules_to_not_convert: list[str] + # Additional HF weight-name substrings to keep in high precision. extra_high_precision_layers_hf: list[str] + # Number of decoder layers at each edge to keep in high precision. num_layers_at_start_in_bf16: int num_layers_at_end_in_bf16: int converted_model_path: str @@ -97,7 +100,7 @@ class SglangSpecificArgs(TypedDict): 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. + # Weight precision and quantized-checkpoint conversion/refit knobs. quantization: NotRequired[SglangQuantizationConfig] sglang_router_config: SGLangRouterConfig diff --git a/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py b/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py index 8108c606193..7c55ed04c14 100644 --- a/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py +++ b/nemo_rl/models/generation/sglang/mxfp8_quantization_core.py @@ -113,21 +113,16 @@ def should_quantize( 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. + Uses flashinfer's swizzle-free MXFP8 kernel + (``is_sf_swizzled_layout=False``). """ 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)." + "installed in the current actor environment. In NeMo-RL this is " + "normally provided by the pinned `mcore` or `sglang` extras." ) from e weight = weight.contiguous() @@ -150,45 +145,3 @@ def source_fp8_to_mxfp8_scale_u8( 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), - ) diff --git a/nemo_rl/models/generation/sglang/mxfp8_setup.py b/nemo_rl/models/generation/sglang/mxfp8_setup.py index daad90f7880..35e3028644c 100644 --- a/nemo_rl/models/generation/sglang/mxfp8_setup.py +++ b/nemo_rl/models/generation/sglang/mxfp8_setup.py @@ -46,10 +46,16 @@ source_fp8_to_mxfp8_scale_u8, strip_weight_suffix, ) +from nemo_rl.models.generation.sglang.quantization_utils import ( + build_dynamic_skip_substrings, + expand_sglang_atomic_high_precision_substrings, + get_dynamic_high_precision_substrings, + validate_checkpoint_high_precision_layout, +) logger = logging.getLogger(__name__) -CONVERTER_VERSION: str = "1" +CONVERTER_VERSION: str = "4" class _ConversionResult: @@ -112,6 +118,7 @@ def _load_source_scale_u8( scale_u8 = None return scale_fp32, scale_u8, scale_key + assert scale_u8 is not None 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] @@ -127,11 +134,8 @@ def _process_file( *, 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, ...], + skip_weight_substrings: tuple[str, ...], source_scale_index: dict[str, str], ) -> None: import safetensors @@ -148,24 +152,6 @@ def _process_file( 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"): @@ -174,7 +160,7 @@ def _process_file( should_quant = should_quantize( key, tensor, - skip_weight_substrings=dynamic_skip_substrings, + skip_weight_substrings=skip_weight_substrings, allow_source_fp8=source_is_block_fp8_ue8m0, ) @@ -251,6 +237,7 @@ def convert_mxfp8( num_layers_at_start_in_bf16: int = 0, num_layers_at_end_in_bf16: int = 0, extra_high_precision_layers_hf: tuple[str, ...] = (), + modules_to_not_convert: tuple[str, ...] = (), ) -> None: """Convert an HF safetensors checkpoint to MXFP8 with UE8M0 scales. @@ -265,7 +252,22 @@ def convert_mxfp8( 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"]) + num_hidden_layers = _get_num_hidden_layers(cfg, config_path=config_path) + selection_config = { + "extra_high_precision_layers_hf": extra_high_precision_layers_hf, + "modules_to_not_convert": modules_to_not_convert, + "num_layers_at_start_in_bf16": num_layers_at_start_in_bf16, + "num_layers_at_end_in_bf16": num_layers_at_end_in_bf16, + } + configured_high_precision_substrings = get_dynamic_high_precision_substrings( + quantization_config=selection_config, + num_hidden_layers=num_hidden_layers, + ) + skip_weight_substrings = build_dynamic_skip_substrings( + quantization_config=selection_config, + num_hidden_layers=num_hidden_layers, + static_skip_substrings=SKIP_WEIGHT_SUBSTRINGS, + ) if is_source_block_fp8_ue8m0_checkpoint(cfg): source_is_block_fp8_ue8m0 = True elif is_bf16_source_checkpoint(cfg): @@ -289,6 +291,16 @@ def convert_mxfp8( index_path = os.path.join(input_path, "model.safetensors.index.json") with open(index_path) as f: weight_map = json.load(f)["weight_map"] + expanded_skip_weight_substrings = expand_sglang_atomic_high_precision_substrings( + weight_names=weight_map, + skip_weight_substrings=skip_weight_substrings, + ) + concrete_atomic_modules = tuple( + substring + for substring in expanded_skip_weight_substrings + if substring not in skip_weight_substrings + ) + skip_weight_substrings = expanded_skip_weight_substrings safetensors_files = sorted(set(weight_map.values())) source_scale_index: dict[str, str] = {} if source_is_block_fp8_ue8m0: @@ -299,6 +311,9 @@ def convert_mxfp8( } result_collector = _ConversionResult() + result_collector.modules_to_not_convert.extend( + (*configured_high_precision_substrings, *concrete_atomic_modules) + ) for filename in safetensors_files: logger.info(f"[mxfp8] Processing {filename}") _process_file( @@ -307,11 +322,8 @@ def convert_mxfp8( 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, + skip_weight_substrings=skip_weight_substrings, source_scale_index=source_scale_index, ) gc.collect() @@ -350,12 +362,151 @@ def _read_source_config(model_dir: str) -> dict[str, Any]: return json.load(f) +def _checkpoint_weight_names(model_dir: str) -> tuple[str, ...]: + """Read tensor names from an indexed or single-file safetensors checkpoint.""" + index_path = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.isfile(index_path): + with open(index_path) as file: + index = json.load(file) + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Missing non-empty weight_map in {index_path}.") + if any( + not isinstance(name, str) or not isinstance(filename, str) + for name, filename in weight_map.items() + ): + raise ValueError(f"{index_path} weight_map must map strings to strings.") + return tuple(weight_map) + + import safetensors + + names: set[str] = set() + shard_names = sorted( + filename + for filename in os.listdir(model_dir) + if filename.endswith(".safetensors") + and os.path.isfile(os.path.join(model_dir, filename)) + ) + if not shard_names: + raise ValueError(f"No safetensors weights found in {model_dir}.") + for filename in shard_names: + with safetensors.safe_open( + os.path.join(model_dir, filename), + framework="pt", + device="cpu", + ) as file: + for name in file.keys(): + if name in names: + raise ValueError( + f"Duplicate tensor {name!r} in checkpoint {model_dir!r}." + ) + names.add(name) + return tuple(names) + + +def _get_num_hidden_layers(cfg: dict[str, Any], *, config_path: str) -> int: + num_hidden_layers = cfg.get("num_hidden_layers") + text_config = cfg.get("text_config") + if num_hidden_layers is None and isinstance(text_config, dict): + num_hidden_layers = text_config.get("num_hidden_layers") + if ( + isinstance(num_hidden_layers, bool) + or not isinstance(num_hidden_layers, int) + or num_hidden_layers <= 0 + ): + raise ValueError( + f"{config_path} must define a positive integer num_hidden_layers." + ) + return num_hidden_layers + + +def _validated_conversion_options( + quantization_cfg: dict[str, Any], + *, + num_hidden_layers: int, +) -> tuple[tuple[str, ...], tuple[str, ...], int, int]: + """Validate and normalize the shared offline/online selection options.""" + get_dynamic_high_precision_substrings( + quantization_config=quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + + extra_value = quantization_cfg.get("extra_high_precision_layers_hf") + modules_value = quantization_cfg.get("modules_to_not_convert") + start_value = quantization_cfg.get("num_layers_at_start_in_bf16") + end_value = quantization_cfg.get("num_layers_at_end_in_bf16") + return ( + () if extra_value is None else tuple(item.strip() for item in extra_value), + () if modules_value is None else tuple(item.strip() for item in modules_value), + 0 if start_value is None else start_value, + 0 if end_value is None else end_value, + ) + + +def _merge_checkpoint_modules_to_not_convert( + *, + checkpoint_path: str, + quantization_cfg: dict[str, Any], +) -> None: + """Merge the checkpoint's concrete ignore list into the refit config.""" + checkpoint_cfg = _read_source_config(checkpoint_path) + checkpoint_quantization_cfg = checkpoint_cfg.get("quantization_config") + if not isinstance(checkpoint_quantization_cfg, dict): + raise ValueError(f"{checkpoint_path}/config.json has no quantization_config.") + + num_hidden_layers = _get_num_hidden_layers( + checkpoint_cfg, + config_path=os.path.join(checkpoint_path, "config.json"), + ) + _, user_modules, _, _ = _validated_conversion_options( + quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + requested_high_precision = get_dynamic_high_precision_substrings( + quantization_config=quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + checkpoint_modules = get_dynamic_high_precision_substrings( + quantization_config={ + "modules_to_not_convert": checkpoint_quantization_cfg.get( + "modules_to_not_convert" + ) + }, + num_hidden_layers=num_hidden_layers, + ) + checkpoint_weight_names = _checkpoint_weight_names(checkpoint_path) + configured_high_precision = tuple( + dict.fromkeys((*requested_high_precision, *checkpoint_modules)) + ) + expanded_high_precision = expand_sglang_atomic_high_precision_substrings( + weight_names=checkpoint_weight_names, + skip_weight_substrings=configured_high_precision, + ) + validate_checkpoint_high_precision_layout( + checkpoint_path=checkpoint_path, + scheme="MXFP8", + weight_names=checkpoint_weight_names, + high_precision_substrings=expanded_high_precision, + quantized_companion_suffixes=(SOURCE_FP8_SCALE_KEY_SUFFIX,), + ) + concrete_atomic_modules = tuple( + substring + for substring in expanded_high_precision + if substring not in configured_high_precision + ) + quantization_cfg["modules_to_not_convert"] = list( + dict.fromkeys((*user_modules, *checkpoint_modules, *concrete_atomic_modules)) + ) + + 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", + # Legacy fingerprint inputs. No longer settable via + # SglangQuantizationConfig; kept so existing cache dirs stay valid. "weight_block_size", "scale_fmt", ) @@ -391,24 +542,20 @@ 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. - """ + """Return an MXFP8 checkpoint path and synchronize its ignore policy.""" if is_existing_mxfp8_checkpoint(model_path): + _merge_checkpoint_modules_to_not_convert( + checkpoint_path=model_path, + quantization_cfg=quantization_cfg, + ) return model_path converted = quantization_cfg.get("converted_model_path") if converted and is_existing_mxfp8_checkpoint(converted): + _merge_checkpoint_modules_to_not_convert( + checkpoint_path=converted, + quantization_cfg=quantization_cfg, + ) return converted cache_root = ( @@ -423,23 +570,33 @@ def ensure_mxfp8_checkpoint( ) if is_existing_mxfp8_checkpoint(save_dir): + _merge_checkpoint_modules_to_not_convert( + checkpoint_path=save_dir, + quantization_cfg=quantization_cfg, + ) return save_dir - extra_high_precision_layers_hf = tuple( - quantization_cfg.get("extra_high_precision_layers_hf", ()) or () + source_cfg = _read_source_config(model_path) + num_hidden_layers = _get_num_hidden_layers( + source_cfg, + config_path=os.path.join(model_path, "config.json"), ) - 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 + ( + extra_high_precision_layers_hf, + modules_to_not_convert, + num_layers_at_start_in_bf16, + num_layers_at_end_in_bf16, + ) = _validated_conversion_options( + quantization_cfg, + num_hidden_layers=num_hidden_layers, ) 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})" + f"extra_hp={extra_high_precision_layers_hf}, " + f"modules_to_not_convert={modules_to_not_convert})" ) convert_mxfp8( model_dir=model_path, @@ -447,5 +604,10 @@ def ensure_mxfp8_checkpoint( 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, + modules_to_not_convert=modules_to_not_convert, + ) + _merge_checkpoint_modules_to_not_convert( + checkpoint_path=save_dir, + quantization_cfg=quantization_cfg, ) return save_dir diff --git a/nemo_rl/models/generation/sglang/nvfp4_quantization_core.py b/nemo_rl/models/generation/sglang/nvfp4_quantization_core.py new file mode 100644 index 00000000000..790793f0d22 --- /dev/null +++ b/nemo_rl/models/generation/sglang/nvfp4_quantization_core.py @@ -0,0 +1,429 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transformer Engine NVFP4 conversion rules for SGLang rollout weights.""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + +from nemo_rl.models.generation.sglang.quantization_utils import ( + HF_MOE_EXPERT_NAME_MARKERS, +) + +FP4_E2M1_MAX = 6.0 +FP8_E4M3_MAX = 448 +NVFP4_GROUP_SIZE = 16 +TE_NVFP4_ROW_ALIGNMENT = 16 + +NVFP4_QUANTIZATION_CONFIG: dict[str, Any] = { + "group_size": NVFP4_GROUP_SIZE, + "ignore": [], + "kv_cache_scheme": {"dynamic": False, "num_bits": 8, "type": "float"}, + "quant_algo": "NVFP4", + "quant_method": "modelopt", +} + +EXPERT_WEIGHT_SUFFIXES = ( + ".w1.weight", + ".w2.weight", + ".w3.weight", + ".gate_proj.weight", + ".up_proj.weight", + ".down_proj.weight", + ".gate_up_proj.weight", +) +GATED_PAIR_SUFFIXES = { + ".gate_proj.weight": "gate", + ".up_proj.weight": "up", + ".w1.weight": "gate", + ".w3.weight": "up", +} +GATED_PAIR_FAMILIES = ( + (".gate_proj.weight", ".up_proj.weight"), + (".w1.weight", ".w3.weight"), +) + + +def strip_weight_suffix(weight_key: str) -> str: + """Remove the trailing ``.weight`` from an HF parameter name.""" + if not weight_key.endswith(".weight"): + raise ValueError(f"Expected key ending with '.weight', got: {weight_key}") + return weight_key[: -len(".weight")] + + +def is_nvfp4_quantization_config(config: dict[str, Any] | None) -> bool: + """Return whether ``config`` describes the ModelOpt NVFP4 layout.""" + if not isinstance(config, dict): + return False + return ( + config.get("quant_algo") == "NVFP4" + and config.get("group_size") == NVFP4_GROUP_SIZE + ) + + +def is_bf16_source_checkpoint(config: dict[str, Any]) -> bool: + """Return whether an HF checkpoint contains ordinary high-precision weights.""" + quantization_config = ( + config.get("quantization_config", {}) if isinstance(config, dict) else {} + ) + if not isinstance(quantization_config, dict) or not quantization_config: + return True + return quantization_config.get("quant_method") in (None, "", "bf16") + + +def is_moe_expert_weight_name(name: str) -> bool: + """Return whether an HF name denotes a quantizable MoE expert GEMM.""" + return ( + name.endswith(".weight") + and any(marker in name for marker in HF_MOE_EXPERT_NAME_MARKERS) + and any(name.endswith(suffix) for suffix in EXPERT_WEIGHT_SUFFIXES) + ) + + +def should_quantize_nvfp4( + name: str, + weight: torch.Tensor, + *, + skip_weight_substrings: tuple[str, ...] = (), +) -> bool: + """Return whether an HF tensor should be emitted in NVFP4 format.""" + if any(substring in name for substring in skip_weight_substrings): + return False + if not is_moe_expert_weight_name(name): + return False + if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return False + if weight.dim() != 2: + return False + if weight.shape[-1] % NVFP4_GROUP_SIZE != 0: + raise ValueError( + f"Last dim {weight.shape[-1]} must be divisible by " + f"{NVFP4_GROUP_SIZE} for NVFP4 ({name})." + ) + return True + + +def split_gated_pair_name(name: str) -> tuple[str | None, str | None]: + """Return the shared base and role for an HF gate/up weight name.""" + for suffix, role in GATED_PAIR_SUFFIXES.items(): + if name.endswith(suffix): + return name[: -len(suffix)], role + return None, None + + +def gated_pair_names(name: str) -> tuple[str, str] | None: + """Return the gate/up names for the pair containing ``name``.""" + for gate_suffix, up_suffix in GATED_PAIR_FAMILIES: + if name.endswith(gate_suffix): + base = name[: -len(gate_suffix)] + return base + gate_suffix, base + up_suffix + if name.endswith(up_suffix): + base = name[: -len(up_suffix)] + return base + gate_suffix, base + up_suffix + return None + + +def should_skip_nvfp4_gated_pair( + name: str, + *, + skip_weight_substrings: tuple[str, ...], +) -> bool: + """Keep both halves high precision when either gate/up name is excluded.""" + pair_names = gated_pair_names(name) + if pair_names is None: + return False + return any( + substring in pair_name + for substring in skip_weight_substrings + for pair_name in pair_names + ) + + +def nvfp4_weight_e4m3_max() -> int: + """Return TE's configured global E4M3 bound for NVFP4 weights.""" + use_4over6 = os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() + use_256 = os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all").strip().lower() + if use_4over6 in ("weights", "all") and use_256 in ("weights", "all"): + return 256 + return FP8_E4M3_MAX + + +def nvfp4_global_encode_scale_te( + global_amax: torch.Tensor, + *, + nvfp4_e4m3_max: int, +) -> torch.Tensor: + """Reproduce Transformer Engine's per-tensor NVFP4 encode scale.""" + fp4_max = torch.tensor( + FP4_E2M1_MAX, + device=global_amax.device, + dtype=torch.float32, + ) + fp8_max = torch.tensor( + float(nvfp4_e4m3_max), + device=global_amax.device, + dtype=torch.float32, + ) + global_encode_scale = torch.div( + fp8_max * fp4_max, + global_amax.to(torch.float32), + ) + global_encode_scale = torch.minimum( + global_encode_scale, + torch.tensor( + torch.finfo(torch.float32).max, + device=global_amax.device, + dtype=torch.float32, + ), + ) + return torch.where( + global_encode_scale == 0.0, + torch.ones_like(global_encode_scale), + global_encode_scale, + ) + + +def nvfp4_global_decode_scale_te( + global_amax: torch.Tensor, + *, + nvfp4_e4m3_max: int, +) -> torch.Tensor: + """Return the FP32 scale stored as ``weight_scale_2`` by ModelOpt.""" + return torch.reciprocal( + nvfp4_global_encode_scale_te( + global_amax, + nvfp4_e4m3_max=nvfp4_e4m3_max, + ) + ) + + +def _nvfp4_4over6_enabled() -> bool: + return os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ( + "weights", + "all", + ) + + +def _pad_rows_for_te_quantizer(weight: torch.Tensor) -> torch.Tensor: + pad_rows = (-weight.shape[0]) % TE_NVFP4_ROW_ALIGNMENT + if pad_rows == 0: + return weight + padding = torch.zeros( + (pad_rows, weight.shape[1]), + device=weight.device, + dtype=weight.dtype, + ) + return torch.cat((weight, padding), dim=0) + + +def _make_nvfp4_quantizer() -> Any: + # Transformer Engine is heavy and optional outside GPU actor environments. + try: + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + except ImportError as error: + raise ImportError( + "Transformer Engine >=2.17 is required for NVFP4 weight " + "quantization. Install NeMo-RL with the `mcore` extra." + ) from error + + try: + return NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_amax_reduction=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=False, + nvfp4_use_4over6=_nvfp4_4over6_enabled(), + nvfp4_e4m3_max=nvfp4_weight_e4m3_max(), + nvfp4_4over6_err_mode=os.getenv( + "NVTE_NVFP4_4OVER6_ERR_MODE", + "MAE", + ) + .strip() + .upper(), + with_random_sign_mask=False, + ) + except TypeError as error: + raise RuntimeError( + "The installed Transformer Engine does not expose the NVFP4 " + "quantizer API required by NeMo-RL. Upgrade to version 2.17 or newer." + ) from error + + +def nvfp4_quantize_2d( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize one 2D weight into ModelOpt's NVFP4 checkpoint layout.""" + if weight.dim() != 2: + raise ValueError(f"nvfp4_quantize_2d expects a 2D tensor, got {weight.dim()}D.") + if weight.shape[1] % NVFP4_GROUP_SIZE != 0: + raise ValueError( + f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {weight.shape[1]}." + ) + + weight = weight.contiguous() + num_rows, num_cols = weight.shape + nvfp4_e4m3_max = nvfp4_weight_e4m3_max() + quantized = _make_nvfp4_quantizer().quantize(_pad_rows_for_te_quantizer(weight)) + + rowwise_data = getattr(quantized, "_rowwise_data", None) + rowwise_scale_inv = getattr(quantized, "_rowwise_scale_inv", None) + amax_rowwise = getattr(quantized, "_amax_rowwise", None) + if rowwise_data is None or rowwise_scale_inv is None or amax_rowwise is None: + raise RuntimeError( + "Transformer Engine returned an incomplete rowwise NVFP4 tensor." + ) + + qweight = rowwise_data[:num_rows, : num_cols // 2].contiguous() + block_scale = rowwise_scale_inv[ + :num_rows, + : num_cols // NVFP4_GROUP_SIZE, + ].contiguous() + weight_scale_2 = nvfp4_global_decode_scale_te( + amax_rowwise.reshape(-1)[0], + nvfp4_e4m3_max=nvfp4_e4m3_max, + ) + return ( + qweight, + block_scale.view(torch.float8_e4m3fn), + weight_scale_2, + ) + + +def _contiguous_pair_view( + first: torch.Tensor, + second: torch.Tensor, +) -> torch.Tensor | None: + if not first.is_contiguous() or not second.is_contiguous(): + return None + if ( + first.device != second.device + or first.dtype != second.dtype + or first.stride() != second.stride() + ): + return None + if first.untyped_storage().data_ptr() != second.untyped_storage().data_ptr(): + return None + if first.storage_offset() + first.numel() != second.storage_offset(): + return None + try: + return first.as_strided( + (first.shape[0] + second.shape[0], first.shape[1]), + first.stride(), + first.storage_offset(), + ) + except RuntimeError: + return None + + +def nvfp4_quantize_2d_pair( + first: torch.Tensor, + second: torch.Tensor, +) -> tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], +]: + """Quantize a gate/up pair together so both weights share global scale.""" + if first.dim() != 2 or second.dim() != 2: + raise ValueError("nvfp4_quantize_2d_pair expects two 2D tensors.") + if first.shape[1] != second.shape[1]: + raise ValueError( + "NVFP4 paired quantization requires matching K dimensions, got " + f"{first.shape[1]} and {second.shape[1]}." + ) + if first.dtype != second.dtype: + raise ValueError( + "NVFP4 paired quantization requires matching dtypes, got " + f"{first.dtype} and {second.dtype}." + ) + if first.device != second.device: + raise ValueError( + "NVFP4 paired quantization requires weights on the same device, got " + f"{first.device} and {second.device}." + ) + + first_rows = first.shape[0] + combined = _contiguous_pair_view(first, second) + if combined is None: + combined = torch.cat((first.contiguous(), second.contiguous()), dim=0) + combined_qweight, combined_block_scale, weight_scale_2 = nvfp4_quantize_2d(combined) + first_result = ( + combined_qweight[:first_rows].contiguous(), + combined_block_scale[:first_rows].contiguous(), + weight_scale_2.clone(), + ) + second_result = ( + combined_qweight[first_rows:].contiguous(), + combined_block_scale[first_rows:].contiguous(), + weight_scale_2.clone(), + ) + return first_result, second_result + + +def quantize_nvfp4( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize a 2D expert weight.""" + if weight.dim() != 2: + raise ValueError( + f"Unsupported weight rank {weight.dim()} for NVFP4 quantization." + ) + return nvfp4_quantize_2d(weight) + + +def quantize_nvfp4_pair( + first: torch.Tensor, + second: torch.Tensor, +) -> tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], +]: + """Quantize a 2D gate/up pair under one shared global scale.""" + if first.dim() != 2 or second.dim() != 2: + raise ValueError( + "NVFP4 paired quantization requires 2D weights, got " + f"{first.dim()}D and {second.dim()}D." + ) + return nvfp4_quantize_2d_pair(first, second) + + +def nvfp4_quantized_entries( + name: str, + quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + include_input_scale: bool, +) -> list[tuple[str, torch.Tensor]]: + """Expand one quantized weight into the ModelOpt tensor names.""" + qweight, block_scale, weight_scale_2 = quantized + base_name = strip_weight_suffix(name) + result = [ + (name, qweight), + (f"{base_name}.weight_scale", block_scale), + (f"{base_name}.weight_scale_2", weight_scale_2), + ] + if include_input_scale: + result.append( + ( + f"{base_name}.input_scale", + torch.ones_like(weight_scale_2, dtype=torch.float32), + ) + ) + return result diff --git a/nemo_rl/models/generation/sglang/nvfp4_setup.py b/nemo_rl/models/generation/sglang/nvfp4_setup.py new file mode 100644 index 00000000000..2c5cad15524 --- /dev/null +++ b/nemo_rl/models/generation/sglang/nvfp4_setup.py @@ -0,0 +1,761 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline HF-to-NVFP4 conversion and SGLang startup checkpoint setup.""" + +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.nvfp4_quantization_core import ( + NVFP4_GROUP_SIZE, + NVFP4_QUANTIZATION_CONFIG, + is_bf16_source_checkpoint, + is_moe_expert_weight_name, + is_nvfp4_quantization_config, + nvfp4_quantized_entries, + quantize_nvfp4, + quantize_nvfp4_pair, + should_quantize_nvfp4, + split_gated_pair_name, + strip_weight_suffix, +) +from nemo_rl.models.generation.sglang.quantization_utils import ( + build_dynamic_skip_substrings, + expand_sglang_atomic_high_precision_substrings, + get_dynamic_high_precision_substrings, + get_hf_moe_expert_container, + validate_checkpoint_high_precision_layout, +) + +logger = logging.getLogger(__name__) + +CONVERTER_VERSION = "1" +_NVFP4_ENV_KEYS = ( + "NVTE_NVFP4_4OVER6", + "NVTE_NVFP4_4OVER6_E4M3_USE_256", + "NVTE_NVFP4_4OVER6_ERR_MODE", + "NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", +) + + +class _ConversionResult: + def __init__(self) -> None: + self.weight_map: dict[str, str] = {} + self.total_size = 0 + self.modules_to_not_convert: set[str] = set() + + def add_tensor(self, filename: str, key: str, tensor: torch.Tensor) -> None: + existing = self.weight_map.get(key) + if existing is not None: + raise ValueError( + f"Duplicate output tensor {key!r} in {existing!r} and {filename!r}." + ) + self.weight_map[key] = filename + self.total_size += tensor.numel() * tensor.element_size() + + +def _read_json(path: str) -> dict[str, Any]: + with open(path) as file: + value = json.load(file) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object in {path}.") + return value + + +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 {} + return _read_json(config_path) + + +def _num_hidden_layers(config: dict[str, Any]) -> int: + value = config.get("num_hidden_layers") + if value is None and isinstance(config.get("text_config"), dict): + value = config["text_config"].get("num_hidden_layers") + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError( + "The source config must define a positive num_hidden_layers, either " + "at the top level or under text_config." + ) + return value + + +def _source_weight_map(model_dir: str) -> dict[str, str]: + """Return a validated tensor-to-shard map for indexed or single-file HF models.""" + import safetensors + + index_path = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.isfile(index_path): + index = _read_json(index_path) + raw_weight_map = index.get("weight_map") + if not isinstance(raw_weight_map, dict) or not raw_weight_map: + raise ValueError(f"Missing non-empty weight_map in {index_path}.") + weight_map = { + str(key): str(filename) for key, filename in raw_weight_map.items() + } + shard_names = sorted(set(weight_map.values())) + else: + shard_names = sorted( + filename + for filename in os.listdir(model_dir) + if filename.endswith(".safetensors") + and os.path.isfile(os.path.join(model_dir, filename)) + ) + if not shard_names: + raise ValueError(f"No safetensors weights found in {model_dir}.") + weight_map: dict[str, str] = {} + + actual_weight_map: dict[str, str] = {} + for filename in shard_names: + shard_path = os.path.join(model_dir, filename) + if not os.path.isfile(shard_path): + raise FileNotFoundError(f"Checkpoint shard not found: {shard_path}") + with safetensors.safe_open( + shard_path, + framework="pt", + device="cpu", + ) as file: + for key in file.keys(): + existing = actual_weight_map.get(key) + if existing is not None: + raise ValueError( + f"Duplicate source tensor {key!r} in {existing!r} and " + f"{filename!r}." + ) + actual_weight_map[key] = filename + + if os.path.isfile(index_path): + if weight_map != actual_weight_map: + missing = sorted(set(weight_map) - set(actual_weight_map)) + unindexed = sorted(set(actual_weight_map) - set(weight_map)) + misplaced = sorted( + key + for key in set(weight_map) & set(actual_weight_map) + if weight_map[key] != actual_weight_map[key] + ) + raise ValueError( + "The safetensors index does not match the checkpoint shards: " + f"missing={missing[:5]}, unindexed={unindexed[:5]}, " + f"misplaced={misplaced[:5]}." + ) + return weight_map + return actual_weight_map + + +def _natural_key(value: str) -> list[Any]: + return [ + int(token) if token.isdigit() else token + for token in re.findall(r"\d+|\D+", value) + ] + + +def _expanded_ignore_modules(modules: set[str]) -> list[str]: + expanded = set(modules) + for module in tuple(expanded): + if module.endswith((".q_proj", ".k_proj", ".v_proj")): + expanded.add(module.rsplit(".", 1)[0] + ".qkv_proj") + container = get_hf_moe_expert_container(module + ".weight") + if container is not None: + expanded.add(container) + + return sorted(expanded, key=_natural_key) + + +def _build_conversion_policy( + *, + weight_map: dict[str, str], + skip_weight_substrings: tuple[str, ...], +) -> tuple[ + set[str], + set[str], + dict[str, tuple[str, str]], +]: + """Select quantized expert tensors and complete gate/up pairs atomically.""" + skipped_containers: set[str] = set() + expert_keys_by_container: dict[str, list[str]] = {} + + for key in weight_map: + if not key.endswith(".weight"): + continue + container = get_hf_moe_expert_container(key) + if container is None: + continue + expert_keys_by_container.setdefault(container, []).append(key) + if not is_moe_expert_weight_name(key) or any( + pattern in key for pattern in skip_weight_substrings + ): + skipped_containers.add(container) + + quantized_keys = { + key + for container, keys in expert_keys_by_container.items() + if container not in skipped_containers + for key in keys + } + + pair_roles: dict[str, dict[str, str]] = {} + for key in sorted(quantized_keys): + pair_base, pair_role = split_gated_pair_name(key) + if pair_base is None or pair_role is None: + continue + roles = pair_roles.setdefault(pair_base, {}) + if pair_role in roles: + raise ValueError( + f"Duplicate NVFP4 {pair_role} weight for pair {pair_base!r}: " + f"{roles[pair_role]!r}, {key!r}." + ) + roles[pair_role] = key + + incomplete = { + base: sorted(roles) + for base, roles in pair_roles.items() + if set(roles) != {"gate", "up"} + } + if incomplete: + raise ValueError( + "NVFP4 gate/up weights must be converted together; incomplete " + f"checkpoint pairs: {incomplete}." + ) + + pair_by_key: dict[str, tuple[str, str]] = {} + for base, roles in pair_roles.items(): + gate_key = roles["gate"] + up_key = roles["up"] + pair_by_key[gate_key] = (base, up_key) + pair_by_key[up_key] = (base, gate_key) + return quantized_keys, skipped_containers, pair_by_key + + +def _load_tensor( + model_dir: str, + filename: str, + key: str, + *, + device: str, +) -> torch.Tensor: + import safetensors + + with safetensors.safe_open( + os.path.join(model_dir, filename), + framework="pt", + device=device, + ) as file: + return file.get_tensor(key) + + +def _add_output_entries( + destination: dict[str, torch.Tensor], + entries: list[tuple[str, torch.Tensor]], +) -> None: + for key, tensor in entries: + if key in destination: + raise ValueError(f"Duplicate NVFP4 output tensor {key!r}.") + destination[key] = tensor + + +def _convert_shards( + *, + model_dir: str, + save_dir: str, + weight_map: dict[str, str], + quantized_keys: set[str], + pair_by_key: dict[str, tuple[str, str]], + skipped_containers: set[str], + device: str, +) -> _ConversionResult: + import safetensors + import safetensors.torch + + result = _ConversionResult() + result.modules_to_not_convert.update(skipped_containers) + deferred: dict[str, dict[str, torch.Tensor]] = {} + processed_pairs: set[str] = set() + processed_source_keys: set[str] = set() + + for filename in sorted(set(weight_map.values())): + logger.info("[nvfp4] Processing %s", filename) + shard_path = os.path.join(model_dir, filename) + with safetensors.safe_open( + shard_path, + framework="pt", + device=device, + ) as file: + weights = {key: file.get_tensor(key) for key in file.keys()} + + output = deferred.pop(filename, {}) + for key, tensor in weights.items(): + if key in processed_source_keys: + continue + + if key in quantized_keys: + if not should_quantize_nvfp4( + key, + tensor, + skip_weight_substrings=(), + ): + raise ValueError( + f"Expert weight {key!r} is not a supported NVFP4 source " + f"(dtype={tensor.dtype}, shape={tuple(tensor.shape)})." + ) + + pair_info = pair_by_key.get(key) + if pair_info is None: + _add_output_entries( + output, + nvfp4_quantized_entries( + key, + quantize_nvfp4(tensor), + include_input_scale=True, + ), + ) + processed_source_keys.add(key) + continue + + pair_base, counterpart_key = pair_info + if pair_base in processed_pairs: + processed_source_keys.add(key) + continue + counterpart_filename = weight_map[counterpart_key] + counterpart = ( + weights[counterpart_key] + if counterpart_filename == filename + else _load_tensor( + model_dir, + counterpart_filename, + counterpart_key, + device=device, + ) + ) + if not should_quantize_nvfp4( + counterpart_key, + counterpart, + skip_weight_substrings=(), + ): + raise ValueError( + f"Expert weight {counterpart_key!r} is not a supported " + f"NVFP4 source (dtype={counterpart.dtype}, " + f"shape={tuple(counterpart.shape)})." + ) + + current_base, current_role = split_gated_pair_name(key) + if current_base != pair_base or current_role is None: + raise RuntimeError(f"Invalid cached NVFP4 pair for {key!r}.") + if current_role == "gate": + gate_key, gate_weight = key, tensor + up_key, up_weight = counterpart_key, counterpart + else: + gate_key, gate_weight = counterpart_key, counterpart + up_key, up_weight = key, tensor + + gate_output, up_output = quantize_nvfp4_pair( + gate_weight, + up_weight, + ) + pair_entries = ( + ( + gate_key, + nvfp4_quantized_entries( + gate_key, + gate_output, + include_input_scale=True, + ), + ), + ( + up_key, + nvfp4_quantized_entries( + up_key, + up_output, + include_input_scale=True, + ), + ), + ) + for source_key, entries in pair_entries: + target_filename = weight_map[source_key] + target = ( + output + if target_filename == filename + else deferred.setdefault(target_filename, {}) + ) + _add_output_entries(target, entries) + + processed_pairs.add(pair_base) + processed_source_keys.update((gate_key, up_key)) + continue + + output[key] = tensor + processed_source_keys.add(key) + if key.endswith(".weight"): + result.modules_to_not_convert.add(strip_weight_suffix(key)) + + output_path = os.path.join(save_dir, filename) + safetensors.torch.save_file( + output, + output_path, + metadata={"format": "pt"}, + ) + for key, tensor in output.items(): + result.add_tensor(filename, key, tensor) + + del output, weights + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + if deferred: + raise RuntimeError( + "Deferred NVFP4 pair outputs were not written because their target " + f"shards were missing: {sorted(deferred)}." + ) + if processed_source_keys != set(weight_map): + missing = sorted(set(weight_map) - processed_source_keys) + raise RuntimeError( + f"NVFP4 conversion did not process source keys: {missing[:10]}." + ) + return result + + +def _copy_checkpoint_metadata(model_dir: str, save_dir: str) -> None: + for filename in os.listdir(model_dir): + source = os.path.join(model_dir, filename) + if os.path.isdir(source) or filename.endswith(".safetensors"): + continue + shutil.copyfile(source, os.path.join(save_dir, filename)) + + +def _write_output_metadata( + *, + model_dir: str, + save_dir: str, + config: dict[str, Any], + result: _ConversionResult, +) -> list[str]: + ignore = _expanded_ignore_modules(result.modules_to_not_convert) + quantization_config = dict(NVFP4_QUANTIZATION_CONFIG) + source_quantization = config.get("quantization_config") + if isinstance(source_quantization, dict): + kv_cache_scheme = source_quantization.get("kv_cache_scheme") + if isinstance(kv_cache_scheme, dict): + quantization_config["kv_cache_scheme"] = kv_cache_scheme + quantization_config["ignore"] = ignore + config["quantization_config"] = quantization_config + + with open(os.path.join(save_dir, "config.json"), "w") as file: + json.dump(config, file, indent=2) + + source_hf_quant_path = os.path.join(model_dir, "hf_quant_config.json") + hf_quant_config = ( + _read_json(source_hf_quant_path) + if os.path.isfile(source_hf_quant_path) + else {"producer": {"name": "modelopt", "version": "nemo-rl"}} + ) + raw_hf_quantization = hf_quant_config.get("quantization") + hf_quantization: dict[str, Any] = ( + dict(raw_hf_quantization) if isinstance(raw_hf_quantization, dict) else {} + ) + hf_quant_config["quantization"] = hf_quantization + hf_quantization.update( + { + "exclude_modules": ignore, + "group_size": NVFP4_GROUP_SIZE, + "kv_cache_quant_algo": "FP8", + "quant_algo": "NVFP4", + } + ) + with open(os.path.join(save_dir, "hf_quant_config.json"), "w") as file: + json.dump(hf_quant_config, file, indent=2) + + index = { + "weight_map": result.weight_map, + "metadata": {"total_size": result.total_size}, + } + with open( + os.path.join(save_dir, "model.safetensors.index.json"), + "w", + ) as file: + json.dump(index, file, indent=2) + return ignore + + +def convert_nvfp4( + 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, ...] = (), + modules_to_not_convert: tuple[str, ...] = (), +) -> list[str]: + """Convert an HF checkpoint to ModelOpt NVFP4 for SGLang. + + Megatron name conversion is intentionally absent: both this converter and + the live refit iterator operate on finalized HF tensor names. + """ + if device.startswith("cuda") and not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available, cannot run NVFP4 quantization.") + + input_path = os.path.abspath(model_dir) + output_path = os.path.abspath(save_dir) + config_path = os.path.join(input_path, "config.json") + if not os.path.isfile(config_path): + raise FileNotFoundError(f"HF config not found: {config_path}") + config = _read_json(config_path) + if not is_bf16_source_checkpoint(config): + raise ValueError( + "NVFP4 conversion only supports BF16/FP16/FP32 source checkpoints." + ) + + num_hidden_layers = _num_hidden_layers(config) + policy_config = { + "extra_high_precision_layers_hf": list(extra_high_precision_layers_hf), + "modules_to_not_convert": list(modules_to_not_convert), + "num_layers_at_start_in_bf16": num_layers_at_start_in_bf16, + "num_layers_at_end_in_bf16": num_layers_at_end_in_bf16, + } + skip_weight_substrings = build_dynamic_skip_substrings( + quantization_config=policy_config, + num_hidden_layers=num_hidden_layers, + ) + weight_map = _source_weight_map(input_path) + quantized_keys, skipped_containers, pair_by_key = _build_conversion_policy( + weight_map=weight_map, + skip_weight_substrings=skip_weight_substrings, + ) + if not quantized_keys: + raise ValueError( + "No eligible MoE expert weights remain after applying NVFP4 skip rules." + ) + + os.makedirs(output_path, exist_ok=True) + _copy_checkpoint_metadata(input_path, output_path) + result = _convert_shards( + model_dir=input_path, + save_dir=output_path, + weight_map=weight_map, + quantized_keys=quantized_keys, + pair_by_key=pair_by_key, + skipped_containers=skipped_containers, + device=device, + ) + ignore = _write_output_metadata( + model_dir=input_path, + save_dir=output_path, + config=config, + result=result, + ) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return ignore + + +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", + ) + payload = { + "config": {key: quantization_cfg.get(key) for key in relevant_keys}, + "environment": {key: os.environ.get(key) for key in _NVFP4_ENV_KEYS}, + } + return hashlib.sha1( + json.dumps(payload, sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:12] + + +def _validated_conversion_options( + quantization_cfg: dict[str, Any], + *, + num_hidden_layers: int, +) -> tuple[tuple[str, ...], tuple[str, ...], int, int]: + """Validate and normalize the shared offline/online selection options.""" + build_dynamic_skip_substrings( + quantization_config=quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + extra_value = quantization_cfg.get("extra_high_precision_layers_hf") + modules_value = quantization_cfg.get("modules_to_not_convert") + start_value = quantization_cfg.get("num_layers_at_start_in_bf16") + end_value = quantization_cfg.get("num_layers_at_end_in_bf16") + return ( + () if extra_value is None else tuple(item.strip() for item in extra_value), + () if modules_value is None else tuple(item.strip() for item in modules_value), + 0 if start_value is None else start_value, + 0 if end_value is None else end_value, + ) + + +def _hash_qualified_save_dir( + *, + model_dir: str, + cache_root: str, + quantization_cfg: dict[str, Any], +) -> str: + absolute_model = os.path.abspath(model_dir) + source_config = _read_source_config(model_dir) + source_fingerprint = hashlib.sha1( + json.dumps(source_config, sort_keys=True, default=str).encode("utf-8") + ).hexdigest()[:12] + quantization_fingerprint = _quantization_fingerprint(quantization_cfg) + payload = ( + f"{absolute_model}|{source_fingerprint}|{quantization_fingerprint}|" + f"v{CONVERTER_VERSION}" + ) + digest = hashlib.sha1(payload.encode("utf-8")).hexdigest()[:16] + base = os.path.basename(os.path.normpath(absolute_model)) or "hf" + return os.path.join(os.path.abspath(cache_root), f"{base}-nvfp4-{digest}") + + +def is_existing_nvfp4_checkpoint(path: str) -> bool: + config = _read_source_config(path) + quantization_config = ( + config.get("quantization_config") if isinstance(config, dict) else None + ) + return is_nvfp4_quantization_config(quantization_config) + + +def _sync_checkpoint_ignore( + checkpoint_path: str, + quantization_cfg: dict[str, Any], +) -> None: + config = _read_source_config(checkpoint_path) + checkpoint_quantization = config.get("quantization_config") + if not isinstance(checkpoint_quantization, dict): + return + ignore = checkpoint_quantization.get("ignore") + if ignore is None: + ignore = [] + if not isinstance(ignore, list): + raise ValueError( + f"{checkpoint_path}/config.json quantization_config.ignore must " + "be a list of strings." + ) + if any(not isinstance(item, str) or not item.strip() for item in ignore): + raise ValueError( + f"{checkpoint_path}/config.json quantization_config.ignore must " + "contain non-empty strings." + ) + num_hidden_layers = _num_hidden_layers(config) + _, configured, _, _ = _validated_conversion_options( + quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + requested_high_precision = get_dynamic_high_precision_substrings( + quantization_config=quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + checkpoint_weight_names = tuple(_source_weight_map(checkpoint_path)) + configured_high_precision = tuple( + dict.fromkeys((*requested_high_precision, *(str(item) for item in ignore))) + ) + expanded_high_precision = expand_sglang_atomic_high_precision_substrings( + weight_names=checkpoint_weight_names, + skip_weight_substrings=configured_high_precision, + ) + validate_checkpoint_high_precision_layout( + checkpoint_path=checkpoint_path, + scheme="NVFP4", + weight_names=checkpoint_weight_names, + high_precision_substrings=expanded_high_precision, + quantized_companion_suffixes=(".weight_scale", ".weight_scale_2"), + ) + concrete_atomic_modules = tuple( + substring + for substring in expanded_high_precision + if substring not in configured_high_precision + ) + merged = list( + dict.fromkeys( + [ + *(str(item) for item in configured), + *(str(item) for item in ignore), + *concrete_atomic_modules, + ] + ) + ) + quantization_cfg["modules_to_not_convert"] = merged + + +def ensure_nvfp4_checkpoint( + *, + model_path: str, + quantization_cfg: dict[str, Any], +) -> str: + """Return an NVFP4 checkpoint path and synchronize its ignore policy.""" + if is_existing_nvfp4_checkpoint(model_path): + _sync_checkpoint_ignore(model_path, quantization_cfg) + return model_path + + converted = quantization_cfg.get("converted_model_path") + if converted and is_existing_nvfp4_checkpoint(converted): + _sync_checkpoint_ignore(converted, quantization_cfg) + return converted + + cache_root = ( + quantization_cfg.get("cache_root") + or os.environ.get("NRL_NVFP4_CACHE") + or os.path.join(os.path.expanduser("~"), ".cache", "nemo_rl", "nvfp4") + ) + save_dir = converted or _hash_qualified_save_dir( + model_dir=model_path, + cache_root=cache_root, + quantization_cfg=quantization_cfg, + ) + if is_existing_nvfp4_checkpoint(save_dir): + _sync_checkpoint_ignore(save_dir, quantization_cfg) + return save_dir + + source_config = _read_source_config(model_path) + num_hidden_layers = _num_hidden_layers(source_config) + ( + extra_high_precision_layers_hf, + modules_to_not_convert, + num_layers_at_start_in_bf16, + num_layers_at_end_in_bf16, + ) = _validated_conversion_options( + quantization_cfg, + num_hidden_layers=num_hidden_layers, + ) + + logger.info( + "[nvfp4] Converting %s -> %s (start_bf16=%s, end_bf16=%s, extra_hp=%s)", + model_path, + save_dir, + num_layers_at_start_in_bf16, + num_layers_at_end_in_bf16, + extra_high_precision_layers_hf, + ) + convert_nvfp4( + 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, + modules_to_not_convert=modules_to_not_convert, + ) + _sync_checkpoint_ignore(save_dir, quantization_cfg) + return save_dir diff --git a/nemo_rl/models/generation/sglang/quantization_utils.py b/nemo_rl/models/generation/sglang/quantization_utils.py new file mode 100644 index 00000000000..f34c68d979f --- /dev/null +++ b/nemo_rl/models/generation/sglang/quantization_utils.py @@ -0,0 +1,344 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared configuration helpers for quantized SGLang checkpoints and refits.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any, Literal, cast + +SglangQuantizationScheme = Literal["bf16", "mxfp8", "nvfp4"] +SUPPORTED_SGLANG_QUANTIZATION_SCHEMES = frozenset({"bf16", "mxfp8", "nvfp4"}) +HF_MOE_EXPERT_NAME_MARKERS = ( + ".experts.", + ".shared_expert.", + ".shared_experts.", + "block_sparse_moe.experts.", + ".moe.experts.", +) +HF_FUSED_LINEAR_WEIGHT_GROUPS = ( + ( + "qkv_proj", + (".q_proj.weight", ".k_proj.weight", ".v_proj.weight"), + ), + ( + "gate_up_proj", + (".gate_proj.weight", ".up_proj.weight"), + ), + ( + "gate_up_proj", + (".w1.weight", ".w3.weight"), + ), +) + + +def get_sglang_quantization_scheme( + quantization_config: dict[str, Any] | None, +) -> SglangQuantizationScheme: + """Return and validate the configured SGLang weight precision. + + An absent or empty quantization block preserves the historical BF16 + behavior. A non-empty block must declare ``scheme`` so a misspelling cannot + silently disable quantization. + """ + if not quantization_config: + return "bf16" + + scheme = quantization_config.get("scheme") + if scheme not in SUPPORTED_SGLANG_QUANTIZATION_SCHEMES: + supported = ", ".join(sorted(SUPPORTED_SGLANG_QUANTIZATION_SCHEMES)) + raise ValueError( + "SGLang quantization.scheme must be one of " + f"{{{supported}}}, got {scheme!r}." + ) + return cast(SglangQuantizationScheme, scheme) + + +def ensure_sglang_quantized_checkpoint( + *, + model_path: str, + quantization_config: dict[str, Any] | None, +) -> str: + """Resolve the startup checkpoint for the configured refit precision.""" + scheme = get_sglang_quantization_scheme(quantization_config) + if scheme == "bf16": + return model_path + if quantization_config is None: + raise RuntimeError(f"{scheme} requires an SGLang quantization config.") + + if scheme == "mxfp8": + from nemo_rl.models.generation.sglang.mxfp8_setup import ( + ensure_mxfp8_checkpoint, + ) + + return ensure_mxfp8_checkpoint( + model_path=model_path, + quantization_cfg=quantization_config, + ) + + from nemo_rl.models.generation.sglang.nvfp4_setup import ( + ensure_nvfp4_checkpoint, + ) + + return ensure_nvfp4_checkpoint( + model_path=model_path, + quantization_cfg=quantization_config, + ) + + +def validate_sglang_quantized_refit_backend( + *, + scheme: SglangQuantizationScheme, + use_megatron: bool, +) -> None: + """Reject quantized refits on backends that cannot emit quantized weights.""" + if scheme != "bf16" and not use_megatron: + raise NotImplementedError( + f"SGLang {scheme} weight refit requires a Megatron policy because " + "the DTensor/FSDP refit path only supports BF16." + ) + + +def validate_checkpoint_high_precision_layout( + *, + checkpoint_path: str, + scheme: str, + weight_names: Iterable[object], + high_precision_substrings: tuple[str, ...], + quantized_companion_suffixes: tuple[str, ...], +) -> None: + """Ensure a reused checkpoint already honors the requested BF16 policy. + + Changing the online refit skip rules cannot change tensors already loaded + by SGLang. A matching weight with quantization companions proves that the + existing checkpoint still stores that tensor in the quantized layout. + """ + names: set[str] = set() + for name in weight_names: + if not isinstance(name, str): + raise TypeError("HF checkpoint weight names must be strings.") + names.add(name) + + conflicts: list[str] = [] + for name in sorted(names): + if not name.endswith(".weight") or not any( + substring in name for substring in high_precision_substrings + ): + continue + base_name = name.removesuffix(".weight") + if any(base_name + suffix in names for suffix in quantized_companion_suffixes): + conflicts.append(name) + + if conflicts: + preview = ", ".join(repr(name) for name in conflicts[:5]) + if len(conflicts) > 5: + preview += f", ... ({len(conflicts)} total)" + raise ValueError( + f"The existing {scheme} checkpoint at {checkpoint_path!r} contains " + "quantized tensors selected for high precision by the current " + f"configuration: {preview}. Reconvert the original HF checkpoint " + "with the requested extra/head/tail skip policy." + ) + + +def get_hf_moe_expert_container(name: str) -> str | None: + """Return the fused SGLang MoE module that owns an HF expert tensor.""" + for marker in sorted(HF_MOE_EXPERT_NAME_MARKERS, key=len, reverse=True): + marker_index = name.find(marker) + if marker_index >= 0: + return name[:marker_index] + marker.rstrip(".") + return None + + +def expand_fused_moe_high_precision_substrings( + *, + weight_names: Iterable[object], + skip_weight_substrings: tuple[str, ...], +) -> tuple[str, ...]: + """Expand any skipped expert tensor to its whole fused MoE container. + + SGLang chooses one quantization method for a complete ``FusedMoE`` module, + so a checkpoint cannot mix quantized and high-precision expert tensors + within that container. + """ + skipped_containers: set[str] = set() + for name in weight_names: + if not isinstance(name, str): + raise TypeError("HF checkpoint weight names must be strings.") + container = get_hf_moe_expert_container(name) + if container is not None and any( + substring in name for substring in skip_weight_substrings + ): + skipped_containers.add(container) + return tuple(dict.fromkeys((*skip_weight_substrings, *sorted(skipped_containers)))) + + +def expand_sglang_atomic_high_precision_substrings( + *, + weight_names: Iterable[object], + skip_weight_substrings: tuple[str, ...], +) -> tuple[str, ...]: + """Expand skip rules across SGLang's fused linear and MoE boundaries.""" + names: set[str] = set() + for name in weight_names: + if not isinstance(name, str): + raise TypeError("HF checkpoint weight names must be strings.") + names.add(name) + + added_modules: set[str] = set() + visited_groups: set[tuple[str, tuple[str, ...]]] = set() + for name in names: + for fused_name, suffixes in HF_FUSED_LINEAR_WEIGHT_GROUPS: + matching_suffix = next( + (suffix for suffix in suffixes if name.endswith(suffix)), + None, + ) + if matching_suffix is None: + continue + prefix = name[: -len(matching_suffix)] + group_key = (prefix, suffixes) + if group_key in visited_groups: + continue + visited_groups.add(group_key) + + members = tuple( + prefix + suffix for suffix in suffixes if prefix + suffix in names + ) + fused_module = f"{prefix}.{fused_name}" + if len(members) >= 2 and ( + any( + substring in member + for member in members + for substring in skip_weight_substrings + ) + or any( + substring in fused_module for substring in skip_weight_substrings + ) + ): + added_modules.update( + member.removesuffix(".weight") for member in members + ) + + expanded = tuple(dict.fromkeys((*skip_weight_substrings, *sorted(added_modules)))) + return expand_fused_moe_high_precision_substrings( + weight_names=names, + skip_weight_substrings=expanded, + ) + + +def _optional_string_tuple( + quantization_config: dict[str, Any], + key: str, +) -> tuple[str, ...]: + value = quantization_config.get(key) + if value is None: + return () + if not isinstance(value, (list, tuple)): + raise TypeError(f"SGLang quantization.{key} must be a list of strings.") + + result: list[str] = [] + for item in value: + if not isinstance(item, str) or not item.strip(): + raise ValueError( + f"SGLang quantization.{key} entries must be non-empty strings." + ) + result.append(item.strip()) + return tuple(result) + + +def _optional_nonnegative_int( + quantization_config: dict[str, Any], + key: str, +) -> int: + value = quantization_config.get(key) + if value is None: + return 0 + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"SGLang quantization.{key} must be an integer.") + if value < 0: + raise ValueError(f"SGLang quantization.{key} must be non-negative.") + return value + + +def get_dynamic_high_precision_substrings( + *, + quantization_config: dict[str, Any], + num_hidden_layers: int, +) -> tuple[str, ...]: + """Build HF-name substrings that must remain in high precision. + + The result combines explicit HF-name patterns, checkpoint loader ignore + rules, and the configured BF16 bands at the start and end of the decoder. + """ + extra_high_precision_layers_hf = _optional_string_tuple( + quantization_config, + "extra_high_precision_layers_hf", + ) + modules_to_not_convert = _optional_string_tuple( + quantization_config, + "modules_to_not_convert", + ) + num_layers_at_start_in_bf16 = _optional_nonnegative_int( + quantization_config, + "num_layers_at_start_in_bf16", + ) + num_layers_at_end_in_bf16 = _optional_nonnegative_int( + quantization_config, + "num_layers_at_end_in_bf16", + ) + + if num_layers_at_start_in_bf16 or num_layers_at_end_in_bf16: + if num_hidden_layers <= 0: + raise ValueError( + "num_hidden_layers must be positive when BF16 head/tail layers " + "are configured." + ) + if num_layers_at_start_in_bf16 + num_layers_at_end_in_bf16 > num_hidden_layers: + raise ValueError( + "The configured BF16 head/tail layer counts exceed the model's " + f"{num_hidden_layers} decoder layers." + ) + + tail_start_idx = num_hidden_layers - num_layers_at_end_in_bf16 + dynamic_layer_prefixes = [ + *(f"model.layers.{i}." for i in range(num_layers_at_start_in_bf16)), + *(f"model.layers.{i}." for i in range(tail_start_idx, num_hidden_layers)), + ] + + # Preserve user order while removing duplicates so matching and cache + # fingerprints remain deterministic. + return tuple( + dict.fromkeys( + ( + *extra_high_precision_layers_hf, + *modules_to_not_convert, + *dynamic_layer_prefixes, + ) + ) + ) + + +def build_dynamic_skip_substrings( + *, + quantization_config: dict[str, Any], + num_hidden_layers: int, + static_skip_substrings: tuple[str, ...] = (), +) -> tuple[str, ...]: + """Combine static exclusions with the configured high-precision rules.""" + dynamic_substrings = get_dynamic_high_precision_substrings( + quantization_config=quantization_config, + num_hidden_layers=num_hidden_layers, + ) + return tuple(dict.fromkeys((*static_skip_substrings, *dynamic_substrings))) diff --git a/nemo_rl/models/generation/sglang/sglang_generation.py b/nemo_rl/models/generation/sglang/sglang_generation.py index 12800bef3a4..d71ca22d675 100644 --- a/nemo_rl/models/generation/sglang/sglang_generation.py +++ b/nemo_rl/models/generation/sglang/sglang_generation.py @@ -410,30 +410,19 @@ def continue_generation(self) -> None: return ray.get([e.continue_generation.remote() for e in engines]) - def post_process_weights( - self, - *, - restore_weights_before_load: bool = False, - post_process_quantization: bool = True, - ) -> None: - """Run SGLang's ``/post_process_weights`` RPC on every node-0 engine. - - Called by the refit dispatch helpers after a colocate IPC or - distributed broadcast refit so SGLang finalizes its weight tables - (e.g. materializes quantized scales, swaps in the fresh buffer). - """ + def begin_weight_update(self) -> None: + """Open a weight-update session on every node-0 engine.""" 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 - ] - ) + ray.get([e.begin_weight_update.remote() for e in engines]) + + def end_weight_update(self) -> None: + """Close a weight-update session and rebuild quantized kernel layouts.""" + engines = [e for e in self.engines if e is not None] + if not engines: + return + ray.get([e.end_weight_update.remote() for e in engines]) def health_monitoring_pause(self) -> None: if self._health_monitor: diff --git a/nemo_rl/models/generation/sglang/sglang_worker.py b/nemo_rl/models/generation/sglang/sglang_worker.py index 80edf3d5bae..2fb4570f07e 100644 --- a/nemo_rl/models/generation/sglang/sglang_worker.py +++ b/nemo_rl/models/generation/sglang/sglang_worker.py @@ -377,23 +377,13 @@ def continue_generation(self): response.raise_for_status() return response - def post_process_weights( - self, - restore_weights_before_load: bool = False, - post_process_quantization: bool = False, - ): - """Finalize engine-side weights after a distributed/IPC refit. + def begin_weight_update(self): + """Open one engine-side session before the first refit bucket.""" + return self._make_request("begin_weight_update", {}) - The HTTP server only posts metadata; the real weights were already - copied on-GPU by the preceding update path. - """ - return self._make_request( - "post_process_weights", - { - "restore_weights_before_load": restore_weights_before_load, - "post_process_quantization": post_process_quantization, - }, - ) + def end_weight_update(self): + """Finalize quantized layouts after the last refit bucket.""" + return self._make_request("end_weight_update", {}) def _simulate_crash(self): """Test-only: tear the engine down to simulate a crash. diff --git a/nemo_rl/models/generation/sglang/utils/patches.py b/nemo_rl/models/generation/sglang/utils/patches.py index cab29239b28..4c306a5588c 100644 --- a/nemo_rl/models/generation/sglang/utils/patches.py +++ b/nemo_rl/models/generation/sglang/utils/patches.py @@ -247,9 +247,347 @@ def _patch_sglang_custom_all_reduce_v2_tms_cudagraph() -> None: ) +def _sglang_weight_update_session_is_complete() -> bool: + """Return whether every component of a native or backported session exists.""" + + def has_all(relative_path: str, sentinels: tuple[str, ...]) -> bool: + with open(_get_sglang_file(relative_path)) as file: + content = file.read() + return all(sentinel in content for sentinel in sentinels) + + common_session_components = ( + has_all( + "srt/managers/io_struct.py", + ( + "class BeginWeightUpdateReqInput", + "class EndWeightUpdateReqInput", + ), + ) + and has_all( + "srt/entrypoints/http_server.py", + ( + '@app.post("/begin_weight_update")', + '@app.post("/end_weight_update")', + ), + ) + and has_all( + "srt/managers/tokenizer_control_mixin.py", + ( + "async def begin_weight_update(", + "async def end_weight_update(", + ), + ) + ) + if not common_session_components: + return False + + # sglang-miles routes lifecycle requests through ``weight_updater``. + native_session = has_all( + "srt/managers/scheduler.py", + ( + "self.weight_updater.begin_weight_update", + "self.weight_updater.end_weight_update", + ), + ) + # The v0.5.12.post1 backport dispatches directly to methods added to + # SchedulerUpdateWeightsMixin. + backported_session = has_all( + "srt/managers/scheduler.py", + ( + "(BeginWeightUpdateReqInput, self.begin_weight_update)", + "(EndWeightUpdateReqInput, self.end_weight_update)", + ), + ) and has_all( + "srt/managers/scheduler_update_weights_mixin.py", + ( + "def begin_weight_update(", + "def end_weight_update(", + "update_weights_from_distributed requires an open session", + "update_weights_from_tensor requires an open session", + ), + ) + return native_session or backported_session + + +def _patch_sglang_weight_update_session() -> None: + """Backport SGLang's begin/end refit session API. + + The pinned release loads tensors but has no lifecycle hook to rerun + quantization post-processing once all buckets arrive, so quantized + weights keep stale derived scales. + """ + # The patch spans five files, so only a complete implementation counts as + # already-patched: a single endpoint may mean a concurrent Ray actor is + # still mid-patch, or died partway through. + if _sglang_weight_update_session_is_complete(): + return + + _patch_sglang_file_replacements( + "srt/managers/io_struct.py", + ( + ( + "class BeginWeightUpdateReqInput(BaseReq):", + "@dataclass\nclass CheckWeightsReqInput(BaseReq):\n", + "@dataclass\n" + "class BeginWeightUpdateReqInput(BaseReq):\n" + " pass\n\n\n" + "@dataclass\n" + "class BeginWeightUpdateReqOutput(BaseReq):\n" + " success: bool\n" + " message: str\n\n\n" + "@dataclass\n" + "class EndWeightUpdateReqInput(BaseReq):\n" + " pass\n\n\n" + "@dataclass\n" + "class EndWeightUpdateReqOutput(BaseReq):\n" + " success: bool\n" + " message: str\n\n\n" + "@dataclass\n" + "class CheckWeightsReqInput(BaseReq):\n", + ), + ), + "weight-update session request types", + ) + _patch_sglang_file_replacements( + "srt/entrypoints/http_server.py", + ( + ( + " BeginWeightUpdateReqInput,\n", + " AttachHiCacheStorageReqInput,\n", + " AttachHiCacheStorageReqInput,\n BeginWeightUpdateReqInput,\n", + ), + ( + " EndWeightUpdateReqInput,\n", + " DumperControlReqInput,\n", + " DumperControlReqInput,\n EndWeightUpdateReqInput,\n", + ), + ( + '@app.post("/begin_weight_update")', + '@app.post("/update_weights_from_tensor")\n', + '@app.post("/begin_weight_update")\n' + "@auth_level(AuthLevel.ADMIN_OPTIONAL)\n" + "async def begin_weight_update(\n" + " obj: BeginWeightUpdateReqInput, request: Request\n" + "):\n" + ' """Open a weight-update session before loading buckets."""\n' + " success, message = (\n" + " await _global_state.tokenizer_manager.begin_weight_update(\n" + " obj, request\n" + " )\n" + " )\n" + ' content = {"success": success, "message": message}\n' + " return ORJSONResponse(\n" + " content,\n" + " status_code=200 if success else HTTPStatus.BAD_REQUEST,\n" + " )\n\n\n" + '@app.post("/end_weight_update")\n' + "@auth_level(AuthLevel.ADMIN_OPTIONAL)\n" + "async def end_weight_update(\n" + " obj: EndWeightUpdateReqInput, request: Request\n" + "):\n" + ' """Finalize quantized weights after all buckets arrive."""\n' + " success, message = (\n" + " await _global_state.tokenizer_manager.end_weight_update(\n" + " obj, request\n" + " )\n" + " )\n" + ' content = {"success": success, "message": message}\n' + " return ORJSONResponse(\n" + " content,\n" + " status_code=200 if success else HTTPStatus.BAD_REQUEST,\n" + " )\n\n\n" + '@app.post("/update_weights_from_tensor")\n', + ), + ), + "weight-update session HTTP endpoints", + ) + _patch_sglang_file_replacements( + "srt/managers/tokenizer_control_mixin.py", + ( + ( + " BeginWeightUpdateReqInput,\n", + " AttachHiCacheStorageReqOutput,\n", + " AttachHiCacheStorageReqOutput,\n" + " BeginWeightUpdateReqInput,\n" + " BeginWeightUpdateReqOutput,\n", + ), + ( + " EndWeightUpdateReqInput,\n", + " DumperControlReqOutput,\n", + " DumperControlReqOutput,\n" + " EndWeightUpdateReqInput,\n" + " EndWeightUpdateReqOutput,\n", + ), + ( + ' ("begin_weight_update", BeginWeightUpdateReqOutput),\n', + ' ("destroy_weights_update_group", DestroyWeightsUpdateGroupReqOutput),\n', + ' ("destroy_weights_update_group", DestroyWeightsUpdateGroupReqOutput),\n' + ' ("begin_weight_update", BeginWeightUpdateReqOutput),\n' + ' ("end_weight_update", EndWeightUpdateReqOutput),\n', + ), + ( + " async def _nemo_rl_weight_update_session_call(", + " async def update_weights_from_distributed(\n", + " async def _nemo_rl_weight_update_session_call(\n" + " self: TokenizerManager, communicator, obj\n" + " ) -> Tuple[bool, str]:\n" + ' """Run a refit lifecycle RPC with pause-aware locking."""\n' + " self.auto_create_handle_loop()\n" + " async with self.is_pause_cond:\n" + " is_paused = self.is_pause\n" + " if is_paused:\n" + " results = await communicator(obj)\n" + " if not is_paused:\n" + " async with self.model_update_lock.writer_lock:\n" + " results = await communicator(obj)\n" + " return FanOutCommunicator.merge_results(results)\n\n" + " async def begin_weight_update(\n" + " self: TokenizerManager,\n" + " obj: BeginWeightUpdateReqInput,\n" + " request: Optional[fastapi.Request] = None,\n" + " ) -> Tuple[bool, str]:\n" + " return await self._nemo_rl_weight_update_session_call(\n" + " self.begin_weight_update_communicator, obj\n" + " )\n\n" + " async def end_weight_update(\n" + " self: TokenizerManager,\n" + " obj: EndWeightUpdateReqInput,\n" + " request: Optional[fastapi.Request] = None,\n" + " ) -> Tuple[bool, str]:\n" + " return await self._nemo_rl_weight_update_session_call(\n" + " self.end_weight_update_communicator, obj\n" + " )\n\n" + " async def update_weights_from_distributed(\n", + ), + ), + "weight-update session tokenizer fanout", + ) + _patch_sglang_file_replacements( + "srt/managers/scheduler.py", + ( + ( + " BeginWeightUpdateReqInput,\n", + " AttachHiCacheStorageReqOutput,\n", + " AttachHiCacheStorageReqOutput,\n BeginWeightUpdateReqInput,\n", + ), + ( + " EndWeightUpdateReqInput,\n", + " DumperControlReqOutput,\n", + " DumperControlReqOutput,\n EndWeightUpdateReqInput,\n", + ), + ( + " (BeginWeightUpdateReqInput, self.begin_weight_update),\n", + " (UpdateWeightFromDiskReqInput, self.update_weights_from_disk),\n", + " (UpdateWeightFromDiskReqInput, self.update_weights_from_disk),\n" + " (BeginWeightUpdateReqInput, self.begin_weight_update),\n" + " (EndWeightUpdateReqInput, self.end_weight_update),\n", + ), + ), + "weight-update session scheduler dispatch", + ) + _patch_sglang_file_replacements( + "srt/managers/scheduler_update_weights_mixin.py", + ( + ( + " BeginWeightUpdateReqInput,\n", + " CheckWeightsReqInput,\n", + " BeginWeightUpdateReqInput,\n" + " BeginWeightUpdateReqOutput,\n" + " CheckWeightsReqInput,\n", + ), + ( + " EndWeightUpdateReqInput,\n", + " GetWeightsByNameReqInput,\n", + " EndWeightUpdateReqInput,\n" + " EndWeightUpdateReqOutput,\n" + " GetWeightsByNameReqInput,\n", + ), + ( + "def _nemo_rl_run_quant_method_hook(", + "logger = logging.getLogger(__name__)\n\n\n" + "class SchedulerUpdateWeightsMixin:\n", + "logger = logging.getLogger(__name__)\n\n\n" + "def _nemo_rl_run_quant_method_hook(model, target_device, hook_name):\n" + " from sglang.srt.lora.layers import BaseLayerWithLoRA\n" + " from sglang.srt.model_loader.loader import device_loading_context\n\n" + " for _, module in model.named_modules():\n" + " if isinstance(module, BaseLayerWithLoRA):\n" + " continue\n" + ' quant_method = getattr(module, "quant_method", None)\n' + " if quant_method is not None and hasattr(quant_method, hook_name):\n" + " with device_loading_context(module, target_device):\n" + " getattr(quant_method, hook_name)(module)\n\n\n" + "class SchedulerUpdateWeightsMixin:\n", + ), + ( + " def begin_weight_update(", + " def update_weights_from_distributed(\n", + " def begin_weight_update(\n" + " self: Scheduler, recv_req: BeginWeightUpdateReqInput\n" + " ):\n" + ' assert not getattr(self, "_weight_update_in_progress", False), (\n' + ' "begin_weight_update called while a session is already open"\n' + " )\n" + " runner = self.tp_worker.model_runner\n" + " _nemo_rl_run_quant_method_hook(\n" + ' runner.model, torch.device(runner.device), "restore_weights_before_loading"\n' + " )\n" + " self._weight_update_in_progress = True\n" + " torch.distributed.barrier(group=self.tp_cpu_group)\n" + ' return BeginWeightUpdateReqOutput(True, "Success")\n\n' + " def end_weight_update(\n" + " self: Scheduler, recv_req: EndWeightUpdateReqInput\n" + " ):\n" + ' assert getattr(self, "_weight_update_in_progress", False), (\n' + ' "end_weight_update called without begin_weight_update"\n' + " )\n" + " runner = self.tp_worker.model_runner\n" + " _nemo_rl_run_quant_method_hook(\n" + ' runner.model, torch.device(runner.device), "process_weights_after_loading"\n' + " )\n" + " self._weight_update_in_progress = False\n" + " torch.distributed.barrier(group=self.tp_cpu_group)\n" + ' return EndWeightUpdateReqOutput(True, "Success")\n\n' + " def update_weights_from_distributed(\n", + ), + ( + ' assert getattr(self, "_weight_update_in_progress", False), (\n' + ' "update_weights_from_distributed requires an open session"\n' + " )\n", + ' """Update the online model parameter."""\n' + " success, message = self.tp_worker.update_weights_from_distributed(recv_req)\n", + ' """Update the online model parameter."""\n' + ' assert getattr(self, "_weight_update_in_progress", False), (\n' + ' "update_weights_from_distributed requires an open session"\n' + " )\n" + " success, message = self.tp_worker.update_weights_from_distributed(recv_req)\n", + ), + ( + ' assert getattr(self, "_weight_update_in_progress", False), (\n' + ' "update_weights_from_tensor requires an open session"\n' + " )\n", + ' """Update the online model parameter from tensors."""\n' + " if recv_req.disable_draft_model:\n", + ' """Update the online model parameter from tensors."""\n' + ' assert getattr(self, "_weight_update_in_progress", False), (\n' + ' "update_weights_from_tensor requires an open session"\n' + " )\n" + " if recv_req.disable_draft_model:\n", + ), + ), + "quantized weight-update session lifecycle", + ) + if not _sglang_weight_update_session_is_complete(): + raise RuntimeError( + "SGLang weight-update session patch did not produce a complete " + "cross-file implementation." + ) + + def _apply_sglang_compat_patches() -> None: _patch_sglang_safe_unpickler() _patch_sglang_custom_all_reduce_v2_tms_cudagraph() + _patch_sglang_weight_update_session() _override_sglang_imbalance_check_env() _patch_megatron_dynamic_context_hook_mode() _patch_megatron_training_hook_mode() diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index c35faa68db3..81c4dbd8ff5 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -2177,6 +2177,17 @@ def _build_sglang_hf_iterator( from nemo_rl.models.policy.workers.megatron_sglang_weight_iterator import ( MegatronSGLangHfWeightIterator, ) + from nemo_rl.models.generation.sglang.quantization_utils import ( + get_sglang_quantization_scheme, + ) + + configured_precision = get_sglang_quantization_scheme(sglang_quantization_cfg) + if target_precision != configured_precision: + raise ValueError( + "SGLang refit target precision does not match its quantization " + f"config: target={target_precision!r}, " + f"configured={configured_precision!r}." + ) if self.refit_conversion_tasks is None: self.refit_conversion_tasks = self.megatron_bridge.get_conversion_tasks( @@ -2184,7 +2195,7 @@ def _build_sglang_hf_iterator( ) num_hidden_layers = 0 - if target_precision == "mxfp8": + if target_precision in ("mxfp8", "nvfp4"): num_hidden_layers = int( getattr(self.megatron_bridge.transformer_config, "num_layers", 0) ) @@ -3395,8 +3406,3 @@ def _percentile(values: list[float], p: float) -> float: ) # pragma: no cover class MegatronPolicyWorker(MegatronPolicyWorkerImpl): pass - - -# --------------------------------------------------------------------------- -# Driver-side SGLang weight-update dispatch (Megatron backend) -# --------------------------------------------------------------------------- diff --git a/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py b/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py index 1e14e8799b2..1ea4a7551c7 100644 --- a/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py +++ b/nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py @@ -15,25 +15,39 @@ """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. +with no vLLM-specific KV/Q scale tensors. Quantized targets apply the same +HF-name selection and tensor conversion used by their offline checkpoint +converters. """ from __future__ import annotations -from typing import Any, Iterator, Literal +from collections.abc import Iterable +from typing import Any, Iterator import torch +from nemo_rl.models.generation.sglang.mxfp8_quantization_core import ( + SKIP_WEIGHT_SUBSTRINGS as MXFP8_SKIP_WEIGHT_SUBSTRINGS, +) from nemo_rl.models.generation.sglang.mxfp8_quantization_core import ( SOURCE_FP8_SCALE_KEY_SUFFIX, - build_dynamic_skip_substrings, quantize_mxfp8, should_quantize, strip_weight_suffix, ) +from nemo_rl.models.generation.sglang.nvfp4_quantization_core import ( + nvfp4_quantized_entries, + quantize_nvfp4, + quantize_nvfp4_pair, + should_quantize_nvfp4, + should_skip_nvfp4_gated_pair, + split_gated_pair_name, +) +from nemo_rl.models.generation.sglang.quantization_utils import ( + SglangQuantizationScheme, + build_dynamic_skip_substrings, +) class MegatronSGLangHfWeightIterator: @@ -42,8 +56,9 @@ class MegatronSGLangHfWeightIterator: 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. + sized by the *post-transformation* tensor footprint. Companion tensors + produced from one source weight (and NVFP4 gate/up pairs) remain in the + same bucket. """ def __init__( @@ -64,7 +79,7 @@ def __init__( def iter_hf_weight_buckets( self, *, - target_precision: Literal["bf16", "mxfp8"] = "bf16", + target_precision: SglangQuantizationScheme = "bf16", buffer_size_bytes: int, ) -> Iterator[list[tuple[str, torch.Tensor]]]: """Yield finalized HF tensor buckets sized by transmitted bytes.""" @@ -73,14 +88,24 @@ def iter_hf_weight_buckets( f"buffer_size_bytes must be positive, got {buffer_size_bytes}" ) - skip_weight_substrings = ( - build_dynamic_skip_substrings( + if target_precision == "mxfp8": + # MXFP8 additionally excludes norms, embeddings, router/gate and + # LM-head weights, which have to stay high precision. NVFP4 has no + # such static list: it only ever targets MoE expert GEMMs. + skip_weight_substrings = build_dynamic_skip_substrings( quantization_config=self._quantization_config, num_hidden_layers=self._num_hidden_layers, + static_skip_substrings=MXFP8_SKIP_WEIGHT_SUBSTRINGS, ) - if target_precision == "mxfp8" - else None - ) + elif target_precision == "nvfp4": + skip_weight_substrings = build_dynamic_skip_substrings( + quantization_config=self._quantization_config, + num_hidden_layers=self._num_hidden_layers, + ) + elif target_precision == "bf16": + skip_weight_substrings = None + else: + raise ValueError(f"Unsupported SGLang target precision: {target_precision}") bucket: list[tuple[str, torch.Tensor]] = [] bucket_size = 0 @@ -89,14 +114,15 @@ def iter_hf_weight_buckets( 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 + finalized_size = sum( + tensor.numel() * tensor.element_size() for _, tensor in finalized + ) + if bucket and bucket_size + finalized_size > buffer_size_bytes: + yield bucket + bucket = [] + bucket_size = 0 + bucket.extend(finalized) + bucket_size += finalized_size if bucket: yield bucket @@ -104,26 +130,32 @@ def iter_hf_weight_buckets( def _iter_finalized_hf_named_tensors( self, *, - target_precision: Literal["bf16", "mxfp8"], + target_precision: SglangQuantizationScheme, 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. + item produces exactly one finalized pair; quantized formats expand a + source weight into the payload required by SGLang. """ - for hf_param_name, tensor in self._bridge.export_hf_weights( + hf_weights = 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 == "nvfp4": + if skip_weight_substrings is None: + raise RuntimeError("NVFP4 refit requires initialized skip rules.") + yield from self._iter_nvfp4_hf_named_tensors( + hf_weights, + skip_weight_substrings=skip_weight_substrings, + ) + return + for hf_param_name, tensor in hf_weights: + # AutoBridge yields plain ``torch.Tensor`` for Megatron (no + # DTensor / async-collective wrapping), so no ``.wait()`` here. if target_precision == "mxfp8" and skip_weight_substrings is not None: if should_quantize( hf_param_name, @@ -139,3 +171,76 @@ def _iter_finalized_hf_named_tensors( continue yield [(hf_param_name, tensor)] + + @staticmethod + def _iter_nvfp4_hf_named_tensors( + hf_weights: Iterable[tuple[str, torch.Tensor]], + *, + skip_weight_substrings: tuple[str, ...], + ) -> Iterator[list[tuple[str, torch.Tensor]]]: + """Quantize NVFP4 expert weights, buffering complete gate/up pairs.""" + pending_pairs: dict[ + str, + dict[str, tuple[str, torch.Tensor]], + ] = {} + + for hf_param_name, tensor in hf_weights: + if should_skip_nvfp4_gated_pair( + hf_param_name, + skip_weight_substrings=skip_weight_substrings, + ): + yield [(hf_param_name, tensor)] + continue + + if not should_quantize_nvfp4( + hf_param_name, + tensor, + skip_weight_substrings=skip_weight_substrings, + ): + yield [(hf_param_name, tensor)] + continue + + pair_base, pair_role = split_gated_pair_name(hf_param_name) + if pair_base is None or pair_role is None: + yield nvfp4_quantized_entries( + hf_param_name, + quantize_nvfp4(tensor), + include_input_scale=False, + ) + continue + + pair = pending_pairs.setdefault(pair_base, {}) + if pair_role in pair: + raise ValueError( + "NVFP4 requires one complete gate/up pair per refit; " + f"found duplicate {pair_role} tensor for {pair_base}." + ) + pair[pair_role] = (hf_param_name, tensor) + if set(pair) != {"gate", "up"}: + continue + + gate_name, gate_weight = pair["gate"] + up_name, up_weight = pair["up"] + gate_output, up_output = quantize_nvfp4_pair(gate_weight, up_weight) + yield [ + *nvfp4_quantized_entries( + gate_name, + gate_output, + include_input_scale=False, + ), + *nvfp4_quantized_entries( + up_name, + up_output, + include_input_scale=False, + ), + ] + del pending_pairs[pair_base] + + if pending_pairs: + incomplete = { + base: sorted(roles) for base, roles in sorted(pending_pairs.items()) + } + raise ValueError( + "NVFP4 gate/up weights must be quantized together; incomplete " + f"pairs: {incomplete}." + ) diff --git a/nemo_rl/weight_sync/dtensor_sglang_refit.py b/nemo_rl/weight_sync/dtensor_sglang_refit.py index 1888e9d20d6..b5b81fa24ce 100644 --- a/nemo_rl/weight_sync/dtensor_sglang_refit.py +++ b/nemo_rl/weight_sync/dtensor_sglang_refit.py @@ -34,11 +34,25 @@ def refit_sglang_colocated( """Refit colocated SGLang engines from the FSDP/DTensor policy. Lifecycle: optional fault-tolerance recover, connect (when new / - recovered engines), pause + KV invalidation, send HF tensor buckets via - Ray IPC, post-process, continue. Mirrors the Megatron colocated driver; - the FSDP path is BF16-only. + recovered engines), pause, open an engine-side update session, send HF + tensor buckets via Ray IPC, close the session, continue. Mirrors the + Megatron colocated driver; the FSDP path is BF16-only. """ - from nemo_rl.models.policy.utils import fetch_updatable_engines_with_recover + from nemo_rl.models.generation.sglang.quantization_utils import ( + get_sglang_quantization_scheme, + ) + from nemo_rl.models.policy.utils import ( + fetch_updatable_engines_with_recover, + get_sglang_quantization_cfg, + ) + + sglang_quantization_cfg = get_sglang_quantization_cfg(policy_generation) + target_precision = get_sglang_quantization_scheme(sglang_quantization_cfg) + if target_precision != "bf16": + raise NotImplementedError( + "The FSDP/DTensor policy only supports BF16 SGLang refits; " + f"got target_precision={target_precision!r}." + ) ( rollout_engines, @@ -58,18 +72,28 @@ def refit_sglang_colocated( "clear_updatable_num_new_engines did not zero num_new_engines" ) - # Pause with the configured mode, then flush: an IPC refit replaces every - # weight in place, so no cached KV entry survives it regardless of mode. - policy_generation.pause_generation(mode=policy_generation.pause_generation_mode) - policy_generation.invalidate_kv_cache() + # Pause with the configured mode, but only invalidate the KV cache when + # the mode actually drops generation state. "in_place" leaves the engine + # paused without dropping its KV cache, so flushing would clobber the + # still-valid in-place state. + pause_mode = policy_generation.pause_generation_mode + policy_generation.pause_generation(mode=pause_mode) + if pause_mode != "in_place": + policy_generation.invalidate_kv_cache() + + policy_generation.begin_weight_update() try: 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_quantization_cfg, ) ray.get(futures) - policy_generation.post_process_weights() finally: + # Close the session and resume on every path, so a failed refit + # leaves the engine usable instead of wedged in the update state. + policy_generation.end_weight_update() policy_generation.continue_generation() return True diff --git a/nemo_rl/weight_sync/megatron_sglang_refit.py b/nemo_rl/weight_sync/megatron_sglang_refit.py index 704b0964f08..bff08981552 100644 --- a/nemo_rl/weight_sync/megatron_sglang_refit.py +++ b/nemo_rl/weight_sync/megatron_sglang_refit.py @@ -34,16 +34,19 @@ def refit_sglang_colocated( """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. + recovered engines), pause, open an engine-side update session, send HF + tensor buckets via Ray IPC, close the session, continue. """ + from nemo_rl.models.generation.sglang.quantization_utils import ( + get_sglang_quantization_scheme, + ) 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") + target_precision = get_sglang_quantization_scheme(sglang_quant) ( rollout_engines, @@ -63,10 +66,16 @@ def refit_sglang_colocated( "clear_updatable_num_new_engines did not zero num_new_engines" ) - # Pause with the configured mode, then flush: an IPC refit replaces every - # weight in place, so no cached KV entry survives it regardless of mode. - policy_generation.pause_generation(mode=policy_generation.pause_generation_mode) - policy_generation.invalidate_kv_cache() + # Pause with the configured mode, but only invalidate the KV cache when + # the mode actually drops generation state. "in_place" leaves the engine + # paused without dropping its KV cache, so flushing would clobber the + # still-valid in-place state. + pause_mode = policy_generation.pause_generation_mode + policy_generation.pause_generation(mode=pause_mode) + if pause_mode != "in_place": + policy_generation.invalidate_kv_cache() + + policy_generation.begin_weight_update() try: # Per-worker actor method is now synchronous (per-chunk ray.get + # lifetime-safe IPC handled inside send_hf_buckets_via_ipc_actor_impl), @@ -79,8 +88,10 @@ def refit_sglang_colocated( sglang_quantization_cfg=sglang_quant, ) ray.get(futures) - policy_generation.post_process_weights() finally: + # Close the session and resume on every path, so a failed refit + # leaves the engine usable instead of wedged in the update state. + policy_generation.end_weight_update() policy_generation.continue_generation() return True @@ -97,13 +108,16 @@ def refit_sglang_distributed( walk the AutoBridge collective inside ``update_weights_to_sglang_distributed`` but do not broadcast. Includes optional fault-tolerance recover prelude. """ + from nemo_rl.models.generation.sglang.quantization_utils import ( + get_sglang_quantization_scheme, + ) 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") + target_precision = get_sglang_quantization_scheme(sglang_quant) ( rollout_engines, @@ -131,6 +145,8 @@ def refit_sglang_distributed( policy_generation.pause_generation(mode=pause_mode) if pause_mode != "in_place": policy_generation.invalidate_kv_cache() + + policy_generation.begin_weight_update() try: futures = policy.update_weights_to_sglang_distributed( rollout_engines=rollout_engines, @@ -140,7 +156,7 @@ def refit_sglang_distributed( sglang_quantization_cfg=sglang_quant, ) ray.get(futures) - policy_generation.post_process_weights() finally: + policy_generation.end_weight_update() policy_generation.continue_generation() return True diff --git a/nemo_rl/weight_sync/sglang_weight_synchronizer.py b/nemo_rl/weight_sync/sglang_weight_synchronizer.py index 4b13701e8af..c1a2d9e4a51 100644 --- a/nemo_rl/weight_sync/sglang_weight_synchronizer.py +++ b/nemo_rl/weight_sync/sglang_weight_synchronizer.py @@ -14,11 +14,12 @@ """Weight synchronizers for the SGLang generation backend. -The refit itself — engine recovery, connect, pause, KV invalidation, bucket -transfer, post-process, continue — lives in the backend-specific driver -modules (``megatron_sglang_refit`` / ``dtensor_sglang_refit``). These -synchronizers only own the GPU phase transitions around that call: which -side gets offloaded/onloaded, and in what order. +The refit itself — engine recovery, connect, pause, conditional KV +invalidation, a begin/end weight-update session around the bucket transfer, +then continue — lives in the backend-specific driver modules +(``megatron_sglang_refit`` / ``dtensor_sglang_refit``). These synchronizers +only own the GPU phase transitions around that call: which side gets +offloaded/onloaded, and in what order. Colocated (``weight_transfer_mode="ipc"``): 1. policy.offload_before_refit() -- free GPU for staging diff --git a/pyproject.toml b/pyproject.toml index 5fd74de6310..5a122cb1ae3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,7 +132,7 @@ automodel = [ "mamba-ssm", "causal-conv1d", "nv-grouped-gemm", - "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.17", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@29d31c095796f3c8ece47ee9cdcc167051bbeed9 ; platform_machine == 'x86_64'", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'", ] @@ -350,7 +350,7 @@ link-mode = "copy" # vllm defaults to opencv for video IO but falls back to torchcodec. # The timm override is needed because sglang requires timm==1.0.16. override-dependencies = [ - "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15", + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.17", "nvidia-cublas==13.5.1.27; sys_platform != 'darwin'", "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", "nvidia-cudnn-frontend==1.23.0", @@ -532,12 +532,12 @@ requires-dist = ["torch", "packaging", "ninja"] [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.15.0+42b8400" +version = "2.17.0+2e559f0" requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[tool.uv.dependency-metadata]] name = "transformer-engine-torch" -version = "2.15.0+42b8400" +version = "2.17.0+2e559f0" requires-dist = ["torch", "transformer-engine"] [[tool.uv.dependency-metadata]] diff --git a/pyrefly.toml b/pyrefly.toml index 30dd49c3ccf..ae7e4f3d1f3 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -8,6 +8,9 @@ replace-imports-with-any = [ "transformers.*", "tensorrt_llm.*", "vllm.*", + "sglang.*", + "flashinfer.*", + "transformer_engine.*", "math_verify.*", "sympy.*", "torchdata.*", @@ -168,6 +171,11 @@ project-includes = [ "nemo_rl/models/generation/megatron/utils.py", "nemo_rl/models/generation/openai_server_utils.py", "nemo_rl/models/generation/sglang/config.py", + "nemo_rl/models/generation/sglang/mxfp8_quantization_core.py", + "nemo_rl/models/generation/sglang/mxfp8_setup.py", + "nemo_rl/models/generation/sglang/nvfp4_quantization_core.py", + "nemo_rl/models/generation/sglang/nvfp4_setup.py", + "nemo_rl/models/generation/sglang/quantization_utils.py", "nemo_rl/models/generation/sglang/utils/http_utils.py", "nemo_rl/models/generation/sglang/utils/ip_port_utils.py", "nemo_rl/models/generation/sglang/utils/patches.py", @@ -200,6 +208,7 @@ project-includes = [ "nemo_rl/models/policy/workers/__init__.py", "nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py", "nemo_rl/models/policy/workers/checkpoint_engine.py", + "nemo_rl/models/policy/workers/megatron_sglang_weight_iterator.py", "nemo_rl/models/policy/workers/patches.py", "nemo_rl/models/value/__init__.py", "nemo_rl/models/value/config.py", diff --git a/tests/unit/models/generation/sglang/test_megatron_sglang_weight_iterator.py b/tests/unit/models/generation/sglang/test_megatron_sglang_weight_iterator.py new file mode 100644 index 00000000000..34dba34a99c --- /dev/null +++ b/tests/unit/models/generation/sglang/test_megatron_sglang_weight_iterator.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Iterable +from typing import Any + +import pytest +import torch + +from nemo_rl.models.policy.workers import ( + megatron_sglang_weight_iterator as weight_iterator, +) + + +class _FakeBridge: + def __init__(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: + self._weights = list(weights) + + def export_hf_weights(self, *_args: Any, **_kwargs: Any): + return iter(self._weights) + + +def _make_iterator( + weights: Iterable[tuple[str, torch.Tensor]], + *, + quantization_config: dict[str, Any] | None = None, + num_hidden_layers: int = 4, +) -> weight_iterator.MegatronSGLangHfWeightIterator: + return weight_iterator.MegatronSGLangHfWeightIterator( + megatron_bridge=_FakeBridge(weights), + models=[object()], + conversion_tasks=object(), + quantization_config=quantization_config, + num_hidden_layers=num_hidden_layers, + ) + + +def _collect_entries( + iterator: weight_iterator.MegatronSGLangHfWeightIterator, + *, + target_precision: str, +) -> list[tuple[str, torch.Tensor]]: + buckets = iterator.iter_hf_weight_buckets( + target_precision=target_precision, # type: ignore[arg-type] + buffer_size_bytes=1 << 30, + ) + return [entry for bucket in buckets for entry in bucket] + + +def _fake_nvfp4_output( + weight: torch.Tensor, + global_scale: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + rows, columns = weight.shape + if global_scale is None: + global_scale = torch.ones((), dtype=torch.float32) + return ( + torch.zeros((rows, columns // 2), dtype=torch.uint8), + torch.zeros( + (rows, columns // 16), + dtype=torch.uint8, + ).view(torch.float8_e4m3fn), + global_scale, + ) + + +def test_mxfp8_iterator_respects_head_tail_and_extra_high_precision( + monkeypatch, +) -> None: + weights = [ + ( + f"model.layers.{layer}.mlp.down_proj.weight", + torch.ones((2, 32), dtype=torch.bfloat16), + ) + for layer in range(4) + ] + monkeypatch.setattr( + weight_iterator, + "quantize_mxfp8", + lambda tensor: ( + torch.zeros_like(tensor, dtype=torch.uint8), + torch.zeros((tensor.shape[0], tensor.shape[1] // 32), dtype=torch.uint8), + ), + ) + iterator = _make_iterator( + weights, + quantization_config={ + "num_layers_at_start_in_bf16": 1, + "num_layers_at_end_in_bf16": 1, + "extra_high_precision_layers_hf": ["model.layers.2."], + }, + ) + + names = [name for name, _ in _collect_entries(iterator, target_precision="mxfp8")] + scale_names = [name for name in names if name.endswith(".weight_scale_inv")] + assert scale_names == ["model.layers.1.mlp.down_proj.weight_scale_inv"] + + +def test_mxfp8_iterator_keeps_synchronized_qkv_group_in_bf16( + monkeypatch, +) -> None: + base = "model.layers.1.self_attn" + names = [ + f"{base}.{projection}.weight" for projection in ("q_proj", "k_proj", "v_proj") + ] + weights = [(name, torch.ones((2, 32), dtype=torch.bfloat16)) for name in names] + + def fail_quantize(_tensor: torch.Tensor): + raise AssertionError("a synchronized high-precision QKV group must stay BF16") + + monkeypatch.setattr(weight_iterator, "quantize_mxfp8", fail_quantize) + entries = _collect_entries( + _make_iterator( + weights, + quantization_config={ + "extra_high_precision_layers_hf": [f"{base}.q_proj"], + "modules_to_not_convert": [ + f"{base}.q_proj", + f"{base}.k_proj", + f"{base}.v_proj", + ], + }, + ), + target_precision="mxfp8", + ) + + assert [name for name, _ in entries] == names + + +def test_nvfp4_iterator_respects_head_tail_and_extra_high_precision( + monkeypatch, +) -> None: + weights = [ + ( + f"model.layers.{layer}.mlp.experts.0.down_proj.weight", + torch.ones((2, 32), dtype=torch.bfloat16), + ) + for layer in range(4) + ] + monkeypatch.setattr( + weight_iterator, + "quantize_nvfp4", + _fake_nvfp4_output, + ) + iterator = _make_iterator( + weights, + quantization_config={ + "num_layers_at_start_in_bf16": 1, + "num_layers_at_end_in_bf16": 1, + "extra_high_precision_layers_hf": ["model.layers.2."], + }, + ) + + names = [name for name, _ in _collect_entries(iterator, target_precision="nvfp4")] + scale_names = [name for name in names if name.endswith(".weight_scale")] + assert scale_names == ["model.layers.1.mlp.experts.0.down_proj.weight_scale"] + + +def test_nvfp4_live_pair_has_shared_scale_and_no_input_scale(monkeypatch) -> None: + gate_name = "model.layers.1.mlp.experts.0.gate_proj.weight" + up_name = "model.layers.1.mlp.experts.0.up_proj.weight" + gate = torch.ones((3, 32), dtype=torch.bfloat16) + up = torch.ones((5, 32), dtype=torch.bfloat16) + pair_calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def fake_pair( + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + ): + pair_calls.append((gate_weight, up_weight)) + shared_scale = torch.tensor(0.25, dtype=torch.float32) + return ( + _fake_nvfp4_output(gate_weight, shared_scale.clone()), + _fake_nvfp4_output(up_weight, shared_scale.clone()), + ) + + monkeypatch.setattr(weight_iterator, "quantize_nvfp4_pair", fake_pair) + entries = _collect_entries( + _make_iterator([(gate_name, gate), (up_name, up)]), + target_precision="nvfp4", + ) + tensors = dict(entries) + + assert pair_calls == [(gate, up)] + assert not any(name.endswith(".input_scale") for name in tensors) + torch.testing.assert_close( + tensors[gate_name.replace(".weight", ".weight_scale_2")], + tensors[up_name.replace(".weight", ".weight_scale_2")], + rtol=0, + atol=0, + ) + + +def test_nvfp4_single_side_skip_keeps_whole_gate_up_pair_in_bf16( + monkeypatch, +) -> None: + gate_name = "model.layers.1.mlp.experts.0.gate_proj.weight" + up_name = "model.layers.1.mlp.experts.0.up_proj.weight" + gate = torch.ones((3, 32), dtype=torch.bfloat16) + up = torch.ones((5, 32), dtype=torch.bfloat16) + + def fail_quantize_pair(*_args: Any, **_kwargs: Any): + raise AssertionError("a partially skipped gate/up pair must not be quantized") + + monkeypatch.setattr( + weight_iterator, + "quantize_nvfp4_pair", + fail_quantize_pair, + ) + entries = _collect_entries( + _make_iterator( + [(gate_name, gate), (up_name, up)], + quantization_config={ + "extra_high_precision_layers_hf": ["gate_proj"], + }, + ), + target_precision="nvfp4", + ) + tensors = dict(entries) + + assert set(tensors) == {gate_name, up_name} + assert tensors[gate_name] is gate + assert tensors[up_name] is up + + +def test_nvfp4_iterator_rejects_incomplete_gate_up_pair() -> None: + gate_name = "model.layers.1.mlp.experts.0.gate_proj.weight" + gate = torch.ones((3, 32), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="incomplete pairs"): + _collect_entries( + _make_iterator([(gate_name, gate)]), + target_precision="nvfp4", + ) + + +def test_nvfp4_iterator_rejects_duplicate_pair_role() -> None: + gate_name = "model.layers.1.mlp.experts.0.gate_proj.weight" + gate = torch.ones((3, 32), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="duplicate gate tensor"): + _collect_entries( + _make_iterator([(gate_name, gate), (gate_name, gate.clone())]), + target_precision="nvfp4", + ) diff --git a/tests/unit/models/generation/sglang/test_mxfp8_setup_selection.py b/tests/unit/models/generation/sglang/test_mxfp8_setup_selection.py new file mode 100644 index 00000000000..e0253758fe6 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_mxfp8_setup_selection.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from nemo_rl.models.generation.sglang.mxfp8_setup import ensure_mxfp8_checkpoint + + +def _write_json(path: Path, value: dict[str, Any]) -> None: + path.write_text(json.dumps(value), encoding="utf-8") + + +def _write_checkpoint_metadata( + checkpoint: Path, + *, + modules_to_not_convert: list[str], + weight_names: list[str], +) -> None: + checkpoint.mkdir() + _write_json( + checkpoint / "config.json", + { + "num_hidden_layers": 2, + "quantization_config": { + "quant_method": "mxfp8", + "weight_block_size": [1, 32], + "scale_fmt": "ue8m0", + "modules_to_not_convert": modules_to_not_convert, + }, + }, + ) + _write_json( + checkpoint / "model.safetensors.index.json", + {"weight_map": {name: "model.safetensors" for name in weight_names}}, + ) + + +def test_existing_mxfp8_checkpoint_preserves_atomic_qkv_skip( + tmp_path: Path, +) -> None: + checkpoint = tmp_path / "mxfp8" + q = "model.layers.0.self_attn.q_proj" + k = "model.layers.0.self_attn.k_proj" + v = "model.layers.0.self_attn.v_proj" + _write_checkpoint_metadata( + checkpoint, + modules_to_not_convert=[q, k, v], + weight_names=[f"{q}.weight", f"{k}.weight", f"{v}.weight"], + ) + config: dict[str, Any] = { + "scheme": "mxfp8", + "extra_high_precision_layers_hf": [q], + } + + result = ensure_mxfp8_checkpoint( + model_path=str(checkpoint), + quantization_cfg=config, + ) + + assert result == str(checkpoint) + assert config["modules_to_not_convert"] == [q, k, v] + + +def test_existing_mxfp8_checkpoint_rejects_new_partial_qkv_skip( + tmp_path: Path, +) -> None: + checkpoint = tmp_path / "mxfp8" + q = "model.layers.0.self_attn.q_proj" + k = "model.layers.0.self_attn.k_proj" + v = "model.layers.0.self_attn.v_proj" + weights = [ + f"{q}.weight", + f"{q}.weight_scale_inv", + f"{k}.weight", + f"{k}.weight_scale_inv", + f"{v}.weight", + f"{v}.weight_scale_inv", + ] + _write_checkpoint_metadata( + checkpoint, + modules_to_not_convert=[], + weight_names=weights, + ) + config: dict[str, Any] = { + "scheme": "mxfp8", + "extra_high_precision_layers_hf": [q], + } + + with pytest.raises(ValueError, match="Reconvert the original HF checkpoint"): + ensure_mxfp8_checkpoint( + model_path=str(checkpoint), + quantization_cfg=config, + ) diff --git a/tests/unit/models/generation/sglang/test_nvfp4_quantization_core.py b/tests/unit/models/generation/sglang/test_nvfp4_quantization_core.py new file mode 100644 index 00000000000..d7725c17cb7 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_nvfp4_quantization_core.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import torch + +from nemo_rl.models.generation.sglang import nvfp4_quantization_core as nvfp4 + + +class _FakeNVFP4Quantizer: + def __init__(self) -> None: + self.quantized_shapes: list[tuple[int, ...]] = [] + + def quantize(self, weight: torch.Tensor) -> SimpleNamespace: + self.quantized_shapes.append(tuple(weight.shape)) + rows, columns = weight.shape + return SimpleNamespace( + _rowwise_data=torch.zeros( + (rows, columns // 2), + dtype=torch.uint8, + device=weight.device, + ), + _rowwise_scale_inv=torch.zeros( + (rows, columns // nvfp4.NVFP4_GROUP_SIZE), + dtype=torch.uint8, + device=weight.device, + ), + _amax_rowwise=weight.abs().max().to(torch.float32).reshape(1), + ) + + +def test_nvfp4_pair_uses_one_quantization_and_independent_shared_scales( + monkeypatch, +) -> None: + quantizer = _FakeNVFP4Quantizer() + monkeypatch.setattr(nvfp4, "_make_nvfp4_quantizer", lambda: quantizer) + + gate = torch.arange(3 * 32, dtype=torch.float32).reshape(3, 32) + up = -torch.arange(5 * 32, dtype=torch.float32).reshape(5, 32) + gate_output, up_output = nvfp4.nvfp4_quantize_2d_pair(gate, up) + + gate_qweight, gate_block_scale, gate_global_scale = gate_output + up_qweight, up_block_scale, up_global_scale = up_output + + # Gate/up are concatenated and padded once, so their ModelOpt global scales + # are bitwise equal while remaining independent tensors. + assert quantizer.quantized_shapes == [(nvfp4.TE_NVFP4_ROW_ALIGNMENT, 32)] + assert gate_qweight.shape == (3, 16) + assert up_qweight.shape == (5, 16) + assert gate_block_scale.shape == (3, 2) + assert up_block_scale.shape == (5, 2) + torch.testing.assert_close(gate_global_scale, up_global_scale, rtol=0, atol=0) + assert ( + gate_global_scale.untyped_storage().data_ptr() + != up_global_scale.untyped_storage().data_ptr() + ) + + +def test_live_nvfp4_entries_do_not_add_static_input_scale() -> None: + name = "model.layers.1.mlp.experts.0.down_proj.weight" + entries = nvfp4.nvfp4_quantized_entries( + name, + ( + torch.zeros((2, 16), dtype=torch.uint8), + torch.zeros((2, 2), dtype=torch.uint8).view(torch.float8_e4m3fn), + torch.ones((), dtype=torch.float32), + ), + include_input_scale=False, + ) + + assert [entry_name for entry_name, _ in entries] == [ + name, + "model.layers.1.mlp.experts.0.down_proj.weight_scale", + "model.layers.1.mlp.experts.0.down_proj.weight_scale_2", + ] diff --git a/tests/unit/models/generation/sglang/test_nvfp4_setup.py b/tests/unit/models/generation/sglang/test_nvfp4_setup.py new file mode 100644 index 00000000000..95f013d2cc8 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_nvfp4_setup.py @@ -0,0 +1,575 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +import safetensors.torch +import torch + +from nemo_rl.models.generation.sglang import nvfp4_setup + + +def _gate(layer: int, expert: int = 0) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.gate_proj.weight" + + +def _up(layer: int, expert: int = 0) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.up_proj.weight" + + +def _down(layer: int, expert: int = 0) -> str: + return f"model.layers.{layer}.mlp.experts.{expert}.down_proj.weight" + + +def _weight(value: float, *, rows: int = 3, columns: int = 32) -> torch.Tensor: + return torch.full((rows, columns), value, dtype=torch.bfloat16) + + +def _write_json(path: Path, value: dict[str, Any]) -> None: + path.write_text(json.dumps(value, indent=2), encoding="utf-8") + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + assert isinstance(value, dict) + return value + + +def _write_source_checkpoint( + model_dir: Path, + shards: dict[str, dict[str, torch.Tensor]], + *, + num_hidden_layers: int, + config_extra: dict[str, Any] | None = None, + indexed: bool = True, +) -> dict[str, str]: + model_dir.mkdir() + config: dict[str, Any] = {"num_hidden_layers": num_hidden_layers} + if config_extra is not None: + config.update(config_extra) + _write_json(model_dir / "config.json", config) + + weight_map: dict[str, str] = {} + total_size = 0 + for filename, tensors in shards.items(): + safetensors.torch.save_file( + tensors, + str(model_dir / filename), + metadata={"format": "pt"}, + ) + for key, tensor in tensors.items(): + assert key not in weight_map + weight_map[key] = filename + total_size += tensor.numel() * tensor.element_size() + + if indexed: + _write_json( + model_dir / "model.safetensors.index.json", + { + "weight_map": weight_map, + "metadata": {"total_size": total_size}, + }, + ) + return weight_map + + +def _fake_quantized( + weight: torch.Tensor, + *, + qvalue: int, + global_scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + qweight = torch.full( + (*weight.shape[:-1], weight.shape[-1] // 2), + qvalue, + dtype=torch.uint8, + ) + block_scale = torch.zeros( + (*weight.shape[:-1], weight.shape[-1] // 16), + dtype=torch.uint8, + ).view(torch.float8_e4m3fn) + weight_scale_2 = torch.full( + weight.shape[:-2], + global_scale, + dtype=torch.float32, + ) + return qweight, block_scale, weight_scale_2 + + +def _quantized_names(weight_name: str) -> set[str]: + base = weight_name.removesuffix(".weight") + return { + weight_name, + f"{base}.weight_scale", + f"{base}.weight_scale_2", + f"{base}.input_scale", + } + + +def _load_shard(path: Path) -> dict[str, torch.Tensor]: + return safetensors.torch.load_file(str(path), device="cpu") + + +def test_same_shard_pair_writes_modelopt_fields_and_metadata( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + gate_name = _gate(0) + up_name = _up(0) + dense_name = "model.layers.0.self_attn.q_proj.weight" + gate = _weight(1.0, rows=3) + up = _weight(2.0, rows=5) + dense = _weight(3.0) + custom_kv_scheme = {"dynamic": True, "num_bits": 8, "type": "float"} + shard_name = "model-00001-of-00001.safetensors" + _write_source_checkpoint( + model_dir, + {shard_name: {gate_name: gate, up_name: up, dense_name: dense}}, + num_hidden_layers=1, + config_extra={"quantization_config": {"kv_cache_scheme": custom_kv_scheme}}, + ) + _write_json( + model_dir / "hf_quant_config.json", + { + "producer": {"name": "source"}, + "quantization": {"calibration": "preserved"}, + }, + ) + _write_json(model_dir / "tokenizer_config.json", {"tokenizer_class": "Test"}) + + pair_calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def fake_pair( + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ]: + pair_calls.append((gate_weight, up_weight)) + return ( + _fake_quantized(gate_weight, qvalue=11, global_scale=0.25), + _fake_quantized(up_weight, qvalue=22, global_scale=0.25), + ) + + monkeypatch.setattr(nvfp4_setup, "quantize_nvfp4_pair", fake_pair) + monkeypatch.setattr( + nvfp4_setup, + "quantize_nvfp4", + lambda _weight: pytest.fail("gate/up must use paired quantization"), + ) + + ignore = nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + ) + + assert len(pair_calls) == 1 + torch.testing.assert_close(pair_calls[0][0], gate) + torch.testing.assert_close(pair_calls[0][1], up) + + output = _load_shard(save_dir / shard_name) + assert set(output) == _quantized_names(gate_name) | _quantized_names(up_name) | { + dense_name + } + assert output[gate_name].dtype == torch.uint8 + assert output[up_name].dtype == torch.uint8 + torch.testing.assert_close(output[dense_name], dense) + for name in (gate_name, up_name): + input_scale = output[name.removesuffix(".weight") + ".input_scale"] + assert input_scale.dtype == torch.float32 + assert input_scale.shape == () + assert input_scale.item() == 1.0 + + output_config = _read_json(save_dir / "config.json") + quantization_config = output_config["quantization_config"] + assert quantization_config["quant_algo"] == "NVFP4" + assert quantization_config["quant_method"] == "modelopt" + assert quantization_config["group_size"] == 16 + assert quantization_config["kv_cache_scheme"] == custom_kv_scheme + assert ignore == quantization_config["ignore"] + assert dense_name.removesuffix(".weight") in ignore + assert "model.layers.0.self_attn.qkv_proj" in ignore + + hf_quant_config = _read_json(save_dir / "hf_quant_config.json") + assert hf_quant_config["producer"] == {"name": "source"} + assert hf_quant_config["quantization"]["calibration"] == "preserved" + assert hf_quant_config["quantization"]["exclude_modules"] == ignore + assert hf_quant_config["quantization"]["quant_algo"] == "NVFP4" + assert hf_quant_config["quantization"]["group_size"] == 16 + assert hf_quant_config["quantization"]["kv_cache_quant_algo"] == "FP8" + assert _read_json(save_dir / "tokenizer_config.json") == {"tokenizer_class": "Test"} + + output_index = _read_json(save_dir / "model.safetensors.index.json") + assert set(output_index["weight_map"]) == set(output) + assert set(output_index["weight_map"].values()) == {shard_name} + expected_total_size = sum( + tensor.numel() * tensor.element_size() for tensor in output.values() + ) + assert output_index["metadata"]["total_size"] == expected_total_size + + +def test_cross_shard_pair_is_quantized_once_and_written_to_source_shards( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + gate_name = _gate(0) + up_name = _up(0) + gate = _weight(1.0, rows=3) + up = _weight(2.0, rows=5) + up_shard = "a-up.safetensors" + gate_shard = "z-gate.safetensors" + _write_source_checkpoint( + model_dir, + { + up_shard: {up_name: up}, + gate_shard: {gate_name: gate}, + }, + num_hidden_layers=1, + ) + + pair_calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def fake_pair( + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ]: + pair_calls.append((gate_weight, up_weight)) + return ( + _fake_quantized(gate_weight, qvalue=31, global_scale=0.5), + _fake_quantized(up_weight, qvalue=47, global_scale=0.5), + ) + + monkeypatch.setattr(nvfp4_setup, "quantize_nvfp4_pair", fake_pair) + nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + ) + + assert len(pair_calls) == 1 + torch.testing.assert_close(pair_calls[0][0], gate) + torch.testing.assert_close(pair_calls[0][1], up) + + gate_output = _load_shard(save_dir / gate_shard) + up_output = _load_shard(save_dir / up_shard) + assert set(gate_output) == _quantized_names(gate_name) + assert set(up_output) == _quantized_names(up_name) + assert torch.all(gate_output[gate_name] == 31) + assert torch.all(up_output[up_name] == 47) + + output_index = _read_json(save_dir / "model.safetensors.index.json") + for name in _quantized_names(gate_name): + assert output_index["weight_map"][name] == gate_shard + for name in _quantized_names(up_name): + assert output_index["weight_map"][name] == up_shard + + +def test_single_side_extra_skip_keeps_entire_expert_container_in_bf16( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + skipped_names = (_gate(0), _up(0), _down(0)) + quantized_name = _down(1) + source_weights = { + skipped_names[0]: _weight(1.0), + skipped_names[1]: _weight(2.0), + skipped_names[2]: _weight(3.0), + quantized_name: _weight(4.0), + } + shard_name = "model.safetensors" + _write_source_checkpoint( + model_dir, + {shard_name: source_weights}, + num_hidden_layers=2, + ) + + quantized_calls: list[torch.Tensor] = [] + + def fake_quantize( + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + quantized_calls.append(weight) + return _fake_quantized(weight, qvalue=61, global_scale=0.75) + + monkeypatch.setattr(nvfp4_setup, "quantize_nvfp4", fake_quantize) + monkeypatch.setattr( + nvfp4_setup, + "quantize_nvfp4_pair", + lambda *_args: pytest.fail("a partially skipped pair must remain BF16"), + ) + + ignore = nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + extra_high_precision_layers_hf=("model.layers.0.mlp.experts.0.gate_proj",), + ) + + assert len(quantized_calls) == 1 + output = _load_shard(save_dir / shard_name) + for name in skipped_names: + torch.testing.assert_close(output[name], source_weights[name]) + assert name.removesuffix(".weight") + ".weight_scale" not in output + assert output[quantized_name].dtype == torch.uint8 + assert "model.layers.0.mlp.experts" in ignore + + +def test_first_and_last_layers_remain_bf16( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + source_weights = {_down(layer): _weight(float(layer + 1)) for layer in range(4)} + shard_name = "model.safetensors" + _write_source_checkpoint( + model_dir, + {shard_name: source_weights}, + num_hidden_layers=4, + ) + + quantized_calls: list[torch.Tensor] = [] + + def fake_quantize( + weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + quantized_calls.append(weight) + return _fake_quantized( + weight, + qvalue=70 + len(quantized_calls), + global_scale=1.0, + ) + + monkeypatch.setattr(nvfp4_setup, "quantize_nvfp4", fake_quantize) + ignore = nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + num_layers_at_start_in_bf16=1, + num_layers_at_end_in_bf16=1, + ) + + assert len(quantized_calls) == 2 + output = _load_shard(save_dir / shard_name) + for layer in (0, 3): + name = _down(layer) + torch.testing.assert_close(output[name], source_weights[name]) + assert name.removesuffix(".weight") + ".weight_scale" not in output + assert f"model.layers.{layer}.mlp.experts" in ignore + for layer in (1, 2): + name = _down(layer) + assert output[name].dtype == torch.uint8 + assert _quantized_names(name) <= set(output) + + +def test_incomplete_gate_up_pair_fails_before_writing_output( + tmp_path: Path, +) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + _write_source_checkpoint( + model_dir, + {"model.safetensors": {_gate(0): _weight(1.0)}}, + num_hidden_layers=1, + ) + + with pytest.raises(ValueError, match="incomplete checkpoint pairs"): + nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + ) + assert not save_dir.exists() + + +def test_duplicate_indexed_tensor_across_shards_fails_loudly( + tmp_path: Path, +) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + model_dir.mkdir() + _write_json(model_dir / "config.json", {"num_hidden_layers": 1}) + gate_name = _gate(0) + up_name = _up(0) + dense_name = "model.embed_tokens.weight" + shard_a = "model-00001-of-00002.safetensors" + shard_b = "model-00002-of-00002.safetensors" + safetensors.torch.save_file( + {gate_name: _weight(1.0), dense_name: _weight(2.0)}, + str(model_dir / shard_a), + ) + safetensors.torch.save_file( + {gate_name: _weight(3.0), up_name: _weight(4.0)}, + str(model_dir / shard_b), + ) + _write_json( + model_dir / "model.safetensors.index.json", + { + "weight_map": { + dense_name: shard_a, + gate_name: shard_b, + up_name: shard_b, + } + }, + ) + + with pytest.raises(ValueError, match="Duplicate source tensor"): + nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + ) + assert not save_dir.exists() + + +def test_index_mismatch_fails_loudly(tmp_path: Path) -> None: + model_dir = tmp_path / "source" + save_dir = tmp_path / "converted" + gate_name = _gate(0) + up_name = _up(0) + _write_source_checkpoint( + model_dir, + {"model.safetensors": {gate_name: _weight(1.0), up_name: _weight(2.0)}}, + num_hidden_layers=1, + ) + index = _read_json(model_dir / "model.safetensors.index.json") + del index["weight_map"][up_name] + _write_json(model_dir / "model.safetensors.index.json", index) + + with pytest.raises(ValueError, match="index does not match"): + nvfp4_setup.convert_nvfp4( + str(model_dir), + str(save_dir), + device="cpu", + ) + assert not save_dir.exists() + + +def test_existing_checkpoint_ignore_is_merged_back_into_refit_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_dir = tmp_path / "source" + source_dir.mkdir() + _write_json(source_dir / "config.json", {"num_hidden_layers": 2}) + converted_dir = tmp_path / "converted" + user_only = "model.layers.0.mlp.experts" + checkpoint_only = "model.layers.1.mlp.experts" + _write_source_checkpoint( + converted_dir, + { + "model.safetensors": { + _down(0): _weight(1.0), + _down(1): _weight(2.0), + } + }, + num_hidden_layers=2, + config_extra={ + "quantization_config": { + "quant_algo": "NVFP4", + "group_size": 16, + "ignore": [checkpoint_only], + } + }, + ) + quantization_cfg: dict[str, Any] = { + "scheme": "nvfp4", + "converted_model_path": str(converted_dir), + "modules_to_not_convert": [user_only], + } + monkeypatch.setattr( + nvfp4_setup, + "convert_nvfp4", + lambda *_args, **_kwargs: pytest.fail("existing checkpoint must be reused"), + ) + + result = nvfp4_setup.ensure_nvfp4_checkpoint( + model_path=str(source_dir), + quantization_cfg=quantization_cfg, + ) + + assert result == str(converted_dir) + assert quantization_cfg["modules_to_not_convert"] == [ + user_only, + checkpoint_only, + ] + + +def test_existing_checkpoint_rejects_new_head_layer_bf16_policy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_dir = tmp_path / "source" + source_dir.mkdir() + _write_json(source_dir / "config.json", {"num_hidden_layers": 2}) + converted_dir = tmp_path / "converted" + gate_name = _gate(0) + gate = _weight(1.0) + qweight, block_scale, global_scale = _fake_quantized( + gate, + qvalue=7, + global_scale=0.5, + ) + gate_base = gate_name.removesuffix(".weight") + _write_source_checkpoint( + converted_dir, + { + "model.safetensors": { + gate_name: qweight, + f"{gate_base}.weight_scale": block_scale, + f"{gate_base}.weight_scale_2": global_scale, + f"{gate_base}.input_scale": torch.ones((), dtype=torch.float32), + } + }, + num_hidden_layers=2, + config_extra={ + "quantization_config": { + "quant_algo": "NVFP4", + "group_size": 16, + "ignore": [], + } + }, + ) + quantization_cfg: dict[str, Any] = { + "scheme": "nvfp4", + "converted_model_path": str(converted_dir), + "num_layers_at_start_in_bf16": 1, + } + monkeypatch.setattr( + nvfp4_setup, + "convert_nvfp4", + lambda *_args, **_kwargs: pytest.fail("existing checkpoint must be reused"), + ) + + with pytest.raises(ValueError, match="Reconvert the original HF checkpoint"): + nvfp4_setup.ensure_nvfp4_checkpoint( + model_path=str(source_dir), + quantization_cfg=quantization_cfg, + ) diff --git a/tests/unit/models/generation/sglang/test_quantization_utils.py b/tests/unit/models/generation/sglang/test_quantization_utils.py new file mode 100644 index 00000000000..2d493ab4ee4 --- /dev/null +++ b/tests/unit/models/generation/sglang/test_quantization_utils.py @@ -0,0 +1,165 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemo_rl.models.generation.sglang.quantization_utils import ( + expand_sglang_atomic_high_precision_substrings, + get_dynamic_high_precision_substrings, + get_sglang_quantization_scheme, + validate_sglang_quantized_refit_backend, +) + + +@pytest.mark.parametrize("scheme", ["bf16", "mxfp8", "nvfp4"]) +def test_get_sglang_quantization_scheme_accepts_supported_values(scheme: str) -> None: + assert get_sglang_quantization_scheme({"scheme": scheme}) == scheme + + +def test_get_sglang_quantization_scheme_defaults_only_when_block_is_absent() -> None: + assert get_sglang_quantization_scheme(None) == "bf16" + assert get_sglang_quantization_scheme({}) == "bf16" + + with pytest.raises(ValueError, match=r"quantization\.scheme"): + get_sglang_quantization_scheme({"modules_to_not_convert": []}) + + +def test_get_sglang_quantization_scheme_rejects_unknown_value() -> None: + with pytest.raises(ValueError, match=r"got 'fp4'"): + get_sglang_quantization_scheme({"scheme": "fp4"}) + + +def test_high_precision_substrings_combine_extra_head_tail_and_deduplicate() -> None: + result = get_dynamic_high_precision_substrings( + quantization_config={ + "extra_high_precision_layers_hf": [ + "model.layers.2.mlp.experts", + "model.layers.0.", + ], + "modules_to_not_convert": [ + "model.layers.2.mlp.experts", + "lm_head", + ], + "num_layers_at_start_in_bf16": 1, + "num_layers_at_end_in_bf16": 1, + }, + num_hidden_layers=4, + ) + + assert result == ( + "model.layers.2.mlp.experts", + "model.layers.0.", + "lm_head", + "model.layers.3.", + ) + + +def test_layer_prefix_for_layer_one_does_not_match_layer_ten() -> None: + substrings = get_dynamic_high_precision_substrings( + quantization_config={"num_layers_at_start_in_bf16": 2}, + num_hidden_layers=12, + ) + + layer_one = "model.layers.1.mlp.experts.0.down_proj.weight" + layer_ten = "model.layers.10.mlp.experts.0.down_proj.weight" + assert any(substring in layer_one for substring in substrings) + assert not any(substring in layer_ten for substring in substrings) + + +def test_atomic_high_precision_expands_fused_linear_and_moe_modules() -> None: + result = expand_sglang_atomic_high_precision_substrings( + weight_names=[ + "model.layers.1.self_attn.q_proj.weight", + "model.layers.1.self_attn.k_proj.weight", + "model.layers.1.self_attn.v_proj.weight", + "model.layers.1.mlp.gate_proj.weight", + "model.layers.1.mlp.up_proj.weight", + "model.layers.2.mlp.experts.0.gate_proj.weight", + "model.layers.2.mlp.experts.0.up_proj.weight", + "model.layers.2.mlp.experts.0.down_proj.weight", + ], + skip_weight_substrings=( + "model.layers.1.self_attn.q_proj", + "model.layers.1.mlp.gate_proj", + "model.layers.2.mlp.experts.0.gate_proj", + ), + ) + + assert "model.layers.1.self_attn.k_proj" in result + assert "model.layers.1.self_attn.v_proj" in result + assert "model.layers.1.mlp.up_proj" in result + assert "model.layers.2.mlp.experts" in result + + +def test_quantized_refit_requires_megatron_backend() -> None: + validate_sglang_quantized_refit_backend(scheme="bf16", use_megatron=False) + validate_sglang_quantized_refit_backend(scheme="mxfp8", use_megatron=True) + validate_sglang_quantized_refit_backend(scheme="nvfp4", use_megatron=True) + + with pytest.raises(NotImplementedError, match="requires a Megatron policy"): + validate_sglang_quantized_refit_backend( + scheme="nvfp4", + use_megatron=False, + ) + + +@pytest.mark.parametrize( + ("quantization_config", "num_hidden_layers", "exception", "match"), + [ + ( + {"num_layers_at_start_in_bf16": -1}, + 4, + ValueError, + "non-negative", + ), + ( + {"num_layers_at_end_in_bf16": True}, + 4, + TypeError, + "must be an integer", + ), + ( + {"num_layers_at_start_in_bf16": 1.5}, + 4, + TypeError, + "must be an integer", + ), + ( + { + "num_layers_at_start_in_bf16": 3, + "num_layers_at_end_in_bf16": 2, + }, + 4, + ValueError, + "exceed", + ), + ( + {"num_layers_at_end_in_bf16": 1}, + 0, + ValueError, + "must be positive", + ), + ], +) +def test_high_precision_substrings_reject_invalid_layer_counts( + quantization_config: dict, + num_hidden_layers: int, + exception: type[Exception], + match: str, +) -> None: + with pytest.raises(exception, match=match): + get_dynamic_high_precision_substrings( + quantization_config=quantization_config, + num_hidden_layers=num_hidden_layers, + ) diff --git a/tests/unit/models/generation/sglang/test_sglang_patches.py b/tests/unit/models/generation/sglang/test_sglang_patches.py new file mode 100644 index 00000000000..873f85c23df --- /dev/null +++ b/tests/unit/models/generation/sglang/test_sglang_patches.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from nemo_rl.models.generation.sglang.utils import patches + + +def test_weight_update_session_completion_requires_cross_file_state( + tmp_path: Path, + monkeypatch, +) -> None: + relative_paths = ( + "srt/managers/io_struct.py", + "srt/entrypoints/http_server.py", + "srt/managers/tokenizer_control_mixin.py", + "srt/managers/scheduler.py", + "srt/managers/scheduler_update_weights_mixin.py", + ) + paths = { + relative_path: tmp_path / relative_path.replace("/", "_") + for relative_path in relative_paths + } + for path in paths.values(): + path.write_text("", encoding="utf-8") + monkeypatch.setattr( + patches, + "_get_sglang_file", + lambda relative_path: str(paths[relative_path]), + ) + + paths["srt/managers/io_struct.py"].write_text( + "class BeginWeightUpdateReqInput:\n" + " pass\n" + "class EndWeightUpdateReqInput:\n" + " pass\n", + encoding="utf-8", + ) + paths["srt/entrypoints/http_server.py"].write_text( + '@app.post("/begin_weight_update")\n@app.post("/end_weight_update")\n', + encoding="utf-8", + ) + assert not patches._sglang_weight_update_session_is_complete() + + paths["srt/managers/tokenizer_control_mixin.py"].write_text( + "async def begin_weight_update():\n" + " pass\n" + "async def end_weight_update():\n" + " pass\n", + encoding="utf-8", + ) + paths["srt/managers/scheduler.py"].write_text( + "self.weight_updater.begin_weight_update\n" + "self.weight_updater.end_weight_update\n", + encoding="utf-8", + ) + assert patches._sglang_weight_update_session_is_complete() + + paths["srt/managers/scheduler.py"].write_text( + "(BeginWeightUpdateReqInput, self.begin_weight_update)\n" + "(EndWeightUpdateReqInput, self.end_weight_update)\n", + encoding="utf-8", + ) + paths["srt/managers/scheduler_update_weights_mixin.py"].write_text( + "def begin_weight_update():\n" + " pass\n" + "def end_weight_update():\n" + " pass\n" + "update_weights_from_distributed requires an open session\n" + "update_weights_from_tensor requires an open session\n", + encoding="utf-8", + ) + assert patches._sglang_weight_update_session_is_complete() diff --git a/uv.lock b/uv.lock index 96543608420..3c90685d993 100644 --- a/uv.lock +++ b/uv.lock @@ -101,7 +101,7 @@ overrides = [ { name = "timm", specifier = "<=1.0.22" }, { name = "torch", marker = "sys_platform != 'darwin'", specifier = "==2.11.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = "==2.11.0", index = "https://pypi.org/simple" }, - { name = "transformer-engine", extras = ["pytorch", "core-cu13"], git = "https://github.com/NVIDIA/TransformerEngine.git?rev=release_v2.15" }, + { name = "transformer-engine", extras = ["pytorch", "core-cu13"], git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.17" }, { name = "xgrammar", specifier = "==0.1.33" }, ] excludes = ["nvidia-cutlass-dsl-libs-base"] @@ -170,12 +170,12 @@ version = "10.16.1.11" [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.15.0+42b8400" +version = "2.17.0+2e559f0" requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[manifest.dependency-metadata]] name = "transformer-engine-torch" -version = "2.15.0+42b8400" +version = "2.17.0+2e559f0" requires-dist = ["torch", "transformer-engine"] [[package]] @@ -4422,7 +4422,7 @@ requires-dist = [ { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = "==0.26.0", index = "https://pypi.org/simple" }, { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.17" }, { name = "transformers", specifier = ">=5.5.0,<5.9.0" }, { name = "transformers", marker = "extra == 'automodel'", specifier = ">=5.5.0,<5.6.0" }, { name = "transformers", marker = "extra == 'sglang'", specifier = "==5.6.0" }, @@ -7890,8 +7890,8 @@ dependencies = [ [[package]] name = "transformer-engine" -version = "2.15.0+42b8400" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=release_v2.15#42b840051647eef89761a16dfdff87e82bb253ab" } +version = "2.17.0+2e559f0" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.17#2e559f062497bef768dfbe9d7e45548fadeca80a" } dependencies = [ { name = "importlib-metadata" }, { name = "packaging" },