[https://nvbugs/6445375][fix] Kept the gated MiniMaxVLLayerNorm fix and added 4 CPU-only meta-init…

tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py

@@ -962,6 +962,31 @@ def from_dict_or_obj(cls, source: Any) -> "CLIPVisionConfig":
         return cls(**filtered)
 
 
+# ---------------------------------------------------------------------------
+# Layer norm.
+# ---------------------------------------------------------------------------
+
+
+class MiniMaxVLLayerNorm(nn.LayerNorm):
+    """``nn.LayerNorm`` whose parameter init is skipped on meta tensors.
+
+    ``reset_parameters`` fills weight/bias via ``aten.fill_.Scalar``, which
+    ``MetaInitMode`` rejects, so a plain ``nn.LayerNorm`` anywhere in the
+    ``__init__`` tree aborts meta-init for the whole model and drops the loader
+    onto its regular-init fallback (hundreds of GB of host allocation per rank
+    for M3). Every layer-norm slot here is covered by the checkpoint, so the
+    skipped values are always overwritten at load time.
+    """
+
+    def reset_parameters(self) -> None:
+        # Skip only on meta: off meta the ones/zeros are still needed, otherwise
+        # a module built without a checkpoint keeps uninitialized storage.
+        # ``weight`` is None when elementwise_affine=False.
+        if self.weight is not None and self.weight.is_meta:
+            return
+        super().reset_parameters()
+
+
 # ---------------------------------------------------------------------------
 # Patch embedding (Conv3d).
 # ---------------------------------------------------------------------------
@@ -1192,9 +1217,13 @@ def __init__(self, config: CLIPVisionConfig, dtype: torch.dtype):
         super().__init__()
         self.embed_dim = config.hidden_size
         self.self_attn = MiniMaxVLEncoderSelfAttention(config, dtype)
-        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps, dtype=dtype)
+        self.layer_norm1 = MiniMaxVLLayerNorm(
+            self.embed_dim, eps=config.layer_norm_eps, dtype=dtype
+        )
         self.mlp = MiniMaxVLEncoderMLP(config, dtype)
-        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps, dtype=dtype)
+        self.layer_norm2 = MiniMaxVLLayerNorm(
+            self.embed_dim, eps=config.layer_norm_eps, dtype=dtype
+        )
 
     def forward(
         self,
@@ -1265,7 +1294,7 @@ def __init__(self, config: CLIPVisionConfig, dtype: torch.dtype):
 
         self.embeddings = MiniMaxVLPatchEmbedding(config, dtype)
         # NOTE: the typo "layrnorm" matches the published checkpoint key.
-        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps, dtype=dtype)
+        self.pre_layrnorm = MiniMaxVLLayerNorm(embed_dim, eps=config.layer_norm_eps, dtype=dtype)
         self.encoder = MiniMaxVLEncoder(config, dtype)
 
         if config.position_embedding_type != "rope" or config.rope_mode != "3d":

tests/integration/test_lists/waives.txt

@@ -185,9 +185,7 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus
 full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4gpus_static_eplb[moe_backend=WIDEEP] SKIP (https://nvbugs/6546609)
 full:B300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346)
 full:B300/accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6525011)
-full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6445375)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188)
-full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6445375)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6529874)
 full:B300/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] SKIP (https://nvbugs/6474894)
 full:B300/accuracy/test_llm_api_pytorch.py::TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] SKIP (https://nvbugs/6539941)

tests/unittest/_torch/models/test_minimax_m3_vl.py

@@ -23,6 +23,7 @@
 import numpy as np
 import pytest
 import torch
+import torch.nn as nn
 from PIL import Image
 from safetensors import safe_open
 from transformers import AutoConfig, AutoProcessor, AutoTokenizer
@@ -31,6 +32,7 @@
 from tensorrt_llm._torch.models.modeling_minimaxm3 import get_text_config
 from tensorrt_llm._torch.models.modeling_minimaxm3_vl import (
     CLIPVisionConfig,
+    MiniMaxVLLayerNorm,
     MiniMaxVLPatchEmbedding,
     MiniMaxVLPatchMerger,
     MiniMaxVLVisionModel,
@@ -47,6 +49,7 @@
     reanchor_multimodal_checkpoint_keys,
     split_multimodal_weights,
 )
+from tensorrt_llm._torch.models.modeling_utils import MetaInitException, MetaInitMode
 
 # ---------------------------------------------------------------------------
 # Shared helpers (mirror the conventions used by test_minimax_m3.py).
@@ -812,6 +815,66 @@ def test_patch_merger_rejects_unaligned_input():
         merge(x)
 
 
+# ---------------------------------------------------------------------------
+# meta-init compatibility (CPU, no checkpoint).
+# ---------------------------------------------------------------------------
+
+
+def _build_tiny_vision_model() -> MiniMaxVLVisionModel:
+    cfg = CLIPVisionConfig.from_dict_or_obj(_tiny_vision_config())
+    return MiniMaxVLVisionModel(
+        config=cfg,
+        text_hidden_size=16,
+        projector_hidden_size=16,
+        dtype=torch.float32,
+    )
+
+
+def test_vision_tower_builds_under_meta_init():
+    """The vision tower must construct inside ``MetaInitMode``.
+
+    See :class:`MiniMaxVLLayerNorm` for why a plain ``nn.LayerNorm`` here aborts
+    meta-init for the whole M3 model.
+    """
+    with MetaInitMode():
+        model = _build_tiny_vision_model()
+
+    layer_norms = [m for m in model.modules() if isinstance(m, nn.LayerNorm)]
+    assert layer_norms, "expected layer norms in the vision tower"
+    # pre_layrnorm + layer_norm1/2 per encoder layer.
+    assert len(layer_norms) == 1 + 2 * len(model.vision_model.encoder.layers)
+    for ln in layer_norms:
+        assert ln.weight.is_meta
+        assert ln.bias.is_meta
+
+
+def test_meta_init_still_rejects_plain_layer_norm_init(monkeypatch):
+    """Control for :func:`test_vision_tower_builds_under_meta_init`.
+
+    Restoring the upstream ``reset_parameters`` must bring the exception back,
+    otherwise that test could pass without the skip doing any work.
+    """
+    monkeypatch.setattr(MiniMaxVLLayerNorm, "reset_parameters", nn.LayerNorm.reset_parameters)
+    with pytest.raises(MetaInitException, match="fill_"):
+        with MetaInitMode():
+            _build_tiny_vision_model()
+
+
+def test_layer_norm_off_meta_init_matches_upstream():
+    """Off meta the skip must not engage, else a checkpoint-free build keeps
+    uninitialized ``torch.empty`` storage instead of ones/zeros."""
+    ln = MiniMaxVLLayerNorm(8, dtype=torch.float32)
+    assert torch.equal(ln.weight, torch.ones(8))
+    assert torch.equal(ln.bias, torch.zeros(8))
+
+
+def test_layer_norm_without_affine_params_builds_under_meta_init():
+    """``elementwise_affine=False`` registers ``weight`` as ``None``."""
+    with MetaInitMode():
+        ln = MiniMaxVLLayerNorm(8, elementwise_affine=False, dtype=torch.float32)
+    assert ln.weight is None
+
+
 # ---------------------------------------------------------------------------
 # full multimodal smoke (CUDA + real checkpoint).
 # ---------------------------------------------------------------------------