diff --git a/tensorrt_llm/_torch/visual_gen/quantization/loader.py b/tensorrt_llm/_torch/visual_gen/quantization/loader.py index 8b042ed25e38..8c1c33a59adc 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/loader.py +++ b/tensorrt_llm/_torch/visual_gen/quantization/loader.py @@ -18,6 +18,18 @@ 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. +# Algos absent from this dict have no verified checkpoint layout in VisualGen; +# the guard fails closed for them (a high-precision weight under any static +# recipe is refused rather than silently cast over uninitialized scales). +_STATIC_SCALE_KEYS = { + QuantAlgo.FP8: ("weight_scale",), + QuantAlgo.FP8_BLOCK_SCALES: ("weight_scale",), + QuantAlgo.NVFP4: ("weight_scale", "weight_scale_2"), +} + class DynamicLinearWeightLoader: """ @@ -135,6 +147,57 @@ 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. + + Algos without an entry in ``_STATIC_SCALE_KEYS`` (e.g. the AWQ + variants, which also allocate ``weight_scale`` uninitialized) fail + closed: a high-precision weight under any static recipe is refused + even when the expected scale names are unknown. + """ + if quant_algo 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 + + expected_scales = _STATIC_SCALE_KEYS.get(quant_algo) + if expected_scales is not None: + missing = [key for key in expected_scales if key not in weight_dict] + if not missing: + return + detail = f"without the expected scale tensor(s) {missing}" + else: + detail = ( + "and no checkpoint scale layout is registered for this algo in " + "_STATIC_SCALE_KEYS, so the scales cannot be verified (the " + "guard fails closed for unverified algos)" + ) + raise ValueError( + f"Static quantization ({quant_algo.name}) is configured for module " + f"'{name}', but the checkpoint provides a {weight.dtype} weight " + f"{detail}. 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 +334,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 diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 4947141a35a9..d7c85557ce8d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -47,6 +47,7 @@ l0_cpu: - unittest/_torch/visual_gen/test_flux_infer.py - unittest/_torch/visual_gen/test_ltx2_pipeline.py - unittest/_torch/visual_gen/test_ltx2_transformer.py + - unittest/_torch/visual_gen/test_quant_static_guard.py - unittest/_torch/visual_gen/test_teacache.py - unittest/_torch/visual_gen/test_tensor_payload.py - unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py diff --git a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py new file mode 100644 index 000000000000..af2a40e2bad9 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py @@ -0,0 +1,115 @@ +# 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 + +# Pure-CPU test (stub Linear, no device): runs in the CPU lane (l0_cpu.yml), +# which selects with `-m cpu_only`; GPU stages deselect it via `not cpu_only`. +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 + + @pytest.mark.parametrize( + "quant_algo", + [QuantAlgo.W4A16_AWQ, QuantAlgo.W4A8_AWQ, QuantAlgo.W8A8_SQ_PER_CHANNEL], + ) + def test_unregistered_static_algo_fails_closed(self, quant_algo): + """Algos accepted by config parsing but absent from _STATIC_SCALE_KEYS + (their VisualGen checkpoint layout is unverified) must also refuse a + high-precision weight instead of silently skipping the check.""" + loader = _make_loader(quant_algo) + module = _StubLinear() + with pytest.raises(ValueError, match="fails closed for unverified algos"): + 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")