A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 07:30 UTC
Fix Nano Nemotron text-only weight loading (#41205)
Signed-off-by: sunghoon.baek <sunghoon.baek@connectfy.cloud> Signed-off-by: Baekpica <35071468+Baekpica@users.noreply.github.com> Signed-off-by: sunghoon.baek <seanbb93@gmail.com> Co-authored-by: sunghoon.baek <sunghoon.baek@connectfy.cloud> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Baekpica committed
May 5, 2026 at 03:43 UTC
8c780943b492a26dc2032f40a18a77f92d1931e6
2 files changed
+129
-3
tests/models/multimodal/test_nano_nemotron_vl.py
new
+114
@@ -0,0 +1,114 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.model_executor.models.nano_nemotron_vl import NemotronH_Nano_VL_V2
7
+
8
+
9
+class _TextOnlyMultiModalConfig:
10
+ def get_limit_per_prompt(self, modality: str) -> int:
11
+ return 0
12
+
13
+
14
+class _ImageOnlyMultiModalConfig:
15
+ def get_limit_per_prompt(self, modality: str) -> int:
16
+ return 1 if modality == "image" else 0
17
+
18
+
19
+class _ModelConfig:
20
+ multimodal_config = _TextOnlyMultiModalConfig()
21
+
22
+
23
+class _ImageOnlyModelConfig:
24
+ multimodal_config = _ImageOnlyMultiModalConfig()
25
+
26
+
27
+class _LanguageModel:
28
+ def __init__(self) -> None:
29
+ self.loaded_weights: list[tuple[str, object]] = []
30
+
31
+ def load_weights(self, weights):
32
+ self.loaded_weights = list(weights)
33
+
34
+
35
+class _MissingMultiModalModule:
36
+ def named_parameters(self):
37
+ raise AssertionError("multimodal weights should not be inspected")
38
+
39
+ def load_weights(self, weights):
40
+ raise AssertionError("multimodal weights should not be loaded")
41
+
42
+
43
+class _AdapterModule:
44
+ def named_parameters(self):
45
+ return []
46
+
47
+
48
+class _VisionModel:
49
+ def __init__(self) -> None:
50
+ self.loaded_weights: list[tuple[str, object]] = []
51
+
52
+ def load_weights(self, weights):
53
+ self.loaded_weights = list(weights)
54
+
55
+
56
+def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
57
+ model = object.__new__(NemotronH_Nano_VL_V2)
58
+ language_model = _LanguageModel()
59
+ object.__setattr__(model, "model_config", _ModelConfig())
60
+ object.__setattr__(model, "language_model", language_model)
61
+ object.__setattr__(model, "mlp1", _AdapterModule())
62
+ object.__setattr__(model, "vision_model", _MissingMultiModalModule())
63
+ object.__setattr__(model, "sound_encoder", None)
64
+
65
+ language_weight = object()
66
+ model.load_weights(
67
+ [
68
+ ("language_model.layers.0.weight", language_weight),
69
+ ("mlp1.0.weight", object()),
70
+ ("vision_model.radio_model.encoder.weight", object()),
71
+ ("sound_encoder.encoder.weight", object()),
72
+ ]
73
+ )
74
+
75
+ assert language_model.loaded_weights == [("layers.0.weight", language_weight)]
76
+
77
+
78
+def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
79
+ model = object.__new__(NemotronH_Nano_VL_V2)
80
+ language_model = _LanguageModel()
81
+ vision_model = _VisionModel()
82
+ object.__setattr__(model, "model_config", _ImageOnlyModelConfig())
83
+ object.__setattr__(model, "language_model", language_model)
84
+ object.__setattr__(model, "mlp1", _AdapterModule())
85
+ object.__setattr__(model, "vision_model", vision_model)
86
+ object.__setattr__(model, "sound_encoder", None)
87
+
88
+ language_weight = object()
89
+ vision_weight = object()
90
+ model.load_weights(
91
+ [
92
+ ("language_model.layers.0.weight", language_weight),
93
+ ("vision_model.radio_model.encoder.weight", vision_weight),
94
+ ]
95
+ )
96
+
97
+ assert language_model.loaded_weights == [("layers.0.weight", language_weight)]
98
+ assert vision_model.loaded_weights == [
99
+ ("radio_model.encoder.weight", vision_weight)
100
+ ]
101
+
102
+
103
+def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights():
104
+ model = object.__new__(NemotronH_Nano_VL_V2)
105
+ language_model = _LanguageModel()
106
+ vision_model = _VisionModel()
107
+ object.__setattr__(model, "model_config", _ImageOnlyModelConfig())
108
+ object.__setattr__(model, "language_model", language_model)
109
+ object.__setattr__(model, "mlp1", _AdapterModule())
110
+ object.__setattr__(model, "vision_model", vision_model)
111
+ object.__setattr__(model, "sound_encoder", None)
112
+
113
+ with pytest.raises(AssertionError):
114
+ model.load_weights([("sound_encoder.encoder.weight", object())])
vllm/model_executor/models/nano_nemotron_vl.py
+15
-3
@@ -1499,6 +1499,11 @@ class NemotronH_Nano_VL_V2(
1499
return self.language_model.compute_logits(hidden_states)
1500
1501
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
1502
+ mm_config = self.model_config.multimodal_config
1503
+ load_multimodal_weights = not all(
1504
+ mm_config.get_limit_per_prompt(modality) == 0
1505
+ for modality in ("image", "video", "audio")
1506
+ )
1507
adapter_dict = dict(self.mlp1.named_parameters())
1508
1509
def is_llm(name: str) -> bool:
@@ -1523,23 +1528,30 @@ class NemotronH_Nano_VL_V2(
1528
# Strip 'language_model.' prefix for LLM weights
1529
llm_weights.append((".".join(name.split(".")[1:]), w))
1530
elif is_adapter_weights((name, w)):
1531
+ if not load_multimodal_weights:
1532
+ continue
1533
# Load vision-language adapter weights directly
1534
trimmed_name = ".".join(name.split(".")[1:])
1535
param = adapter_dict[trimmed_name]
1536
with torch.no_grad():
1537
default_weight_loader(param, w)
1538
elif is_vision_weights(name):
1539
+ if not load_multimodal_weights:
1540
+ continue
1541
# Convert: vision_model.radio_model.* → radio_model.*
1542
hf_key = name[len("vision_model.") :] # Remove "vision_model." prefix
1543
vision_weights.append((hf_key, w))
1544
elif is_sound_weights(name):
1545
+ if not load_multimodal_weights:
1546
+ continue
1547
assert self.sound_encoder is not None
1548
sound_weights.append((name, w))
1549
1550
self.language_model.load_weights(llm_weights)
1540
- self.vision_model.load_weights(vision_weights)
1541
- if self.sound_encoder is not None and len(sound_weights) > 0:
1542
- self.sound_encoder.load_weights(sound_weights)
1551
+ if load_multimodal_weights:
1552
+ self.vision_model.load_weights(vision_weights)
1553
+ if self.sound_encoder is not None and len(sound_weights) > 0:
1554
+ self.sound_encoder.load_weights(sound_weights)
1555
1556
def get_vit_model_from_radio_config(self, hf_config):
1557
hf_config_vision = hf_config.vision_config