diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3e10c786753..da0c237f239 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Changelog **Bug Fixes** +- Fix FP8 MHA quantization for diffusion attention modules that delegate to a module-level helper, and preserve fusion-ready Q/DQ placement around attention scaling, key transposes, and softmax outputs in exported ONNX graphs. - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..be84c77bba4 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -15,6 +15,7 @@ """Support quantization of diffusers layers.""" +import inspect from collections.abc import Callable, Iterator from functools import partial from types import ModuleType @@ -25,6 +26,7 @@ import torch from diffusers.models.attention_processor import Attention from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from diffusers.models.modeling_utils import ModelMixin from packaging.version import parse as parse_version if parse_version(diffusers.__version__) >= parse_version("0.35.0"): @@ -64,7 +66,7 @@ TensorQuantizer, ) from ...nn.modules.quant_conv import _QuantConv3d -from ..custom import _QuantFunctionalMixin +from ..custom import CUSTOM_MODEL_PLUGINS, _QuantFunctionalMixin onnx_dtype_map = { "BFloat16": onnx.TensorProto.BFLOAT16, @@ -75,6 +77,7 @@ "UINT8": onnx.TensorProto.UINT8, } mha_valid_precisions = {"Half", "BFloat16"} +_QKV_PARAMETER_NAMES = (("q", "k", "v"), ("query", "key", "value")) class _QuantLoRACompatibleLinearConvBase(QuantLinearConvBase): @@ -108,6 +111,26 @@ def _quantized_baddbmm(self, input, batch1, batch2, *args, **kwargs): return torch._baddbmm(input, self.q_bmm_quantizer(q), self.k_bmm_quantizer(k), *args, **kwargs) +def _fp8_mha_disabled(self, query, key, value): + if hasattr(self, "_disable_fp8_mha") or not getattr(self, "_auto_fp8_mha", False): + return getattr(self, "_disable_fp8_mha", True) + names = ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer", "softmax_quantizer") + quantizers = tuple(getattr(self, name) for name in names) + if any(not quantizer.is_enabled or not quantizer.is_fp8 for quantizer in quantizers): + return True + high_precision = self.q_bmm_quantizer.trt_high_precision_dtype + if high_precision not in mha_valid_precisions or { + quantizer.trt_high_precision_dtype for quantizer in quantizers + } != {high_precision}: + return True + if any( + tensor.dtype != getattr(torch, high_precision.lower()) for tensor in (query, key, value) + ): + return True + head_dims = tuple(int(tensor.shape[-1]) for tensor in (query, key, value)) + return len(set(head_dims)) != 1 or head_dims[0] % 16 != 0 + + def _quantized_sdpa(self, *args, **kwargs): fp8_sdpa = FP8SDPA.apply parameters = [ @@ -132,11 +155,14 @@ def _quantized_sdpa(self, *args, **kwargs): while fp8_sdpa_args and fp8_sdpa_args[-1] is None: fp8_sdpa_args.pop() query, key, value = fp8_sdpa_args[:3] - - if not torch.onnx.is_in_onnx_export(): + exporting = torch.onnx.is_in_onnx_export() + if not exporting: query = self.q_bmm_quantizer(query) key = self.k_bmm_quantizer(key) value = self.v_bmm_quantizer(value) + disable_fp8_mha = _fp8_mha_disabled(self, query, key, value) + if exporting and getattr(self, "_auto_fp8_mha", False) and not disable_fp8_mha: + self._delegated_sdpa_export_calls = getattr(self, "_delegated_sdpa_export_calls", 0) + 1 q_quantized_scale = self.q_bmm_quantizer._get_amax(query) k_quantized_scale = self.k_bmm_quantizer._get_amax(key) @@ -155,7 +181,7 @@ def _quantized_sdpa(self, *args, **kwargs): self.q_bmm_quantizer.trt_high_precision_dtype if hasattr(self.q_bmm_quantizer, "trt_high_precision_dtype") else "Half", - self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, + disable_fp8_mha, ) ) @@ -211,6 +237,126 @@ def forward(self, *args, **kwargs): original_scaled_dot_product_attention = F.scaled_dot_product_attention +def _call_delegated_attention(self, helper, parameter_names, *args, **kwargs): + bound = inspect.signature(helper).bind(*args, **kwargs) + qkv = tuple(bound.arguments[name] for name in parameter_names) + if torch.onnx.is_in_onnx_export(): + before = self._delegated_sdpa_export_calls + eligible = not _fp8_mha_disabled(self, *qkv) + output = helper(*bound.args, **bound.kwargs) + self._delegated_missed_sdpa |= eligible and before == self._delegated_sdpa_export_calls + return output + + quantizers = (self.q_bmm_quantizer, self.k_bmm_quantizer, self.v_bmm_quantizer) + for name, quantizer in zip(parameter_names, quantizers): + bound.arguments[name] = quantizer(bound.arguments[name]) + return self.bmm2_output_quantizer(helper(*bound.args, **bound.kwargs)) + + +class _QuantDelegatedAttention(_QuantAttention): + def forward(self, *args, **kwargs): + self._delegated_missed_sdpa = self._delegated_sdpa_export_calls = 0 + output = super().forward(*args, **kwargs) + if torch.onnx.is_in_onnx_export() and self._delegated_missed_sdpa: + raise RuntimeError("Eligible delegated FP8 MHA export did not reach SDPA.") + return output + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + module, name, parameter_names = self._delegated_attention_helper + helper = getattr(module, name, None) + if not callable(helper): + return + quantized_sdpa = partial(_quantized_sdpa, self) + if helper is original_scaled_dot_product_attention: + replacement = quantized_sdpa + else: + replacement = partial(_call_delegated_attention, self, helper, parameter_names) + yield module, name, replacement + if not torch.onnx.is_in_onnx_export(): + return + + yield F, "scaled_dot_product_attention", quantized_sdpa + if helper is original_scaled_dot_product_attention: + return + helper = inspect.unwrap(helper) + helper_globals = getattr(helper, "__globals__", {}) + helper_module = inspect.getmodule(helper) + code = getattr(helper, "__code__", None) + if code is None or not isinstance(helper_module, ModuleType): + return + if helper_module.__dict__ is helper_globals: + for alias in code.co_names: + if helper_globals.get(alias) is original_scaled_dot_product_attention: + yield helper_module, alias, quantized_sdpa + + +def _try_register_attention(attention_cls): + try: + forward = inspect.unwrap(attention_cls.forward) + except (TypeError, ValueError): + return + code = getattr(forward, "__code__", None) + function_globals = getattr(forward, "__globals__", {}) + if ( + code is None + or not isinstance(module := inspect.getmodule(forward), ModuleType) + or module.__dict__ is not function_globals + ): + return + candidates = [] + for name in dict.fromkeys(code.co_names): + helper = function_globals.get(name) + if not callable(helper): + continue + try: + parameter_names = tuple(inspect.signature(helper).parameters)[:3] + except (TypeError, ValueError): + if helper is not original_scaled_dot_product_attention: + continue + parameter_names = _QKV_PARAMETER_NAMES[1] + helper_name = f"{name} {getattr(helper, '__name__', '')}".lower() + if parameter_names in _QKV_PARAMETER_NAMES and ( + "attention" in helper_name or "attn" in helper_name + ): + candidates.append((module, name, parameter_names)) + if len(candidates) != 1: + return + quantized_cls = type( + f"_Quant{attention_cls.__name__}", + (_QuantDelegatedAttention,), + {"_auto_fp8_mha": True, "_delegated_attention_helper": candidates[0]}, + ) + QuantModuleRegistry.register({attention_cls: attention_cls.__name__})(quantized_cls) + + +def _register_diffusers_attentions_on_the_fly(model): + if not isinstance(model, ModelMixin): + return + + seen_classes = set() + for module in model.modules(): + module_type = type(module) + if ( + not module_type.__name__.endswith("Attention") + or module_type in seen_classes + or module_type in QuantModuleRegistry + or module_type.__module__.startswith("transformers.") + or hasattr(module_type, "_setup") + or not all(isinstance(getattr(module, name, None), torch.nn.Module) for name in "qkv") + or any( + child is not module and type(child).__name__.endswith("Attention") + for child in module.modules() + ) + ): + continue + seen_classes.add(module_type) + _try_register_attention(module_type) + + +CUSTOM_MODEL_PLUGINS.add(_register_diffusers_attentions_on_the_fly) + + class FP8SDPA(Function): """A customized FP8 SDPA op for the onnx export.""" diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_attention.py b/tests/unit/torch/quantization/plugins/test_diffusers_attention.py new file mode 100644 index 00000000000..4c7e0be308f --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_diffusers_attention.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import io + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +ModelMixin = pytest.importorskip("diffusers").ModelMixin + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.plugins.diffusion.diffusers import _fp8_mha_disabled + +_SDPA_ALIAS = F.scaled_dot_product_attention + + +def _fake_attention(q, k, v): + return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v + + +def _wan_attention(q, k, v): + if torch.onnx.is_in_onnx_export(): + return F.scaled_dot_product_attention(q, k, v) + return _fake_attention(q, k, v) + + +class DelegatedAttention(nn.Module): + def __init__(self, style="keyword"): + super().__init__() + self.style = style + self.q, self.k, self.v, self.o = (nn.Linear(32, 32) for _ in range(4)) + + def _qkv(self, x): + qkv = (self.q(x), self.k(x), self.v(x)) + return tuple(tensor.view(2, 4, 2, 16).transpose(1, 2) for tensor in qkv) + + def forward(self, x): + q, k, v = self._qkv(x) + if self.style == "keyword": + output = _wan_attention(q=q, k=k, v=v) + else: + output = _wan_attention(q, k, v) + if self.style == "repeated": + output = _wan_attention(output, k, v) + return self.o(output.transpose(1, 2).flatten(2)) + + +class AliasedAttention(DelegatedAttention): + def forward(self, x): + return self.o(_SDPA_ALIAS(*self._qkv(x)).transpose(1, 2).flatten(2)) + + +class NonSDPAAttention(DelegatedAttention): + def forward(self, x): + return self.o(_fake_attention(*self._qkv(x)).transpose(1, 2).flatten(2)) + + +class DelegatedModel(ModelMixin): + def __init__(self, attention_cls=DelegatedAttention, **kwargs): + super().__init__() + self.attn = attention_cls(**kwargs) + + def forward(self, x): + return self.attn(x) + + +@pytest.fixture(autouse=True) +def _clean_registrations(): + yield + for cls in (DelegatedAttention, AliasedAttention, NonSDPAAttention): + if cls in mtq.QuantModuleRegistry: + mtq.unregister(cls) + + +def _quantize(model, enabled=True): + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + attrs = {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"} + config["quant_cfg"].extend( + {"quantizer_name": name, "cfg": attrs, "enable": enabled} + for name in ("*[qkv]_bmm_quantizer", "*softmax_quantizer") + ) + inputs = torch.randn(2, 4, 32) + mtq.quantize(model, config, lambda quant_model: quant_model(inputs)) + return inputs + + +def _export(attention_cls): + onnx = pytest.importorskip("onnx") + model = DelegatedModel(attention_cls).eval() + inputs, buffer = _quantize(model), io.BytesIO() + torch.onnx.export(model.half(), inputs.half(), buffer, opset_version=20, dynamo=False) + graph = onnx.load_model_from_string(buffer.getvalue()) + onnx.checker.check_model(graph) + return graph + + +def _attention_qdq(graph): + producers = {output: node for node in graph.graph.node for output in node.output} + + def source(name): + node = producers.get(name) + if node is not None and node.op_type == "Cast": + node = producers.get(node.input[0]) + if node is None or node.op_type != "TRT_FP8DequantizeLinear": + return None + quantize = producers.get(node.input[0]) + if quantize is None or quantize.op_type != "TRT_FP8QuantizeLinear": + return None + return quantize, producers.get(quantize.input[0]) + + return producers, [ + (node, inputs) + for node in graph.graph.node + if node.op_type == "MatMul" and all(inputs := [source(name) for name in node.input]) + ] + + +@pytest.mark.parametrize(("style", "calls"), [("keyword", 1), ("positional", 1), ("repeated", 2)]) +def test_delegated_helper_call_styles(style, calls): + model = DelegatedModel(style=style) + inputs = _quantize(model) + seen = [] + with model.attn.q_bmm_quantizer.register_forward_hook(lambda *_: seen.append(None)): + model(inputs) + assert len(seen) == calls + assert all(getattr(model.attn, f"{name}_bmm_quantizer").amax is not None for name in "qkv") + + +@pytest.mark.parametrize("attention_cls", [DelegatedAttention, AliasedAttention]) +def test_export_uses_fp8_sdpa_symbolic(attention_cls): + graph = _export(attention_cls) + producers, attention = _attention_qdq(graph) + qk = next(x for _, x in attention if all(source.op_type == "Mul" for _, source in x)) + assert any(producers[source.input[0]].op_type == "Transpose" for _, source in qk) + pv, inputs = next( + x for x in attention if any(source.op_type == "Softmax" for _, source in x[1]) + ) + softmax_q = next(q for q, source in inputs if source.op_type == "Softmax") + scale = producers[softmax_q.input[1]].attribute[0].t.raw_data + assert scale == torch.tensor(1 / 448, dtype=torch.float16).numpy().tobytes() + reachable = set(pv.output) + for node in graph.graph.node: + if reachable.intersection(node.input): + reachable.update(node.output) + qdq = {node.input[0] for node in graph.graph.node if node.op_type == "TRT_FP8QuantizeLinear"} + assert reachable & qdq + + +def test_export_fails_when_helper_does_not_reach_sdpa(): + helper = _fake_attention + with pytest.raises(RuntimeError, match="did not reach SDPA"): + _export(NonSDPAAttention) + assert _fake_attention is helper and F.scaled_dot_product_attention is _SDPA_ALIAS + + +@pytest.mark.parametrize("case", ["eligible", "fp32", "int8", "misaligned", "explicit"]) +def test_fp8_mha_eligibility(case): + model = DelegatedModel() + _quantize(model) + model.attn.q_bmm_quantizer.num_bits = 8 if case == "int8" else (4, 3) + if case == "explicit": + model.attn._disable_fp8_mha = True + dtype = torch.float32 if case == "fp32" else torch.float16 + head_dim = 15 if case == "misaligned" else 16 + qkv = (torch.randn(2, 2, 4, head_dim, dtype=dtype) for _ in range(3)) + assert _fp8_mha_disabled(model.attn, *qkv) is (case != "eligible") + + +@pytest.mark.parametrize(("model_mixin", "enabled"), [(True, True), (True, False), (False, True)]) +def test_registration_and_enablement_guards(model_mixin, enabled): + model = DelegatedModel() if model_mixin else nn.Sequential(DelegatedAttention()) + _quantize(model, enabled) + attention = model.attn if model_mixin else model[0] + assert bool(getattr(attention, "_auto_fp8_mha", False)) is model_mixin + if model_mixin: + assert attention.q_bmm_quantizer.is_enabled is enabled