Skip to main content
`Ctrl`+`K`

#

# Source code for quark.onnx.quantization.config.config

    #
    # Copyright (C) 2023 - 2025 Advanced Micro Devices, Inc. All rights reserved.
    # SPDX-License-Identifier: MIT
    #
    """Quark Quantization Config API for ONNX"""

    from __future__ import annotations

    from dataclasses import dataclass
    from typing import Any

    from .algorithm import AlgoConfig
    from .data_type import DataType
    from .legacy import QuantizationConfig
    from .spec import Int8Spec, QLayerConfig

    @dataclass(eq=True)
    class Config:
        """
        A class that encapsulates comprehensive quantization configurations for a machine learning model, allowing for detailed and hierarchical control over quantization parameters across different model components.

        :param QuantizationConfig global_quant_config: Global quantization configuration applied to the entire model unless overridden at the layer level.
        """

        # Global quantization configuration applied to the entire model unless overridden at the layer level.
        global_quant_config: QuantizationConfig

    # TODO: Move QConfig into quark/shares

    @dataclass(eq=True, init=False)
    class QConfig:
        """
        A class that defines quantization configuration at multiple levels (global, specific layers, specific operation types),
        and provides flexibility for specifying algorithm settings.

        :param QLayerConfig global_config: Global quantization configuration applied to all layers unless overridden.
        :param Dict[DataType, List[str]] specific_layer_config: Dictionary mapping specific layer names to their quantization
            configuration. Overrides ``global_config`` for those layers. Default is ``None``.
        :param Dict[Optional[DataType], List[str]] layer_type_config: Dictionary mapping layer types (e.g., Conv, Gemm) to
            quantization configurations. Overrides ``global_config`` for those operation types. Default is ``None``.
        :param List[Union[str, List[Tuple[List[str]]]]] exclude: List of nodes or subgraphs excluded from quantization. Default is ``None``.
        :param List[AlgoConfig] algo_config: Algorithm configuration(s), such as CLE, SmoothQuant,
            or AdaRound. Can be a list of algorithm configurations. Default is ``None``.
        :param bool use_external_data_format: Whether to use ONNX external data format when saving the quantized model.
            Default is ``False``.
            advanced customization and extension.
        :param Dict[str, Any] extra_options: Dictionary for additional options. Default is ``None``.
        """

        global_config: QLayerConfig = QLayerConfig(activation=Int8Spec(), weight=Int8Spec())
        specific_layer_config: dict[DataType, list[str]] | None
        layer_type_config: dict[DataType | None, list[str]] | None
        exclude: list[str | list[tuple[list[str]]]] | None
        algo_config: list[AlgoConfig] | None
        use_external_data_format: bool

        def __init__(
            self,
            global_config: QLayerConfig,
            specific_layer_config: dict[DataType, list[str]] | None = None,
            layer_type_config: dict[DataType | None, list[str]] | None = None,
            exclude: list[str | list[tuple[list[str]]]] | None = None,
            algo_config: list[AlgoConfig] | None = None,
            use_external_data_format: bool = False,
            **kwargs: dict[str, Any],
        ):
            self.global_config = global_config
            self.specific_layer_config = specific_layer_config or {}
            self.layer_type_config = layer_type_config or {}
            self.exclude = exclude or []
            self.algo_config = algo_config or []  # type: ignore
            self.use_external_data_format = use_external_data_format
            self.extra_options = kwargs

        @staticmethod
        def get_default_config(config_name: str) -> Config:
            """
            Retrieve the default quantization configuration by name.

            This function looks up the provided `config_name` in the
            `DefaultConfigMapping`. If a match is found, it returns a
            `Config` object with the corresponding global quantization
            configuration. Otherwise, it raises a ValueError.

            Args:
                config_name (str): The name of the default configuration
                to look up like XINT8.

            Returns:
                Config: A configuration object containing the default
                quantization settings.

            Raises:
                ValueError: If the provided `config_name` is not found
                in `DefaultConfigMapping`.
            """
            from . import DefaultConfigMapping

            if config_name in DefaultConfigMapping:
                return Config(global_quant_config=DefaultConfigMapping[config_name])
            else:
                raise ValueError("The quantization config is invalid.")