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 3, 22:45 UTC
Avoid silent weights corruption when loading Nemotron Nano VL with reusable-buffer loaders like runai distributed streaming (#42244)
Signed-off-by: Noa Neria <nneria@nvidia.com>
Noa Neria committed
May 11, 2026 at 15:03 UTC
ac062147fae75e482e1fd0b2930746678253d7bb
2 files changed
+55
-28
tests/models/multimodal/test_nano_nemotron_vl.py
+14
-1
@@ -53,6 +53,19 @@ class _VisionModel:
53
self.loaded_weights = list(weights)
54
55
56
+class _FakeTensor:
57
+ """Sentinel stand-in for torch.Tensor in load_weights tests. Supports the
58
+ .detach().clone() chain used by load_weights for buffered mm weights;
59
+ both methods return self so identity (and the existing equality
60
+ assertions) are preserved through cloning."""
61
+
62
+ def detach(self):
63
+ return self
64
+
65
+ def clone(self):
66
+ return self
67
+
68
+
69
def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
70
model = object.__new__(NemotronH_Nano_VL_V2)
71
language_model = _LanguageModel()
@@ -86,7 +99,7 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
99
object.__setattr__(model, "sound_encoder", None)
100
101
language_weight = object()
89
- vision_weight = object()
102
+ vision_weight = _FakeTensor()
103
model.load_weights(
104
[
105
("language_model.layers.0.weight", language_weight),
vllm/model_executor/models/nano_nemotron_vl.py
+41
-27
@@ -1518,37 +1518,51 @@ class NemotronH_Nano_VL_V2(
1518
def is_sound_weights(name: str) -> bool:
1519
return name.startswith("sound")
1520
1521
- # Separate weights by component
1522
- llm_weights = []
1523
- vision_weights = []
1524
- sound_weights = []
1525
-
1526
- for name, w in weights:
1527
- if is_llm(name):
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
1521
+ # LLM weights (the bulk of the model) are streamed lazily through a
1522
+ # generator so each tensor is copied into its parameter before the
1523
+ # iterator advances, avoiding stale-reference corruption with
1524
+ # reusable-buffer streamers. The smaller mm components (mlp1, vision,
1525
+ # sound) are detach+cloned on append so they are independent of any
1526
+ # reusable buffer the streamer may use, then loaded after the LLM.
1527
+ adapter_weights: list[tuple[str, torch.Tensor]] = []
1528
+ vision_weights: list[tuple[str, torch.Tensor]] = []
1529
+ sound_weights: list[tuple[str, torch.Tensor]] = []
1530
+
1531
+ def llm_weights_gen():
1532
+ for name, w in weights:
1533
+ if is_llm(name):
1534
+ # Strip 'language_model.' prefix for LLM weights
1535
+ yield ".".join(name.split(".")[1:]), w
1536
+ elif is_adapter_weights((name, w)):
1537
+ if not load_multimodal_weights:
1538
+ continue
1539
+ trimmed_name = ".".join(name.split(".")[1:])
1540
+ adapter_weights.append((trimmed_name, w.detach().clone()))
1541
+ elif is_vision_weights(name):
1542
+ if not load_multimodal_weights:
1543
+ continue
1544
+ # Convert: vision_model.radio_model.* → radio_model.*
1545
+ hf_key = name[len("vision_model.") :]
1546
+ vision_weights.append((hf_key, w.detach().clone()))
1547
+ elif is_sound_weights(name):
1548
+ if not load_multimodal_weights:
1549
+ continue
1550
+ assert self.sound_encoder is not None
1551
+ sound_weights.append((name, w.detach().clone()))
1552
+
1553
+ # Fully drain the generator so every mm tensor is buffered, even if
1554
+ # the LLM loader stops iterating early.
1555
+ llm_weights_iter = llm_weights_gen()
1556
+ self.language_model.load_weights(llm_weights_iter)
1557
+ for _ in llm_weights_iter:
1558
+ pass
1559
+
1560
+ if load_multimodal_weights:
1561
+ for trimmed_name, w in adapter_weights:
1562
# Load vision-language adapter weights directly
1534
- trimmed_name = ".".join(name.split(".")[1:])
1563
param = adapter_dict[trimmed_name]
1564
with torch.no_grad():
1565
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)
1551
- if load_multimodal_weights:
1566
self.vision_model.load_weights(vision_weights)
1567
if self.sound_encoder is not None and len(sound_weights) > 0:
1568
self.sound_encoder.load_weights(sound_weights)