Model-Optimizer/modelopt/torch/quantization/plugins/diffusers.py

# 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.

"""Support quantization of diffusers layers."""

import functools
from collections.abc import Callable, Iterator
from functools import partial
from types import ModuleType

import onnx
import torch
from diffusers.models.attention_processor import Attention
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
from torch.autograd import Function
from torch.nn import functional as F
from torch.onnx import symbolic_helper
from torch.onnx._internal import jit_utils, registration

from ..export_onnx import export_fp8_mha
from ..nn import (
    QuantConv2d,
    QuantInputBase,
    QuantLinear,
    QuantLinearConvBase,
    QuantModuleRegistry,
    TensorQuantizer,
)
from .custom import _QuantFunctionalMixin

_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=18)

onnx_dtype_map = {
    "BFloat16": onnx.TensorProto.BFLOAT16,
    "Float": onnx.TensorProto.FLOAT,
    "Float8": onnx.TensorProto.FLOAT8E4M3FN,
    "Half": onnx.TensorProto.FLOAT16,
    "INT8": onnx.TensorProto.INT8,
    "UINT8": onnx.TensorProto.UINT8,
}
mha_valid_precisions = {"Half", "BFloat16"}


class _QuantLoRACompatibleLinearConvBase(QuantLinearConvBase):
    def _setup(self):
        assert self.lora_layer is None, (
            f"To quantize {self}, lora_layer should be None. Please fuse the LoRA layer before"
            " quantization."
        )
        return super()._setup()


@QuantModuleRegistry.register({LoRACompatibleConv: "LoRACompatibleConv"})
class _QuantLoRACompatibleConv(_QuantLoRACompatibleLinearConvBase):
    default_quant_desc_weight = QuantConv2d.default_quant_desc_weight


@QuantModuleRegistry.register({LoRACompatibleLinear: "LoRACompatibleLinear"})
class _QuantLoRACompatibleLinear(_QuantLoRACompatibleLinearConvBase):
    default_quant_desc_weight = QuantLinear.default_quant_desc_weight


def _quantized_bmm(self, input, mat2, *args, **kwargs):
    attn, v = input, mat2
    return self.bmm2_output_quantizer(
        torch._bmm(self.softmax_quantizer(attn), self.v_bmm_quantizer(v), *args, **kwargs)
    )


def _quantized_baddbmm(self, input, batch1, batch2, *args, **kwargs):
    q, k = batch1, batch2
    return torch._baddbmm(input, self.q_bmm_quantizer(q), self.k_bmm_quantizer(k), *args, **kwargs)


def _quantized_sdpa(self, *args, **kwargs):
    fp8_sdpa = FP8SDPA.apply
    parameters = [
        "query",
        "key",
        "value",
        "attn_mask",
        "dropout_p",
        "is_causal",
        "scale",
        "q_quantized_scale",
        "k_quantized_scale",
        "v_quantized_scale",
        "high_precision_flag",
    ]
    default_values = [None, None, None, None, 0.0, False, None, None, None, None, "Half"]
    param_dict = dict(zip(parameters, default_values))
    for i, arg in enumerate(args):
        param_dict[parameters[i]] = arg
    param_dict.update(kwargs)
    fp8_sdpa_args = [param_dict[param] for param in parameters]
    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():
        query = self.q_bmm_quantizer(query)
        key = self.k_bmm_quantizer(key)
        value = self.v_bmm_quantizer(value)

    q_quantized_scale = self.q_bmm_quantizer._get_amax(query)
    k_quantized_scale = self.k_bmm_quantizer._get_amax(key)
    v_quantized_scale = self.v_bmm_quantizer._get_amax(value)

    # We don't need to calibrate the output of softmax
    return self.bmm2_output_quantizer(
        fp8_sdpa(
            query,
            key,
            value,
            *fp8_sdpa_args[3:7],
            q_quantized_scale,
            k_quantized_scale,
            v_quantized_scale,
            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,
        )
    )


class _QuantAttention(_QuantFunctionalMixin):
    """FP8 processor for performing attention-related computations."""

    _functionals_to_replace = [
        (torch, "bmm", _quantized_bmm),
        (torch, "baddbmm", _quantized_baddbmm),
        (F, "scaled_dot_product_attention", _quantized_sdpa),
    ]

    @property
    def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]:
        for package, func_name, quantized_func in self._functionals_to_replace:
            if not hasattr(package, func_name):
                continue
            quantized_func = partial(quantized_func, self)
            yield package, func_name, quantized_func

    def _setup(self):
        self.q_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input)
        self.k_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input)
        self.v_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input)
        self.softmax_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input)
        self.bmm2_output_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input)


QuantModuleRegistry.register({Attention: "Attention"})(_QuantAttention)


original_scaled_dot_product_attention = F.scaled_dot_product_attention


class FP8SDPA(Function):
    """A customized FP8 SDPA op for the onnx export."""

    @staticmethod
    def forward(
        ctx,
        query,
        key,
        value,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=False,
        scale=None,
        q_quantized_scale=None,
        k_quantized_scale=None,
        v_quantized_scale=None,
        high_precision_flag=None,
        disable_fp8_mha=True,
    ):
        """Forward method."""
        ctx.save_for_backward(query, key, value, attn_mask)
        ctx.q_quantized_scale = q_quantized_scale
        ctx.k_quantized_scale = k_quantized_scale
        ctx.v_quantized_scale = v_quantized_scale
        # During runtime, ignore x or use it as needed
        return original_scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=attn_mask,
            dropout_p=dropout_p,
            is_causal=is_causal,
            scale=scale,
        )

    @staticmethod
    @symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "t", "t", "t", "s", "b")
    def symbolic(
        g: jit_utils.GraphContext,
        query: torch._C.Value,
        key: torch._C.Value,
        value: torch._C.Value,
        attn_mask: torch._C.Value | None = None,
        dropout_p: float = 0.0,
        is_causal: bool = False,
        scale: torch._C.Value | None = None,
        q_quantized_scale: float = 1.0,
        k_quantized_scale: float = 1.0,
        v_quantized_scale: float = 1.0,
        high_precision_flag: str = "Half",
        disable_fp8_mha: bool = True,
    ):
        """Symbolic method."""
        return export_fp8_mha(
            g,
            query,
            key,
            value,
            attn_mask,
            dropout_p,
            is_causal,
            scale,
            q_quantized_scale,
            k_quantized_scale,
            v_quantized_scale,
            high_precision_flag,
            disable_fp8_mha,
        )