scripts/quantize.py

from enum import Enum

from tqdm import tqdm
from typing import Set, List, Optional
import onnx
import os

from dataclasses import dataclass, field

from transformers import HfArgumentParser

from onnxruntime.quantization import QuantType, QuantizationMode
from onnxruntime.quantization.onnx_quantizer import ONNXQuantizer
from onnxruntime.quantization.registry import IntegerOpsRegistry
from onnxruntime.quantization.matmul_4bits_quantizer import MatMul4BitsQuantizer
from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer

from . import float16
from .utils import check_and_save_model

class QuantMode(Enum):
    # F32 = 'fp32'
    FP16 = "fp16"
    Q8 = "q8"
    QI8 = "int8"
    QU8 = "uint8"
    Q4 = "q4"
    Q4F16 = "q4f16"
    BNB4 = "bnb4"


QUANTIZE_SUFFIX_MAPPING = {
    QuantMode.Q8: "quantized",
}

QUANTIZE_OPTIONS = tuple(x.value for x in QuantMode)

# A list of operators that, when detected in a model, should select QUInt8 as the weight type for 8-bit quantization.
QUINT8_OPS = (
    # NOTE:
    # As of 2024/11/29, the latest version of onnxruntime-web is 1.20.1, and does not support INT8 weights for Conv layers.
    # If you attempt to run a model with INT8 weights for Conv layers, you will get an error like:
    # `Can't create a session. ERROR_CODE: 9, ERROR_MESSAGE: Could not find an implementation for ConvInteger(10) node with name '/.../Conv_quant'`
    #
    # For this reason, we choose model weight types to ensure compatibility with onnxruntime-web.
    #
    # As per docs, signed weight type (QInt8) is faster on most CPUs, so, we use that unless the model contains a Conv layer.
    # For more information, see:
    #  - https://github.com/microsoft/onnxruntime/issues/3130#issuecomment-1105200621
    #  - https://github.com/microsoft/onnxruntime/issues/2339
    "Conv",

    # Models produced by onnxruntime-genai contain optimized operators that perform better with QUInt8 weights.
    "GroupQueryAttention",
    "MultiHeadAttention",

    # TODO: "SimplifiedLayerNormalization", "SkipSimplifiedLayerNormalization"
)

@dataclass
class IOArguments:
    """
    Arguments to specify input and output folders
    """
    input_folder: str = field(
        metadata={
            "help": "Path of the input folder containing the .onnx models to quantize"
        }
    )
    output_folder: str = field(
        metadata={
            "help": "Path of the output folder where the quantized .onnx models will be saved"
        }
    )

@dataclass
class QuantizationArguments:
    """
    Arguments for quantizing ONNX models
    """

    modes: QuantMode = field(
        default=QUANTIZE_OPTIONS,
        metadata={
            "help": "Quantization mode to use.",
            "choices": QUANTIZE_OPTIONS,
            "nargs": "+",
        },
    )

    # 8-bit quantization
    per_channel: bool = field(
        default=None, metadata={"help": "Whether to quantize weights per channel"}
    )
    reduce_range: bool = field(
        default=None,
        metadata={
            "help": "Whether to quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode"
        },
    )

    # 4-bit quantization
    block_size: int = field(
        default=None,
        metadata={
            "help": "Block size for blockwise quantization. Note: bnb.nn.Linear4bit only uses block_size=64"
        },
    )

    # MatMul4BitsQuantizer
    is_symmetric: bool = field(
        default=True,
        metadata={"help": "Indicate whether to quantize the model symmetrically"},
    )
    accuracy_level: int = field(
        default=None,
        metadata={
            "help": "Accuracy level of the 4-bit quantized MatMul computation. "
            "Refer to the MatMulNBits contrib op's 'accuracy_level' attribute for details "
            "(https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftmatmulnbits)."
        },
    )

    # MatMulBnb4Quantizer
    quant_type: int = field(
        default=MatMulBnb4Quantizer.NF4,
        metadata={
            "help": "Quantization data type. 0: FP4, 1: NF4",
            "choices": [MatMulBnb4Quantizer.FP4, MatMulBnb4Quantizer.NF4],
        },
    )

    op_block_list: List[str] = field(
        default=None,
        metadata={
            "help": "List of operators to exclude from quantization."
            "Can be any standard ONNX operator (see https://onnx.ai/onnx/operators/)"
            "or your custom implemented operators.",
            "nargs": "+",
        },
    )


def get_operators(model: onnx.ModelProto) -> Set[str]:
    operators = set()

    def traverse_graph(graph):
        for node in graph.node:
            operators.add(node.op_type)
            for attr in node.attribute:
                if attr.type == onnx.AttributeProto.GRAPH:
                    traverse_graph(attr.g)

    traverse_graph(model.graph)
    return operators


def quantize_q8(
    model: onnx.ModelProto,
    save_path: str,
    per_channel: bool,
    reduce_range: bool,
    weight_type: QuantType,
    op_block_list: Optional[List[str]]
):
    """
    Quantize the weights of the model from float32 to int8/uint8

    Uses unsigned ints for activation values, signed ints for weights, per
    https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection
    it is faster on most CPU architectures
    """

    op_types_to_quantize = set(IntegerOpsRegistry.keys())
    if op_block_list is not None:
        op_types_to_quantize.difference_update(op_block_list)

    quantizer = ONNXQuantizer(
        model,
        per_channel,
        reduce_range,
        mode=QuantizationMode.IntegerOps,
        static=False,
        weight_qType=weight_type,
        activation_qType=QuantType.QUInt8,  # dynamic activation only supports uint8
        tensors_range=None,
        nodes_to_quantize=[],
        nodes_to_exclude=[],
        op_types_to_quantize=op_types_to_quantize,
        extra_options=dict(
            EnableSubgraph=True,
            MatMulConstBOnly=True,
        ),
    )

    quantizer.quantize_model()
    check_and_save_model(quantizer.model.model, save_path)


