diff --git a/CHANGELOG.rst b/CHANGELOG.rst index be05c1c441e..880378ae39b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -63,6 +63,8 @@ Experimental **Bug Fixes** +- Add FP8, INT8, and AutoQuantize recipes that quantize timm ResNet shortcut inputs + immediately before residual adds. - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export. diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index f479bab3ae1..cb185c50b16 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -54,6 +54,7 @@ The `torch_quant_to_onnx.py` script quantizes [timm](https://github.com/huggingf - Loads a pretrained timm torch model (default: ViT-Base). - Quantizes the torch model to FP8, MXFP8, INT8, NVFP4, or INT4_AWQ using ModelOpt. - For models with Conv2d layers (e.g., SwinTransformer), automatically overrides Conv2d quantization to FP8 (for MXFP8/NVFP4 modes) or INT8 (for INT4_AWQ mode) for TensorRT compatibility. +- Uses ResNet FP8, INT8, and AutoQuantize recipes to quantize shortcut inputs immediately before residual adds. - Exports the quantized model to ONNX. - Postprocesses the ONNX model to be compatible with TensorRT. - Saves the final ONNX model. @@ -65,7 +66,7 @@ The `torch_quant_to_onnx.py` script quantizes [timm](https://github.com/huggingf ```bash python torch_quant_to_onnx.py \ --timm_model_name= \ - --quantize_mode= \ + --quantize_mode= \ --onnx_save_path= ``` @@ -74,6 +75,9 @@ Quantization configs are loaded from the YAML preset recipes under `--recipe=` to use a different recipe (e.g. `--recipe=nvfp4_awq_lite` or `--recipe=/path/to/my_quant_cfg.yaml`). +ResNet AutoQuantize searches FP8 and INT8 at eight effective bits and keeps its +residual connections in FP8. + ### Conv2d Quantization Override TensorRT only supports FP8 and INT8 for convolution operations. When quantizing models with Conv2d layers (like SwinTransformer), the script automatically applies the following overrides: diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index e0ffc75a294..4f615d8c5d5 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -34,9 +34,11 @@ from evaluation import evaluate import modelopt.torch.quantization as mtq -from modelopt.recipe import load_config +from modelopt.recipe import load_config, load_recipe from modelopt.recipe.presets import MODEL_QUANT_PRESET_DIR from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.custom import CUSTOM_POST_CONVERSION_PLUGINS """ Quantize a timm vision model and export to ONNX for TensorRT deployment. @@ -70,6 +72,13 @@ def load_quant_config(recipe: str) -> dict: return load_config(recipe, schema_type=QuantizeConfig).model_dump() +_RESNET_RECIPES = { + "fp8": "timm/resnet/ptq/fp8", + "int8": "timm/resnet/ptq/int8", + "auto": "timm/resnet/auto_quantize/fp8_int8_at_8p0bits", +} + + _FP8_CONV_OVERRIDE: list = [ { "parent_class": "nn.Conv2d", @@ -151,6 +160,27 @@ def get_quant_config(quantize_mode, recipe=None): return config +def _get_resnet_recipe(model, quantize_mode): + if not isinstance(model, timm.models.resnet.ResNet): + return None + recipe_path = _RESNET_RECIPES.get(quantize_mode) + return load_recipe(recipe_path) if recipe_path is not None else None + + +def _add_resnet_residual_quantizers(model): + block_types = (timm.models.resnet.BasicBlock, timm.models.resnet.Bottleneck) + for block in (module for module in model.modules() if isinstance(module, block_types)): + if block.downsample is None: + block.downsample = torch.nn.Sequential() + elif not isinstance(block.downsample, torch.nn.Sequential): + block.downsample = torch.nn.Sequential(block.downsample) + if "residual_quantizer" in block.downsample._modules: + continue + residual_quantizer = TensorQuantizer() + residual_quantizer.disable() + block.downsample.add_module("residual_quantizer", residual_quantizer) + + def filter_func(name): """Filter function to exclude certain layers from quantization. @@ -366,9 +396,10 @@ def auto_quantize_model( model, data_loader, quantization_formats, - effective_bits=4.8, + effective_bits=None, num_calib_steps=512, num_score_steps=128, + recipe=None, ): """Auto-quantize the model using optimal per-layer quantization search. @@ -385,26 +416,48 @@ def auto_quantize_model( Tuple of (quantized_model, search_state_dict) """ _disable_inplace_relu(model) - constraints = {"effective_bits": effective_bits} - - # Convert string format names to config objects, incorporating Conv2d TRT overrides. - # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. - # By including the overrides in the format configs, the auto_quantize search - # correctly accounts for Conv2d being FP8/INT8 in the effective_bits budget. - format_configs: list[dict[str, Any] | str] = [] - for fmt in quantization_formats: - if isinstance(fmt, str): - config = load_quant_config(fmt) - if fmt in _NEEDS_FP8_CONV_OVERRIDE: - config["quant_cfg"].extend(_FP8_CONV_OVERRIDE) - elif fmt in _NEEDS_INT8_CONV_OVERRIDE: - config["quant_cfg"].extend(_INT8_CONV_OVERRIDE) - format_configs.append(config) - else: - format_configs.append(fmt) - - print(f"Starting auto-quantization search with {len(format_configs)} formats...") - print(f"Effective bits constraint: {effective_bits}") + if recipe is None: + constraints = {"effective_bits": 4.8 if effective_bits is None else effective_bits} + format_configs: list[dict[str, Any] | str] = [] + for fmt in quantization_formats: + if isinstance(fmt, str): + config = load_quant_config(fmt) + if fmt in _NEEDS_FP8_CONV_OVERRIDE: + config["quant_cfg"].extend(_FP8_CONV_OVERRIDE) + elif fmt in _NEEDS_INT8_CONV_OVERRIDE: + config["quant_cfg"].extend(_INT8_CONV_OVERRIDE) + format_configs.append(config) + else: + format_configs.append(fmt) + fixed_quantization_config = None + module_search_spaces = None + disabled_layers = None + method = "gradient" + else: + auto_config = recipe.auto_quantize + constraints = auto_config.constraints.model_dump(exclude_none=True) + if effective_bits is not None: + constraints["effective_bits"] = effective_bits + format_configs = [] + fixed_quantization_config = recipe.quantize.model_dump() + module_search_spaces = [ + { + "module_name_patterns": search_space.module_name_patterns, + "quantization_formats": [ + candidate.model_dump() for candidate in search_space.candidate_formats + ], + "allow_no_quant": search_space.allow_no_quant, + } + for search_space in auto_config.module_search_spaces + ] + disabled_layers = auto_config.disabled_layers + method = auto_config.auto_quantize_method + + format_count = len(format_configs) or sum( + len(search_space["quantization_formats"]) for search_space in module_search_spaces or [] + ) + print(f"Starting auto-quantization search with {format_count} formats...") + print(f"Effective bits constraint: {constraints['effective_bits']}") print(f"Calibration steps: {num_calib_steps}, Scoring steps: {num_score_steps}") quantized_model, search_state = mtq.auto_quantize( @@ -417,6 +470,10 @@ def auto_quantize_model( num_calib_steps=num_calib_steps, num_score_steps=num_score_steps, verbose=True, + fixed_quantization_config=fixed_quantization_config, + module_search_spaces=module_search_spaces, + disabled_layers=disabled_layers, + method=method, ) # Disable quantization for specified layers @@ -513,8 +570,11 @@ def main(): parser.add_argument( "--effective_bits", type=float, - default=4.8, - help="Target effective bits for auto quantization constraint. Default is 4.8.", + default=None, + help=( + "Target effective bits for auto quantization. Defaults to 4.8 without a recipe " + "and overrides the ResNet recipe when provided." + ), ) parser.add_argument( "--num_score_steps", @@ -571,6 +631,10 @@ def main(): ) print(f"Base Model - Top-1 Accuracy: {top1:.2f}%, Top-5 Accuracy: {top5:.2f}%") + resnet_recipe = None if args.recipe else _get_resnet_recipe(model, args.quantize_mode) + if resnet_recipe is not None: + CUSTOM_POST_CONVERSION_PLUGINS.add(_add_resnet_residual_quantizers) + # Quantize model based on mode if args.quantize_mode == "auto": # Auto quantization requires labels for loss computation @@ -589,13 +653,18 @@ def main(): args.effective_bits, args.calibration_data_size, args.num_score_steps, + recipe=resnet_recipe, ) else: # Standard quantization - load calibration data # Note: MXFP8 is dynamic and does not need calibration itself, but when # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8 # quantizers require calibration data. - config = get_quant_config(args.quantize_mode, args.recipe) + config = ( + resnet_recipe.quantize.model_dump() + if resnet_recipe is not None + else get_quant_config(args.quantize_mode, args.recipe) + ) data_loader = load_calibration_data( model, @@ -612,6 +681,7 @@ def main(): # Auto mode also needs this when an MXFP8/NVFP4 candidate format is in the search set. uses_dynamic_quantize = args.quantize_mode in ("mxfp8", "nvfp4") or ( args.quantize_mode == "auto" + and resnet_recipe is None and any(fmt in _NEEDS_FP8_CONV_OVERRIDE for fmt in args.auto_quantization_formats) ) if uses_dynamic_quantize: @@ -621,7 +691,10 @@ def main(): # no tactic for that 3-channel Q→Conv fusion. Skip for pure INT8 (unaffected). uses_fp8_conv_input = args.quantize_mode in ("fp8", "mxfp8", "nvfp4") or ( args.quantize_mode == "auto" - and any(fmt not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats) + and ( + resnet_recipe is not None + or any(fmt not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats) + ) ) if uses_fp8_conv_input: _disable_low_channel_conv_input_quantizers(quantized_model) diff --git a/modelopt_recipes/timm/resnet/auto_quantize/fp8_int8_at_8p0bits.yaml b/modelopt_recipes/timm/resnet/auto_quantize/fp8_int8_at_8p0bits.yaml new file mode 100644 index 00000000000..233e8f88fe8 --- /dev/null +++ b/modelopt_recipes/timm/resnet/auto_quantize/fp8_int8_at_8p0bits.yaml @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + fp8: configs/numerics/fp8 + fp8_model: configs/ptq/presets/model/fp8 + int8_model: configs/ptq/presets/model/int8 + +metadata: + recipe_type: auto_quantize + description: FP8 and INT8 ResNet AutoQuantize with FP8 residual connections. + +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*residual_quantizer' + cfg: + $import: fp8 + +auto_quantize: + constraints: + effective_bits: 8.0 + module_search_spaces: + - module_name_patterns: + - '*' + candidate_formats: + - $import: fp8_model + - $import: int8_model + allow_no_quant: false + auto_quantize_method: gradient + score_size: 128 diff --git a/modelopt_recipes/timm/resnet/ptq/fp8.yaml b/modelopt_recipes/timm/resnet/ptq/fp8.yaml new file mode 100644 index 00000000000..9de05bd98a9 --- /dev/null +++ b/modelopt_recipes/timm/resnet/ptq/fp8.yaml @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + fp8: configs/numerics/fp8 + w8a8_fp8_fp8: configs/ptq/units/w8a8_fp8_fp8 + +metadata: + recipe_type: ptq + description: FP8 ResNet PTQ with quantized residual connections. + +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - $import: w8a8_fp8_fp8 + - $import: default_disabled_quantizers + - quantizer_name: '*residual_quantizer' + cfg: + $import: fp8 diff --git a/modelopt_recipes/timm/resnet/ptq/int8.yaml b/modelopt_recipes/timm/resnet/ptq/int8.yaml new file mode 100644 index 00000000000..68c10f055d6 --- /dev/null +++ b/modelopt_recipes/timm/resnet/ptq/int8.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + int8: configs/numerics/int8 + int8_per_channel: configs/numerics/int8_per_channel + +metadata: + recipe_type: ptq + description: INT8 ResNet PTQ with quantized residual connections. + +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: int8_per_channel + - quantizer_name: '*input_quantizer' + cfg: + $import: int8 + - $import: default_disabled_quantizers + - quantizer_name: '*residual_quantizer' + cfg: + $import: int8 diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index fe99cae9bc3..430100478aa 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -14,6 +14,8 @@ # limitations under the License. +from collections import defaultdict + import onnx import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command @@ -21,6 +23,7 @@ # TODO: Add int4_awq once the INT4 exporter supports non-MatMul/Gemm consumer patterns # (e.g., DQ -> Reshape -> Slice in small ViT / SwinTransformer ONNX graphs). _QUANT_MODES = ["fp8", "int8", "mxfp8", "nvfp4", "auto"] +_RESNET_QUANT_MODES = {"fp8", "int8", "auto"} _MODELS = { "vit_tiny": ("vit_tiny_patch16_224", '{"depth": 1}'), @@ -30,6 +33,36 @@ } +def _assert_residual_inputs_are_quantized(onnx_save_path): + model = onnx.load(onnx_save_path) + consumers = defaultdict(list) + producers = {} + for node in model.graph.node: + for input_name in node.input: + consumers[input_name].append(node) + for output_name in node.output: + producers[output_name] = node + + residual_adds = [ + node + for node in model.graph.node + if node.op_type == "Add" + and [consumer.op_type for consumer in consumers[node.output[0]]] == ["Relu"] + ] + assert len(residual_adds) == 16 + for add in residual_adds: + input_producers = [producers[input_name] for input_name in add.input] + input_producers = [ + producers[node.input[0]] if node.op_type == "Cast" else node for node in input_producers + ] + assert any( + node.op_type.endswith("DequantizeLinear") + and "/downsample/residual_quantizer/" in node.name + and producers[node.input[0]].op_type.endswith("QuantizeLinear") + for node in input_producers + ) + + @pytest.mark.parametrize("quantize_mode", _QUANT_MODES) @pytest.mark.parametrize("model_key", list(_MODELS)) def test_torch_onnx(model_key, quantize_mode): @@ -48,6 +81,9 @@ def test_torch_onnx(model_key, quantize_mode): cmd_parts.extend(["--no_pretrained", "--trt_build"]) run_example_command(cmd_parts, "torch_onnx") + if model_key == "resnet50" and quantize_mode in _RESNET_QUANT_MODES: + _assert_residual_inputs_are_quantized(onnx_save_path) + def test_torch_onnx_recipe_flag(tmp_path): timm_model_name, model_kwargs = _MODELS["vit_tiny"]