[TRTLLM-15404][fix] VisualGen: refuse static quant recipes against unquantized checkpoints (silent weight corruption)

tensorrt_llm/_torch/visual_gen/quantization/loader.py

@@ -18,6 +18,15 @@
 from tensorrt_llm.quantization.mode import QuantAlgo
 from tensorrt_llm.quantization.utils import fp4_utils
 
+# Checkpoint scale tensors each static (pre-quantized) recipe expects
+# alongside the quantized weight (e.g. as exported by ModelOpt). Used to
+# detect unquantized checkpoints when dynamic weight quantization is off.
+_STATIC_SCALE_KEYS = {
+    QuantAlgo.FP8: ("weight_scale",),
+    QuantAlgo.FP8_BLOCK_SCALES: ("weight_scale",),
+    QuantAlgo.NVFP4: ("weight_scale", "weight_scale_2"),
+}
+
 
 class DynamicLinearWeightLoader:
     """
@@ -135,6 +144,43 @@ def _get_quant_algo_for_layer(self, name: str) -> Optional[QuantAlgo]:
 
         return None
 
+    def _check_static_quant_scales(
+        self, weight_dict: Dict[str, torch.Tensor], quant_algo: Optional[QuantAlgo], name: str
+    ) -> None:
+        """Refuse static quant recipes against checkpoints without scales.
+
+        With ``dynamic_weight_quant=False`` the loader expects the checkpoint
+        to provide pre-quantized weights plus their scale tensors. If a module
+        was built for a quantized recipe but the checkpoint holds a
+        high-precision weight with no scales (an unquantized checkpoint),
+        ``Linear.load_weights`` would silently cast the weight into the
+        quantized buffer while the scales keep their default or uninitialized
+        values, corrupting the model without any error. Fail fast instead.
+        """
+        expected_scales = _STATIC_SCALE_KEYS.get(quant_algo)
+        if expected_scales is None or self.dynamic_weight_quant:
+            return
+
+        if self.quant_config is not None:
+            if self.quant_config.is_module_excluded_from_quantization(name):
+                return
+
+        weight = weight_dict.get("weight")
+        if weight is None or weight.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+            return
+
+        missing = [key for key in expected_scales if key not in weight_dict]
+        if missing:
+            raise ValueError(
+                f"Static quantization ({quant_algo.name}) is configured for module "
+                f"'{name}', but the checkpoint provides a {weight.dtype} weight "
+                f"without the expected scale tensor(s) {missing}. The checkpoint "
+                "appears to be unquantized; loading it would silently corrupt the "
+                "weights. Use a checkpoint that carries quantized weights and "
+                "scales (e.g. exported by ModelOpt), or enable load-time "
+                "quantization by setting 'dynamic': true in quant_config."
+            )
+
     def _should_dynamic_quantize(
         self, weight_dict: Dict[str, torch.Tensor], quant_algo: Optional[QuantAlgo], name: str
     ) -> bool:
@@ -271,6 +317,11 @@ def load_linear_weights(
         else:
             quant_algo = self._get_quant_algo_for_layer(name)
 
+        # Static (pre-quantized) recipes must not be loaded from checkpoints
+        # that do not carry the expected scale tensors.
+        for weight_dict in weight_dicts:
+            self._check_static_quant_scales(weight_dict, quant_algo, name)
+
         # Special handling for fused NVFP4 dynamic quantization
         # Fused weights (Q,K,V or gate,up) must be quantized TOGETHER
         # to ensure consistent global scale

tests/unittest/_torch/visual_gen/test_quant_static_guard.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Static quant recipes must be refused against unquantized checkpoints.
+
+With ``dynamic_weight_quant=False`` the loader expects pre-quantized weights
+plus their scale tensors from the checkpoint. Loading a high-precision
+checkpoint through a static recipe silently corrupts the weights (scales stay
+at their default or uninitialized values), so the loader must raise instead.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig
+from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+pytestmark = pytest.mark.cpu_only
+
+
+class _StubLinear:
+    """Records load_weights calls; quant_algo resolves via the loader's global config."""
+
+    quant_config = None
+
+    def __init__(self):
+        self.loaded = None
+
+    def load_weights(self, weight_dicts):
+        self.loaded = weight_dicts
+
+
+def _make_loader(quant_algo, dynamic_weight_quant=False, exclude_modules=None):
+    model_config = DiffusionModelConfig(
+        quant_config=QuantConfig(quant_algo=quant_algo, exclude_modules=exclude_modules),
+        dynamic_weight_quant=dynamic_weight_quant,
+    )
+    return DynamicLinearWeightLoader(model_config)
+
+
+def _bf16_weights():
+    return {"weight": torch.zeros(8, 16, dtype=torch.bfloat16)}
+
+
+class TestStaticQuantGuard:
+    @pytest.mark.parametrize(
+        "quant_algo",
+        [QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.NVFP4],
+    )
+    def test_static_recipe_vs_unquantized_checkpoint_raises(self, quant_algo):
+        loader = _make_loader(quant_algo)
+        module = _StubLinear()
+        with pytest.raises(ValueError, match="appears to be unquantized"):
+            loader.load_linear_weights(module, "blocks.0.attn1.to_q", [_bf16_weights()])
+        assert module.loaded is None
+
+    def test_static_fp8_checkpoint_with_scales_loads(self):
+        loader = _make_loader(QuantAlgo.FP8)
+        module = _StubLinear()
+        weight_dict = {
+            "weight": torch.zeros(8, 16, dtype=torch.float8_e4m3fn),
+            "weight_scale": torch.ones(1, dtype=torch.float32),
+            "input_scale": torch.ones(1, dtype=torch.float32),
+        }
+        loader.load_linear_weights(module, "blocks.0.attn1.to_q", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_static_nvfp4_checkpoint_with_scales_loads(self):
+        loader = _make_loader(QuantAlgo.NVFP4)
+        module = _StubLinear()
+        weight_dict = {
+            "weight": torch.zeros(8, 8, dtype=torch.uint8),
+            "weight_scale": torch.zeros(8, 1, dtype=torch.float8_e4m3fn),
+            "weight_scale_2": torch.ones(1, dtype=torch.float32),
+        }
+        loader.load_linear_weights(module, "blocks.0.attn1.to_q", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_excluded_module_keeps_high_precision_weights(self):
+        loader = _make_loader(QuantAlgo.FP8, exclude_modules=["proj_out"])
+        module = _StubLinear()
+        weight_dict = _bf16_weights()
+        loader.load_linear_weights(module, "proj_out", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_unquantized_recipe_is_unaffected(self):
+        loader = _make_loader(None)
+        module = _StubLinear()
+        weight_dict = _bf16_weights()
+        loader.load_linear_weights(module, "blocks.0.attn1.to_q", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_dynamic_recipe_skips_the_guard(self):
+        loader = _make_loader(QuantAlgo.FP8, dynamic_weight_quant=True)
+        # The guard is a no-op for dynamic recipes; quantization happens later
+        # in _maybe_dynamic_quantize (GPU path, not exercised here).
+        loader._check_static_quant_scales(_bf16_weights(), QuantAlgo.FP8, "blocks.0.attn1.to_q")