def quantize_fp16(
    model: onnx.ModelProto,
    save_path: str,
    op_block_list: Optional[List[str]]
):
    """
    Quantize the weights of the model from float32 to float16
    """

    # Check whether we should disable shape infer:
    # ValueError: Message onnx.ModelProto exceeds maximum protobuf size of 2GB: 2338583841
    disable_shape_infer = model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF

    blocked_ops = set(float16.DEFAULT_OP_BLOCK_LIST)
    if op_block_list is not None:
        blocked_ops.update(op_block_list)

    model_fp16 = float16.convert_float_to_float16(
        model,
        keep_io_types=True,
        disable_shape_infer=disable_shape_infer,
        op_block_list=blocked_ops,
    )
    check_and_save_model(model_fp16, save_path)


def quantize_q4(
    model: onnx.ModelProto,
    save_path: str | None,
    block_size: int,
    is_symmetric: bool,
    accuracy_level: int,
):
    """
    Quantize the weights of the model from float32 to 4-bit int
    """

    quantizer = MatMul4BitsQuantizer(
        model=model,
        block_size=block_size,
        is_symmetric=is_symmetric,
        accuracy_level=accuracy_level,
    )
    quantizer.process()
    if save_path:
        check_and_save_model(quantizer.model.model, save_path)
    return quantizer.model.model


def quantize_bnb4(
    model: onnx.ModelProto,
    save_path: str,
    block_size: int,
    quant_type: int,
):
    """
    Quantize the weights of the model from float32 to 4-bit int using MatMulBnb4Quantizer
    """

    quantizer = MatMulBnb4Quantizer(
        model=model,
        block_size=block_size,
        quant_type=quant_type,
    )
    quantizer.process()
    check_and_save_model(quantizer.model.model, save_path)
    return quantizer.model.model


def quantize(input_folder, output_folder, quantization_args: QuantizationArguments):

    # (Step 1) Validate the arguments
    if not quantization_args.modes:
        raise ValueError("At least one quantization mode must be specified")

    if not os.path.exists(input_folder):
        raise ValueError(f"Input folder {input_folder} does not exist")

    model_names_or_paths = [
        os.path.join(input_folder, file)
        for file in os.listdir(input_folder)
        if file.endswith(".onnx")
    ]
    if not model_names_or_paths:
        raise ValueError(f"No .onnx models found in {input_folder}")

    os.makedirs(output_folder, exist_ok=True)

    # (Step 2) Quantize the models
    for model_path in (progress_models := tqdm(model_names_or_paths)):
        progress_models.set_description(f"Processing {model_path}")

        file_name_without_extension = os.path.splitext(os.path.basename(model_path))[0]

        for mode in (progress := tqdm(quantization_args.modes)):
            progress.set_description(f" - Quantizing to {mode}")
            mode = QuantMode(mode)
            suffix = QUANTIZE_SUFFIX_MAPPING.get(mode, mode.value)
            save_path = os.path.join(
                output_folder,
                f"{file_name_without_extension}_{suffix}.onnx",
            )

            # NOTE: Unfortunately, we need to reload the model for each quantization mode,
            # which is memory inefficient. This is because the quantization functions
            # modify the model in-place, and we need to keep the original model for each mode.
            model = onnx.load_model(model_path)

            if mode == QuantMode.FP16:
                quantize_fp16(
                    model,
                    save_path,
                    quantization_args.op_block_list
                )

            elif mode in (QuantMode.Q4, QuantMode.Q4F16):
                block_size = quantization_args.block_size or 32

                q4_model = quantize_q4(
                    model,
                    save_path=None if mode == QuantMode.Q4F16 else save_path,
                    block_size=block_size,
                    is_symmetric=quantization_args.is_symmetric,
                    accuracy_level=quantization_args.accuracy_level,
                )
                if mode == QuantMode.Q4F16:
                    quantize_fp16(
                        q4_model,
                        save_path,
                        quantization_args.op_block_list,
                    )

            elif mode == QuantMode.BNB4:
                quantize_bnb4(
                    model,
                    save_path,
                    block_size=quantization_args.block_size or 64,
                    quant_type=(
                        quantization_args.quant_type
                        if quantization_args.quant_type is not None
                        else MatMulBnb4Quantizer.NF4
                    ),
                )

            elif mode in (QuantMode.Q8, QuantMode.QI8, QuantMode.QU8):
                if mode == QuantMode.Q8:
                    op_types = get_operators(model)
                    weight_type = (
                        QuantType.QUInt8
                        if any(x in QUINT8_OPS for x in op_types)
                        else QuantType.QInt8
                    )

                elif mode == QuantMode.QI8:
                    weight_type = QuantType.QInt8

                else:  # mode == QuantMode.QU8:
                    weight_type = QuantType.QUInt8

                quantize_q8(
                    model,
                    save_path,
                    per_channel=quantization_args.per_channel,
                    reduce_range=quantization_args.reduce_range,
                    weight_type=weight_type,
                    op_block_list=quantization_args.op_block_list,
                )


def main():
    parser = HfArgumentParser((IOArguments, QuantizationArguments))
    io_args, quantization_args = parser.parse_args_into_dataclasses()
    input_folder = io_args.input_folder
    output_folder = io_args.output_folder
    quantize(input_folder, output_folder, quantization_args)

if __name__ == "__main__":
    main()