diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9ee5a31b096..bad70d89b25 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add an NVFP4 tied-LSQ QAD recipe that quantizes Qwen3-VL visual attention, MLP, and merger linear layers and exports restored LSQ post-amax scales. - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections. - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 8726dc5c719..083ab07b62f 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -158,7 +158,11 @@ def get_weights_scaling_factor_from_quantizer( if cls._is_static_quantizer(weight_quantizer): # Static path: use pre-computed per-block amax values from quantizer global_amax = cls._get_static_global_amax(weight_quantizer).float() - per_block_amax = weight_quantizer._amax.float() + per_block_amax = ( + weight_quantizer.amax_post + if getattr(weight_quantizer, "_lsq", False) + else weight_quantizer.amax + ).float() # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX diff --git a/modelopt_recipes/general/qad/nvfp4_lsq-mse_init-fp8_kv-visual-linear.yaml b/modelopt_recipes/general/qad/nvfp4_lsq-mse_init-fp8_kv-visual-linear.yaml new file mode 100644 index 00000000000..93077e1b82f --- /dev/null +++ b/modelopt_recipes/general/qad/nvfp4_lsq-mse_init-fp8_kv-visual-linear.yaml @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +metadata: + description: >- + QAD recipe using static NVFP4 E2M1 weights, dynamic NVFP4 E2M1 + activations and FP8 KV cache. Unlike the default multimodal-safe recipe, + this recipe explicitly re-enables Qwen3-VL visual attention, MLP and + merger Linear modules after importing the default visual exclusions. +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + kv_fp8: configs/ptq/units/kv_fp8 + nvfp4: configs/numerics/nvfp4 + nvfp4_static: configs/numerics/nvfp4_static +quantize: + algorithm: + method: lsq + learnable_amax: + - pre + - post + tied_amax: true + quantize_pre_scale: true + scale_algorithm: + method: mse + fp8_scale_sweep: true + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers + - quantizer_name: '*visual*.attn.*.weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*visual.*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*visual*.linear_fc*.weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*visual.*patch_embed*' + enable: false + - quantizer_name: '*visual.*_bmm_quantizer' + enable: false diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py index d37d3aba3b1..e592e7e5e0f 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -167,6 +167,7 @@ def quantize(model, quant_cfg, forward_loop): @pytest.mark.parametrize( "recipe", [ + "general/qad/nvfp4_lsq-mse_init-fp8_kv-visual-linear", "model_type/qwen3_vl/ptq/fp8_vision-kv_none", "model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast", "model_type/qwen3_5/ptq/fp8_vision-kv_none", diff --git a/tests/gpu/torch/export/test_qwen_vision_recipe_export.py b/tests/gpu/torch/export/test_qwen_vision_recipe_export.py index 389375428b1..dc54f876234 100644 --- a/tests/gpu/torch/export/test_qwen_vision_recipe_export.py +++ b/tests/gpu/torch/export/test_qwen_vision_recipe_export.py @@ -27,6 +27,7 @@ _VISION_RECIPE = "fp8_vision-kv_none" _JOINT_RECIPE = "fp8_vision_lm-kv_fp8_cast" +_NVFP4_VISUAL_QAD_RECIPE = "general/qad/nvfp4_lsq-mse_init-fp8_kv-visual-linear" def _get_tiny_qwen_vlm(model_type): @@ -102,3 +103,57 @@ def calibration_forward(_model): assert ( any(name.startswith("model.language_model.") for name in scale_names) is quantizes_language ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9), + reason="NVFP4 visual QAD export requires compute capability 8.9 or newer", +) +@pytest.mark.parametrize( + ("model_type", "architecture"), + [ + ("qwen3_vl", "Qwen3VLForConditionalGeneration"), + ("qwen3_5", "Qwen3_5ForConditionalGeneration"), + ], +) +def test_nvfp4_lsq_visual_recipe_calibrates_and_exports(tmp_path, model_type, architecture): + model = _get_tiny_qwen_vlm(model_type).to("cuda").eval() + model.config.architectures = [architecture] + quant_cfg = load_recipe(_NVFP4_VISUAL_QAD_RECIPE).quantize.model_dump() + vision_config = model.config.vision_config + pixel_width = ( + vision_config.in_channels * vision_config.temporal_patch_size * vision_config.patch_size**2 + ) + pixel_values = torch.randn(4, pixel_width, dtype=torch.bfloat16, device="cuda") + image_grid_thw = torch.tensor([[1, 2, 2]], device="cuda") + input_ids = torch.randint(5, 16, (1, 4), device="cuda") + + def calibration_forward(_model): + model.model.visual(pixel_values, image_grid_thw) + model(input_ids=input_ids) + + mtq.quantize(model, quant_cfg, forward_loop=calibration_forward) + + export_path = tmp_path / f"{model_type}-nvfp4-visual-qad" + export_hf_checkpoint(model, export_dir=export_path) + + with (export_path / "hf_quant_config.json").open() as config_file: + quantization = json.load(config_file)["quantization"] + assert quantization["quant_algo"] == "NVFP4" + assert quantization["kv_cache_quant_algo"] == "FP8" + assert "quantized_layers" not in quantization + + tensor_names = set() + for checkpoint_path in export_path.glob("*.safetensors"): + with safe_open(str(checkpoint_path), framework="pt") as checkpoint: + tensor_names.update(checkpoint.keys()) + scale_names = { + name + for name in tensor_names + if name.endswith(("input_scale", "weight_scale", "weight_scale_2")) + } + assert any(name.startswith("model.visual.") for name in scale_names) + assert any(".attn.qkv.weight_scale" in name for name in scale_names) + assert any(".merger.linear_fc1.weight_scale" in name for name in scale_names) + assert any(name.startswith("model.language_model.") for name in scale_names) + assert not any("patch_embed" in name for name in scale_names) diff --git a/tests/unit/recipe/test_lsq_recipes.py b/tests/unit/recipe/test_lsq_recipes.py index dfc864a4009..86adbd2ed1f 100644 --- a/tests/unit/recipe/test_lsq_recipes.py +++ b/tests/unit/recipe/test_lsq_recipes.py @@ -27,6 +27,7 @@ _LSQ_RECIPES = { "nvfp4_dual_lsq-mse_init-fp8_kv.yaml": (False, False), "nvfp4_lsq-mse_init-fp8_kv.yaml": (True, True), + "nvfp4_lsq-mse_init-fp8_kv-visual-linear.yaml": (True, True), } @@ -76,3 +77,21 @@ def test_lsq_recipe_resolves_modular_quant_cfg(filename): assert input_cfg["block_sizes"]["type"] == "dynamic" assert input_cfg["num_bits"] == (2, 1) assert kv_cfg["num_bits"] == (4, 3) + + +def test_visual_linear_recipe_reenables_visual_quantizers(): + quantize = _load_lsq_recipe("nvfp4_lsq-mse_init-fp8_kv-visual-linear.yaml") + entries = {entry.quantizer_name: entry for entry in quantize.quant_cfg} + + expected = { + "*visual*.attn.*.weight_quantizer": "static", + "*visual.*input_quantizer": "dynamic", + "*visual*.linear_fc*.weight_quantizer": "static", + } + for name, block_type in expected.items(): + cfg = entries[name].cfg.model_dump() + assert cfg["num_bits"] == (2, 1) + assert cfg["block_sizes"]["type"] == block_type + + assert entries["*visual.*patch_embed*"].enable is False + assert entries["*visual.*_bmm_quantizer"].enable is False diff --git a/tests/unit/recipe/test_qwen_vision_recipe.py b/tests/unit/recipe/test_qwen_vision_recipe.py index 264e3b73239..77f2faffcf3 100644 --- a/tests/unit/recipe/test_qwen_vision_recipe.py +++ b/tests/unit/recipe/test_qwen_vision_recipe.py @@ -82,3 +82,52 @@ def test_qwen_vision_recipes_select_expected_quantizers( enabled_bmm = {name for name in enabled if name.endswith("_bmm_quantizer")} assert bool(enabled_bmm) is quantizes_kv assert all(name.startswith("model.language_model.") for name in enabled_bmm) + + +@pytest.mark.parametrize("model_type", ["qwen3_vl", "qwen3_5"]) +def test_nvfp4_lsq_visual_recipe_selects_expected_quantizers(model_type): + model = _get_tiny_qwen_vlm(model_type) + quant_cfg = load_recipe( + "general/qad/nvfp4_lsq-mse_init-fp8_kv-visual-linear" + ).quantize.model_dump() + quant_cfg["algorithm"] = None + + mtq.quantize(model, quant_cfg, forward_loop=None) + modules = dict(model.named_modules()) + enabled = {name for name, module in modules.items() if getattr(module, "is_enabled", False)} + + expected_visual_suffixes = { + "blocks.0.attn.qkv.weight_quantizer", + "blocks.0.attn.qkv.input_quantizer", + "blocks.0.attn.proj.weight_quantizer", + "blocks.0.attn.proj.input_quantizer", + "blocks.0.mlp.linear_fc1.weight_quantizer", + "blocks.0.mlp.linear_fc1.input_quantizer", + "blocks.0.mlp.linear_fc2.weight_quantizer", + "blocks.0.mlp.linear_fc2.input_quantizer", + "merger.linear_fc1.weight_quantizer", + "merger.linear_fc1.input_quantizer", + "merger.linear_fc2.weight_quantizer", + "merger.linear_fc2.input_quantizer", + } + if model_type == "qwen3_vl": + expected_visual_suffixes.update( + { + "deepstack_merger_list.0.linear_fc1.weight_quantizer", + "deepstack_merger_list.0.linear_fc1.input_quantizer", + "deepstack_merger_list.0.linear_fc2.weight_quantizer", + "deepstack_merger_list.0.linear_fc2.input_quantizer", + } + ) + assert {f"model.visual.{suffix}" for suffix in expected_visual_suffixes} <= enabled + assert not any("patch_embed" in name for name in enabled) + assert not any( + name.startswith("model.visual.") and name.endswith("_bmm_quantizer") for name in enabled + ) + + vision_weight = modules["model.visual.blocks.0.attn.qkv.weight_quantizer"] + vision_input = modules["model.visual.blocks.0.attn.qkv.input_quantizer"] + assert vision_weight.num_bits == (2, 1) + assert vision_weight.block_sizes["type"] == "static" + assert vision_input.num_bits == (2, 1) + assert vision_input.block_sizes["type"] == "dynamic" diff --git a/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py b/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py index ded38fa9dc2..72562b0d0dd 100644 --- a/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py +++ b/tests/unit/torch/quantization/test_nvfp4_static_export_cpu.py @@ -239,6 +239,31 @@ def test_static_matches_dynamic_when_amax_is_block_max(self): assert torch.allclose(static_ws_2, dynamic_ws_2) +def test_lsq_export_uses_post_amax_for_tied_and_dual_lsq(): + weight = torch.ones(2, BLOCK_SIZE, dtype=torch.bfloat16) + + for tied_amax in (False, True): + quantizer = _make_static_quantizer( + torch.ones(2, 1), + torch.tensor(6.0), + ) + quantizer.enable_lsq( + learnable_amax=["pre", "post"], + tied_amax=tied_amax, + ) + quantizer._amax_post.data.fill_(3.0) + if not tied_amax: + quantizer._amax_pre.data.fill_(1.0) + + scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer( + quantizer, + weight, + keep_high_precision=True, + ) + + assert torch.equal(scale, torch.full((2, 1), 0.5)) + + class TestNVFP4StaticManualRoundTrip: """Cross-check the export path against a manual per-block computation.