diff --git a/python/sglang/srt/arg_groups/model_overrides/nemotron_h.py b/python/sglang/srt/arg_groups/model_overrides/nemotron_h.py index d316db88f..d9fa5bf3c 100644 --- a/python/sglang/srt/arg_groups/model_overrides/nemotron_h.py +++ b/python/sglang/srt/arg_groups/model_overrides/nemotron_h.py @@ -1,6 +1,7 @@ """Config-time override declarations for nemotron_h. -Architectures: NemotronHForCausalLM, NemotronHPuzzleForCausalLM. +Architectures: NemotronHForCausalLM, NemotronHPuzzleForCausalLM, +NemotronH_Omni_Reasoning_V3. """ import logging @@ -17,7 +18,11 @@ from sglang.srt.runtime_context import get_platform logger = logging.getLogger(__name__) -@_register_for("NemotronHForCausalLM", "NemotronHPuzzleForCausalLM") +@_register_for( + "NemotronHForCausalLM", + "NemotronHPuzzleForCausalLM", + "NemotronH_Omni_Reasoning_V3", +) def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict: """NemotronH quantization / MoE runner / attention backend defaults (absorbed from the retired arg_groups/nemotron_h_hook.py; the mamba radix @@ -35,7 +40,12 @@ def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict: ] quantization = cfg.quantization if is_modelopt: - assert model_config.hf_config.mlp_hidden_act == "relu2" + language_config = ( + model_config.hf_text_config + if model_arch == "NemotronH_Omni_Reasoning_V3" + else hf_config + ) + assert language_config.mlp_hidden_act == "relu2" if model_config.quantization == "modelopt": quant_algo = model_config.hf_config.quantization_config["quant_algo"] if quant_algo == "MIXED_PRECISION": diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index b49b490a6..cd6b64db7 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -456,6 +456,7 @@ _MAMBA_RADIX_CACHE_ARCHS = frozenset( "MiniCPMV4_6ForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", + "NemotronH_Omni_Reasoning_V3", "FalconH1ForCausalLM", "JetNemotronForCausalLM", "JetVLMForConditionalGeneration", @@ -488,6 +489,7 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset( "Glm5NextForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", + "NemotronH_Omni_Reasoning_V3", # KDA-based: same MambaPool ping-pong machinery as GDN; requires the # KDA backend's track-snapshot writes (decode + extend) so donated # slots hold real states for prefix-cache restores. @@ -1001,6 +1003,7 @@ _FLASHINFER_ALLREDUCE_FUSION_ARCHS = frozenset( "Qwen3_5ForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", + "NemotronH_Omni_Reasoning_V3", } ) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 7278a8d08..0f2a97ebb 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -57,6 +57,7 @@ from sglang.srt.configs.nanbeige import NanbeigeConfig from sglang.srt.configs.nano_nemotron_vl import ( NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, + NemotronH_Omni_Reasoning_V3_Config, ) from sglang.srt.configs.nemotron_h import NemotronHConfig, NemotronHPuzzleConfig from sglang.srt.configs.olmo3 import Olmo3Config @@ -137,6 +138,7 @@ __all__ = [ "NemotronHPuzzleConfig", "NemotronH_Nano_VL_V2_Config", "NemotronH_Nano_Omni_Reasoning_V3_Config", + "NemotronH_Omni_Reasoning_V3_Config", "NanbeigeConfig", "JetNemotronConfig", "JetVLMConfig", diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 0f8874b63..4de15464b 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -901,6 +901,14 @@ class ModelConfig: self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 + if ( + is_draft_model + and self.hf_config.architectures[0] == "NemotronH_Omni_Reasoning_V3" + ): + self.hf_config = self.hf_text_config + self.hf_config.architectures = ["NemotronHForCausalLMMTP"] + self.hf_config.num_nextn_predict_layers = 1 + if is_draft_model and self.hf_config.architectures[0] in [ "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", @@ -2086,6 +2094,7 @@ multimodal_model_archs = [ "MossVLForConditionalGeneration", "NemotronH_Nano_VL_V2", "NemotronH_Nano_Omni_Reasoning_V3", + "NemotronH_Omni_Reasoning_V3", "MuseGlimmerForConditionalGeneration", "PixtralForConditionalGeneration", "Qwen2AudioForConditionalGeneration", diff --git a/python/sglang/srt/configs/nano_nemotron_vl.py b/python/sglang/srt/configs/nano_nemotron_vl.py index 30c6b4390..2e3798783 100644 --- a/python/sglang/srt/configs/nano_nemotron_vl.py +++ b/python/sglang/srt/configs/nano_nemotron_vl.py @@ -30,6 +30,15 @@ def float_triplet(seq: Any): return a, b, c +def _nemotron_h_compatible_config(config: dict) -> dict: + config = dict(config) + aliases = {"linear_attention": "mamba", "full_attention": "attention"} + for field in ("layers_block_type", "mtp_layers_block_type"): + if config.get(field) is not None: + config[field] = [aliases.get(value, value) for value in config[field]] + return config + + class NemotronH_Nano_VL_V2_Config(PretrainedConfig): model_type = "NemotronH_Nano_VL_V2" is_composition = True @@ -166,3 +175,15 @@ class NemotronH_Nano_Omni_Reasoning_V3_Config(NemotronH_Nano_VL_V2_Config): # Explicit __init__ prevents PretrainedConfig.__init_subclass__ from # replacing the parent's custom __init__ with a dataclass-generated one. super().__init__(*args, **kwargs) + + +class NemotronH_Omni_Reasoning_V3_Config(NemotronH_Nano_Omni_Reasoning_V3_Config): + model_type = "nemotron_h_omni" + + def __init__(self, *args, **kwargs): + args = list(args) + if len(args) > 1 and args[1] is not None: + args[1] = _nemotron_h_compatible_config(args[1]) + elif kwargs.get("llm_config") is not None: + kwargs["llm_config"] = _nemotron_h_compatible_config(kwargs["llm_config"]) + super().__init__(*args, **kwargs) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 19672d8ca..f0914b18d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -972,10 +972,11 @@ class Scheduler( initialize_mamba_selective_state_update_backend(self.server_args) def init_moe_gemm_config(self): - # For the MM models, check the text_config for MoE settings - config_to_check = getattr( - self.model_config.hf_config, "text_config", self.model_config.hf_config - ) + config_to_check = self.model_config.hf_config + if hasattr(self.model_config.hf_config, "text_config"): + config_to_check = self.model_config.hf_config.text_config + elif hasattr(self.model_config, "hf_text_config"): + config_to_check = self.model_config.hf_text_config # Different MoE architectures expose the per-token expert count under # different attribute names (e.g. Gemma4 uses ``top_k_experts``, diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index 021d33ad8..f35109ae1 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/python/sglang/srt/models/nano_nemotron_vl.py @@ -16,6 +16,7 @@ # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/nano_nemotron_vl.py import logging +from collections import deque from typing import Iterable import torch @@ -178,6 +179,18 @@ class NemotronH_Nano_VL_V2(EVS): x = x.permute(0, 2, 1, 3).contiguous() return x + def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor: + return features + + def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None: + return + + def get_embed_and_head(self): + return self.language_model.get_embed_and_head() + + def set_embed_and_head(self, embed, head): + self.language_model.set_embed_and_head(embed, head) + def extract_feature_dynamic(self, pixel_values_list: list[torch.Tensor]): """Extract features from variable-size images (dynamic resolution). @@ -190,6 +203,7 @@ class NemotronH_Nano_VL_V2(EVS): offset = 0 for i, num_patches in enumerate(num_patches_list): img_feats = features[0, offset : offset + num_patches] + img_feats = self._normalize_vision_features(img_feats) h_patches = pixel_values_list[i].shape[-2] // patch_size w_patches = pixel_values_list[i].shape[-1] // patch_size img_feats = img_feats.reshape(1, h_patches, w_patches, -1) @@ -203,6 +217,7 @@ class NemotronH_Nano_VL_V2(EVS): def extract_video_feature_temporal(self, pixel_values, num_frames): """Extract video features with temporal compression (tubelet grouping).""" vit_embeds = self.vision_model(pixel_values, num_frames=num_frames) + vit_embeds = self._normalize_vision_features(vit_embeds) num_tubelets = vit_embeds.shape[0] patch_size = self.config.patch_size h_patches = pixel_values.shape[-2] // patch_size @@ -217,6 +232,9 @@ class NemotronH_Nano_VL_V2(EVS): def get_input_embeddings(self): return self.language_model.get_input_embeddings() + def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None: + self.language_model.set_dflash_layers_to_capture(layer_ids) + def extract_feature(self, pixel_values): micro_batch_size = 128 n = pixel_values.shape[0] @@ -229,6 +247,7 @@ class NemotronH_Nano_VL_V2(EVS): batch_size = chunk.shape[0] vit_embeds = self.vision_model(chunk) vit_embeds = vit_embeds.to(dtype=self.model_dtype) + vit_embeds = self._normalize_vision_features(vit_embeds) vit_embeds = vit_embeds.reshape(batch_size, h_patches, w_patches, -1) vit_embeds = self.pixel_shuffle( vit_embeds, scale_factor=self.downsample_ratio @@ -358,36 +377,96 @@ class NemotronH_Nano_VL_V2(EVS): def is_sound_weights(name: str) -> bool: return name.startswith("sound") - # Separate weights by component - llm_weights = [] - vision_weights = [] - sound_weights = [] - - for name, w in weights: - if is_llm(name): - # Strip 'language_model.' prefix for LLM weights - llm_weights.append((".".join(name.split(".")[1:]), w)) - elif is_adapter_weights((name, w)): - # Load vision-language adapter weights directly - trimmed_name = ".".join(name.split(".")[1:]) - param = adapter_dict[trimmed_name] - with torch.no_grad(): - default_weight_loader(param, w) - elif is_vision_weights(name): - # Convert: vision_model.radio_model.* → radio_model.* - hf_key = name[len("vision_model.") :] - vision_weights.append((hf_key, w)) - elif is_sound_weights(name): - sound_weights.append((name, w)) + def iter_llm_weights(): + for name, w in weights: + if is_llm(name): + yield (".".join(name.split(".")[1:]), w) + elif is_adapter_weights((name, w)): + trimmed_name = ".".join(name.split(".")[1:]) + param = adapter_dict[trimmed_name] + with torch.no_grad(): + default_weight_loader(param, w) + elif is_vision_weights(name): + hf_key = name[len("vision_model.") :] + self.vision_model.load_weights([(hf_key, w)]) + elif is_sound_weights(name): + if self.sound_encoder is not None: + self.sound_encoder.load_weights([(name, w)]) + else: + self._load_extra_weight(name, w) + llm_weights = iter_llm_weights() self.language_model.load_weights(llm_weights) - self.vision_model.load_weights(vision_weights) - if self.sound_encoder is not None and len(sound_weights) > 0: - self.sound_encoder.load_weights(sound_weights) + deque(llm_weights, maxlen=0) class NemotronH_Nano_Omni_Reasoning_V3(NemotronH_Nano_VL_V2): pass -EntryClass = [NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3] +class NemotronH_Omni_Reasoning_V3(NemotronH_Nano_VL_V2): + packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping + _hf_projector_weight_names = { + "vision_projector.mlp1.norm.": "mlp1.0.", + "vision_projector.mlp1.linear1.": "mlp1.1.", + "vision_projector.mlp1.linear2.": "mlp1.3.", + } + + def __init__(self, config, quant_config=None, prefix: str = ""): + super().__init__(config, quant_config, prefix) + self.vision_final_layernorm = ( + nn.LayerNorm( + config.vit_hidden_size, + eps=config.raw_vision_config.get("layer_norm_eps", 1e-6), + ).to(self.model_dtype) + if (config.llm_config.num_nextn_predict_layers or 0) > 0 + else None + ) + + @property + def lm_head(self): + return self.language_model.lm_head + + def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor: + if self.vision_final_layernorm is None: + return features + return self.vision_final_layernorm(features) + + def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None: + prefix = "vision_projector.vision_final_layernorm." + if not name.startswith(prefix): + raise ValueError(f"Unexpected Nemotron-H Omni weight: {name}") + if self.vision_final_layernorm is None: + raise ValueError(f"Unexpected vision projector weight: {name}") + parameter_name = name.removeprefix(prefix) + parameters = dict(self.vision_final_layernorm.named_parameters()) + if parameter_name not in parameters: + raise ValueError(f"Unexpected vision projector weight: {name}") + parameter = parameters[parameter_name] + default_weight_loader(parameter, weight) + + @classmethod + def _remap_checkpoint_weight_name(cls, name: str) -> str: + for source, target in cls._hf_projector_weight_names.items(): + if name.startswith(source): + return name.replace(source, target, 1) + if name.startswith("vision_model.") and not name.startswith( + "vision_model.radio_model." + ): + return name.replace( + "vision_model.", "vision_model.radio_model.hf_model.", 1 + ) + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + super().load_weights( + (self._remap_checkpoint_weight_name(name), weight) + for name, weight in weights + ) + + +EntryClass = [ + NemotronH_Nano_VL_V2, + NemotronH_Nano_Omni_Reasoning_V3, + NemotronH_Omni_Reasoning_V3, +] diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 5df677188..d64caae1f 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -27,6 +27,7 @@ from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ColumnParallelLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4A16LinearMethod from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -38,6 +39,7 @@ from sglang.srt.models.nemotron_h import ( NemotronHMoEDecoderLayer, ) from sglang.srt.models.nemotron_h_utils import is_attn_layer +from sglang.srt.models.utils import WeightsMapper from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix @@ -293,7 +295,21 @@ class NemotronHMultiTokenPredictor(nn.Module): inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor: if inputs_embeds is None: - inputs_embeds = self.get_input_embeddings(input_ids) + inputs_embeds = forward_batch.mm_input_embeds + if ( + forward_batch.forward_mode.is_extend() + and forward_batch.contains_mm_inputs() + and not forward_batch.forward_mode.is_draft_extend_v2() + ): + assert inputs_embeds is not None + last_indices = ( + forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1 + ).long() + inputs_embeds[last_indices] = self.get_input_embeddings( + input_ids[last_indices] + ) + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings(input_ids) hidden_states = forward_batch.spec_info.hidden_states residual = None @@ -309,6 +325,10 @@ class NemotronHMultiTokenPredictor(nn.Module): class NemotronHForCausalLMMTP(NemotronHForCausalLM): + hf_to_sglang_mapper = NemotronHForCausalLM.hf_to_sglang_mapper | WeightsMapper( + orig_to_new_prefix={"language_model.mtp.": "mtp."} + ) + def __init__( self, config: NemotronHConfig, @@ -319,6 +339,7 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM): config = config.get_mtp_config() self.config = config self.quant_config = quant_config + self._owns_lm_head = False # Required for parent's load_weights self.pp_group = get_pp_group() @@ -366,10 +387,56 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM): def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False ): - super().load_weights(weights, is_mtp=True) + has_mtp_layers = False + has_target_layers = False + head_weights = set() + + def normalized_weights(): + nonlocal has_mtp_layers, has_target_layers + for name, weight in weights: + name = name.removeprefix("language_model.") + has_mtp_layers |= name.startswith("mtp.layers.") + has_target_layers |= name.startswith( + ("backbone.layers.", "model.layers.") + ) + if name.startswith("lm_head."): + head_weights.add(name) + yield name, weight + + # Inspect names while streaming: buffering a full target checkpoint here + # would double its host-memory footprint during embedded MTP loading. + super().load_weights(normalized_weights(), is_mtp=True) + self._owns_lm_head = bool( + has_mtp_layers and not has_target_layers and head_weights + ) + if self._owns_lm_head: + expected = { + name + for name, _ in self.named_parameters() + if name.startswith("lm_head.") + } + if "lm_head.input_scale" in expected and isinstance( + self.lm_head.quant_method, ModelOptNvFp4A16LinearMethod + ): + # NVFP4A16 accepts this loader placeholder but never uses it. + expected.remove("lm_head.input_scale") + missing = (expected | {"lm_head.weight"}) - head_weights + if missing: + raise ValueError( + f"Incomplete standalone MTP lm_head: missing {sorted(missing)}" + ) + + def set_embed_and_head(self, embed, head): + if not self._owns_lm_head: + return super().set_embed_and_head(embed, head) + # Standalone MTP checkpoints can supply a differently quantized head. + # Share only the input embeddings; retain the entire loaded head module. + self.model.embed_tokens.weight = embed + torch.cuda.empty_cache() + torch.cuda.synchronize() def set_lm_head_from_target(self, target_lm_head: nn.Module) -> None: - if self.config.tie_word_embeddings: + if self.config.tie_word_embeddings or self._owns_lm_head: return self.lm_head = target_lm_head diff --git a/python/sglang/srt/models/radio.py b/python/sglang/srt/models/radio.py index 20edc607b..bdb2cc79c 100644 --- a/python/sglang/srt/models/radio.py +++ b/python/sglang/srt/models/radio.py @@ -42,6 +42,41 @@ input_dim_t: TypeAlias = int | tuple[int, int] norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor +def _map_hf_radio_weight_name(name: str) -> tuple[str, str | None] | None: + prefix = "radio_model.hf_model." + if not name.startswith(prefix): + return None + + name = name.removeprefix(prefix) + if name == "summary_idxs": + return None + + embedding_names = { + "embeddings.cls_register_token": "model.patch_generator.cls_token.token", + "embeddings.patch_projection": "model.patch_generator.embedder", + "embeddings.position_embedding": "model.patch_generator.pos_embed", + "embeddings.video_patch_projection": "model.patch_generator.video_embedder", + } + for source, target in embedding_names.items(): + if name == source or name.startswith(f"{source}."): + return name.replace(source, target, 1), None + + name = name.replace("encoder.layer.", "model.encoder.layers.", 1) + attention_names = { + ".attention.attention.query.": (".attn.attn.qkv_proj.", "q"), + ".attention.attention.key.": (".attn.attn.qkv_proj.", "k"), + ".attention.attention.value.": (".attn.attn.qkv_proj.", "v"), + ".attention.output.dense.": (".attn.attn.proj.", None), + } + for source, (target, shard_id) in attention_names.items(): + if source in name: + return name.replace(source, target, 1), shard_id + + name = name.replace(".layer_scale1.lambda1", ".ls1") + name = name.replace(".layer_scale2.lambda1", ".ls2") + return name, None + + def _ntuple(n): def parse(x): if isinstance(x, Iterable) and not isinstance(x, str): @@ -588,18 +623,34 @@ class RadioModel(nn.Module): weights_list = list(weights) for name, weight in weights_list: - if not name.startswith("radio_model."): - # Skip non-radio weights - continue - name = replace_substrings(name, remap_substrings) - name = replace_prefix(name, remap_prefixes) + source_name = name + is_hf_export = name.startswith("radio_model.hf_model.") + loaded_shard_id = None + if is_hf_export: + mapped_weight = _map_hf_radio_weight_name(name) + if mapped_weight is None: + continue + name, loaded_shard_id = mapped_weight + else: + if not name.startswith("radio_model."): + # Skip non-radio weights + continue + name = replace_substrings(name, remap_substrings) + name = replace_prefix(name, remap_prefixes) if name and name in params_dict: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, weight) + if loaded_shard_id is None: + weight_loader(param, weight) + else: + weight_loader(param, weight, loaded_shard_id) loaded_params.add(name) if "video_embedder" in name: self.model.patch_generator._video_embedder_loaded = True + elif is_hf_export: + raise ValueError( + f"Unexpected HF RADIO weight: {source_name} (mapped to {name})" + ) return loaded_params diff --git a/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py b/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py index 08f0b9aeb..cd3e62c2a 100644 --- a/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py +++ b/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py @@ -22,6 +22,7 @@ from PIL import Image from sglang.srt.configs.nano_nemotron_vl import ( NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, + NemotronH_Omni_Reasoning_V3_Config, ) from sglang.srt.managers.schedule_batch import ( Modality, @@ -31,6 +32,7 @@ from sglang.srt.managers.schedule_batch import ( from sglang.srt.models.nano_nemotron_vl import ( NemotronH_Nano_Omni_Reasoning_V3, NemotronH_Nano_VL_V2, + NemotronH_Omni_Reasoning_V3, ) from sglang.srt.models.parakeet import ParakeetExtractor from sglang.srt.multimodal.audio_from_video import extract_audio_from_video_bytes @@ -58,7 +60,11 @@ MAX_FRAMES = 128 class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): - models = [NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3] + models = [ + NemotronH_Nano_VL_V2, + NemotronH_Nano_Omni_Reasoning_V3, + NemotronH_Omni_Reasoning_V3, + ] gpu_image_decode = ( False # NanoNemotronVL processes loaded image as PIL image explicitly ) @@ -70,6 +76,7 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): { NemotronH_Nano_VL_V2_Config: NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3_Config: NemotronH_Nano_Omni_Reasoning_V3, + NemotronH_Omni_Reasoning_V3_Config: NemotronH_Omni_Reasoning_V3, }, ) Image.MAX_IMAGE_PIXELS = None diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index d61d2407c..dc00e67f8 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -68,6 +68,7 @@ from sglang.srt.configs import ( NanbeigeConfig, NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, + NemotronH_Omni_Reasoning_V3_Config, NemotronHConfig, NemotronHPuzzleConfig, Olmo3Config, @@ -138,6 +139,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { Dots3Config, NemotronH_Nano_VL_V2_Config, NemotronH_Nano_Omni_Reasoning_V3_Config, + NemotronH_Omni_Reasoning_V3_Config, NemotronHConfig, NemotronHPuzzleConfig, NanbeigeConfig, diff --git a/test/registered/unit/configs/test_model_config.py b/test/registered/unit/configs/test_model_config.py index bdc43661f..e077075ae 100644 --- a/test/registered/unit/configs/test_model_config.py +++ b/test/registered/unit/configs/test_model_config.py @@ -7,6 +7,7 @@ from sglang.srt.configs.model_config import ( ModelConfig, get_hybrid_layer_ids, is_embedding_gemma, + is_multimodal_model, resolve_spec_hidden_size, ) from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig @@ -56,6 +57,9 @@ class TestEmbeddingGemmaConfig(CustomTestCase): class TestDraftModelConfig(CustomTestCase): + def test_nemotron_h_omni_is_multimodal(self): + self.assertTrue(is_multimodal_model(["NemotronH_Omni_Reasoning_V3"])) + def test_qwen35_mtp_depth_is_synced_to_text_config(self): config = object.__new__(ModelConfig) config.is_draft_model = True @@ -71,6 +75,21 @@ class TestDraftModelConfig(CustomTestCase): self.assertEqual(config.hf_config.num_nextn_predict_layers, 1) self.assertEqual(config.hf_text_config.num_nextn_predict_layers, 1) + def test_nemotron_h_omni_mtp_uses_language_model_config(self): + config = object.__new__(ModelConfig) + config.is_draft_model = True + config.speculative_algorithm = "EAGLE" + config.hf_config = SimpleNamespace( + architectures=["NemotronH_Omni_Reasoning_V3"] + ) + config.hf_text_config = SimpleNamespace(architectures=["NemotronHForCausalLM"]) + + config._config_draft_model() + + self.assertIs(config.hf_config, config.hf_text_config) + self.assertEqual(config.hf_config.architectures, ["NemotronHForCausalLMMTP"]) + self.assertEqual(config.hf_config.num_nextn_predict_layers, 1) + def test_qwen4_exp_spec_hidden_size_keeps_hc_width(self): """Qwen4-Exp's MTP draft consumes the hc-flattened target stream, so spec_hidden_size must stay hidden_size * hc_mult; hy_v4 collapses first.""" diff --git a/test/registered/unit/configs/test_nano_nemotron_vl_config.py b/test/registered/unit/configs/test_nano_nemotron_vl_config.py new file mode 100644 index 000000000..a8612ccb2 --- /dev/null +++ b/test/registered/unit/configs/test_nano_nemotron_vl_config.py @@ -0,0 +1,51 @@ +"""Unit tests for Nano Nemotron VL configuration compatibility.""" + +import unittest + +from sglang.srt.configs.nano_nemotron_vl import ( + NemotronH_Omni_Reasoning_V3_Config, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestNemotronHOmniConfig(CustomTestCase): + def test_uses_checkpoint_model_type(self): + config = NemotronH_Omni_Reasoning_V3_Config( + vision_config={"args": {"model": "radio"}}, + llm_config={}, + architectures=["NemotronH_Omni_Reasoning_V3"], + ) + + self.assertEqual(config.model_type, "nemotron_h_omni") + + def test_normalizes_current_nemotron_h_layer_names(self): + llm_config = { + "layers_block_type": ["linear_attention", "moe", "full_attention"], + "num_nextn_predict_layers": 1, + "mtp_layers_block_type": ["full_attention", "moe"], + } + + config = NemotronH_Omni_Reasoning_V3_Config( + vision_config={"args": {"model": "radio"}}, + llm_config=llm_config, + ) + + self.assertEqual( + config.llm_config.layers_block_type, + ["mamba", "moe", "attention"], + ) + self.assertEqual( + config.llm_config.mtp_layers_block_type, + ["attention", "moe"], + ) + self.assertEqual( + llm_config["layers_block_type"], + ["linear_attention", "moe", "full_attention"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 97c6011ab..1e8f8e374 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -41,6 +41,7 @@ from sglang.srt.model_loader.weight_utils import ( ) from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM from sglang.srt.models.muse_glimmer import MuseGlimmerForConditionalGeneration +from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci @@ -716,6 +717,29 @@ class TestModelOptFp4LoaderSelection(CustomTestCase): class TestModelOptMixedPrecisionConfig(CustomTestCase): + def test_nemotron_h_omni_resolves_fused_qkv_from_split_layers(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + f"language_model.model.layers.7.mixer.{projection}": { + "quant_algo": "FP8" + } + for projection in ("q_proj", "k_proj", "v_proj") + }, + "packed_modules_mapping": ( + NemotronH_Omni_Reasoning_V3.packed_modules_mapping + ), + } + ) + + self.assertEqual( + quant_config._resolve_quant_algo( + "language_model.model.layers.7.mixer.qkv_proj" + ), + "FP8", + ) + def test_fp8_pb_wo_dispatches_to_native_block_fp8(self): quant_config = ModelOptMixedPrecisionConfig.from_config( { diff --git a/test/registered/unit/models/test_nano_nemotron_vl.py b/test/registered/unit/models/test_nano_nemotron_vl.py new file mode 100644 index 000000000..5a31c3536 --- /dev/null +++ b/test/registered/unit/models/test_nano_nemotron_vl.py @@ -0,0 +1,182 @@ +"""Unit tests for native Nemotron-H Omni model integration.""" + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from sglang.srt.models.nano_nemotron_vl import ( + NemotronH_Nano_VL_V2, + NemotronH_Omni_Reasoning_V3, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestNemotronHOmniModel(CustomTestCase): + def test_existing_nano_model_keeps_ignoring_unrecognized_weights(self): + model = object.__new__(NemotronH_Nano_VL_V2) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.language_model = SimpleNamespace( + load_weights=lambda weights: list(weights) + ) + model.vision_model = SimpleNamespace(load_weights=lambda weights: None) + model.sound_encoder = None + + model.load_weights([("unrecognized.weight", torch.ones(1))]) + + def test_model_registry_resolves_new_architecture(self): + from sglang.srt.models.registry import ModelRegistry + + model_class, architecture = ModelRegistry.resolve_model_cls( + "NemotronH_Omni_Reasoning_V3" + ) + + self.assertIs(model_class, NemotronH_Omni_Reasoning_V3) + self.assertEqual(architecture, "NemotronH_Omni_Reasoning_V3") + + def test_exposes_language_embed_and_head(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + embed = object() + head = object() + model.language_model = SimpleNamespace( + get_embed_and_head=lambda: (embed, head), + lm_head=head, + ) + + self.assertEqual(model.get_embed_and_head(), (embed, head)) + self.assertIs(model.lm_head, head) + + def test_delegates_dflash_capture_to_language_model(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + captured_layer_ids = [] + model.language_model = SimpleNamespace( + set_dflash_layers_to_capture=captured_layer_ids.extend + ) + + model.set_dflash_layers_to_capture([1, 22, 43, 64, 85]) + + self.assertEqual(captured_layer_ids, [1, 22, 43, 64, 85]) + + def test_vision_final_layernorm_is_loaded_and_applied(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = nn.LayerNorm(2) + model.language_model = SimpleNamespace(load_weights=lambda weights: None) + model.vision_model = SimpleNamespace(load_weights=lambda weights: None) + model.sound_encoder = None + + weight = torch.tensor([2.0, 3.0]) + bias = torch.tensor([0.5, -0.5]) + model.load_weights( + [ + ("vision_projector.vision_final_layernorm.weight", weight), + ("vision_projector.vision_final_layernorm.bias", bias), + ] + ) + + features = torch.tensor([[1.0, 3.0]]) + expected = nn.functional.layer_norm(features, (2,), weight, bias) + torch.testing.assert_close(model._normalize_vision_features(features), expected) + + def test_hf_vision_and_projector_names_are_remapped(self): + remap = NemotronH_Omni_Reasoning_V3._remap_checkpoint_weight_name + + self.assertEqual( + remap("vision_model.embeddings.position_embedding"), + "vision_model.radio_model.hf_model.embeddings.position_embedding", + ) + self.assertEqual( + remap("vision_model.embeddings.video_patch_projection.weight"), + ( + "vision_model.radio_model.hf_model.embeddings." + "video_patch_projection.weight" + ), + ) + self.assertEqual( + remap("vision_projector.mlp1.linear1.weight"), + "mlp1.1.weight", + ) + self.assertEqual( + remap("vision_model.radio_model.model.patch_generator.pos_embed"), + "vision_model.radio_model.model.patch_generator.pos_embed", + ) + + def test_unexpected_checkpoint_weight_raises(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = nn.LayerNorm(2) + model.language_model = SimpleNamespace(load_weights=lambda weights: None) + model.vision_model = SimpleNamespace(load_weights=lambda weights: None) + model.sound_encoder = None + + cases = ( + ("vision_projector.unknown.weight", "Unexpected Nemotron-H Omni"), + ( + "vision_projector.vision_final_layernorm.running_mean", + "Unexpected vision projector weight", + ), + ) + for name, message in cases: + with self.subTest(name=name), self.assertRaisesRegex(ValueError, message): + model.load_weights([(name, torch.ones(1))]) + + def test_language_weights_are_streamed_and_remaining_components_are_routed(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = None + source_exhausted = False + loaded_language_weights = [] + loaded_vision_weights = [] + loaded_sound_weights = [] + + def source_weights(): + nonlocal source_exhausted + yield "language_model.model.layer.weight", torch.ones(1) + yield "vision_model.radio_model.encoder.weight", torch.ones(1) + yield "sound_encoder.projection.weight", torch.ones(1) + source_exhausted = True + + def load_language_weights(weights): + self.assertFalse(source_exhausted) + loaded_language_weights.append(next(weights)) + + def load_vision_weights(weights): + self.assertFalse(source_exhausted) + loaded_vision_weights.extend(weights) + + def load_sound_weights(weights): + self.assertFalse(source_exhausted) + loaded_sound_weights.extend(weights) + + model.language_model = SimpleNamespace(load_weights=load_language_weights) + model.vision_model = SimpleNamespace(load_weights=load_vision_weights) + model.sound_encoder = SimpleNamespace(load_weights=load_sound_weights) + + model.load_weights(source_weights()) + + self.assertTrue(source_exhausted) + self.assertEqual( + [name for name, _ in loaded_language_weights], ["model.layer.weight"] + ) + self.assertEqual( + [name for name, _ in loaded_vision_weights], + ["radio_model.encoder.weight"], + ) + self.assertEqual( + [name for name, _ in loaded_sound_weights], + ["sound_encoder.projection.weight"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_nemotron_h_mtp.py b/test/registered/unit/models/test_nemotron_h_mtp.py new file mode 100644 index 000000000..d6de48d81 --- /dev/null +++ b/test/registered/unit/models/test_nemotron_h_mtp.py @@ -0,0 +1,226 @@ +"""Unit tests for Nemotron-H MTP model behavior.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.nn as nn + +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptMixedPrecisionConfig, + ModelOptNvFp4A16LinearMethod, +) +from sglang.srt.models.nemotron_h_mtp import ( + NemotronHForCausalLMMTP, + NemotronHMultiTokenPredictor, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _RecordingLayer(nn.Module): + def __init__(self): + super().__init__() + self.inputs_embeds = None + + def forward(self, *, inputs_embeds, hidden_states, residual, forward_batch): + self.inputs_embeds = inputs_embeds + return hidden_states, residual + + +class TestNemotronHMultiTokenPredictor(CustomTestCase): + def test_text_only_forward_uses_model_embeddings(self): + model = object.__new__(NemotronHMultiTokenPredictor) + nn.Module.__init__(model) + model.embed_tokens = nn.Embedding(8, 2) + model.embed_tokens.weight.data.copy_(torch.arange(16).reshape(8, 2)) + model.pattern_len = 1 + layer = _RecordingLayer() + model.layers = nn.ModuleDict({"0": layer}) + input_ids = torch.tensor([1, 2, 3]) + forward_batch = SimpleNamespace( + mm_input_embeds=None, + forward_mode=SimpleNamespace(is_extend=lambda: False), + contains_mm_inputs=lambda: False, + spec_info=SimpleNamespace(hidden_states=torch.zeros(3, 2)), + ) + + model( + input_ids=input_ids, + positions=torch.arange(3), + forward_batch=forward_batch, + ) + + torch.testing.assert_close( + layer.inputs_embeds, + model.embed_tokens(input_ids), + ) + + def test_multimodal_prefill_reuses_target_embeddings(self): + model = object.__new__(NemotronHMultiTokenPredictor) + nn.Module.__init__(model) + model.embed_tokens = nn.Embedding(8, 2) + model.embed_tokens.weight.data.copy_(torch.arange(16).reshape(8, 2)) + model.pattern_len = 1 + layer = _RecordingLayer() + model.layers = nn.ModuleDict({"0": layer}) + + target_embeddings = torch.tensor( + [[101.0, 102.0], [103.0, 104.0], [105.0, 106.0]] + ) + forward_batch = SimpleNamespace( + mm_input_embeds=target_embeddings.clone(), + forward_mode=SimpleNamespace( + is_extend=lambda: True, + is_draft_extend_v2=lambda: False, + ), + contains_mm_inputs=lambda: True, + extend_start_loc=torch.tensor([0]), + extend_seq_lens=torch.tensor([3]), + spec_info=SimpleNamespace(hidden_states=torch.zeros(3, 2)), + ) + + model( + input_ids=torch.tensor([100, 101, 2]), + positions=torch.arange(3), + forward_batch=forward_batch, + ) + + expected = target_embeddings.clone() + expected[-1] = model.embed_tokens(torch.tensor(2)) + torch.testing.assert_close(layer.inputs_embeds, expected) + + +class TestNemotronHForCausalLMMTP(CustomTestCase): + def _make_head_model(self): + model = object.__new__(NemotronHForCausalLMMTP) + nn.Module.__init__(model) + model.config = SimpleNamespace( + max_n_routed_experts=0, tie_word_embeddings=False + ) + model.pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + model.model = nn.Module() + model.model.embed_tokens = nn.Embedding(4, 2) + model.model.layers = nn.ModuleList([nn.Linear(2, 2, bias=False)]) + model.lm_head = nn.Linear(2, 4, bias=False) + model.lm_head.quant_method = None + model.lm_head.register_parameter( + "weight_scale", nn.Parameter(torch.zeros(1), requires_grad=False) + ) + return model + + def test_standalone_mtp_head_survives_both_target_sharing_calls(self): + # Replacing either the head weight or its module silently discards the + # external checkpoint's output projection (including quantization scales). + for prefix in ("", "language_model."): + with self.subTest(prefix=prefix): + model = self._make_head_model() + model.load_weights( + iter( + [ + (prefix + "mtp.layers.0.weight", torch.ones(2, 2)), + ( + prefix + "lm_head.weight", + torch.arange(8.0).reshape(4, 2), + ), + (prefix + "lm_head.weight_scale", torch.tensor([0.5])), + ] + ) + ) + draft_head = model.lm_head + draft_weight = draft_head.weight + target_embed = nn.Parameter(torch.ones(4, 2)) + target_head = nn.Linear(2, 4, bias=False) + with patch("torch.cuda.synchronize"), patch("torch.cuda.empty_cache"): + model.set_embed_and_head(target_embed, target_head.weight) + self.assertIs(model.lm_head.weight, draft_weight) + model.set_lm_head_from_target(target_head) + self.assertIs(model.lm_head, draft_head) + self.assertIs(model.model.embed_tokens.weight, target_embed) + torch.testing.assert_close( + model.lm_head(torch.ones(1, 2)), + torch.tensor([[1.0, 5.0, 9.0, 13.0]]), + ) + torch.testing.assert_close( + model.lm_head.weight_scale, torch.tensor([0.5]) + ) + + def test_embedded_and_headless_mtp_share_complete_target_head(self): + for embedded in (False, True): + with self.subTest(embedded=embedded): + model = self._make_head_model() + weights = [("mtp.layers.0.weight", torch.ones(2, 2))] + if embedded: + # Full checkpoints also contain lm_head tensors; their + # presence alone must not opt out of embedded head sharing. + weights += [ + ("lm_head.weight", torch.ones(4, 2)), + ("lm_head.weight_scale", torch.ones(1)), + ("backbone.layers.0.weight", torch.ones(2, 2)), + ] + model.load_weights(iter(weights)) + target_head = nn.Linear(2, 4, bias=False) + target_embed = nn.Parameter(torch.ones(4, 2)) + with patch("torch.cuda.synchronize"), patch("torch.cuda.empty_cache"): + model.set_embed_and_head(target_embed, target_head.weight) + model.set_lm_head_from_target(target_head) + self.assertIs(model.lm_head, target_head) + self.assertIs(model.model.embed_tokens.weight, target_embed) + + def test_incomplete_standalone_head_is_rejected(self): + for missing in ("weight", "weight_scale"): + with self.subTest(missing=missing): + model = self._make_head_model() + weights = { + "mtp.layers.0.weight": torch.ones(2, 2), + "lm_head.weight": torch.ones(4, 2), + "lm_head.weight_scale": torch.ones(1), + } + del weights["lm_head." + missing] + with self.assertRaisesRegex( + ValueError, "Incomplete standalone MTP lm_head" + ): + model.load_weights(iter(weights.items())) + + def test_w4a16_head_does_not_require_unused_input_scale(self): + model = self._make_head_model() + model.lm_head.quant_method = ModelOptNvFp4A16LinearMethod(quant_config=None) + model.lm_head.register_parameter( + "input_scale", nn.Parameter(torch.zeros(1), requires_grad=False) + ) + # NVFP4A16 registers this loader placeholder but discards it before + # inference. Requiring it would reject valid standalone W4A16 heads. + model.load_weights( + iter( + [ + ("mtp.layers.0.weight", torch.ones(2, 2)), + ("lm_head.weight", torch.ones(4, 2)), + ("lm_head.weight_scale", torch.ones(1)), + ] + ) + ) + + def test_maps_quantized_mtp_metadata(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "language_model.mtp.layers.0.mixer.q_proj": {"quant_algo": "FP8"} + }, + } + ) + quant_config.apply_weight_name_mapper( + NemotronHForCausalLMMTP.hf_to_sglang_mapper + ) + + self.assertEqual( + quant_config._resolve_quant_algo("mtp.layers.0.mixer.q_proj"), + "FP8", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_nemotron_h_weight_loading.py b/test/registered/unit/models/test_nemotron_h_weight_loading.py index 0a4309992..1943cd1fc 100644 --- a/test/registered/unit/models/test_nemotron_h_weight_loading.py +++ b/test/registered/unit/models/test_nemotron_h_weight_loading.py @@ -1,13 +1,4 @@ -""" -Unit tests for NemotronHForCausalLM.load_weights. - -Regression test for Nemotron-H expert scale checkpoint tensors that map to -parameters absent from the current runtime model. -""" - -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=12, suite="base-a-test-cpu") +"""Unit tests for Nemotron-H target and MTP checkpoint weight loading.""" import unittest from types import SimpleNamespace @@ -15,6 +6,11 @@ from types import SimpleNamespace import torch from sglang.srt.models.nemotron_h import NemotronHForCausalLM +from sglang.srt.models.nemotron_h_mtp import NemotronHForCausalLMMTP +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=12, suite="base-a-test-cpu") class _FakePPGroup: @@ -43,9 +39,19 @@ class _RecordingParam: self.loaded_weight = loaded_weight -class TestNemotronHWeightLoading(unittest.TestCase): - def _make_minimal_model(self, named_parameters=()): - model = object.__new__(NemotronHForCausalLM) +class _RecordingStackedParam: + def __init__(self): + self.loads = [] + + def weight_loader(self, param, loaded_weight, shard_id): + self.loads.append((param, loaded_weight, shard_id)) + + +class TestNemotronHWeightLoading(CustomTestCase): + def _make_minimal_model( + self, named_parameters=(), model_class=NemotronHForCausalLM + ): + model = object.__new__(model_class) model.config = SimpleNamespace(n_routed_experts=2, max_n_routed_experts=2) model.model = SimpleNamespace() model.pp_group = _FakePPGroup() @@ -134,6 +140,59 @@ class TestNemotronHWeightLoading(unittest.TestCase): skipped.loaded_weight, "non-MTP target weight should be skipped" ) + def test_mtp_strips_multimodal_language_model_prefix(self): + embed = _RecordingParam() + head = _RecordingParam() + mtp_layer = _RecordingParam() + model = self._make_minimal_model( + [ + ("model.embed_tokens.weight", embed), + ("lm_head.weight", head), + ("model.layers.0.norm.weight", mtp_layer), + ], + model_class=NemotronHForCausalLMMTP, + ) + model.remap_prefix = {"backbone": "model"} + model.remap_substr = {"embeddings": "embed_tokens"} + + w_embed, w_head, w_mtp = (torch.ones(1) for _ in range(3)) + model.load_weights( + [ + ("language_model.backbone.embeddings.weight", w_embed), + ("language_model.lm_head.weight", w_head), + ("language_model.mtp.layers.0.norm.weight", w_mtp), + ] + ) + + self.assertIs(embed.loaded_weight, w_embed) + self.assertIs(head.loaded_weight, w_head) + self.assertIs(mtp_layer.loaded_weight, w_mtp) + + def test_split_qkv_fp8_scales_load_into_fused_parameter(self): + input_scale = _RecordingStackedParam() + model = self._make_minimal_model( + [("model.layers.7.mixer.qkv_proj.input_scale", input_scale)] + ) + model.stacked_params_mapping = NemotronHForCausalLM.stacked_params_mapping + + q_scale, k_scale, v_scale = (torch.tensor(value) for value in (1, 2, 3)) + model.load_weights( + [ + ("model.layers.7.mixer.q_proj.input_scale", q_scale), + ("model.layers.7.mixer.k_proj.input_scale", k_scale), + ("model.layers.7.mixer.v_proj.input_scale", v_scale), + ] + ) + + self.assertEqual( + input_scale.loads, + [ + (input_scale, q_scale, "q"), + (input_scale, k_scale, "k"), + (input_scale, v_scale, "v"), + ], + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_radio.py b/test/registered/unit/models/test_radio.py new file mode 100644 index 000000000..3ef21eccb --- /dev/null +++ b/test/registered/unit/models/test_radio.py @@ -0,0 +1,130 @@ +"""Unit tests for RADIO checkpoint weight loading.""" + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from sglang.srt.models.radio import RadioModel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _RecordingWeight: + def __init__(self): + self.loads = [] + + def weight_loader(self, param, weight, shard_id=None): + self.loads.append((param, weight, shard_id)) + + +class TestRadioWeightLoading(CustomTestCase): + def _make_model(self, named_parameters=()): + model = object.__new__(RadioModel) + nn.Module.__init__(model) + model.named_parameters = lambda: iter(named_parameters) + model.model = SimpleNamespace( + patch_generator=SimpleNamespace(_video_embedder_loaded=False) + ) + return model + + def test_hf_export_maps_embeddings_and_split_qkv(self): + position_embedding = _RecordingWeight() + qkv_weight = _RecordingWeight() + model = self._make_model( + [ + ("model.patch_generator.pos_embed", position_embedding), + ("model.encoder.layers.0.attn.attn.qkv_proj.weight", qkv_weight), + ] + ) + + position = torch.ones(1) + query, key, value = (torch.full((1,), value) for value in (2, 3, 4)) + loaded = model.load_weights( + [ + ("radio_model.hf_model.embeddings.position_embedding", position), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention." + "query.weight", + query, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention." + "key.weight", + key, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention." + "value.weight", + value, + ), + ("radio_model.hf_model.summary_idxs", torch.tensor([0, 1])), + ] + ) + + self.assertEqual( + loaded, + { + "model.patch_generator.pos_embed", + "model.encoder.layers.0.attn.attn.qkv_proj.weight", + }, + ) + self.assertEqual( + position_embedding.loads, [(position_embedding, position, None)] + ) + self.assertEqual( + qkv_weight.loads, + [ + (qkv_weight, query, "q"), + (qkv_weight, key, "k"), + (qkv_weight, value, "v"), + ], + ) + + def test_hf_export_loads_encoder_parameters(self): + cases = { + "embeddings.video_patch_projection.weight": ( + "model.patch_generator.video_embedder.weight" + ), + "encoder.layer.1.attention.output.dense.weight": ( + "model.encoder.layers.1.attn.attn.proj.weight" + ), + "encoder.layer.2.layer_scale1.lambda1": "model.encoder.layers.2.ls1", + "encoder.layer.3.layer_scale2.lambda1": "model.encoder.layers.3.ls2", + "encoder.layer.4.mlp.fc1.bias": "model.encoder.layers.4.mlp.fc1.bias", + "encoder.layer.5.norm2.weight": "model.encoder.layers.5.norm2.weight", + } + for source, target in cases.items(): + with self.subTest(source=source): + parameter = _RecordingWeight() + model = self._make_model([(target, parameter)]) + weight = torch.ones(1) + + self.assertEqual( + model.load_weights([(f"radio_model.hf_model.{source}", weight)]), + {target}, + ) + self.assertEqual(parameter.loads, [(parameter, weight, None)]) + + def test_unmapped_hf_export_weight_raises(self): + model = self._make_model() + + with self.assertRaisesRegex(ValueError, "Unexpected HF RADIO weight"): + model.load_weights( + [("radio_model.hf_model.encoder.layer.0.unknown.weight", torch.ones(1))] + ) + + def test_legacy_unknown_weight_remains_ignored(self): + model = self._make_model() + + self.assertEqual( + model.load_weights([("radio_model.unknown.weight", torch.ones(1))]), + set(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py b/test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py new file mode 100644 index 000000000..cee894743 --- /dev/null +++ b/test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py @@ -0,0 +1,24 @@ +"""Unit tests for the Nano Nemotron VL processor registry.""" + +import unittest + +from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 +from sglang.srt.multimodal.processors.nano_nemotron_vl import ( + NanoNemotronVLImageProcessor, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestNanoNemotronVLProcessor(CustomTestCase): + def test_supports_nemotron_h_omni(self): + self.assertIn( + NemotronH_Omni_Reasoning_V3, + NanoNemotronVLImageProcessor.models, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 913088836..7f146151f 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1124,6 +1124,46 @@ class TestGoldenModelOverrides(_IsolatedPublish): self.assertNotIn("attention_backend", overrides) self.assertNotIn("speculative_draft_attention_backend", overrides) + def test_nemotron_h_omni_uses_inner_text_config(self): + outer_config = SimpleNamespace( + architectures=["NemotronH_Omni_Reasoning_V3"], + quantization_config={"quant_algo": "NVFP4"}, + ) + model_config = SimpleNamespace( + quantization="modelopt", + hf_config=outer_config, + hf_text_config=SimpleNamespace(mlp_hidden_act="relu2"), + ) + server_args = SimpleNamespace( + quantization=None, + moe_runner_backend="auto", + moe_a2a_backend="none", + attention_backend=None, + _model_config=model_config, + ) + + with ( + override_platform(is_blackwell=False), + override_platform(is_sm100=False), + override_platform(is_cuda=False), + ): + self.assertEqual( + collect_model_override_declarations( + "NemotronH_Omni_Reasoning_V3", + server_args, + outer_config, + ), + [ + ( + "_nemotron_h_overrides", + { + "quantization": "modelopt_fp4", + "moe_runner_backend": "flashinfer_cutlass", + }, + ) + ], + ) + def test_nemotron_h_w4a16_moe_rejects_a2a_backend(self): from sglang.srt.arg_groups.model_overrides.nemotron_h import ( _nemotron_h_overrides, @@ -1979,6 +2019,12 @@ class TestGoldenModelOverrides(_IsolatedPublish): _flashinfer_allreduce_fusion_auto_enable(_view()), {"flashinfer_allreduce_fusion_backend": "auto"}, ) + self.assertEqual( + _flashinfer_allreduce_fusion_auto_enable( + _view(arch="NemotronH_Omni_Reasoning_V3") + ), + {"flashinfer_allreduce_fusion_backend": "auto"}, + ) # guards: unsupported arch / tp==1 / dp attention / a2a backend self.assertEqual( _flashinfer_allreduce_fusion_auto_enable( @@ -2341,13 +2387,18 @@ class TestGoldenModelOverrides(_IsolatedPublish): ) # NemotronH routes through the pass (covered by the guard union, # not the branch chain — its hook invokes the handler) - self.assertEqual( - _mamba_radix_cache_resolution(_view("NemotronHForCausalLM")), - { - "uses_mamba_radix_cache": True, - "mamba_radix_cache_strategy": "extra_buffer", - }, - ) + for architecture in ( + "NemotronHForCausalLM", + "NemotronH_Omni_Reasoning_V3", + ): + with self.subTest(architecture=architecture): + self.assertEqual( + _mamba_radix_cache_resolution(_view(architecture)), + { + "uses_mamba_radix_cache": True, + "mamba_radix_cache_strategy": "extra_buffer", + }, + ) # GraniteMoeHybrid is guarded on mamba layer types self.assertEqual( _mamba_radix_cache_resolution(