minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support NemotronH_Omni_Reasoning_V3 in SGLang (#35599)

Browse Source 
Signed-off-by: Ryan Stewart <rystewart@nvidia.com>
Signed-off-by: rystewart-nvidia <rystewart@nvidia.com>
Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
rystewart-nvidia
2026-09-10 16:57:22 -07:00
committed by GitHub
co-authored by elvischenv Po-Han Huang
parent 203d7e812c
commit fae8cd84cb
20 changed files with 1080 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/configs/test_model_config.py
+19
View File
@@ -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."""
test/registered/unit/configs/test_nano_nemotron_vl_config.py
+51
View File
@@ -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()
test/registered/unit/model_loader/test_modelopt_loader.py
+24
View File
@@ -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(
{
test/registered/unit/models/test_nano_nemotron_vl.py
+182
View File
@@ -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()
test/registered/unit/models/test_nemotron_h_mtp.py
+226
View File
@@ -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()
test/registered/unit/models/test_nemotron_h_weight_loading.py
+72 -13
View File
@@ -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()
test/registered/unit/models/test_radio.py
+130
View File
@@ -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()
test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py
+24
View File
@@ -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()
test/registered/unit/test_model_overrides.py
+58 -7
View File
@@ -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(