[Bugfix] Load and apply Nemotron Super final vision LayerNorm (#61401)

Signed-off-by: Shaun Kotek <skotek@nvidia.com>
Co-authored-by: Shaun Kotek - Nvidia <skotek@nvidia.com>

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:

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.