From 7fcba5cb0a3f67d42cb8369669ddac4f1511a0df Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:02:49 -0700 Subject: [PATCH 1/2] [nvbugs/6445375][fix] Skip MiniMax-M3 VL LayerNorm init on meta tensors Plain nn.LayerNorm.reset_parameters() fills weight/bias with aten.fill_.Scalar, which MetaInitMode does not allow, so the vision tower's LayerNorm construction raised MetaInitException and aborted meta-init for the whole model. The loader fell back to regular init, really allocating every weight on the host. Subclass nn.LayerNorm and short-circuit reset_parameters() when the weight is on the meta device, following NemotronLayerNormPlus1. The gate matters: an unconditional skip leaves uninitialized storage for modules built without a checkpoint and NaNs the vision-tower forward tests. All 65 layer-norm slots are present in the checkpoint, so the skipped values are overwritten at load time. Unwaive the two B300 MiniMax-M3 tests filed under this bug. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3_vl.py | 34 +++++++++++++++++-- tests/integration/test_lists/waives.txt | 2 -- 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index 55ce331d176c..3166336476cd 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py @@ -962,6 +962,30 @@ 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. + + ``nn.LayerNorm.reset_parameters`` fills weight/bias via ``aten.fill_.Scalar``, + which ``MetaInitMode`` does not allow, so a plain ``nn.LayerNorm`` anywhere in + the ``__init__`` tree aborts meta-init for the whole model and forces 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: + # Only skip on meta: real tensors still need ones/zeros, otherwise a + # module built without a checkpoint keeps uninitialized storage. + if self.weight is not None and self.weight.is_meta: + return + super().reset_parameters() + + # --------------------------------------------------------------------------- # Patch embedding (Conv3d). # --------------------------------------------------------------------------- @@ -1192,9 +1216,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 +1293,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": diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 71b5595714c3..d50290564212 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) From c27da3bcf90d142508450723c0f10a6736a95de9 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:01:10 -0700 Subject: [PATCH 2/2] [nvbugs/6445375][test] Add meta-init regression tests for MiniMax-M3 VL LayerNorm The MiniMaxVLLayerNorm meta-init skip had no in-tree coverage, so a regression would only surface as a silent fallback to regular model init (hundreds of GB of host allocation per rank) rather than a test failure. Covers the skip working under MetaInitMode, a control that restores the upstream reset_parameters and asserts the exception returns, and the two edge cases the gate depends on: off-meta init keeping ones/zeros, and elementwise_affine=False registering weight as None. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3_vl.py | 15 ++--- .../_torch/models/test_minimax_m3_vl.py | 63 +++++++++++++++++++ 2 files changed, 71 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index 3166336476cd..727b3326ca49 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py @@ -970,17 +970,18 @@ def from_dict_or_obj(cls, source: Any) -> "CLIPVisionConfig": class MiniMaxVLLayerNorm(nn.LayerNorm): """``nn.LayerNorm`` whose parameter init is skipped on meta tensors. - ``nn.LayerNorm.reset_parameters`` fills weight/bias via ``aten.fill_.Scalar``, - which ``MetaInitMode`` does not allow, so a plain ``nn.LayerNorm`` anywhere in - the ``__init__`` tree aborts meta-init for the whole model and forces 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 + ``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: - # Only skip on meta: real tensors still need ones/zeros, otherwise a - # module built without a checkpoint keeps uninitialized storage. + # 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() diff --git a/tests/unittest/_torch/models/test_minimax_m3_vl.py b/tests/unittest/_torch/models/test_minimax_m3_vl.py index 4000e954185c..7c42e2e076da 100644 --- a/tests/unittest/_torch/models/test_minimax_m3_vl.py +++ b/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). # ---------------------------------------------------------------------------