From 0838cd6593ba751aa215c9ecd44048bb29dcff1b Mon Sep 17 00:00:00 2001 From: Nathan Sala <33957625+natsala13@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:42:29 +0300 Subject: [PATCH] [Bugfix] Load and apply Nemotron Super final vision LayerNorm (#61401) Signed-off-by: Shaun Kotek Co-authored-by: Shaun Kotek - Nvidia --- .../multimodal/test_nano_nemotron_vl.py | 57 +++++++++++++++++++ .../model_executor/models/nano_nemotron_vl.py | 35 +++++++++++- 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index 75900cde53..41f9428e6d 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -97,6 +97,7 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", _MissingMultiModalModule()) + object.__setattr__(model, "vision_final_layernorm", _MissingMultiModalModule()) object.__setattr__(model, "sound_encoder", None) language_weight = object() @@ -106,6 +107,8 @@ def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): ("mlp1.0.weight", object()), ("vision_model.radio_model.encoder.weight", object()), ("sound_encoder.encoder.weight", object()), + ("vision_projector.vision_final_layernorm.weight", object()), + ("vision_projector.vision_final_layernorm.bias", object()), ] ) @@ -120,6 +123,7 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "vision_final_layernorm", None) object.__setattr__(model, "sound_encoder", None) language_weight = object() @@ -145,12 +149,65 @@ def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights(): object.__setattr__(model, "language_model", language_model) object.__setattr__(model, "mlp1", _AdapterModule()) object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "vision_final_layernorm", None) object.__setattr__(model, "sound_encoder", None) with pytest.raises(AssertionError): model.load_weights([("sound_encoder.encoder.weight", object())]) +def _vision_norm_model(use_layernorm=True): + model = object.__new__(NemotronH_Nano_VL_V2) + torch.nn.Module.__init__(model) + model.model_config = _ImageOnlyModelConfig() + model.language_model = _LanguageModel() + model.vision_model = _VisionModel() + model.sound_encoder = None + model.mlp1 = torch.nn.Identity() + model.vision_final_layernorm = ( + torch.nn.LayerNorm(4, eps=1e-6, dtype=torch.bfloat16) if use_layernorm else None + ) + model.patch_size = 2 + model.downsample_ratio = 0.5 + model.ps_version = "v2" + model.video_temporal_patch_size = 2 + return model + + +@pytest.mark.parametrize("mode", ["image", "dynamic", "video"]) +@pytest.mark.parametrize("use_layernorm", [True, False]) +def test_vision_layernorm_is_loaded_and_applied_before_pixel_shuffle( + mode, use_layernorm +): + model = _vision_norm_model(use_layernorm) + weight = torch.tensor([2.0, 3.0, 4.0, 5.0], dtype=torch.bfloat16) + bias = torch.tensor([0.5, -0.5, 1.0, -1.0], dtype=torch.bfloat16) + model.load_weights( + [ + ("vision_projector.vision_final_layernorm.weight", weight), + ("vision_projector.vision_final_layernorm.bias", bias), + ] + if use_layernorm + else [] + ) + features = torch.arange(16, dtype=torch.bfloat16).reshape(1, 4, 4) + expected = ( + torch.nn.functional.layer_norm(features, (4,), weight, bias, eps=1e-6) + if use_layernorm + else features + ) + expected = model.pixel_shuffle(expected.reshape(1, 2, 2, 4)).reshape(1, 1, 16) + model.vision_model = lambda *args, **kwargs: (None, features) + pixels = torch.zeros(2 if mode == "video" else 1, 3, 4, 4) + if mode == "dynamic": + output = model.extract_feature_dynamic(pixels, imgs_sizes=[(4, 4)]) + else: + output = model.extract_feature( + pixels, num_frames=2 if mode == "video" else None + ) + torch.testing.assert_close(output, expected) + + def _make_mm_items_with_video_bytes( video_bytes: bytes, ) -> MultiModalDataItems: diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 074a89abb9..b77206295d 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1001,6 +1001,16 @@ class NemotronH_Nano_VL_V2( vision_projection_hidden_size = config.projector_hidden_size llm_hidden_size = config.text_config.hidden_size + # Match the HF vision projector for checkpoints trained with MTP. + self.vision_final_layernorm = ( + nn.LayerNorm( + vit_hidden_size, + eps=getattr(vision_config, "layer_norm_eps", 1e-6), + ).to(llm_dtype) + if (getattr(config.text_config, "num_nextn_predict_layers", 0) or 0) > 0 + else None + ) + mlp1 = nn.Sequential( RMSNorm( hidden_size=vit_hidden_size @@ -1095,6 +1105,8 @@ class NemotronH_Nano_VL_V2( """Dynamic resolution extract_feature for images.""" _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes) vit_embeds = vit_embeds.to(dtype=torch.bfloat16) + if self.vision_final_layernorm is not None: + vit_embeds = self.vision_final_layernorm(vit_embeds) vit_embeds = self.pixel_shuffle_dynamic_res(vit_embeds, imgs_sizes=imgs_sizes) vit_embeds = self.mlp1(vit_embeds) return vit_embeds @@ -1127,6 +1139,8 @@ class NemotronH_Nano_VL_V2( else: _, vit_embeds = self.vision_model(chunk) vit_embeds = vit_embeds.to(dtype=torch.bfloat16) + if self.vision_final_layernorm is not None: + vit_embeds = self.vision_final_layernorm(vit_embeds) vit_embeds = vit_embeds.reshape( vit_embeds.shape[0], H_patches, W_patches, -1 ) @@ -1575,7 +1589,7 @@ class NemotronH_Nano_VL_V2( """Get the module prefix in multimodal models""" return MultiModelKeys.from_string_field( language_model="language_model", - connector=["mlp1", "sound_encoder.projection"], + connector=["mlp1", "sound_encoder.projection", "vision_final_layernorm"], tower_model=["vision_model", "sound_encoder.encoder"], ) @@ -1593,6 +1607,11 @@ class NemotronH_Nano_VL_V2( for modality in ("image", "video", "audio") ) adapter_dict = dict(self.mlp1.named_parameters()) + vision_norm_params = ( + dict(self.vision_final_layernorm.named_parameters()) + if load_multimodal_weights and self.vision_final_layernorm is not None + else {} + ) def is_llm(name: str) -> bool: return name.startswith("language_model") @@ -1606,6 +1625,11 @@ class NemotronH_Nano_VL_V2( def is_sound_weights(name: str) -> bool: return name.startswith("sound") + vision_norm_prefix = "vision_projector.vision_final_layernorm." + + def is_vision_norm_weights(name: str) -> bool: + return name.startswith(vision_norm_prefix) + # LLM weights (the bulk of the model) are streamed lazily through a # generator so each tensor is copied into its parameter before the # iterator advances, avoiding stale-reference corruption with @@ -1637,6 +1661,15 @@ class NemotronH_Nano_VL_V2( continue assert self.sound_encoder is not None sound_weights.append((name, w.detach().clone())) + elif is_vision_norm_weights(name): + if not load_multimodal_weights: + continue + param_name = name.removeprefix(vision_norm_prefix) + assert param_name in vision_norm_params, ( + f"Unexpected vision LayerNorm weight: {name}" + ) + with torch.no_grad(): + default_weight_loader(vision_norm_params[param_name], w) # Fully drain the generator so every mm tensor is buffered, even if # the LLM loader stops iterating early.