diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 20a677d0a0b..2b595515f58 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,12 +21,14 @@ Changelog **Backward Breaking Changes** - The ``quant_cfg`` field in quantization configs is now an **ordered list** of ``QuantizerCfgEntry`` dicts instead of a flat dictionary. Each entry specifies a ``quantizer_name`` wildcard, an optional ``parent_class`` filter, a ``cfg`` dict of quantizer attributes, and/or an ``enable`` flag. Entries are applied in list order with later entries overriding earlier ones. The old dict-based format is still accepted and automatically converted via ``normalize_quant_cfg_list()``, but now emits a ``DeprecationWarning``; new code should use the list format. All built-in configs (e.g. ``FP8_DEFAULT_CFG``, ``INT4_AWQ_CFG``, ``NVFP4_DEFAULT_CFG``), examples, and YAML recipes have been updated. See the :ref:`quant-cfg` documentation for the new format reference and migration guide. +- Deprecated Mllama (Llama 3.2 Vision) support in the ``llm_ptq`` and ``vlm_ptq`` examples. The ``model_type == "mllama"`` branches and ``MllamaImageProcessor`` usage have been removed from ``hf_ptq.py`` and ``example_utils.py``. For image-text calibration of VLMs, use ``--calib_with_images`` with a supported VLM (see Nemotron VL section in ``examples/llm_ptq/README.md``). **Bug Fixes** - Fix Megatron utility functions for generation (with pipeline parallelism) and ~10x speedup in MMLU score evaluation (by batching prefill passes). - Fix Minitron pruning (``mcore_minitron``) for MoE models. Importance estimation hooks were incorrectly registered for MoE modules and NAS step was hanging before this. - Fix TRT support for remote autotuning in ONNX Autotune from 10.16+ to 10.15+ and fix TRT versioning check to the ``trtexec`` version instead of the TRT Python API when using ``trtexec`` backend. +- Exclude MatMul/Gemm nodes with K or N < 16 from ONNX INT8 and FP8 quantization. Such small-dimension GEMMs cannot efficiently use INT8/FP8 Tensor Cores and the added Q/DQ layers cause perf regressions in TensorRT. Honors Gemm ``transB`` when deriving K. **Misc** diff --git a/LICENSE_HEADER b/LICENSE_HEADER index 7b0d9f06b06..9031d2eca55 100644 --- a/LICENSE_HEADER +++ b/LICENSE_HEADER @@ -1,4 +1,4 @@ -SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +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"); diff --git a/README.md b/README.md index df46a075215..f16d69813f7 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,8 @@ To install stable release packages for Model Optimizer with `pip` from [PyPI](ht pip install -U nvidia-modelopt[all] ``` +Model Optimizer will download and install additional third-party open source software projects. Review the license terms of these open source projects before use. + To install from source in editable mode with all development dependencies or to use the latest features, run: ```bash @@ -79,8 +81,14 @@ cd Model-Optimizer pip install -e .[dev] ``` -You can also directly use the [TensorRT-LLM docker images](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release/tags) -(e.g., `nvcr.io/nvidia/tensorrt-llm/release:`), which have Model Optimizer pre-installed. +You can also directly use NVIDIA container images, which have Model Optimizer pre-installed: + +- `nvcr.io/nvidia/pytorch:-py3` +- `nvcr.io/nvidia/nemo:` +- `nvcr.io/nvidia/tensorrt-llm/release:` +- `nvcr.io/nvidia/tensorrt:-py3` + +Before pulling and using the container images, please review their respective license terms. Make sure to upgrade Model Optimizer to the latest version as described above. Visit our [installation guide](https://nvidia.github.io/Model-Optimizer/getting_started/2_installation.html) for more fine-grained control on installed dependencies or for alternative docker images and environment variables to setup. diff --git a/docs/source/getting_started/_installation_for_Linux.rst b/docs/source/getting_started/_installation_for_Linux.rst index 2b2d4d8219b..a18b45ee7c4 100644 --- a/docs/source/getting_started/_installation_for_Linux.rst +++ b/docs/source/getting_started/_installation_for_Linux.rst @@ -32,11 +32,11 @@ Environment setup To use Model Optimizer with full dependencies (e.g. TensorRT/TensorRT-LLM deployment), we recommend using the `TensorRT-LLM docker image `_, - e.g., ``nvcr.io/nvidia/tensorrt-llm/release:``. + e.g., ``nvcr.io/nvidia/tensorrt-llm/release:`` (Model Optimizer pre-installed). Make sure to upgrade Model Optimizer to the latest version using ``pip`` as described in the next section. - You would also need to setup appropriate environment variables for the TensorRT binaries as follows: + If relevant, you would also need to setup appropriate environment variables for the TensorRT binaries as follows: .. code-block:: shell @@ -48,11 +48,16 @@ Environment setup **Alternative NVIDIA docker images** For PyTorch, you can also use `NVIDIA NGC PyTorch container `_ - and for NVIDIA Megatron-Bridge or Megatron-LM framework, you can use the `NeMo container `_. - Both of these containers come with Model Optimizer pre-installed. Make sure to update the Model Optimizer to the latest version if not already. + (``nvcr.io/nvidia/pytorch:-py3``, Model Optimizer pre-installed) + and for NVIDIA Megatron-Bridge or Megatron-LM framework, you can use the `NeMo container `_ + (``nvcr.io/nvidia/nemo:``, Model Optimizer pre-installed). + Make sure to update the Model Optimizer to the latest version if not already. For ONNX / TensorRT use cases, you can also use the `TensorRT container `_ - which provides superior performance to the PyTorch container. + (``nvcr.io/nvidia/tensorrt:-py3``), which provides superior performance to the PyTorch container. + + .. note:: + Before pulling and using the container images, please review their respective license terms. .. tab:: Local environment (PIP / Conda) @@ -82,8 +87,8 @@ Environment setup Install Model Optimizer ======================= -ModelOpt including its dependencies can be installed via ``pip``. Please review the license terms of ModelOpt and any -dependencies before use. +Model Optimizer will download and install additional third-party open source software projects. Review the license +terms of these open source projects before use. If you build and use ModelOpt's docker image, you can skip this step as the image already contains ModelOpt and all optional dependencies pre-installed. diff --git a/docs/source/getting_started/windows/_installation_for_Windows.rst b/docs/source/getting_started/windows/_installation_for_Windows.rst index f68ee90b5dd..1925f610f68 100644 --- a/docs/source/getting_started/windows/_installation_for_Windows.rst +++ b/docs/source/getting_started/windows/_installation_for_Windows.rst @@ -30,6 +30,10 @@ The following system requirements are necessary to install and use Model Optimiz The Model Optimizer - Windows can be used in following ways: +.. note:: + Model Optimizer will download and install additional third-party open source software projects. + Review the license terms of these open source projects before use. + .. toctree:: :glob: :maxdepth: 1 diff --git a/examples/llm_eval/lm_eval_hf.py b/examples/llm_eval/lm_eval_hf.py index dbdf22d8689..7d1f9f19935 100755 --- a/examples/llm_eval/lm_eval_hf.py +++ b/examples/llm_eval/lm_eval_hf.py @@ -38,15 +38,15 @@ # limitations under the License. import contextlib import warnings +from importlib.metadata import version import datasets -import lm_eval from lm_eval import utils from lm_eval.__main__ import cli_evaluate, parse_eval_args, setup_parser -if not lm_eval.__version__.startswith("0.4.8"): +if not version("lm_eval").startswith("0.4.8"): warnings.warn( - f"lm_eval_hf.py is tested with lm-eval 0.4.8; found {lm_eval.__version__}. " + f"lm_eval_hf.py is tested with lm-eval 0.4.8; found {version('lm_eval')}. " "Later versions may have incompatible API changes." ) from lm_eval.api.model import T diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index 90532efe38d..9455157645c 100755 --- a/examples/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -46,8 +46,6 @@ except ImportError: snapshot_download = None -from modelopt.torch.utils.image_processor import BaseImageProcessor, MllamaImageProcessor - logger = logging.getLogger(__name__) SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] @@ -285,13 +283,10 @@ def get_tokenizer(ckpt_path, trust_remote_code=False, **kwargs) -> PreTrainedTok def get_processor( ckpt_path, model_type, - device: torch.device = "auto", trust_remote_code=False, attn_implementation=None, -) -> BaseImageProcessor | ProcessorMixin | None: - """ - Returns a :class:`modelopt.torch.utils.image_processor.MllamaImageProcessor` object. - """ +) -> ProcessorMixin | None: + """Load a processor appropriate for the given model type.""" model_kwargs = {"trust_remote_code": trust_remote_code} if attn_implementation is not None: model_kwargs["attn_implementation"] = attn_implementation @@ -309,19 +304,6 @@ def get_processor( ) return processor - elif model_type == "mllama": - processor = AutoProcessor.from_pretrained( - ckpt_path, - padding_side="left", - **model_kwargs, - ) - if processor.tokenizer.pad_token is None: - processor.tokenizer.pad_token = processor.tokenizer.eos_token - assert processor.tokenizer.pad_token is not None, ( - f"Pad token for {ckpt_path} cannot be set!" - ) - - return MllamaImageProcessor(processor, device) else: # Try to load AutoProcessor for other VL models (e.g., Nemotron-Parse) try: diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 831d230a672..81649b81284 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -77,7 +77,6 @@ get_max_batch_size, get_supported_datasets, ) -from modelopt.torch.utils.image_processor import BaseImageProcessor, MllamaImageProcessor from modelopt.torch.utils.memory_monitor import launch_memory_monitor from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader @@ -202,7 +201,7 @@ def _to_device(value): def make_calib_dataloader( args: argparse.Namespace, language_model: torch.nn.Module, - processor: BaseImageProcessor | ProcessorMixin | None, + processor: ProcessorMixin | None, tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, @@ -250,19 +249,6 @@ def make_calib_dataloader( use_media_shards=True, max_shards=1, ) - elif model_type == "mllama": - assert processor is not None and isinstance(processor, MllamaImageProcessor), ( - "The MllamaImageProcessor must be set." - ) - assert len(args.calib_size) == 1, ( - "mllama only supports one dataset for calibration, can extend this in the future" - ) - calib_dataloader = get_vlm_dataset_dataloader( - dataset_name=args.dataset[0] if args.dataset else "scienceqa", - processor=processor, - batch_size=args.batch_size, - num_samples=args.calib_size[0], - ) elif model_type == "whisper": assert processor is not None and isinstance(processor, WhisperProcessor), ( "The AutoProcessor must be set." @@ -292,6 +278,7 @@ def make_calib_dataloader( tokenizer=tokenizer, batch_size=args.batch_size, num_samples=args.calib_size, + max_sample_length=args.calib_seq, device=device, include_labels=include_labels, ) @@ -472,23 +459,14 @@ def load_model(args: argparse.Namespace): print("Nemotron VL model detected. Enabling image-text calibration by default.") args.calib_with_images = True - if model_type == "mllama": + if model_type == "whisper": processor = get_processor( args.pyt_ckpt_path, model_type, - device, trust_remote_code=args.trust_remote_code, - attn_implementation=args.attn_implementation, ) - elif model_type == "whisper": - processor = get_processor( - args.pyt_ckpt_path, - model_type, - device, - trust_remote_code=args.trust_remote_code, - ) - elif is_nemotron_vl_model and args.calib_with_images: - # For Nemotron VL image calibration, we need an AutoProcessor to build multimodal inputs. + elif args.calib_with_images: + # For VLM image calibration, we need an AutoProcessor to build multimodal inputs. processor = AutoProcessor.from_pretrained( args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code, @@ -715,13 +693,6 @@ def export_quantized( print(f"Warning: Could not save processor config: {e}") print("This is normal for some VLM architectures that don't use AutoProcessor") - if model_type == "mllama": - full_model_config = full_model.config - # TRT-LLM expects both the vision_config and text_config to be set for export. - setattr(full_model.config, "vision_config", full_model_config.vision_config) - setattr(full_model.config, "text_config", full_model_config.text_config) - setattr(full_model.config, "architectures", full_model_config.architectures) - start_time = time.time() if ( model_type in ["t5", "bart", "whisper"] @@ -858,7 +829,7 @@ def post_quantize( language_model: torch.nn.Module, model_type: str | None, tokenizer: PreTrainedTokenizerBase | None, - processor: BaseImageProcessor | ProcessorMixin | None, + processor: ProcessorMixin | None, preview_input_ids, generated_ids_before_ptq, is_nemotron_vl_model, @@ -921,9 +892,7 @@ def post_quantize( ) def input_decode(input_ids): - if processor is not None and isinstance(processor, MllamaImageProcessor): - return processor.tokenizer.batch_decode(input_ids) - elif processor is not None and isinstance(processor, WhisperProcessor): + if processor is not None and isinstance(processor, WhisperProcessor): return first_text_speech_dataset elif tokenizer is not None: return tokenizer.batch_decode(input_ids) @@ -936,8 +905,6 @@ def output_decode(generated_ids, input_shape): return processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] elif tokenizer is not None: return tokenizer.batch_decode(generated_ids, skip_special_tokens=True) - elif processor is not None and isinstance(processor, MllamaImageProcessor): - return processor.tokenizer.batch_decode(generated_ids[:, input_shape:]) elif tokenizer is not None: return tokenizer.batch_decode(generated_ids[:, input_shape:]) else: @@ -982,7 +949,7 @@ def quantize_main( language_model: torch.nn.Module, model_type: str | None, calibration_only: bool, - processor: BaseImageProcessor | ProcessorMixin | None, + processor: ProcessorMixin | None, tokenizer: PreTrainedTokenizerBase | None, default_padding_side, default_pad_token, diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 5bd7c650648..89183073399 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -270,19 +270,59 @@ For a quick smoke test, add `--limit 10`. > **Alternative:** For server-based evaluation via an OpenAI-compatible endpoint, > see [evaluation/nemo_evaluator_instructions.md](./evaluation/nemo_evaluator_instructions.md). -## Inference Performance Benchmarking +## Deploy compressed model in vLLM -Now let's evaluate how much speedup we get with the compressed model in terms of throughput and latency. +To deploy a compressed model in vLLM, install vLLM fork with AnyModel enabled: + +```bash +git clone https://github.com/askliar/vllm.git +cd vllm +git checkout feature/add_anymodel_to_vllm +VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto +``` + +See [vLLM documentation](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source) for more details on installation. + +**NOTE:** This is a temporary workaround pending official vLLM integration. You can track merge status [here](https://github.com/vllm-project/vllm/pull/36512). + +Then, add the following to the model's `config.json` file (here we use Llama as an example): -- Install [vLLM from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source). -- Rearrange the model safetensors to be used for vLLM. +```json +{ + ... + "architectures": ["AnyModel"], + "base_architecture": "LlamaForCausalLM", + ... +} +``` + +For new architectures that are not supported by vLLM, you additionally need to add the following to the `config.json` file (using Llama3 as an example): + +```json +{ + ... + "anymodel_arch_info": { + "decoder_layer_module": ".", + "decoder_layer_class": "", + "base_model_module": ".", + "layers_path": "", + "init_prefix": "model", + "Layer_hf_config": "" + } + ... +} +``` + +With these changes it is now possible to load the compressed model in vLLM for inference: ```bash -cd path/to/model -mv subblocks_safetensors/* . -sed -i 's+subblocks_safetensors/++g' model.safetensors.index.json +vllm serve ``` +### Inference Performance Benchmarking + +Now let's evaluate how much speedup we get with the compressed model in terms of throughput and latency. + - Benchmark latency ```bash diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 98daa2c13d4..97bd0b60c18 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -88,6 +88,22 @@ }, ] +# FP8 MHA-aware config entries: quantize LayerNorm output so TRT can fuse the shared +# Q/DQ across all downstream Q/K/V/FC consumers. Softmax-output Q/DQ is handled by the +# FP8 ONNX exporter's post-processing pass (fixed 1/448 scale, data-independent). +_FP8_MHA_OVERRIDE: list = [ + { + "parent_class": "nn.LayerNorm", + "quantizer_name": "*output_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + }, + { + "parent_class": "nn.LayerNorm", + "quantizer_name": "*input_quantizer", + "enable": False, + }, +] + # Auto-quantize format configs that use block quantization and need Conv2d overrides for TRT. # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. _NEEDS_FP8_CONV_OVERRIDE: set[str] = { @@ -102,11 +118,16 @@ def get_quant_config(quantize_mode): """Get quantization config, overriding Conv2d for TRT compatibility. 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 = copy.deepcopy(QUANT_CONFIG_DICT[quantize_mode]) - if quantize_mode in ("mxfp8", "nvfp4"): + if quantize_mode == "fp8": + config["quant_cfg"].extend(_FP8_MHA_OVERRIDE) + elif quantize_mode in ("mxfp8", "nvfp4"): warnings.warn( f"TensorRT only supports FP8/INT8 for Conv layers. " f"Overriding Conv2d quantization to FP8 for '{quantize_mode}' mode." @@ -126,6 +147,9 @@ def filter_func(name): ``downsample.reduction`` (Swin/SwinV2) is excluded because it operates on 4D tensors and TRT's DynamicQuantize layer (used for MXFP8/NVFP4) requires 2D/3D input. + Other 4D-input layers (e.g. Swin's ``norm1``, ``downsample.norm``, top-level ``norm``) + are handled dynamically by ``_disable_high_rank_input_quantizers`` via a forward-pass + rank probe — that avoids false positives on ViT, whose same-named ``norm`` sees 3D input. """ pattern = re.compile( r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|" @@ -135,6 +159,65 @@ def filter_func(name): return pattern.match(name) is not None +def _disable_high_rank_input_quantizers(model, input_shape, device): + """Disable quantizers on Linear/LayerNorm modules that receive 4D+ input. + + TRT's MXFP8/NVFP4 ``DynamicQuantize`` op only supports 2D/3D input, so Swin's + per-block ``norm1``, ``downsample.norm``, and top-level ``norm`` (all 4D in Swin + but 3D in ViT) must be skipped. A forward pass with hooks identifies them at + runtime, so this works across architectures without hardcoded paths. + """ + high_rank: set[str] = set() + handles = [] + for name, mod in model.named_modules(): + if isinstance(mod, (torch.nn.Linear, torch.nn.LayerNorm)): + + def hook(m, inp, out, _n=name): + if inp and hasattr(inp[0], "ndim") and inp[0].ndim > 3: + high_rank.add(_n) + + handles.append(mod.register_forward_hook(hook)) + + was_training = model.training + model.eval() + try: + with torch.no_grad(): + model(torch.randn(input_shape, device=device)) + finally: + for h in handles: + 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)) + + +def _disable_low_channel_conv_input_quantizers(model): + """Disable ``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 + fails to find an FP8/MXFP8/NVFP4 tactic for this first-layer Q→Conv fusion: + + Error Code 10: Could not find any implementation for node + /conv1/input_quantizer/TRT_FP8QuantizeLinear ... [ElementWise] + + Ada (8.9) happens to have a tactic, which is why local runs pass. Disabling the + input quantizer on the raw-RGB conv is also standard quantization practice — + first/last layers are typically left in higher precision. Weight quantization + still applies. Swin/ViT's ``patch_embed.proj`` is already excluded via + ``filter_func``'s ``patch_embed`` pattern, so this helper is effectively the + ResNet-shaped analogue. + """ + 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: + q.disable() + + def load_calibration_data(model, data_size, batch_size, device, with_labels=False): """Load and prepare calibration data. @@ -169,6 +252,28 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals ) +def _disable_dead_quantizers(model): + """Disable quantizers whose calibrated ``amax`` is non-positive or NaN. + + ``export_fp8`` computes ``scale = 448 / amax`` and blows up on ``amax == 0``. + This shows up on SwinV2 with ``--no_pretrained``: timm's ``res-post-norm`` scheme + zero-inits each block's ``norm1``/``norm2`` weight and bias, so those LayerNorm + outputs are exactly zero at init and the MHA override's output_quantizer + calibrates to ``amax == 0``. Disable such dead quantizers — they have nothing + meaningful to quantize and would otherwise break ONNX export. + """ + for _, mod in model.named_modules(): + for attr in ("input_quantizer", "output_quantizer", "weight_quantizer"): + q = getattr(mod, attr, None) + if q is None or not q.is_enabled: + continue + amax = q.amax + if amax is None or not torch.is_tensor(amax): + continue + if torch.any(torch.isnan(amax)) or torch.all(amax <= 0): + q.disable() + + def _calibrate_uncalibrated_quantizers(model, data_loader): """Calibrate FP8 quantizers that weren't calibrated by mtq.quantize(). @@ -219,6 +324,10 @@ def forward_loop(model): 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) + return quantized_model @@ -303,6 +412,8 @@ def auto_quantize_model( # Disable quantization for specified layers mtq.disable_quantizer(quantized_model, filter_func) + _disable_dead_quantizers(quantized_model) + return quantized_model, search_state @@ -458,6 +569,7 @@ def main(): # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8 # quantizers require calibration data. config = get_quant_config(args.quantize_mode) + data_loader = load_calibration_data( model, args.calibration_data_size, @@ -468,6 +580,25 @@ 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 != "INT8_DEFAULT_CFG" for fmt in args.auto_quantization_formats) + ) + if uses_fp8_conv_input: + _disable_low_channel_conv_input_quantizers(quantized_model) + # Print quantization summary print("\nQuantization Summary:") mtq.print_quant_summary(quantized_model) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index aa8d3a5388b..6b658551f15 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -572,7 +572,7 @@ def load_state_dict_from_path( saved_quant_dict = { key.replace("quantizer_", "quantizer._"): value for key, value in saved_quant_dict.items() - if "quantizer_" in key + if "quantizer" in key } saved_quant_dict = convert_dict_to_vllm(saved_quant_dict) diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index dcae618dd0a..427a7791f3b 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -17,6 +17,7 @@ import time +import numpy as np import onnx import onnx_graphsurgeon as gs import torch @@ -26,6 +27,11 @@ from .base_exporter import ONNXQuantExporter +# FP8 E4M3 max representable magnitude; softmax output in [0, 1] saturates exactly at 1.0 +# when using 1/448 as the Q scale (single fixed value — softmax range is data-independent). +_FP8_E4M3_MAX = 448.0 +_FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -62,6 +68,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: # Fold constants is required since the scale is not constant yet. graph.cleanup().toposort().fold_constants().cleanup() + n_t_folded = 0 + for node in graph.nodes: if node.op == "TRT_FP8QuantizeLinear": # Should not remove input QDQ (only process weight quantization) @@ -74,9 +82,44 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: torch_scale = torch.from_numpy(scale.values) quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] - assert dq_op.op == "TRT_FP8DequantizeLinear", ( - f"QDQ does not occur in pairs. You reached {dq_op.op}" - ) + if dq_op.op != "TRT_FP8DequantizeLinear": + raise RuntimeError(f"QDQ does not occur in pairs. You reached {dq_op.op}") + + # Pre-transpose constant weights if DQ feeds ``Transpose → MatMul`` (or + # ``Cast → Transpose → MatMul`` after fp16 conversion) so TRT sees DQ→MatMul. + # Control flow: scan candidates; a Cast-wrapped candidate is accepted only if it + # leads to a Transpose; a bare Transpose whose all consumers are MatMul wins and + # breaks the loop. Any other shape defaults `cast_to_remove` back to None and + # continues scanning. + transpose_to_remove = None + cast_to_remove = None + for candidate in list(dq_op.outputs[0].outputs): + if candidate.op == "Cast": + cast_to_remove = candidate + candidate = next( + (c for c in candidate.outputs[0].outputs if c.op == "Transpose"), + None, + ) + if candidate is None: + cast_to_remove = None + continue + if candidate.op != "Transpose": + cast_to_remove = None + continue + t_consumers = list(candidate.outputs[0].outputs) + # Only fold the transpose when every downstream consumer is MatMul; otherwise + # non-MatMul consumers would observe the un-transposed weights. + if t_consumers and all(c.op == "MatMul" for c in t_consumers): + perm = candidate.attrs.get("perm", None) + torch_weights = ( + torch_weights.permute(*perm).contiguous() + if perm is not None + else torch_weights.T.contiguous() + ) + transpose_to_remove = candidate + else: + cast_to_remove = None + break # Replace it with Dequantize with FP8 weights. This is a WAR because numpy does not support fp8. numpy_weights = ( @@ -94,9 +137,23 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: dq_op.inputs[0] = onnx_weights_fp8 dq_op.op = "DequantizeLinear" dq_op.outputs[0].dtype = dq_op.inputs[1].dtype + dq_op.outputs[0].shape = list(numpy_weights.shape) + + if transpose_to_remove is not None: + t_out = transpose_to_remove.outputs[0] + for consumer in list(t_out.outputs): + for i, inp in enumerate(consumer.inputs): + if inp is t_out: + consumer.inputs[i] = dq_op.outputs[0] + transpose_to_remove.outputs.clear() + if cast_to_remove is not None: + cast_to_remove.outputs.clear() + n_t_folded += 1 graph.cleanup().toposort() end_time = time.time() + if n_t_folded > 0: + logger.info(f"Folded {n_t_folded} weight Transpose nodes during weight compression") print(f"fp8 qdq replaced with only dq completed in {end_time - start_time}s.") return gs.export_onnx(graph) @@ -121,7 +178,6 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: Returns: Number of Conv weight DQ nodes inserted. """ - fp8_max = 448.0 count = 0 for node in list(graph.nodes): @@ -142,7 +198,7 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: amax = torch_weights.abs().max().float() if amax == 0: continue - scale_val = (amax / fp8_max).item() + scale_val = (amax / _FP8_E4M3_MAX).item() # Quantize weights to FP8 (WAR: numpy doesn't support fp8) fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() @@ -155,8 +211,6 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: ) # Scale in FP16 — DQ output type matches scale dtype, must match activation type - import numpy as np - scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", np.array(scale_val, dtype=np.float16), @@ -175,13 +229,220 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: return count + @staticmethod + def _move_mul_before_qdq(graph: gs.Graph) -> int: + """Move attention-scaling Mul(const) from after DQ to before Q for TRT MatMul fusion. + + Handles both ``DQ → Mul → MatMul`` and ``DQ → Transpose → Mul → MatMul`` (K path). + """ + count = 0 + for mul_node in list(graph.nodes): + if mul_node.op != "Mul": + continue + + const_input = next( + (i for i in mul_node.inputs if isinstance(i, gs.Constant) and i.values.size == 1), + None, + ) + tensor_input = next( + (i for i in mul_node.inputs if not isinstance(i, gs.Constant)), None + ) + if const_input is None or tensor_input is None: + continue + if not (isinstance(tensor_input, gs.Variable) and len(tensor_input.inputs) == 1): + continue + + producer = tensor_input.inputs[0] + transpose_node = producer if producer.op == "Transpose" else None + dq_node = producer if producer.op == "DequantizeLinear" else None + if transpose_node is not None: + t_input = transpose_node.inputs[0] + if ( + isinstance(t_input, gs.Variable) + and len(t_input.inputs) == 1 + and t_input.inputs[0].op == "DequantizeLinear" + ): + dq_node = t_input.inputs[0] + if dq_node is None: + continue + + q_output = dq_node.inputs[0] + if ( + not isinstance(q_output, gs.Variable) + or len(q_output.inputs) != 1 + or q_output.inputs[0].op != "QuantizeLinear" + ): + continue + q_node = q_output.inputs[0] + q_input = q_node.inputs[0] + if not isinstance(q_input, gs.Variable): + continue + + mul_output = mul_node.outputs[0] + mul_consumers = list(mul_output.outputs) + # Require every consumer to be MatMul: rewiring all consumers to bypass the Mul + # would silently drop the scale for any non-MatMul branch. + if not mul_consumers or not all(c.op == "MatMul" for c in mul_consumers): + continue + + new_mul_output = gs.Variable( + q_input.name + "_scaled", dtype=q_input.dtype, shape=q_input.shape + ) + graph.nodes.append( + gs.Node( + op="Mul", + name=mul_node.name + "_moved", + inputs=[q_input, const_input], + outputs=[new_mul_output], + ) + ) + q_node.inputs[0] = new_mul_output + + replacement = ( + transpose_node.outputs[0] if transpose_node is not None else dq_node.outputs[0] + ) + for consumer in mul_consumers: + for i, inp in enumerate(consumer.inputs): + if inp is mul_output: + consumer.inputs[i] = replacement + mul_node.outputs.clear() + count += 1 + + graph.cleanup().toposort() + return count + + @staticmethod + def _move_transpose_before_qdq(graph: gs.Graph) -> int: + """Move Transpose from ``DQ → Transpose → MatMul`` to ``Transpose → Q → DQ → MatMul`` (K path).""" + count = 0 + for transpose_node in list(graph.nodes): + if transpose_node.op != "Transpose": + continue + + t_input = transpose_node.inputs[0] + if ( + not isinstance(t_input, gs.Variable) + or len(t_input.inputs) != 1 + or t_input.inputs[0].op != "DequantizeLinear" + ): + continue + dq_node = t_input.inputs[0] + + dq_input = dq_node.inputs[0] + if ( + not isinstance(dq_input, gs.Variable) + or len(dq_input.inputs) != 1 + or dq_input.inputs[0].op != "QuantizeLinear" + ): + continue + q_node = dq_input.inputs[0] + q_input = q_node.inputs[0] + if not isinstance(q_input, gs.Variable): + continue + + t_output = transpose_node.outputs[0] + t_consumers = list(t_output.outputs) + # Require every consumer to be MatMul: rewiring to dq_node.outputs[0] would drop + # the transpose for any non-MatMul branch, producing a wrong-shape tensor. + if not t_consumers or not all(c.op == "MatMul" for c in t_consumers): + continue + + new_t_output = gs.Variable(q_input.name + "_transposed", dtype=q_input.dtype) + graph.nodes.append( + gs.Node( + op="Transpose", + name=transpose_node.name + "_moved", + inputs=[q_input], + outputs=[new_t_output], + attrs=transpose_node.attrs, + ) + ) + q_node.inputs[0] = new_t_output + + for consumer in t_consumers: + for i, inp in enumerate(consumer.inputs): + if inp is t_output: + consumer.inputs[i] = dq_node.outputs[0] + transpose_node.outputs.clear() + count += 1 + + graph.cleanup().toposort() + return count + + @staticmethod + def _insert_qdq_after_softmax(graph: gs.Graph) -> int: + """Insert FP8 Q→DQ on Softmax outputs feeding MatMul (required by TRT MHA fusion). + + Softmax output is data-independently bounded to [0, 1], so we use a fixed scale + ``_FP8_E4M3_SOFTMAX_SCALE`` (1/448) that saturates exactly at 1.0 while covering + the full FP8 E4M3 representable range. No calibration is required. Only applied + when every Softmax consumer is a MatMul so we do not insert quantization error + on unrelated branches. + """ + count = 0 + for softmax_node in list(graph.nodes): + if softmax_node.op != "Softmax": + continue + softmax_output = softmax_node.outputs[0] + consumers = list(softmax_output.outputs) + if not consumers or not all(c.op == "MatMul" for c in consumers): + continue + if any(c.op == "QuantizeLinear" for c in consumers): + continue + + # Match scale dtype to the graph's current float dtype so TRT stronglyTyped + # sees consistent Q/DQ types with the surrounding compute. + scale_dtype = softmax_output.dtype if softmax_output.dtype is not None else np.float32 + scale_val = np.array(_FP8_E4M3_SOFTMAX_SCALE, dtype=scale_dtype) + scale_constant = gs.Constant(softmax_node.name + "/softmax_q_scale", scale_val) + dq_scale_constant = gs.Constant( + softmax_node.name + "/softmax_dq_scale", scale_val.copy() + ) + + zp_tensor = onnx.TensorProto() + zp_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN + zp_tensor.dims.extend([1]) + zp_tensor.raw_data = b"\x00" + zp_constant = gs.Constant( + softmax_node.name + "/softmax_q_zero_point", LazyValues(zp_tensor) + ) + + q_output = gs.Variable(softmax_node.name + "/q_output") + dq_output = gs.Variable(softmax_node.name + "/dq_output", dtype=softmax_output.dtype) + q_node = gs.Node( + op="QuantizeLinear", + name=softmax_node.name + "/QuantizeLinear", + inputs=[softmax_output, scale_constant, zp_constant], + outputs=[q_output], + attrs={"saturate": 1}, + ) + dq_node = gs.Node( + op="DequantizeLinear", + name=softmax_node.name + "/DequantizeLinear", + inputs=[q_output, dq_scale_constant], + outputs=[dq_output], + ) + graph.nodes.extend([q_node, dq_node]) + + for consumer in consumers: + if consumer is q_node: + continue + for i, inp in enumerate(consumer.inputs): + if inp is softmax_output: + consumer.inputs[i] = dq_output + count += 1 + + graph.cleanup().toposort() + return count + @staticmethod def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for FP8 quantization. - Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear and + Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear, adds FP8 weight DQ for Conv layers whose weight quantizers were disabled during - TorchScript export. + TorchScript export, and rewrites attention scaling / K-transpose / softmax-output + patterns so TRT can fuse DQ into the attention MatMul kernels. Args: onnx_model: The ONNX model containing TRT_FP8 quantization nodes. @@ -223,5 +484,15 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: if count > 0: logger.info(f"Inserted FP8 weight DequantizeLinear for {count} Conv nodes") + # Attention-aware rewrites so TRT can fuse DQ into the attention MatMuls. + n_mul = FP8QuantExporter._move_mul_before_qdq(graph) + n_t = FP8QuantExporter._move_transpose_before_qdq(graph) + n_sm = FP8QuantExporter._insert_qdq_after_softmax(graph) + if n_mul or n_t or n_sm: + logger.info( + f"Attention QDQ rewrites: moved {n_mul} Mul, {n_t} Transpose; " + f"inserted QDQ on {n_sm} Softmax outputs" + ) + graph.cleanup().toposort() return gs.export_onnx(graph) diff --git a/modelopt/onnx/quantization/graph_utils.py b/modelopt/onnx/quantization/graph_utils.py index e8b15e30595..164af24839b 100755 --- a/modelopt/onnx/quantization/graph_utils.py +++ b/modelopt/onnx/quantization/graph_utils.py @@ -1089,11 +1089,14 @@ def find_nodes_from_matmul_to_exclude( calibration_eps: list[str] = ["cpu", "cuda:0", "trt"], calibration_shapes: str | dict | None = None, ) -> list[str]: - """Find MatMul nodes that meets gemv condition to exclude. + """Find MatMul nodes that meet gemv or small-gemm conditions and should be excluded. - Either of m or n in matmul is 1, this matmul cannot utilize - TensorCores. The perf of adding Q/DQ layers is not good in - TRT. Thus, in this case, do not add Q/DQ layers to this matmul. + A MatMul is excluded if either: + + - m or n in the output is 1 (GEMV): cannot utilize TensorCores; or + - K or N is smaller than ``_MIN_MATMUL_DIM`` (16): both INT8 and FP8 Tensor Core + kernels need K/N >= 16 to be efficient, and adding Q/DQ layers on such small + GEMMs causes TRT perf regressions. Args: onnx_path: Path to the onnx model. @@ -1143,6 +1146,10 @@ def find_nodes_from_matmul_to_exclude( _MIN_CHANNELS_FP8 = 16 +# Minimum K/N dim for MatMul/Gemm under INT8 or FP8 quantization. Both INT8 and FP8 +# Tensor Core kernels need K/N >= 16 to be efficient; adding Q/DQ layers on smaller +# GEMMs causes TRT perf regressions. +_MIN_MATMUL_DIM = 16 def find_nodes_from_convs_to_exclude(graph: Graph, quantize_mode: str = "int8"): @@ -1231,10 +1238,47 @@ def find_nodes_from_convs_to_exclude(graph: Graph, quantize_mode: str = "int8"): return unsupported_conv_nodes +def _get_inp_b_k_dim( + matmul_node, value_info_map: dict | None = None, output_map: dict | None = None +): + """Get the K dimension from the second input of a MatMul/Gemm node. + + Tries Constant shape first, then falls back to shape inference (value_info_map) + or runtime inference (output_map). For Gemm nodes, honors the ``transB`` attribute: + when ``transB=1``, B has shape ``[N, K]`` so K lives at axis -1; otherwise B is + ``[..., K, N]`` and K is at axis -2. + + Returns: + The K dimension value, or None if it cannot be determined. + """ + # For Gemm, transB=1 means B is [N, K] (K is last axis); default/MatMul is [K, N]. + trans_b = bool(matmul_node.attrs.get("transB", 0)) if matmul_node.op == "Gemm" else False + k_axis = -1 if trans_b else -2 + + inp_b = matmul_node.inputs[1] + if hasattr(inp_b, "values") and inp_b.values is not None: + inp_b_shape = inp_b.values.shape + if len(inp_b_shape) >= 2: + return inp_b_shape[k_axis] + if value_info_map is not None: + inp_b_info = value_info_map.get(inp_b.name) + if inp_b_info: + inp_b_dims = inp_b_info.type.tensor_type.shape.dim + if len(inp_b_dims) >= 2: + return inp_b_dims[k_axis].dim_value + if output_map is not None and inp_b.name in output_map: + inp_b_out = output_map[inp_b.name] + if len(inp_b_out.shape) >= 2: + return inp_b_out.shape[k_axis] + return None + + def _exclude_matmuls_by_shape_inference( - model: onnx.ModelProto, matmul_nodes: list, calibration_shapes: str | dict | None = None + model: onnx.ModelProto, + matmul_nodes: list, + calibration_shapes: str | dict | None = None, ) -> list[str]: - """Use shape inference to find MatMuls with dimension 1.""" + """Use shape inference to find MatMuls with dimension 1 or small K/N.""" # Prepare model for symbolic inference for graph_input in model.graph.input: for dim in graph_input.type.tensor_type.shape.dim: @@ -1263,7 +1307,10 @@ def _exclude_matmuls_by_shape_inference( dim.dim_value = new_dim_value model = infer_shapes(model) - value_info_map = {vi.name: vi for vi in model.graph.value_info} + # Include graph inputs, value_info, and outputs so B that comes from a graph input + # is visible when deriving K. + value_info_map = {vi.name: vi for vi in model.graph.input} + value_info_map.update({vi.name: vi for vi in model.graph.value_info}) value_info_map.update({vi.name: vi for vi in model.graph.output}) nodes_to_exclude = [] @@ -1280,8 +1327,23 @@ def _exclude_matmuls_by_shape_inference( if dims[-1].dim_value == 1 or dims[-2].dim_value == 1: nodes_to_exclude.append(matmul_node.name) + continue elif len(dims) < 3 and any(out.dim_value == 1 for out in dims): nodes_to_exclude.append(matmul_node.name) + continue + + # Small-gemm check: applies to both INT8 and FP8 quantization. + n_dim = dims[-1].dim_value if len(dims) >= 2 else 0 + k_dim = _get_inp_b_k_dim(matmul_node, value_info_map=value_info_map) + small_n = 0 < n_dim < _MIN_MATMUL_DIM + small_k = k_dim is not None and 0 < k_dim < _MIN_MATMUL_DIM + + if small_n or small_k: + logger.debug( + f"Excluding small-dim MatMul from quantization: {matmul_node.name} " + f"(N={n_dim}, K={k_dim}, threshold={_MIN_MATMUL_DIM})" + ) + nodes_to_exclude.append(matmul_node.name) return nodes_to_exclude @@ -1295,10 +1357,20 @@ def _exclude_matmuls_by_inference( calibration_data_reader: CalibrationDataReader, calibration_eps: list[str], ) -> list[str]: - """Use actual inference to find MatMuls with dimension 1.""" - # Add matmul outputs to model outputs + """Use actual inference to find MatMuls with dimension 1 or small K/N.""" + # Add matmul outputs and second-input outputs to model outputs + existing_output_names = {out.name for out in model.graph.output} for matmul_node in matmul_nodes: - model.graph.output.extend([onnx.ValueInfoProto(name=matmul_node.outputs[0].name)]) + out_name = matmul_node.outputs[0].name + if out_name not in existing_output_names: + model.graph.output.extend([onnx.ValueInfoProto(name=out_name)]) + existing_output_names.add(out_name) + # Also add second input for K-dimension check (only if it's a Variable, not a Constant) + if isinstance(matmul_node.inputs[1], Variable): + inp_b_name = matmul_node.inputs[1].name + if inp_b_name not in existing_output_names: + model.graph.output.extend([onnx.ValueInfoProto(name=inp_b_name)]) + existing_output_names.add(inp_b_name) output_map = get_extended_model_outputs( onnx_path, @@ -1319,8 +1391,23 @@ def _exclude_matmuls_by_inference( or matmul_output.shape[-2] == 1 ): nodes_to_exclude.append(matmul_node.name) + continue elif len(matmul_output.shape) < 3 and any(out == 1 for out in matmul_output.shape): nodes_to_exclude.append(matmul_node.name) + continue + + # Small-gemm check: applies to both INT8 and FP8 quantization. + n_dim = matmul_output.shape[-1] if len(matmul_output.shape) >= 2 else 0 + k_dim = _get_inp_b_k_dim(matmul_node, output_map=output_map) + small_n = 0 < n_dim < _MIN_MATMUL_DIM + small_k = k_dim is not None and 0 < k_dim < _MIN_MATMUL_DIM + + if small_n or small_k: + logger.debug( + f"Excluding small-dim MatMul from quantization: {matmul_node.name} " + f"(N={n_dim}, K={k_dim}, threshold={_MIN_MATMUL_DIM})" + ) + nodes_to_exclude.append(matmul_node.name) return nodes_to_exclude diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index ac93bc2a26c..7b1d7903c1a 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1415,6 +1415,70 @@ def _bypass_cast_node(model: onnx.ModelProto, node: onnx.NodeProto) -> None: consumer.input[i] = input_tensor +_DQ_OPS = {"DequantizeLinear", "TRT_FP8DequantizeLinear"} +_Q_OPS = {"QuantizeLinear", "TRT_FP8QuantizeLinear"} + + +def _scale_fp32_to_fp16(scale_init: onnx.TensorProto) -> None: + """Convert a scalar Q/DQ scale initializer in-place from FP32 to FP16. + + Warns if any non-zero scale saturates to 0/inf in FP16 (out of FP16 representable range). + """ + if scale_init.data_type != onnx.TensorProto.FLOAT: + return + scale_data = np.frombuffer(scale_init.raw_data, dtype=np.float32) + if not scale_data.size: + scale_data = np.array(scale_init.float_data, dtype=np.float32) + fp16_data = scale_data.astype(np.float16) + if np.any(np.isinf(fp16_data)) or (np.any(fp16_data == 0) and np.any(scale_data != 0)): + logger.warning(f"Q/DQ scale '{scale_init.name}' overflows or underflows when cast to FP16") + scale_init.data_type = onnx.TensorProto.FLOAT16 + scale_init.raw_data = fp16_data.tobytes() + del scale_init.float_data[:] + + +def fold_q_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + """Remove ``Cast(FP16→FP32) → Q`` patterns inserted by ``convert_float_to_float16``. + + The Q scale is rewritten to FP16 so Q consumes the FP16 graph directly. Skipped for + opsets below ``BASE_MIN_OPSET`` since FP16 Q scales require opset >= 19. + """ + if get_opset_version(onnx_model) < BASE_MIN_OPSET: + logger.debug( + f"Skipping fold_q_fp16_to_fp32_casts: opset < {BASE_MIN_OPSET} (FP16 Q scale unsupported)" + ) + return onnx_model + + consumer_map: dict[str, list[onnx.NodeProto]] = {} + for node in onnx_model.graph.node: + for inp in node.input: + consumer_map.setdefault(inp, []).append(node) + initializers = {init.name: init for init in onnx_model.graph.initializer} + + to_remove = [] + for node in onnx_model.graph.node: + if node.op_type != "Cast": + continue + cast_to = next((a.i for a in node.attribute if a.name == "to"), None) + if cast_to != onnx.TensorProto.FLOAT: + continue + consumers = consumer_map.get(node.output[0], []) + if not consumers or not all(c.op_type in _Q_OPS for c in consumers): + continue + + for q_node in consumers: + if len(q_node.input) >= 2 and q_node.input[1] in initializers: + _scale_fp32_to_fp16(initializers[q_node.input[1]]) + + _bypass_cast_node(onnx_model, node) + to_remove.append(node) + + logger.debug(f"Folded {len(to_remove)} Cast(FP16->FP32) -> Q patterns") + for node in to_remove: + onnx_model.graph.node.remove(node) + return onnx_model + + def _is_foldable_constant_cast_pattern(model: onnx.ModelProto, node: onnx.NodeProto) -> bool: """Check if a Constant -> Cast pattern can be folded.""" assert node.op_type == "Cast" @@ -1523,7 +1587,12 @@ def fold_dq_fp32_to_fp16_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: Returns: The ONNX model with Cast nodes removed and DQ outputs set to FP16. """ - import numpy as np + if get_opset_version(onnx_model) < BASE_MIN_OPSET: + logger.debug( + f"Skipping fold_dq_fp32_to_fp16_casts: opset < {BASE_MIN_OPSET} " + "(FP16 DQ scale unsupported)" + ) + return onnx_model dq_ops = {"DequantizeLinear", "TRT_FP8DequantizeLinear"} @@ -1623,6 +1692,13 @@ def fold_qdq_scale_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.Model Returns: The ONNX model with redundant scale-path casts removed. """ + if get_opset_version(onnx_model) < BASE_MIN_OPSET: + logger.debug( + f"Skipping fold_qdq_scale_fp16_to_fp32_casts: opset < {BASE_MIN_OPSET} " + "(FP16 Q/DQ scale unsupported)" + ) + return onnx_model + qdq_ops = { "QuantizeLinear", "DequantizeLinear", diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 9ec110b7887..01fb754bbae 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -48,6 +48,7 @@ change_casts_to_fp16, check_model_uses_external_data, fold_dq_fp32_to_fp16_casts, + fold_q_fp16_to_fp32_casts, fold_qdq_scale_fp16_to_fp32_casts, get_input_names, get_input_shapes, @@ -663,6 +664,11 @@ def get_onnx_bytes_and_metadata( onnx_opt_graph = remove_redundant_casts(onnx_opt_graph) + # Remove Cast nodes around Q/DQ for optimal TRT fusion + if is_fp8_quantized(model): + onnx_opt_graph = fold_q_fp16_to_fp32_casts(onnx_opt_graph) + onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) + # TensorRT expects all scales to be postive onnx_opt_graph = replace_zero_scale_with_smallest_nonzero(onnx_opt_graph) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index c8e45be3650..85952f8dc58 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -135,8 +135,10 @@ def _get_quantized_state( # string then it usually ends with "." which needs to be removed. self.exclude_modules.append(prefix.removesuffix(".")) block_size = 0 - - if hasattr(module, "weight") and module.weight is not None: + name_to_value = self._get_weight_bias(module, dtype, name_to_value) + if "weight" in name_to_value: + # Use the original device (avoid the CPU round-trip introduced by _get_weight_bias; + # fake-quantization runs on CUDA and the result is moved to CPU below). weight = module.weight.to(dtype) # Fold the weight_quantizer into the weight by applying fake-quantization # (quantize then dequantize). The weight_quantizer amax is not exported; @@ -171,9 +173,6 @@ def _get_quantized_state( else: return name_to_value, qformat, block_size - if hasattr(module, "bias") and module.bias is not None: - name_to_value["bias"] = module.bias.to(dtype).cpu() - # Only save input/output quantizer state; weight_quantizer amax is not exported # since it has been folded into the weight above. for name, param in get_quantizer_state_dict(module).items(): diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 89b718623da..62053e549c8 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -743,6 +743,44 @@ def _custom_mapping_to_lambda(mapping): return all_rules + def _get_weight_bias( + self, + module: torch.nn.Module, + dtype: torch.dtype = torch.float16, + name_to_value: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: + """Get the weight and bias of the module. + + Args: + module: The target module to get the weight and bias. + dtype: The data type of the weight and bias. + name_to_value: The dictionary to store the weight and bias. A new dict is created + if not provided. + + Returns: + The dictionary containing the weight and bias. + """ + if name_to_value is None: + name_to_value = {} + # numel() > 0 intentionally excludes zero-element weight tensors (e.g. MoE routing + # layers whose weight is a placeholder) so callers can use "weight" in name_to_value + # as a reliable guard without re-inspecting module.weight. + if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0: + weight = module.weight.to(dtype).cpu() + name_to_value["weight"] = weight + + if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0: + name_to_value["bias"] = module.bias.to(dtype).cpu() + + if ( + hasattr(module, "expert_bias") + and module.expert_bias is not None + and module.expert_bias.numel() > 0 + ): + name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu() + + return name_to_value + def _get_quantized_state( self, module: torch.nn.Module, @@ -767,21 +805,10 @@ def _get_quantized_state( self.exclude_modules.append(prefix.removesuffix(".")) block_size = get_weight_block_size(module) - if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0: - weight = module.weight.to(dtype).cpu() - name_to_value["weight"] = weight - else: - return name_to_value, qformat, block_size + name_to_value = self._get_weight_bias(module, dtype, name_to_value) - if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0: - name_to_value["bias"] = module.bias.to(dtype).cpu() - - if ( - hasattr(module, "expert_bias") - and module.expert_bias is not None - and module.expert_bias.numel() > 0 - ): - name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu() + if "weight" not in name_to_value: + return name_to_value, qformat, block_size if qformat == QUANTIZATION_NONE: return name_to_value, qformat, block_size diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index 05efe48842f..7b42dce5782 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -216,56 +216,36 @@ def _fp8_quantize( g: "GraphContext", inputs: torch.Value, scale_inv: float, - trt_high_precision_dtype: str, ): """Helper Function for Quantization.""" + # Emit the scale in the native input dtype so no Cast is inserted between the + # graph and Q/DQ (Cast nodes block TRT from fusing DQ into the MatMul kernel). output_shape = sym_help._get_tensor_sizes(inputs) - - # TRT StronglyType only supports FP16 QDQs - # custom ops, so cast the input if needed. - input_type = inputs.type().scalarType() - assert trt_high_precision_dtype in (input_type, "Float"), ( - "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." - ) - if trt_high_precision_dtype != input_type: - inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) - scale = g.op( "Constant", - value_t=torch.tensor(scale_inv).to(torch_dtype_map[trt_high_precision_dtype]), + value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]), ) - q_op = g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( + return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) ) - return q_op def _fp8_dequantize( g: "GraphContext", inputs: torch.Value, scale_inv: float, - trt_high_precision_dtype: str, otype: str | None = None, ): """Helper Function for Dequantization.""" output_shape = sym_help._get_tensor_sizes(inputs) - assert trt_high_precision_dtype in (otype, "Float"), ( - "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." - ) scale = g.op( "Constant", value_t=torch.tensor(scale_inv, dtype=torch_dtype_map[otype]), # type: ignore[index] ) - out = g.op("trt::TRT_FP8DequantizeLinear", inputs, scale).setType( - inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) + return g.op("trt::TRT_FP8DequantizeLinear", inputs, scale).setType( + inputs.type().with_dtype(torch_dtype_map[otype]).with_sizes(output_shape) # type: ignore[index] ) - # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT - # custom ops, so cast the output if needed. - if trt_high_precision_dtype != otype: - out = g.op("Cast", out, to_i=onnx_dtype_map[otype]) # type: ignore[index] - return out - def export_fp8( g: "GraphContext", @@ -273,14 +253,17 @@ def export_fp8( amax: float, trt_high_precision_dtype: str | None, ): - """Export quantized model to FP8 ONNX.""" + """Export quantized model to FP8 ONNX. + + ``trt_high_precision_dtype`` is accepted for API compatibility but unused: Q/DQ now + emit scales in the native input dtype, so no intermediate Cast is required. + """ + del trt_high_precision_dtype scale = 1.0 if amax is None else 448.0 / float(amax) otype = inputs.type().scalarType() - if trt_high_precision_dtype is None: - trt_high_precision_dtype = otype - q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype) - return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype, otype) + q_tensor = _fp8_quantize(g, inputs, 1.0 / scale) + return _fp8_dequantize(g, q_tensor, 1.0 / scale, otype) def scaled_dot_product_attention( diff --git a/modelopt/torch/quantization/nn/__init__.py b/modelopt/torch/quantization/nn/__init__.py index ca7082eb1cb..af9490c8311 100644 --- a/modelopt/torch/quantization/nn/__init__.py +++ b/modelopt/torch/quantization/nn/__init__.py @@ -19,6 +19,7 @@ from .modules.quant_batchnorm import * from .modules.quant_conv import * from .modules.quant_instancenorm import * +from .modules.quant_layernorm import * from .modules.quant_linear import * from .modules.quant_module import * from .modules.quant_pooling import * diff --git a/modelopt/torch/quantization/nn/modules/quant_layernorm.py b/modelopt/torch/quantization/nn/modules/quant_layernorm.py new file mode 100644 index 00000000000..10e9ba47582 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_layernorm.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Registers ``torch.nn.LayerNorm`` with ``QuantInputBase``. + +Enables LayerNorm output quantizers to be honored during quantization. Required for FP8 +attention fusion where a single LayerNorm output QDQ is shared across all downstream +Q/K/V/FC consumers (instead of repeating it on each input), which enables TRT to fuse DQ +into the attention MatMul kernels. +""" + +import torch.nn as nn + +from .quant_module import QuantInputBase, QuantModuleRegistry + +QuantModuleRegistry.register({nn.LayerNorm: "nn.LayerNorm"})(QuantInputBase) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 59bcd215bbc..990d0c0348d 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -274,6 +274,22 @@ def forward(self, *args, **kwargs): return super().forward(*args, **kwargs) +def _wraps_nested_attention(module): + """Return True when ``module`` contains another Attention child on this specific instance. + + Checked per-instance (not by class) so an attention class reused as both wrapper and + leaf is not dropped everywhere. In a 3-level hierarchy (Outer → Middle → Inner), both + Outer and Middle are treated as wrappers and only Inner is registered for KV-cache + quantization. Used to avoid double-patching ``eager_attention_forward`` when a wrapper + attention module (e.g. ``ViTAttention``) delegates to a nested self-attention child + (e.g. ``ViTSelfAttention``). + """ + return any( + child is not module and type(child).__name__.endswith("Attention") + for _, child in module.named_modules() + ) + + def register_hf_attentions_on_the_fly(model): """Find HF Attention modules in the model and register them for KV Cache quantization. @@ -286,9 +302,12 @@ def register_hf_attentions_on_the_fly(model): attention_cls = set() registered_attn_module = False + for name, module in model.named_modules(): # Only register attention classes that are from Huggingface transformers if type(module).__name__.endswith("Attention"): + if _wraps_nested_attention(module): + continue attention_type = _QuantAttention.get_attn_type(module) # Add modules to be registered only if they arent already registered if ( diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 6721a8f798f..95ca3240b73 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -385,14 +385,19 @@ def _invoke_fused_moe_quantized_function( # First layer of expert A = self.w13_input_quantizer(A) # noqa: N806 if self.w13_weight_quantizer.is_enabled: # pragma: no cover - original_weight, self.w13_weight = ( - self.w13_weight, - self.w13_weight_quantizer(self.w13_weight), - ) - # In case the weight quantizer isn't folded yet in vllm_serve_fakequant, pass the - # quantized weight to the kernel. - B = self.w13_weight # noqa: N806 + # Same pattern as FakeQuantMethod.apply: wrap as nn.Parameter if needed, swap + # w13_weight, call kernel, restore (tensor cannot stay assigned to nn.Parameter slot). + original_weight = self.w13_weight + quantized_tensor = self.w13_weight_quantizer(original_weight) try: + if isinstance(original_weight, torch.nn.Parameter) and not isinstance( + quantized_tensor, torch.nn.Parameter + ): + quantized_tensor = torch.nn.Parameter( + quantized_tensor, requires_grad=original_weight.requires_grad + ) + self.w13_weight = quantized_tensor + B = quantized_tensor # noqa: N806 original_kernel(A, B, C, *args, **kwargs) finally: self.w13_weight = original_weight @@ -403,14 +408,17 @@ def _invoke_fused_moe_quantized_function( elif B is self.w2_weight: A = self.w2_input_quantizer(A) # noqa: N806 if self.w2_weight_quantizer.is_enabled: # pragma: no cover - original_weight, self.w2_weight = ( - self.w2_weight, - self.w2_weight_quantizer(self.w2_weight), - ) - # In case the weight quantizer isn't folded yet in vllm_serve_fakequant, pass the - # quantized weight to the kernel. - B = self.w2_weight # noqa: N806 + original_weight = self.w2_weight + quantized_tensor = self.w2_weight_quantizer(original_weight) try: + if isinstance(original_weight, torch.nn.Parameter) and not isinstance( + quantized_tensor, torch.nn.Parameter + ): + quantized_tensor = torch.nn.Parameter( + quantized_tensor, requires_grad=original_weight.requires_grad + ) + self.w2_weight = quantized_tensor + B = quantized_tensor # noqa: N806 original_kernel(A, B, C, *args, **kwargs) finally: self.w2_weight = original_weight diff --git a/modelopt/torch/speculative/plugins/transformers.py b/modelopt/torch/speculative/plugins/transformers.py index e2133931917..3e862a838cc 100644 --- a/modelopt/torch/speculative/plugins/transformers.py +++ b/modelopt/torch/speculative/plugins/transformers.py @@ -1080,6 +1080,8 @@ def forward( batch_size, seq_len_s, device=eagle_input_hiddens.device ).argsort(dim=1)[:, :num_to_replace] + # Clone to avoid inplace modification that breaks autograd + eagle_input_hiddens = eagle_input_hiddens.clone() batch_indices = torch.arange(batch_size)[:, None] eagle_input_hiddens[batch_indices, rand_indices] = eagle_output_hiddens[ batch_indices, rand_indices diff --git a/modelopt/torch/utils/image_processor.py b/modelopt/torch/utils/image_processor.py deleted file mode 100644 index 6374642e3de..00000000000 --- a/modelopt/torch/utils/image_processor.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -# Adapted from tensorrt_llm/quantization/image_processing.py -"""Utility classes for image processing.""" - -import torch - - -class BaseImageProcessor: - """Base class for image processors.""" - - def __init__(self, tokenizer, device="cuda"): - """Constructor.""" - self.tokenizer = tokenizer - self.device = device - - def __call__(self, **kwargs): - """Call the tokenizer.""" - return self.tokenizer(**kwargs) - - def preprocess_function(self, examples): - """Preprocess function.""" - raise NotImplementedError("Each image processor must implement its own preprocess method") - - def collate_function(self, examples): - """Collate function to process images during data loading.""" - raise NotImplementedError("Each image processor must implement its own collate method") - - -# A light Encapsulation for Huggingface MllamaImageProcessor - - -class MllamaImageProcessor(BaseImageProcessor): - """Image processor for Mllama.""" - - def preprocess_function(self, examples): - """Preprocess function.""" - # Prepare prompts in a generic chat format - question = examples.get("question", "Describe this image.") - - if examples["image"] is not None: - if self.tokenizer.chat_template is not None: - prompt = self.tokenizer.apply_chat_template( - [ - { - "role": "user", - "content": [{"type": "image"}, {"type": "text", "text": question}], - } - ], - add_generation_prompt=True, - ) - else: - prompt = f"<|image|><|begin_of_text|>{question}" - - # Process images using the processor's image processor - values = self.tokenizer(text=prompt, images=examples["image"], return_tensors="pt").to( - self.device - ) - else: - if self.tokenizer.chat_template is not None: - prompt = self.tokenizer.apply_chat_template( - [ - { - "role": "user", - "content": [{"type": "text", "text": question}], - } - ], - add_generation_prompt=True, - ) - else: - prompt = question - - values = self.tokenizer(text=prompt, images=None, return_tensors="pt").to(self.device) - - values["pixel_values"] = None - values["aspect_ratio_ids"] = None - values["aspect_ratio_mask"] = None - values["cross_attention_mask"] = None - - return values - - def collate_function(self, batch): - """Collate function to process images during data loading.""" - batch[0]["input_ids"] = torch.LongTensor(batch[0]["input_ids"]).to(self.device) - batch[0]["attention_mask"] = torch.LongTensor(batch[0]["attention_mask"]).to(self.device) - - if batch[0]["pixel_values"] is not None: - batch[0]["pixel_values"] = torch.Tensor(batch[0]["pixel_values"]).to(self.device) - batch[0]["aspect_ratio_ids"] = torch.LongTensor(batch[0]["aspect_ratio_ids"]).to( - self.device - ) - batch[0]["aspect_ratio_mask"] = torch.LongTensor(batch[0]["aspect_ratio_mask"]).to( - self.device - ) - batch[0]["cross_attention_mask"] = torch.LongTensor( - batch[0]["cross_attention_mask"] - ).to(self.device) - - return batch[0] diff --git a/modelopt/torch/utils/vlm_dataset_utils.py b/modelopt/torch/utils/vlm_dataset_utils.py index 3f07c57715b..9de40792e4b 100644 --- a/modelopt/torch/utils/vlm_dataset_utils.py +++ b/modelopt/torch/utils/vlm_dataset_utils.py @@ -30,7 +30,6 @@ import torch from torch.utils.data import DataLoader -from .image_processor import MllamaImageProcessor from .nemotron_vlm_dataset_utils import NemotronTarPlusJsonlIterable, list_repo_files_cached # Use dict to store the config for each dataset. @@ -379,18 +378,6 @@ def get_vlm_dataset_dataloader( max_shards=max_shards, ) - # Legacy path: our internal image processor wrapper (e.g., Mllama). - if isinstance(processor, MllamaImageProcessor): - processed_dataset = dataset.map( - processor.preprocess_function, batched=False, remove_columns=dataset.column_names - ) - return DataLoader( - processed_dataset, - batch_size=batch_size, - shuffle=False, - collate_fn=processor.collate_function, - ) - # Generic HF ProcessorMixin / AutoProcessor path: tokenize & process images at collate-time. # For Nemotron VLM datasets, we prefer to follow the model-card flow: # prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) @@ -457,7 +444,7 @@ def _collate_fn(examples: list[dict[str, Any]]) -> dict[str, torch.Tensor] | dic "return_tensors": "pt", "padding": True, } - if max_length is not None: + if max_length is not None and "images" not in kwargs: kwargs.update({"truncation": True, "max_length": max_length}) enc = processor(**kwargs) diff --git a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py new file mode 100644 index 00000000000..1f7251a9ad9 --- /dev/null +++ b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for the attention-aware FP8 ONNX graph rewrites in ``FP8QuantExporter``.""" + +import numpy as np +import onnx_graphsurgeon as gs +import pytest + +from modelopt.onnx.export.fp8_exporter import FP8QuantExporter + + +def _var(name): + return gs.Variable(name, dtype=np.float32) + + +def _qdq(src): + """Build ``QuantizeLinear → DequantizeLinear`` and return [Q, DQ], dq_out.""" + scale = gs.Constant("scale", np.array(0.1, dtype=np.float32)) + q_out, dq_out = _var("q_out"), _var("dq_out") + return [ + gs.Node(op="QuantizeLinear", inputs=[src, scale], outputs=[q_out]), + gs.Node(op="DequantizeLinear", inputs=[q_out, scale], outputs=[dq_out]), + ], dq_out + + +def _graph(nodes, inputs, outputs): + return gs.Graph(nodes=nodes, inputs=inputs, outputs=outputs, opset=19) + + +def test_move_mul_before_qdq_rewrites_dq_mul_matmul_pattern(): + """``DQ → Mul(const) → MatMul`` collapses to ``Mul → Q → DQ → MatMul``.""" + x, k, y, mul_out = _var("x"), _var("k"), _var("y"), _var("mul_out") + qdq_nodes, dq_out = _qdq(x) + mul = gs.Node( + op="Mul", + inputs=[dq_out, gs.Constant("c", np.array(0.5, dtype=np.float32))], + outputs=[mul_out], + ) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph([*qdq_nodes, mul, mm], [x, k], [y]) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert q.inputs[0].inputs[0].op == "Mul" + + +def test_move_transpose_before_qdq_rewrites_dq_transpose_matmul_pattern(): + """``DQ → Transpose → MatMul`` collapses to ``Transpose → Q → DQ → MatMul``.""" + k_in, q_in, scores, t_out = _var("k_in"), _var("q_in"), _var("scores"), _var("t_out") + qdq_nodes, dq_out = _qdq(k_in) + t = gs.Node(op="Transpose", inputs=[dq_out], outputs=[t_out], attrs={"perm": [0, 2, 1]}) + mm = gs.Node(op="MatMul", inputs=[q_in, t_out], outputs=[scores]) + graph = _graph([*qdq_nodes, t, mm], [k_in, q_in], [scores]) + + assert FP8QuantExporter._move_transpose_before_qdq(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert q.inputs[0].inputs[0].op == "Transpose" + + +def test_insert_qdq_after_softmax_adds_fixed_scale_q_dq(): + """Softmax → MatMul picks up ``Q → DQ`` with the fixed ``1/448`` scale.""" + scores, v, y, sm_out = _var("scores"), _var("v"), _var("y"), _var("sm_out") + sm = gs.Node(op="Softmax", inputs=[scores], outputs=[sm_out], attrs={"axis": -1}) + mm = gs.Node(op="MatMul", inputs=[sm_out, v], outputs=[y]) + graph = _graph([sm, mm], [scores, v], [y]) + + assert FP8QuantExporter._insert_qdq_after_softmax(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert np.isclose(float(q.inputs[1].values), 1.0 / 448.0) + + +@pytest.mark.parametrize( + "rewrite", ["_move_mul_before_qdq", "_move_transpose_before_qdq", "_insert_qdq_after_softmax"] +) +def test_rewrites_skip_when_non_matmul_consumer_exists(rewrite): + """Every MHA rewrite must skip when the candidate tensor fans out to a non-MatMul branch.""" + x, k, y_mm, y_side, shared = _var("x"), _var("k"), _var("y_mm"), _var("y_side"), _var("shared") + + if rewrite == "_move_mul_before_qdq": + qdq_nodes, dq_out = _qdq(x) + producer = gs.Node( + op="Mul", + inputs=[dq_out, gs.Constant("c", np.array(0.5, dtype=np.float32))], + outputs=[shared], + ) + prelude = [*qdq_nodes, producer] + elif rewrite == "_move_transpose_before_qdq": + qdq_nodes, dq_out = _qdq(x) + producer = gs.Node( + op="Transpose", inputs=[dq_out], outputs=[shared], attrs={"perm": [1, 0]} + ) + prelude = [*qdq_nodes, producer] + else: + prelude = [gs.Node(op="Softmax", inputs=[x], outputs=[shared], attrs={"axis": -1})] + + graph = _graph( + [ + *prelude, + gs.Node(op="MatMul", inputs=[shared, k], outputs=[y_mm]), + gs.Node(op="Relu", inputs=[shared], outputs=[y_side]), + ], + [x, k], + [y_mm, y_side], + ) + assert getattr(FP8QuantExporter, rewrite)(graph) == 0 diff --git a/tests/unit/onnx/quantization/test_graph_utils.py b/tests/unit/onnx/quantization/test_graph_utils.py index 1deaa1b8d31..d72e99b8cc1 100644 --- a/tests/unit/onnx/quantization/test_graph_utils.py +++ b/tests/unit/onnx/quantization/test_graph_utils.py @@ -13,11 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + import numpy as np import onnx_graphsurgeon as gs import pytest +from onnx import TensorProto, helper -from modelopt.onnx.quantization.graph_utils import find_nodes_from_convs_to_exclude +from modelopt.onnx.quantization.graph_utils import ( + _exclude_matmuls_by_inference, + _exclude_matmuls_by_shape_inference, + _get_inp_b_k_dim, + find_nodes_from_convs_to_exclude, +) def _make_conv_graph(output_channels, input_channels, kernel_shape=(3, 3), name="Conv_0"): @@ -85,3 +93,318 @@ def test_fp8_channels_below_16_excluded_by_general_check(oc, ic): graph = _make_conv_graph(output_channels=oc, input_channels=ic, kernel_shape=(3, 3)) excluded = find_nodes_from_convs_to_exclude(graph, quantize_mode="fp8") assert "Conv_0" in excluded + + +def _make_matmul_model(m, k, n, name="MatMul_0", inp_b_constant=True): + """Build a minimal ONNX model with a single MatMul: [M, K] x [K, N] -> [M, N].""" + inp_a = helper.make_tensor_value_info("A", TensorProto.FLOAT, [m, k]) + out = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [m, n]) + + if inp_b_constant: + b_init = helper.make_tensor("B", TensorProto.FLOAT, [k, n], np.ones(k * n).tolist()) + matmul = helper.make_node("MatMul", ["A", "B"], ["Y"], name=name) + graph = helper.make_graph([matmul], "test", [inp_a], [out], initializer=[b_init]) + else: + inp_b = helper.make_tensor_value_info("B", TensorProto.FLOAT, [k, n]) + matmul = helper.make_node("MatMul", ["A", "B"], ["Y"], name=name) + graph = helper.make_graph([matmul], "test", [inp_a, inp_b], [out]) + + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + return model + + +def _get_nodes_by_op(model, op): + """Import an ONNX model and return its gs.Nodes whose op matches ``op``.""" + graph = gs.import_onnx(model) + return [n for n in graph.nodes if n.op == op] + + +def test_get_inp_b_k_dim_constant(): + """K dimension should be read from the Constant weight shape.""" + model = _make_matmul_model(m=32, k=8, n=64) + nodes = _get_nodes_by_op(model, "MatMul") + assert _get_inp_b_k_dim(nodes[0]) == 8 + + +def test_get_inp_b_k_dim_variable_with_output_map(): + """K dimension should be read from output_map for Variable inputs.""" + model = _make_matmul_model(m=32, k=10, n=64, inp_b_constant=False) + nodes = _get_nodes_by_op(model, "MatMul") + output_map = {"B": np.zeros((10, 64))} + assert _get_inp_b_k_dim(nodes[0], output_map=output_map) == 10 + + +def test_get_inp_b_k_dim_returns_none_when_unknown(): + """Should return None if K cannot be determined.""" + model = _make_matmul_model(m=32, k=8, n=64, inp_b_constant=False) + nodes = _get_nodes_by_op(model, "MatMul") + assert _get_inp_b_k_dim(nodes[0]) is None + + +@pytest.mark.parametrize( + ("m", "k", "n", "expected_excluded"), + [ + (32, 64, 8, True), + (32, 64, 15, True), + (32, 8, 64, True), + (32, 15, 64, True), + (32, 8, 8, True), + (32, 64, 16, False), + (32, 16, 64, False), + (32, 64, 64, False), + (32, 32, 32, False), + ], +) +def test_matmul_small_gemm_exclusion(m, k, n, expected_excluded): + """MatMuls with N or K < 16 should be excluded by shape inference.""" + model = _make_matmul_model(m=m, k=k, n=n) + nodes = _get_nodes_by_op(model, "MatMul") + calibration_shapes = {"A": [m, k]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + if expected_excluded: + assert "MatMul_0" in excluded + else: + assert "MatMul_0" not in excluded + + +def test_matmul_gemv_excluded(): + """MatMul with N=1 (GEMV) should be excluded regardless of other dims.""" + model = _make_matmul_model(m=32, k=64, n=1) + nodes = _get_nodes_by_op(model, "MatMul") + calibration_shapes = {"A": [32, 64]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + assert "MatMul_0" in excluded + + +def test_matmul_gemv_variable_b_excluded(): + """All-Variable MatMul with N=1 should be excluded via the all-Variable GEMV branch.""" + # inp_b_constant=False makes B a graph input (Variable) so the all-Variable path is taken. + model = _make_matmul_model(m=32, k=64, n=1, inp_b_constant=False) + nodes = _get_nodes_by_op(model, "MatMul") + calibration_shapes = {"A": [32, 64], "B": [64, 1]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + assert "MatMul_0" in excluded + + +def test_matmul_large_dims_not_excluded(): + """MatMul with all large dims should not be excluded.""" + model = _make_matmul_model(m=128, k=256, n=64) + nodes = _get_nodes_by_op(model, "MatMul") + calibration_shapes = {"A": [128, 256]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + assert "MatMul_0" not in excluded + + +def _make_gemm_model(m, k, n, trans_b, name="Gemm_0"): + """Build a minimal ONNX model with a single Gemm node and a constant B. + + If trans_b is 1, B has shape [N, K] (K is last axis). + Otherwise B has shape [K, N]. + """ + inp_a = helper.make_tensor_value_info("A", TensorProto.FLOAT, [m, k]) + out = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [m, n]) + + b_shape = [n, k] if trans_b else [k, n] + b_init = helper.make_tensor( + "B", TensorProto.FLOAT, b_shape, np.ones(b_shape[0] * b_shape[1]).tolist() + ) + gemm = helper.make_node("Gemm", ["A", "B"], ["Y"], name=name, transB=trans_b) + graph = helper.make_graph([gemm], "test", [inp_a], [out], initializer=[b_init]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + return model + + +@pytest.mark.parametrize("trans_b", [0, 1]) +def test_get_inp_b_k_dim_gemm_transb_constant(trans_b): + """Gemm should honor transB when deriving K from a Constant B.""" + model = _make_gemm_model(m=32, k=10, n=64, trans_b=trans_b) + nodes = _get_nodes_by_op(model, "Gemm") + assert _get_inp_b_k_dim(nodes[0]) == 10 + + +@pytest.mark.parametrize("trans_b", [0, 1]) +def test_get_inp_b_k_dim_gemm_transb_output_map(trans_b): + """Gemm should honor transB when deriving K from an output_map.""" + # Build with a Variable B so the node's input is not a Constant. + inp_a = helper.make_tensor_value_info("A", TensorProto.FLOAT, [32, 10]) + inp_b = helper.make_tensor_value_info("B", TensorProto.FLOAT, [64, 10] if trans_b else [10, 64]) + out = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [32, 64]) + gemm = helper.make_node("Gemm", ["A", "B"], ["Y"], name="Gemm_0", transB=trans_b) + graph = helper.make_graph([gemm], "test", [inp_a, inp_b], [out]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + nodes = _get_nodes_by_op(model, "Gemm") + + b_runtime_shape = (64, 10) if trans_b else (10, 64) + output_map = {"B": np.zeros(b_runtime_shape)} + assert _get_inp_b_k_dim(nodes[0], output_map=output_map) == 10 + + +def test_gemm_small_k_excluded_with_transb(): + """Gemm with transB=1 and small K should be excluded (regression: prior code read N).""" + # N=64 is large; K=8 is small. With transB=1, B=[N,K]=[64,8], K axis is -1. + # If _get_inp_b_k_dim ignored transB it would read 64 (N) and not exclude. + model = _make_gemm_model(m=32, k=8, n=64, trans_b=1) + nodes = _get_nodes_by_op(model, "Gemm") + calibration_shapes = {"A": [32, 8]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + assert "Gemm_0" in excluded + + +def test_gemm_large_dims_not_excluded_with_transb(): + """Gemm with transB=1 and all large dims should NOT be excluded.""" + model = _make_gemm_model(m=32, k=64, n=64, trans_b=1) + nodes = _get_nodes_by_op(model, "Gemm") + calibration_shapes = {"A": [32, 64]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + assert "Gemm_0" not in excluded + + +def _make_matmul_model_graph_input_b(m, k, n, name="MatMul_0"): + """MatMul where B is a graph input (its shape lives in model.graph.input only).""" + inp_a = helper.make_tensor_value_info("A", TensorProto.FLOAT, [m, k]) + inp_b = helper.make_tensor_value_info("B", TensorProto.FLOAT, [k, n]) + out = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [m, n]) + matmul = helper.make_node("MatMul", ["A", "B"], ["Y"], name=name) + graph = helper.make_graph([matmul], "test", [inp_a, inp_b], [out]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + +def test_matmul_small_k_graph_input_b_excluded(): + """Small-K MatMul whose B is a graph input should still be excluded. + + Regression: previous value_info_map only covered model.graph.value_info/output, + missing graph inputs, so K was undetectable and the MatMul wasn't excluded. + """ + model = _make_matmul_model_graph_input_b(m=32, k=8, n=64) + nodes = _get_nodes_by_op(model, "MatMul") + calibration_shapes = {"A": [32, 8], "B": [8, 64]} + excluded = _exclude_matmuls_by_shape_inference(model, nodes, calibration_shapes) + assert "MatMul_0" in excluded + + +@pytest.mark.parametrize( + ("k", "n", "expected_excluded"), + [ + (8, 64, True), + (64, 8, True), + (64, 64, False), + ], +) +def test_exclude_matmuls_by_inference_runtime_path(k, n, expected_excluded): + """Exercise the runtime-inference path with B as a graph input (read from output_map).""" + m = 32 + model = _make_matmul_model_graph_input_b(m=m, k=k, n=n) + nodes = _get_nodes_by_op(model, "MatMul") + + # Mock get_extended_model_outputs to return a synthetic output_map so we don't + # need an actual ORT session. + fake_output_map = { + "Y": np.zeros((m, n), dtype=np.float32), + "B": np.zeros((k, n), dtype=np.float32), + } + with mock.patch( + "modelopt.onnx.quantization.graph_utils.get_extended_model_outputs", + return_value=fake_output_map, + ): + excluded = _exclude_matmuls_by_inference( + onnx_path="unused.onnx", + model=model, + matmul_nodes=nodes, + use_external_data_format=False, + intermediate_generated_files=[], + calibration_data_reader=None, + calibration_eps=["cpu"], + ) + if expected_excluded: + assert "MatMul_0" in excluded + else: + assert "MatMul_0" not in excluded + + +def test_exclude_matmuls_by_inference_gemv_variable_b(): + """All-Variable GEMV (N=1) should be excluded via the runtime-inference all-Variable path.""" + m, k, n = 32, 64, 1 + model = _make_matmul_model_graph_input_b(m=m, k=k, n=n) + nodes = _get_nodes_by_op(model, "MatMul") + fake_output_map = { + "Y": np.zeros((m, n), dtype=np.float32), + "B": np.zeros((k, n), dtype=np.float32), + } + with mock.patch( + "modelopt.onnx.quantization.graph_utils.get_extended_model_outputs", + return_value=fake_output_map, + ): + excluded = _exclude_matmuls_by_inference( + onnx_path="unused.onnx", + model=model, + matmul_nodes=nodes, + use_external_data_format=False, + intermediate_generated_files=[], + calibration_data_reader=None, + calibration_eps=["cpu"], + ) + assert "MatMul_0" in excluded + + +def test_exclude_matmuls_by_inference_gemv_constant_b(): + """Constant-B GEMV (N=1) should be excluded via the runtime-inference elif path.""" + m, k, n = 32, 64, 1 + model = _make_matmul_model(m=m, k=k, n=n, inp_b_constant=True) + nodes = _get_nodes_by_op(model, "MatMul") + # B is a Constant (initializer) so only the matmul output is added to graph outputs. + fake_output_map = {"Y": np.zeros((m, n), dtype=np.float32)} + with mock.patch( + "modelopt.onnx.quantization.graph_utils.get_extended_model_outputs", + return_value=fake_output_map, + ): + excluded = _exclude_matmuls_by_inference( + onnx_path="unused.onnx", + model=model, + matmul_nodes=nodes, + use_external_data_format=False, + intermediate_generated_files=[], + calibration_data_reader=None, + calibration_eps=["cpu"], + ) + assert "MatMul_0" in excluded + + +def test_exclude_matmuls_by_inference_dedupes_added_outputs(): + """Two MatMuls sharing the same Variable B must not create duplicate graph outputs.""" + # Build two MatMuls sharing B as a graph input. + m, k, n = 32, 8, 64 + inp_a1 = helper.make_tensor_value_info("A1", TensorProto.FLOAT, [m, k]) + inp_a2 = helper.make_tensor_value_info("A2", TensorProto.FLOAT, [m, k]) + inp_b = helper.make_tensor_value_info("B", TensorProto.FLOAT, [k, n]) + out1 = helper.make_tensor_value_info("Y1", TensorProto.FLOAT, [m, n]) + out2 = helper.make_tensor_value_info("Y2", TensorProto.FLOAT, [m, n]) + mm1 = helper.make_node("MatMul", ["A1", "B"], ["Y1"], name="MatMul_0") + mm2 = helper.make_node("MatMul", ["A2", "B"], ["Y2"], name="MatMul_1") + graph = helper.make_graph([mm1, mm2], "test", [inp_a1, inp_a2, inp_b], [out1, out2]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + nodes = _get_nodes_by_op(model, "MatMul") + + fake_output_map = { + "Y1": np.zeros((m, n), dtype=np.float32), + "Y2": np.zeros((m, n), dtype=np.float32), + "B": np.zeros((k, n), dtype=np.float32), + } + with mock.patch( + "modelopt.onnx.quantization.graph_utils.get_extended_model_outputs", + return_value=fake_output_map, + ): + excluded = _exclude_matmuls_by_inference( + onnx_path="unused.onnx", + model=model, + matmul_nodes=nodes, + use_external_data_format=False, + intermediate_generated_files=[], + calibration_data_reader=None, + calibration_eps=["cpu"], + ) + output_names = [o.name for o in model.graph.output] + # B should appear only once in the graph outputs. + assert output_names.count("B") == 1 + # Both MatMuls should be excluded (small K). + assert "MatMul_0" in excluded + assert "MatMul_1" in excluded diff --git a/tests/unit/onnx/test_fold_casts.py b/tests/unit/onnx/test_fold_casts.py new file mode 100644 index 00000000000..59a434d1206 --- /dev/null +++ b/tests/unit/onnx/test_fold_casts.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for the FP16 Q/DQ scale cast-folding helpers in ``modelopt.onnx.utils``.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.utils import fold_dq_fp32_to_fp16_casts, fold_q_fp16_to_fp32_casts + + +def _dq_cast_model(opset): + """``DQ → Cast(FP32→FP16) → MatMul(x)`` with FP32 scale.""" + nodes = [ + helper.make_node("DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["dq_out"], "dq"), + helper.make_node("Cast", ["dq_out"], ["cast_out"], "cast", to=TensorProto.FLOAT16), + helper.make_node("MatMul", ["x", "cast_out"], ["y"], "matmul"), + ] + inits = [ + numpy_helper.from_array(np.ones((4, 4), dtype=np.int8), "w_q"), + numpy_helper.from_array(np.array(0.1, dtype=np.float32), "w_scale"), + numpy_helper.from_array(np.array(0, dtype=np.int8), "w_zp"), + ] + return helper.make_model( + helper.make_graph( + nodes, + "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [None, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT16, [None, 4])], + initializer=inits, + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + + +def _cast_q_model(opset): + """``Cast(FP16→FP32) → Q → DQ → MatMul`` with FP32 scale.""" + nodes = [ + helper.make_node("Cast", ["x"], ["c_out"], "cast", to=TensorProto.FLOAT), + helper.make_node("QuantizeLinear", ["c_out", "scale", "zp"], ["q_out"], "q"), + helper.make_node("DequantizeLinear", ["q_out", "scale", "zp"], ["dq_out"], "dq"), + helper.make_node("MatMul", ["dq_out", "w"], ["y"], "matmul"), + ] + inits = [ + numpy_helper.from_array(np.ones((4, 4), dtype=np.float16), "w"), + numpy_helper.from_array(np.array(0.1, dtype=np.float32), "scale"), + numpy_helper.from_array(np.array(0, dtype=np.int8), "zp"), + ] + return helper.make_model( + helper.make_graph( + nodes, + "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [None, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [None, 4])], + initializer=inits, + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + + +@pytest.mark.parametrize( + ("fold_fn", "build_model", "scale_name"), + [ + (fold_dq_fp32_to_fp16_casts, _dq_cast_model, "w_scale"), + (fold_q_fp16_to_fp32_casts, _cast_q_model, "scale"), + ], +) +def test_fold_rewrites_cast_and_scale_at_opset_19(fold_fn, build_model, scale_name): + folded = fold_fn(build_model(opset=19)) + assert "Cast" not in {n.op_type for n in folded.graph.node} + scale = next(i for i in folded.graph.initializer if i.name == scale_name) + assert scale.data_type == TensorProto.FLOAT16 + + +@pytest.mark.parametrize( + ("fold_fn", "build_model", "scale_name"), + [ + (fold_dq_fp32_to_fp16_casts, _dq_cast_model, "w_scale"), + (fold_q_fp16_to_fp32_casts, _cast_q_model, "scale"), + ], +) +def test_fold_is_noop_below_min_opset(fold_fn, build_model, scale_name): + folded = fold_fn(build_model(opset=18)) + assert "Cast" in {n.op_type for n in folded.graph.node} + scale = next(i for i in folded.graph.initializer if i.name == scale_name) + assert scale.data_type == TensorProto.FLOAT diff --git a/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py b/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py new file mode 100644 index 00000000000..7a919799b86 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for the per-instance nested-attention skip in the HF quantization plugin.""" + +import pytest +import torch.nn as nn + +pytest.importorskip("transformers") + +from modelopt.torch.quantization.plugins.huggingface import _wraps_nested_attention + + +def _attn(name, child=None): + """Build a module whose class name ends with ``Attention`` and optionally wraps ``child``.""" + cls = type(name, (nn.Module,), {"__init__": lambda self: nn.Module.__init__(self)}) + m = cls() + if child is not None: + m.inner = child + return m + + +def test_wraps_nested_attention_flags_only_wrappers_per_instance(): + """Leaf attention is not a wrapper; wrappers (any level) are; same class reused is + checked per-instance.""" + leaf = _attn("SelfAttention") + wrapper = _attn("ViTAttention", child=_attn("ViTSelfAttention")) + outer = _attn("OuterAttention", child=wrapper) + reused_wrapper = _attn("ReusedAttention", child=_attn("ReusedAttention")) + + assert not _wraps_nested_attention(leaf) + assert _wraps_nested_attention(wrapper) + assert _wraps_nested_attention(outer) and _wraps_nested_attention(outer.inner) + assert _wraps_nested_attention(reused_wrapper) + assert not _wraps_nested_attention(reused_wrapper.inner) diff --git a/tests/unit/torch/quantization/test_quant_layernorm.py b/tests/unit/torch/quantization/test_quant_layernorm.py new file mode 100644 index 00000000000..1e20662524a --- /dev/null +++ b/tests/unit/torch/quantization/test_quant_layernorm.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for ``torch.nn.LayerNorm`` being registered in ``QuantModuleRegistry``.""" + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.nn import QuantModuleRegistry + + +def test_layernorm_quant_wrapper_is_identity_when_quantizers_disabled(): + qln = QuantModuleRegistry.convert(nn.LayerNorm(8)) + qln.input_quantizer.disable() + qln.output_quantizer.disable() + + x = torch.randn(2, 8) + ref = nn.functional.layer_norm(x, (8,), qln.weight, qln.bias, eps=qln.eps) + assert torch.allclose(qln(x), ref, rtol=0, atol=0) diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative.py b/tests/unit/torch/speculative/plugins/test_hf_speculative.py index b41b7fae2b0..f5a14a08976 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative.py +++ b/tests/unit/torch/speculative/plugins/test_hf_speculative.py @@ -17,6 +17,7 @@ from copy import deepcopy import pytest +import torch from _test_utils.torch.transformers_models import ( get_tiny_llama, tf_modelopt_state_and_output_tester, @@ -48,3 +49,39 @@ def test_eagle_model_convert_save_and_restore(tmp_path, eagle_config): model_test = AutoModelForCausalLM.from_pretrained(tmp_path / "modelopt_model") assert isinstance(model_test, mtsp.plugins.HFEagleModel) tf_modelopt_state_and_output_tester(model_ref, model_test) + + +@pytest.mark.parametrize("eagle_config", [EAGLE3_DEFAULT_CFG]) +@pytest.mark.parametrize("eagle_ttt_steps", [1, 2]) +def test_eagle_mix_hidden_states_backward(eagle_config, eagle_ttt_steps): + """Regression test for GitHub issue #1088. + + Verifies that the EAGLE training forward+backward pass does not crash with + ``eagle_mix_hidden_states=True`` due to an in-place tensor modification + breaking autograd. + """ + model = get_tiny_llama(num_hidden_layers=8) + + config = deepcopy(eagle_config["config"]) + config["eagle_architecture_config"].update( + { + "draft_vocab_size": model.config.vocab_size, + "hidden_size": model.config.hidden_size, + } + ) + config["eagle_mix_hidden_states"] = True + config["eagle_ttt_steps"] = eagle_ttt_steps + config["eagle_use_torch_compile"] = False + + mtsp.convert(model, mode=[("eagle", config)]) + model.train() + + input_ids = torch.randint(0, model.config.vocab_size, (2, 16)) + labels = input_ids.clone() + + outputs = model(input_ids=input_ids, labels=labels) + assert outputs.loss is not None + outputs.loss.backward() + + eagle_grads = [p.grad for p in model.eagle_module.parameters() if p.grad is not None] + assert len(eagle_grads) > 0, "Expected gradients to flow to eagle_module"