diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 632b5532ebf..7262ddc0526 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,6 +21,7 @@ Changelog **Bug Fixes** +- Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. 0.46 (2026-08-17) diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index f479bab3ae1..bf50cfd2c69 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -54,6 +54,9 @@ 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. +- Supports FP8 and INT8 recipes for convolutional architectures such as ResNet. Other formats are + not supported for convolutional models because of limited TensorRT kernel support. +- ResNet FP8 and INT8 recipes 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,25 +68,31 @@ 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= \ + --qformat= \ --onnx_save_path= ``` -Quantization configs are loaded from the YAML preset recipes under -`modelopt_recipes/configs/ptq/presets/model/`, selected by `--quantize_mode`. Pass -`--recipe=` to use a different -recipe (e.g. `--recipe=nvfp4_awq_lite` or `--recipe=/path/to/my_quant_cfg.yaml`). +Without `--recipe`, `--qformat` selects a quantization preset. Pass a built-in recipe name or YAML +path to `--recipe` to use a PTQ or AutoQuantize recipe instead. The recipe is authoritative when +provided, so `--qformat` is ignored. + +Convolutional architectures such as ResNet support only FP8 and INT8 quantization. MXFP8, NVFP4, +INT4_AWQ, and AutoQuantize are not supported for these models because TensorRT does not provide +the required convolution kernels. ### 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: -| Quantize Mode | Conv2d Override | Reason | +| Qformat | Conv2d Override | Reason | | :---: | :---: | :--- | | FP8, INT8 | None (already compatible) | Native TRT support | | MXFP8, NVFP4 | Conv2d -> FP8 | TRT Conv limitation | | INT4_AWQ | Conv2d -> INT8 | TRT Conv limitation | +These overrides support transformer architectures that contain individual Conv2d layers; they do +not make MXFP8, NVFP4, INT4_AWQ, or AutoQuantize supported for convolutional architectures. + ### Evaluation If the input model is of type image classification, use the following script to evaluate it. The script automatically downloads and uses the [ILSVRC/imagenet-1k](https://huggingface.co/datasets/ILSVRC/imagenet-1k) dataset from Hugging Face. This gated repository requires authentication via Hugging Face access token. See for details. @@ -332,7 +341,7 @@ For full documentation, see the [TensorRT-Edge-LLM Developer Guide](https://nvid ## Mixed Precision Quantization (Auto Mode) -The `auto` mode enables mixed precision quantization by searching for the optimal quantization format per layer. This approach balances model accuracy and compression by assigning different precision formats (e.g., NVFP4, FP8) to different layers based on their sensitivity. +AutoQuantize recipes enable mixed precision quantization by searching for the optimal quantization format per layer. This approach balances model accuracy and compression by assigning different precision formats (e.g., NVFP4, FP8) to different layers based on their sensitivity. The `--qformat=auto` CLI mode remains available for configuring the search with individual flags. ### How it works @@ -353,7 +362,18 @@ The `auto` mode enables mixed precision quantization by searching for the optima ```bash python torch_quant_to_onnx.py \ --timm_model_name=vit_base_patch16_224 \ - --quantize_mode=auto \ + --recipe=general/auto_quantize/nvfp4_fp8_at_5p4bits \ + --calibration_data_size=512 \ + --evaluate \ + --onnx_save_path=vit_base_patch16_224.auto_quant.onnx +``` + +The equivalent flag-based form is: + +```bash +python torch_quant_to_onnx.py \ + --timm_model_name=vit_base_patch16_224 \ + --qformat=auto \ --auto_quantization_formats nvfp4_awq_lite fp8 \ --effective_bits=4.8 \ --num_score_steps=128 \ @@ -369,7 +389,7 @@ python torch_quant_to_onnx.py \ | [vit_base_patch16_224](https://huggingface.co/timm/vit_base_patch16_224.augreg_in21k_ft_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [swin_tiny_patch4_window7_224](https://huggingface.co/timm/swin_tiny_patch4_window7_224.ms_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [swinv2_tiny_window8_256](https://huggingface.co/timm/swinv2_tiny_window8_256.ms_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | -| [resnet50](https://huggingface.co/timm/resnet50.a1_in1k) | ✅ | ✅ | ✅ | ✅ | | ✅ | +| [resnet50](https://huggingface.co/timm/resnet50.a1_in1k) | ✅ | ✅ | N/A | N/A | N/A | N/A | ## Resources diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index e0ffc75a294..7450f124274 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -19,8 +19,8 @@ import subprocess import sys import warnings +from copy import deepcopy from pathlib import Path -from typing import Any # Add onnx_ptq to path for shared modules sys.path.insert(0, str(Path(__file__).parent.parent / "onnx_ptq")) @@ -34,9 +34,10 @@ from evaluation import evaluate import modelopt.torch.quantization as mtq -from modelopt.recipe import load_config -from modelopt.recipe.presets import MODEL_QUANT_PRESET_DIR -from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.recipe import ModelOptAutoQuantizeRecipe, ModelOptPTQRecipe, load_recipe +from modelopt.recipe.presets import QUANT_CFG_CHOICES +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. @@ -58,18 +59,6 @@ mp.set_start_method("spawn", force=True) # Needed for data loader with multiple workers -def load_quant_config(recipe: str) -> dict: - """Load a quantization config from a recipe YAML. - - ``recipe`` is either a preset basename under - ``modelopt_recipes/configs/ptq/presets/model/`` (e.g. ``nvfp4``) or a path to a - ``QuantizeConfig`` YAML file (filesystem or built-in recipe library). - """ - if "/" not in recipe and not recipe.endswith((".yml", ".yaml")): - recipe = f"{MODEL_QUANT_PRESET_DIR}/{recipe}" - return load_config(recipe, schema_type=QuantizeConfig).model_dump() - - _FP8_CONV_OVERRIDE: list = [ { "parent_class": "nn.Conv2d", @@ -122,27 +111,28 @@ def load_quant_config(recipe: str) -> dict: _NEEDS_INT8_CONV_OVERRIDE: set[str] = {"int4_awq"} -def get_quant_config(quantize_mode, recipe=None): +def get_quant_config(qformat): """Get quantization config, overriding Conv2d for TRT compatibility. - The config is loaded from ``recipe`` when given, else from the preset YAML - matching ``quantize_mode``. TensorRT only supports FP8 and INT8 for Conv layers. + The config is loaded from the preset YAML matching ``qformat``. TensorRT only supports FP8 + and INT8 for Conv layers. - For FP8: add MHA-aware LayerNorm output quantizer so TRT fuses shared Q/DQ into downstream attention matmuls. Softmax-output Q/DQ is inserted by the FP8 ONNX exporter's post-processing (fixed 1/448 scale, no calibration needed). - For MXFP8, NVFP4: override Conv2d to FP8 - For INT4_AWQ: override Conv2d to INT8 """ - config: dict = load_quant_config(recipe or quantize_mode) - if quantize_mode == "fp8": + config = deepcopy(QUANT_CFG_CHOICES[qformat]) + if qformat == "fp8": config["quant_cfg"].extend(_FP8_MHA_OVERRIDE) - elif quantize_mode in ("mxfp8", "nvfp4"): + elif qformat in ("mxfp8", "nvfp4"): warnings.warn( f"TensorRT only supports FP8/INT8 for Conv layers. " - f"Overriding Conv2d quantization to FP8 for '{quantize_mode}' mode." + f"Overriding Conv2d quantization to FP8 for '{qformat}' format." ) config["quant_cfg"].extend(_FP8_CONV_OVERRIDE) - elif quantize_mode == "int4_awq": + config["algorithm"] = "max" + elif qformat == "int4_awq": warnings.warn( "TensorRT only supports FP8/INT8 for Conv layers. " "Overriding Conv2d quantization to INT8 for 'int4_awq' mode." @@ -151,6 +141,57 @@ def get_quant_config(quantize_mode, recipe=None): return config +def _prepare_auto_quantize_format(fmt): + config = deepcopy(QUANT_CFG_CHOICES[fmt]) if isinstance(fmt, str) else fmt.model_dump() + block_num_bits = { + entry["cfg"]["num_bits"] + for entry in config["quant_cfg"] + if isinstance(entry.get("cfg"), dict) and entry["cfg"].get("block_sizes") + } + if (isinstance(fmt, str) and fmt in _NEEDS_FP8_CONV_OVERRIDE) or block_num_bits & { + (4, 3), + (2, 1), + }: + config["quant_cfg"].extend(_FP8_CONV_OVERRIDE) + elif (isinstance(fmt, str) and fmt in _NEEDS_INT8_CONV_OVERRIDE) or 4 in block_num_bits: + config["quant_cfg"].extend(_INT8_CONV_OVERRIDE) + return config + + +def _add_resnet_residual_quantizers(model): + """Add disabled quantizers immediately before each ResNet residual addition. + + Appending to ``downsample`` places the quantizer on the shortcut immediately before the + residual addition. Identity shortcuts use an empty ``Sequential`` so the placement is the + same for every block. Quantizers start disabled and are enabled only by an explicit recipe. + """ + 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 hasattr(block.downsample, "residual_quantizer"): + continue + residual_quantizer = TensorQuantizer() + residual_quantizer.disable() + block.downsample.add_module("residual_quantizer", residual_quantizer) + + +def _enables_resnet_residual_quantization(recipe): + if not isinstance(recipe, ModelOptPTQRecipe): + return False + for entry in recipe.quantize.quant_cfg: + config = entry.model_dump(exclude_unset=True) + patterns = config.get("quantizer_name", []) + patterns = [patterns] if isinstance(patterns, str) else patterns + if any("residual_quantizer" in pattern for pattern in patterns) and config.get( + "enable", True + ): + return True + return False + + def filter_func(name): """Filter function to exclude certain layers from quantization. @@ -176,6 +217,15 @@ def _disable_high_rank_input_quantizers(model, input_shape, device): but 3D in ViT) must be skipped. A forward pass with hooks identifies them at runtime, so this works across architectures without hardcoded paths. """ + if not any( + isinstance(quantizer, TensorQuantizer) + and quantizer.is_enabled + and quantizer.block_sizes + and name.endswith("input_quantizer") + for name, quantizer in model.named_modules() + ): + return + high_rank: set[str] = set() handles = [] for name, mod in model.named_modules(): @@ -197,14 +247,15 @@ def hook(m, inp, out, _n=name): h.remove() model.train(was_training) - if not high_rank: - return - prefixes = tuple(n + "." for n in high_rank) - mtq.disable_quantizer(model, lambda n: n.startswith(prefixes)) + modules = dict(model.named_modules()) + for name in high_rank: + quantizer = getattr(modules[name], "input_quantizer", None) + if quantizer is not None and quantizer.is_enabled and quantizer.block_sizes: + quantizer.disable() -def _disable_low_channel_conv_input_quantizers(model): - """Disable ``input_quantizer`` on Conv2d modules whose ``in_channels <= 3``. +def _disable_low_channel_fp8_conv_input_quantizers(model): + """Disable FP8 ``input_quantizer`` on Conv2d modules whose ``in_channels <= 3``. The first Conv2d of an image backbone (e.g. ResNet50's ``conv1``) consumes raw RGB input, so ``in_channels == 3``. On Blackwell (compute capability 12.0) TRT @@ -223,10 +274,22 @@ def _disable_low_channel_conv_input_quantizers(model): for _, mod in model.named_modules(): if isinstance(mod, torch.nn.Conv2d) and mod.in_channels <= 3: q = getattr(mod, "input_quantizer", None) - if q is not None and q.is_enabled: + if q is not None and q.is_enabled and q.num_bits == (4, 3): q.disable() +def _validate_resnet_quantizers(model): + """Reject enabled ResNet quantizers other than per-tensor FP8 or INT8.""" + for name, quantizer in model.named_modules(): + if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled: + continue + if quantizer.num_bits not in ((4, 3), 8) or quantizer.block_sizes: + raise ValueError( + f"ResNet quantizer '{name}' uses an unsupported format; only FP8 and INT8 " + "are supported for convolutional models." + ) + + def load_calibration_data(model, data_size, batch_size, device, with_labels=False): """Load and prepare calibration data. @@ -272,7 +335,12 @@ def _disable_dead_quantizers(model): meaningful to quantize and would otherwise break ONNX export. """ for _, mod in model.named_modules(): - for attr in ("input_quantizer", "output_quantizer", "weight_quantizer"): + for attr in ( + "input_quantizer", + "output_quantizer", + "weight_quantizer", + "residual_quantizer", + ): q = getattr(mod, attr, None) if q is None or not q.is_enabled: continue @@ -283,36 +351,6 @@ def _disable_dead_quantizers(model): q.disable() -def _calibrate_uncalibrated_quantizers(model, data_loader): - """Calibrate FP8 quantizers that weren't calibrated by mtq.quantize(). - - When MXFP8/NVFP4 modes override Conv2d to FP8, the FP8 quantizers may not - be calibrated because the MXFP8/NVFP4 quantization pipeline skips standard - calibration. This function explicitly calibrates those uncalibrated quantizers. - """ - uncalibrated = [] - for _, module in model.named_modules(): - for attr_name in ("input_quantizer", "weight_quantizer"): - if not hasattr(module, attr_name): - continue - quantizer = getattr(module, attr_name) - if quantizer.is_enabled and not quantizer.block_sizes and quantizer.amax is None: - quantizer.enable_calib() - uncalibrated.append(quantizer) - - if not uncalibrated: - return - - model.eval() - with torch.no_grad(): - for batch in data_loader: - model(batch) - - for quantizer in uncalibrated: - quantizer.disable_calib() - quantizer.load_calib_amax(strict=False) - - def quantize_model(model, config, data_loader=None): """Quantize the model using the given config and calibration data.""" if data_loader is not None: @@ -325,14 +363,8 @@ def forward_loop(model): else: quantized_model = mtq.quantize(model, config) - # Disable filtered quantizers BEFORE calibrating override quantizers so we don't - # waste time calibrating quantizers that are about to be turned off. mtq.disable_quantizer(quantized_model, filter_func) - # Calibrate any FP8 override quantizers that weren't calibrated by mtq.quantize(). - if data_loader is not None: - _calibrate_uncalibrated_quantizers(quantized_model, data_loader) - # Drop quantizers whose calibration saw only zeros (e.g. SwinV2 zero-init norm1/norm2) # so ``export_fp8`` doesn't divide by zero. _disable_dead_quantizers(quantized_model) @@ -362,13 +394,48 @@ def _disable_inplace_relu(model): module.inplace = False +def _mtq_inputs_from_auto_quantize_config(auto_config, fixed_quantize_config=None): + """Map a resolved AutoQuantizeConfig to mtq.auto_quantize inputs.""" + constraints = auto_config.constraints.model_dump(exclude_none=True) + if auto_config.cost_excluded_layers: + constraints.setdefault("cost", {})["excluded_module_name_patterns"] = ( + auto_config.cost_excluded_layers + ) + return { + "constraints": constraints, + "quantization_formats": [ + _prepare_auto_quantize_format(fmt) for fmt in auto_config.candidate_formats + ], + "fixed_quantization_config": ( + _prepare_auto_quantize_format(fixed_quantize_config) + if fixed_quantize_config is not None + else None + ), + "module_search_spaces": [ + { + "module_name_patterns": search_space.module_name_patterns, + "quantization_formats": [ + _prepare_auto_quantize_format(candidate) + 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, + "num_score_steps": auto_config.score_size, + } + + 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,38 +452,37 @@ 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}") - print(f"Calibration steps: {num_calib_steps}, Scoring steps: {num_score_steps}") + if recipe is None: + inputs = { + "constraints": {"effective_bits": 4.8 if effective_bits is None else effective_bits}, + "quantization_formats": [ + _prepare_auto_quantize_format(fmt) for fmt in quantization_formats + ], + "fixed_quantization_config": None, + "module_search_spaces": None, + "disabled_layers": None, + "method": "gradient", + "num_score_steps": num_score_steps, + } + else: + inputs = _mtq_inputs_from_auto_quantize_config(recipe.auto_quantize, recipe.quantize) + + format_count = len(inputs["quantization_formats"]) or sum( + len(search_space["quantization_formats"]) + for search_space in inputs["module_search_spaces"] or [] + ) + print(f"Starting auto-quantization search with {format_count} formats...") + print(f"Effective bits constraint: {inputs['constraints']['effective_bits']}") + print(f"Calibration steps: {num_calib_steps}, Scoring steps: {inputs['num_score_steps']}") quantized_model, search_state = mtq.auto_quantize( model, - constraints=constraints, - quantization_formats=format_configs, data_loader=data_loader, forward_step=forward_step, loss_func=loss_func, num_calib_steps=num_calib_steps, - num_score_steps=num_score_steps, verbose=True, + **inputs, ) # Disable quantization for specified layers @@ -450,20 +516,18 @@ def main(): type=str, ) parser.add_argument( - "--quantize_mode", + "--qformat", choices=["fp8", "mxfp8", "int8", "nvfp4", "int4_awq", "auto"], default="mxfp8", - help="Type of quantization to apply. Default is MXFP8.", + help="Quantization format to apply when --recipe is not provided. Default is MXFP8.", ) parser.add_argument( "--recipe", type=str, default=None, help=( - "Quantization config recipe: a preset basename under " - "modelopt_recipes/configs/ptq/presets/model/ or a path to a QuantizeConfig " - "YAML. Defaults to the preset matching --quantize_mode. Not supported with " - "--quantize_mode=auto." + "PTQ or AutoQuantize recipe YAML file or built-in recipe name. The recipe is " + "authoritative when provided; --qformat is used only without a recipe." ), ) parser.add_argument( @@ -513,8 +577,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 --qformat=auto. Defaults to 4.8 and is ignored when " + "an AutoQuantize recipe is provided." + ), ) parser.add_argument( "--num_score_steps", @@ -541,10 +608,12 @@ def main(): args = parser.parse_args() - if args.recipe and args.quantize_mode == "auto": + recipe = load_recipe(args.recipe) if args.recipe is not None else None + if recipe is not None and not isinstance( + recipe, (ModelOptPTQRecipe, ModelOptAutoQuantizeRecipe) + ): parser.error( - "--recipe is not supported with --quantize_mode=auto; " - "use --auto_quantization_formats instead." + f"Expected a PTQ or AutoQuantize recipe, got {type(recipe).__name__} from {args.recipe}." ) # Create model and move to appropriate device @@ -571,8 +640,20 @@ def main(): ) print(f"Base Model - Top-1 Accuracy: {top1:.2f}%, Top-5 Accuracy: {top5:.2f}%") - # Quantize model based on mode - if args.quantize_mode == "auto": + is_resnet = isinstance(model, timm.models.resnet.ResNet) + run_auto_quantize = isinstance(recipe, ModelOptAutoQuantizeRecipe) or ( + recipe is None and args.qformat == "auto" + ) + if is_resnet and run_auto_quantize: + raise ValueError("AutoQuantize is not supported for convolutional models such as ResNet.") + if is_resnet and recipe is None and args.qformat not in {"fp8", "int8"}: + raise ValueError( + f"ResNet does not support qformat '{args.qformat}'; only FP8 and INT8 are supported." + ) + if is_resnet and _enables_resnet_residual_quantization(recipe): + CUSTOM_POST_CONVERSION_PLUGINS.add(_add_resnet_residual_quantizers) + + if run_auto_quantize: # Auto quantization requires labels for loss computation data_loader = load_calibration_data( model, @@ -589,13 +670,18 @@ def main(): args.effective_bits, args.calibration_data_size, args.num_score_steps, + recipe=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 = ( + recipe.quantize.model_dump() + if isinstance(recipe, ModelOptPTQRecipe) + else get_quant_config(args.qformat) + ) data_loader = load_calibration_data( model, @@ -607,24 +693,14 @@ def main(): quantized_model = quantize_model(model, config, data_loader) - # MXFP8/NVFP4 lower their input quantizers to TRT DynamicQuantize (2D/3D only). - # Disable quantizers on 4D-input layers (Swin's norm1 / downsample.norm / top-level norm). - # 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 any(fmt in _NEEDS_FP8_CONV_OVERRIDE for fmt in args.auto_quantization_formats) - ) - if uses_dynamic_quantize: - _disable_high_rank_input_quantizers(quantized_model, input_shape, device) - - # FP8-family modes emit TRT_FP8QuantizeLinear on the first-layer conv; Blackwell has - # 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) - ) - if uses_fp8_conv_input: - _disable_low_channel_conv_input_quantizers(quantized_model) + if is_resnet: + _validate_resnet_quantizers(quantized_model) + + # Disable block quantizers on 4D-input layers, which TRT DynamicQuantize does not support. + _disable_high_rank_input_quantizers(quantized_model, input_shape, device) + + # Blackwell has no tactic for an FP8 Q→Conv fusion on the first RGB layer. + _disable_low_channel_fp8_conv_input_quantizers(quantized_model) # Print quantization summary print("\nQuantization Summary:") diff --git a/modelopt_recipes/README.md b/modelopt_recipes/README.md index b366e4cc670..1937d181361 100644 --- a/modelopt_recipes/README.md +++ b/modelopt_recipes/README.md @@ -43,12 +43,14 @@ huggingface/qwen3_5/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast`. |-----------|-----------------| | `general/` | **Model-agnostic** recipes — a good starting point for any model. PTQ combos, speculative-decoding training, and distillation. | | `huggingface//` | **Model-specific** recipes keyed by a HF `model_type`, optionally nested by released checkpoint. Use these first if your model has an entry. | +| `timm//` | **Architecture-specific** recipes for timm models. | | `models//` | **Instance-specific** recipes that mirror a particular published checkpoint's quantization config. | | `configs/` | Shared building blocks (`numerics/`, `ptq/units/`, `ptq/presets/`) that recipes compose from via `$import`. Not run directly. | **Choosing where to look:** check `huggingface//` (then any nested -`/`) for your model first; if there's no entry, fall back to -`general/`. The presence of a model folder signals a recommended, tuned recipe. +`/`) for a Hugging Face model or `timm//` for a timm model. If there's +no entry, fall back to `general/`. The presence of a model folder signals a recommended, tuned +recipe. --- @@ -90,6 +92,7 @@ a per-component mixed-precision scheme tuned to match a specific release. Browse - **Tuned for a HF architecture** → `huggingface///`, with a `README.md` documenting the delta from the generic preset. Verify the exact `model_type` against the checkpoint's `config.json` before placing it. +- **Tuned for a timm architecture** → `timm///`. - **Mirrors a specific released checkpoint** → `models//`. - Share reused bodies via a `# modelopt-schema:`-tagged snippet and `$import` it; keep recipe wrappers thin. diff --git a/modelopt_recipes/timm/resnet/ptq/README.md b/modelopt_recipes/timm/resnet/ptq/README.md new file mode 100644 index 00000000000..14941cdbc51 --- /dev/null +++ b/modelopt_recipes/timm/resnet/ptq/README.md @@ -0,0 +1,17 @@ +# ResNet PTQ recipes + +These recipes use the shared FP8 and INT8 numerics and standard disabled-quantizer units. Their +ResNet-specific change is the `*residual_quantizer` entry, which quantizes the shortcut immediately +before each residual addition. + +The residual quantizer modules are inserted by +`examples/torch_onnx/torch_quant_to_onnx.py`; these recipes require that integration and do not add +the modules themselves. + +| Recipe | ResNet-specific behavior | +|--------|--------------------------| +| `fp8.yaml` | Enables per-tensor FP8 shortcut quantization. | +| `int8.yaml` | Enables per-tensor INT8 shortcut quantization. | + +Only FP8 and INT8 are supported for convolutional architectures because TensorRT has limited +convolution kernel support. diff --git a/modelopt_recipes/timm/resnet/ptq/fp8.yaml b/modelopt_recipes/timm/resnet/ptq/fp8.yaml new file mode 100644 index 00000000000..5040246348c --- /dev/null +++ b/modelopt_recipes/timm/resnet/ptq/fp8.yaml @@ -0,0 +1,35 @@ +# 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 +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..8ebb6151a1d --- /dev/null +++ b/modelopt_recipes/timm/resnet/ptq/int8.yaml @@ -0,0 +1,40 @@ +# 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 +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..d2748942972 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -14,13 +14,18 @@ # 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 +from modelopt.recipe import load_recipe + # 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"] +_QFORMATS = ["fp8", "int8", "mxfp8", "nvfp4", "auto"] +_RESNET_RECIPE_QFORMATS = {"fp8", "int8"} _MODELS = { "vit_tiny": ("vit_tiny_patch16_224", '{"depth": 1}'), @@ -30,36 +35,86 @@ } -@pytest.mark.parametrize("quantize_mode", _QUANT_MODES) +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("qformat", _QFORMATS) @pytest.mark.parametrize("model_key", list(_MODELS)) -def test_torch_onnx(model_key, quantize_mode): +def test_torch_onnx(tmp_path, model_key, qformat): + if model_key == "resnet50" and qformat not in _RESNET_RECIPE_QFORMATS: + pytest.skip("Only FP8 and INT8 quantization are supported for ResNet") + timm_model_name, model_kwargs = _MODELS[model_key] - onnx_save_path = f"{model_key}.{quantize_mode}.onnx" + onnx_save_path = tmp_path / f"{model_key}.{qformat}.onnx" cmd_parts = extend_cmd_parts( ["python", "torch_quant_to_onnx.py"], timm_model_name=timm_model_name, model_kwargs=model_kwargs, - quantize_mode=quantize_mode, - onnx_save_path=onnx_save_path, + qformat=qformat, + recipe=( + f"timm/resnet/ptq/{qformat}" + if model_key == "resnet50" and qformat in _RESNET_RECIPE_QFORMATS + else None + ), + onnx_save_path=str(onnx_save_path), calibration_data_size="1", num_score_steps="1", ) cmd_parts.extend(["--no_pretrained", "--trt_build"]) run_example_command(cmd_parts, "torch_onnx") + if model_key == "resnet50" and qformat in _RESNET_RECIPE_QFORMATS: + _assert_residual_inputs_are_quantized(onnx_save_path) + def test_torch_onnx_recipe_flag(tmp_path): timm_model_name, model_kwargs = _MODELS["vit_tiny"] onnx_save_path = tmp_path / "vit_tiny.recipe.onnx" recipe_path = tmp_path / "disable_all.yaml" - recipe_path.write_text("quant_cfg:\n - quantizer_name: '*'\n enable: false\n") + recipe_path.write_text( + "metadata:\n" + " recipe_type: ptq\n" + " description: Disable all quantizers.\n" + "quantize:\n" + " algorithm: max\n" + " quant_cfg:\n" + " - quantizer_name: '*'\n" + " enable: false\n" + ) cmd_parts = extend_cmd_parts( ["python", "torch_quant_to_onnx.py"], timm_model_name=timm_model_name, model_kwargs=model_kwargs, - quantize_mode="int8", + qformat="int8", recipe=str(recipe_path), onnx_save_path=str(onnx_save_path), calibration_data_size="1", @@ -75,3 +130,60 @@ def test_torch_onnx_recipe_flag(tmp_path): "TRT_FP8QuantizeLinear", } assert not quantize_ops & {node.op_type for node in onnx.load(onnx_save_path).graph.node} + + +def test_torch_onnx_auto_quantize_recipe(tmp_path): + timm_model_name, model_kwargs = _MODELS["vit_tiny"] + onnx_save_path = tmp_path / "vit_tiny.auto_recipe.onnx" + recipe_path = tmp_path / "auto.yaml" + recipe_path.write_text( + "imports:\n" + " fp8: configs/ptq/presets/model/fp8\n" + " int8: configs/ptq/presets/model/int8\n" + "metadata:\n" + " recipe_type: auto_quantize\n" + " description: Test AutoQuantize recipe.\n" + "auto_quantize:\n" + " constraints:\n" + " effective_bits: 8.0\n" + " candidate_formats:\n" + " - $import: fp8\n" + " - $import: int8\n" + " auto_quantize_method: gradient\n" + " score_size: 1\n" + ) + + cmd_parts = extend_cmd_parts( + ["python", "torch_quant_to_onnx.py"], + timm_model_name=timm_model_name, + model_kwargs=model_kwargs, + qformat="int8", + recipe=str(recipe_path), + onnx_save_path=str(onnx_save_path), + calibration_data_size="1", + ) + cmd_parts.append("--no_pretrained") + run_example_command(cmd_parts, "torch_onnx") + + +def test_auto_quantize_recipe_mapping(): + from examples.torch_onnx.torch_quant_to_onnx import ( + _enables_resnet_residual_quantization, + _mtq_inputs_from_auto_quantize_config, + ) + + recipe = load_recipe("general/auto_quantize/nvfp4_fp8_at_5p4bits") + inputs = _mtq_inputs_from_auto_quantize_config(recipe.auto_quantize, recipe.quantize) + + assert inputs["num_score_steps"] == recipe.auto_quantize.score_size + assert len(inputs["quantization_formats"]) == 2 + block_format = next( + config + for config in inputs["quantization_formats"] + if any( + isinstance(entry.get("cfg"), dict) and entry["cfg"].get("block_sizes") + for entry in config["quant_cfg"] + ) + ) + assert any(entry.get("parent_class") == "nn.Conv2d" for entry in block_format["quant_cfg"]) + assert _enables_resnet_residual_quantization(load_recipe("timm/resnet/ptq/fp8"))