diff --git a/docker/quark-quantization/amd/Dockerfile b/docker/quark-quantization/amd/Dockerfile new file mode 100644 index 00000000..218fffb5 --- /dev/null +++ b/docker/quark-quantization/amd/Dockerfile @@ -0,0 +1,39 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# Licensed under the MIT License. + +FROM rocm/dev-ubuntu-22.04:6.1 + +LABEL maintainer="Hugging Face" + +ARG DEBIAN_FRONTEND=noninteractive + +RUN apt-get update && apt-get install -y --no-install-recommends \ + sudo \ + python3.10 \ + python3.10-dev \ + python3-pip \ + git \ + wget \ + unzip \ + libsndfile1-dev \ + tesseract-ocr \ + espeak-ng \ + rocthrust-dev \ + hipsparse-dev \ + hipblas-dev && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* && \ + update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1 && \ + python -m pip install -U pip + +RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.0 --no-cache-dir + +WORKDIR /quark +RUN wget -O quark.zip https://www.xilinx.com/bin/public/openDownload?filename=quark-0.1.0+a9827f5.zip && \ + unzip quark.zip -d quark && \ + pip install quark/quark-0.1.0+a9827f5/quark-0.1.0+a9827f5-py39-none-any.whl && \ + rm -rf quark.zip quark + +RUN python -c "import quark.torch.kernel" + +ENTRYPOINT ["bash"] diff --git a/optimum/amd/quantization/__init__.py b/optimum/amd/quantization/__init__.py new file mode 100644 index 00000000..afddd61b --- /dev/null +++ b/optimum/amd/quantization/__init__.py @@ -0,0 +1,8 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# Licensed under the MIT License. + +from .quark import ( + AutoQuantizationConfig, + QuarkQuantizationConfig, + QuarkQuantizer, +) diff --git a/optimum/amd/quantization/quark/algo_config_constants.py b/optimum/amd/quantization/quark/algo_config_constants.py new file mode 100644 index 00000000..24283625 --- /dev/null +++ b/optimum/amd/quantization/quark/algo_config_constants.py @@ -0,0 +1,195 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# Licensed under the MIT License. + + +ALGO_CONFIG_PARAMS = { + "llama": { + "scaling_layers": [ + { + "prev_op": "input_layernorm", + "layers": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + "inp": "self_attn.q_proj", + "module2inspect": "self_attn", + "has_kwargs": True, + "help": "attention input", + }, + { + "prev_op": "self_attn.v_proj", + "layers": ["self_attn.o_proj"], + "inp": "self_attn.o_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "attention out, Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696, if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape", + "condition": "module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape", + }, + { + "prev_op": "post_attention_layernorm", + "layers": ["mlp.gate_proj", "mlp.up_proj"], + "inp": "mlp.gate_proj", + "module2inspect": "mlp", + "has_kwargs": False, + "help": "linear 1", + }, + { + "prev_op": "mlp.up_proj", + "layers": ["mlp.down_proj"], + "inp": "mlp.down_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "linear 2", + }, + ], + "inside_layer_modules": [ + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.q_proj", + "self_attn.o_proj", + "mlp.up_proj", + "mlp.gate_proj", + "mlp.down_proj", + ], + "model_decoder_layers": "model.layers", + "embedding_layers": ["model.embed_tokens"], + }, + "mistral": { + "scaling_layers": [ + { + "prev_op": "input_layernorm", + "layers": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + "inp": "self_attn.q_proj", + "module2inspect": "self_attn", + "has_kwargs": True, + "help": "attention input", + }, + { + "prev_op": "self_attn.v_proj", + "layers": ["self_attn.o_proj"], + "inp": "self_attn.o_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "attention out, Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696, if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape", + "condition": "module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape", + }, + { + "prev_op": "post_attention_layernorm", + "layers": ["mlp.gate_proj", "mlp.up_proj"], + "inp": "mlp.gate_proj", + "module2inspect": "mlp", + "has_kwargs": False, + "help": "linear 1", + }, + { + "prev_op": "mlp.up_proj", + "layers": ["mlp.down_proj"], + "inp": "mlp.down_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "linear 2", + }, + ], + "inside_layer_modules": [ + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.q_proj", + "self_attn.o_proj", + "mlp.up_proj", + "mlp.gate_proj", + "mlp.down_proj", + ], + "model_decoder_layers": "model.layers", + "embedding_layers": ["model.embed_tokens"], + }, + "opt": { + "scaling_layers": [ + { + "prev_op": "self_attn_layer_norm", + "layers": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + "inp": "self_attn.q_proj", + "module2inspect": "self_attn", + "has_kwargs": True, + "help": "attention input", + }, + { + "prev_op": "self_attn.v_proj", + "layers": ["self_attn.out_proj"], + "inp": "self_attn.out_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "attention out", + }, + { + "prev_op": "final_layer_norm", + "layers": ["fc1"], + "inp": "fc1", + "module2inspect": None, + "has_kwargs": False, + "help": "linear 1", + }, + { + "prev_op": "fc1", + "layers": ["fc2"], + "inp": "fc2", + "module2inspect": None, + "has_kwargs": False, + "help": "linear 2", + }, + ], + "inside_layer_modules": [ + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.q_proj", + "self_attn.out_proj", + "fc1", + "fc2", + ], + "model_decoder_layers": "model.decoder.layers", + "embedding_layers": ["model.decoder.embed_tokens", "model.decoder.embed_positions"], + }, + "qwen2": { + "scaling_layers": [ + { + "prev_op": "input_layernorm", + "layers": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + "inp": "self_attn.q_proj", + "module2inspect": "self_attn", + "has_kwargs": True, + "help": "attention input", + }, + { + "prev_op": "self_attn.v_proj", + "layers": ["self_attn.o_proj"], + "inp": "self_attn.o_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "attention out, Please refer to https://github.com/mit-han-lab/llm-awq/pull/67#issue-1850622696, if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape", + }, + { + "prev_op": "post_attention_layernorm", + "layers": ["mlp.gate_proj", "mlp.up_proj"], + "inp": "mlp.gate_proj", + "module2inspect": "mlp", + "has_kwargs": False, + "help": "linear 1", + }, + { + "prev_op": "mlp.up_proj", + "layers": ["mlp.down_proj"], + "inp": "mlp.down_proj", + "module2inspect": None, + "has_kwargs": False, + "help": "linear 2", + }, + ], + "inside_layer_modules": [ + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.q_proj", + "self_attn.o_proj", + "mlp.up_proj", + "mlp.gate_proj", + "mlp.down_proj", + ], + "model_decoder_layers": "model.layers", + "embedding_layers": ["model.embed_tokens"], + }, +} diff --git a/optimum/amd/quantization/quark/configuration.py b/optimum/amd/quantization/quark/configuration.py new file mode 100644 index 00000000..b93c7787 --- /dev/null +++ b/optimum/amd/quantization/quark/configuration.py @@ -0,0 +1,367 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# Licensed under the MIT License. +"""Configuration classes for quantization with AMD Quark.""" + +from enum import Enum +from typing import Dict, List, Optional + +from quark.torch.quantization.config.config import ( + AlgoConfig, + AWQConfig, + Config, + GPTQConfig, + QuantizationConfig, + QuantizationSpec, + SmoothQuantConfig, +) +from quark.torch.quantization.config.custom_config import ( + DEFAULT_AWQ_CONFIG, + DEFAULT_FLOAT16_CONFIG, + DEFAULT_GPTQ_CONFIG, + DEFAULT_SMOOTH_QUANT_CONFIG, + DEFAULT_W_FP8_A_FP8_OFP8_PER_TENSOR_CONFIG, + DEFAULT_W_FP8_A_FP8_PER_TENSOR_CONFIG, + DEFAULT_W_INT4_PER_CHANNEL_CONFIG, + DEFAULT_W_INT4_PER_GROUP_SYM_CONFIG, + DEFAULT_W_INT4_PER_TENSOR_CONFIG, + DEFAULT_W_INT8_A_INT8_PER_TENSOR_CONFIG, + DEFAULT_W_INT8_A_INT8_PER_TENSOR_DYNAMIC_CONFIG, + DEFAULT_W_UINT4_A_BFLOAT16_PER_GROUP_CONFIG, + DEFAULT_W_UINT4_PER_GROUP_CONFIG, + FP8_PER_TENSOR_SPEC, +) + +from .algo_config_constants import ALGO_CONFIG_PARAMS + + +QuarkQuantizationConfig = Config + +__all__ = [ + "QuarkQuantizationConfig", + "QuantizationSpec", + "SmoothQuantConfig", + "AWQConfig", + "GPTQConfig", + "AlgoConfig", +] + + +class KVCacheDType(Enum): + FP8 = "fp8" + + +class AutoQuantizationConfig: + @staticmethod + def _apply_fp8_config(layer_quant_config): + KV_CACHE_CFG = { + "*.v_proj": QuantizationConfig( + input_tensors=FP8_PER_TENSOR_SPEC, output_tensors=FP8_PER_TENSOR_SPEC, weight=FP8_PER_TENSOR_SPEC + ), + "*.k_proj": QuantizationConfig( + input_tensors=FP8_PER_TENSOR_SPEC, output_tensors=FP8_PER_TENSOR_SPEC, weight=FP8_PER_TENSOR_SPEC + ), + } + return {**layer_quant_config, **KV_CACHE_CFG} + + @staticmethod + def _validate_kv_cache_dtype(kv_cache_dtype: str): + if kv_cache_dtype.lower() != "fp8": + raise ValueError(f"Invalid value for kv_cache_dtype: {kv_cache_dtype}. Expected 'fp8'.") + + @staticmethod + def _validate_model_type(model_type: str): + SUPPORTED_MODEL_TYPES = ["llama", "mistral", "opt", "qwen2"] + + if model_type not in SUPPORTED_MODEL_TYPES: + raise ValueError( + f"Invalid value for model_type for AutoQuantizationConfig: {model_type}. Expected one of {SUPPORTED_MODEL_TYPES}." + ) + + return model_type + + @staticmethod + def w_fp8_a_fp8( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_FP8_A_FP8_PER_TENSOR_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_fp8_a_fp8_o_fp8( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_FP8_A_FP8_OFP8_PER_TENSOR_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_int4_per_tensor( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_INT4_PER_TENSOR_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_int4_per_channel_sym( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_INT4_PER_CHANNEL_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_int4_per_group_sym( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_INT4_PER_GROUP_SYM_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_uint4_per_group_asym_awq( + model_type: str, + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + model_type = AutoQuantizationConfig._validate_model_type(model_type) + algo_config_info = ALGO_CONFIG_PARAMS[model_type] + + quant_config = DEFAULT_AWQ_CONFIG + quant_config.algo_config.scaling_layers = algo_config_info["scaling_layers"] + quant_config.algo_config.model_decoder_layers = algo_config_info["model_decoder_layers"] + quant_config.algo_config.embedding_layers = algo_config_info["embedding_layers"] + return QuarkQuantizationConfig( + global_quant_config=quant_config.global_quant_config, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + algo_config=quant_config.algo_config, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_uint4_per_group_asym_smoothquant( + model_type: str, + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + model_type = AutoQuantizationConfig._validate_model_type(model_type) + algo_config_info = ALGO_CONFIG_PARAMS[model_type] + + quant_config = DEFAULT_SMOOTH_QUANT_CONFIG + quant_config.algo_config.scaling_layers = algo_config_info["scaling_layers"] + quant_config.algo_config.model_decoder_layers = algo_config_info["model_decoder_layers"] + quant_config.algo_config.embedding_layers = algo_config_info["embedding_layers"] + return QuarkQuantizationConfig( + global_quant_config=quant_config.global_quant_config, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + algo_config=quant_config.algo_config, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_uint4_per_group_asym_gptq( + model_type: str, + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + model_type = AutoQuantizationConfig._validate_model_type(model_type) + algo_config_info = ALGO_CONFIG_PARAMS[model_type] + + quant_config = DEFAULT_GPTQ_CONFIG + quant_config.algo_config.inside_layer_modules = algo_config_info["inside_layer_modules"] + quant_config.algo_config.model_decoder_layers = algo_config_info["model_decoder_layers"] + quant_config.algo_config.embedding_layers = algo_config_info["embedding_layers"] + return QuarkQuantizationConfig( + global_quant_config=quant_config.global_quant_config, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + algo_config=quant_config.algo_config, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_uint4_per_group_asym( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_UINT4_PER_GROUP_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_uint4_a_bfloat16_per_group_asym( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_UINT4_A_BFLOAT16_PER_GROUP_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_int8_a_int8_per_tensor_sym( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_INT8_A_INT8_PER_TENSOR_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def w_int8_a_int8_per_tensor_sym_dynamic( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_W_INT8_A_INT8_PER_TENSOR_DYNAMIC_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) + + @staticmethod + def float16( + layer_type_quant_config: Dict[str, QuantizationConfig] = {}, + layer_quant_config: Dict[str, QuantizationConfig] = {}, + exclude: List[str] = [], + pre_quant_opt_config: Optional[QuantizationConfig] = None, + kv_cache_dtype: Optional[KVCacheDType] = None, + ): + if kv_cache_dtype: + AutoQuantizationConfig._validate_kv_cache_dtype(kv_cache_dtype) + layer_quant_config = AutoQuantizationConfig._apply_fp8_config(layer_quant_config) + + return QuarkQuantizationConfig( + global_quant_config=DEFAULT_FLOAT16_CONFIG, + layer_type_quant_config=layer_type_quant_config, + layer_quant_config=layer_quant_config, + exclude=exclude, + pre_quant_opt_config=pre_quant_opt_config, + ) diff --git a/optimum/amd/quantization/quark/quantizer.py b/optimum/amd/quantization/quark/quantizer.py new file mode 100644 index 00000000..27a0740b --- /dev/null +++ b/optimum/amd/quantization/quark/quantizer.py @@ -0,0 +1,219 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# Licensed under the MIT License. +"""AMD Quark Quantizer""" + +import logging +from typing import Callable, Dict, List, Optional, Union + +import torch +from datasets import Dataset, load_dataset +from quark.torch import ModelExporter, ModelQuantizer +from quark.torch.export.config.custom_config import DEFAULT_EXPORTER_CONFIG, EMPTY_EXPORTER_CONFIG +from torch.utils.data import DataLoader + +from optimum.exporters import TasksManager +from optimum.quantization_base import OptimumQuantizer + +from .configuration import QuarkQuantizationConfig + + +logger = logging.getLogger(__name__) + + +class QuarkQuantizer(OptimumQuantizer): + """ + Handles the quantization process for models shared on huggingface.co/models. + """ + + def __init__( + self, + model: torch.nn.Module, + quantization_config: QuarkQuantizationConfig, + model_name_or_path: Optional[str], + ): + super().__init__() + self.model_name_or_path = model_name_or_path + + self.model = model + self.model_type = model.config.model_type + self.model_dtype = next(model.parameters()).dtype + + # Initialize the quantizer + self.quantizer = ModelQuantizer(quantization_config) + + @classmethod + def from_pretrained( + cls, + model_name_or_path: str, + quantization_config: QuarkQuantizationConfig, + subfolder: str = "", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + trust_remote_code: bool = False, + force_download: bool = False, + local_files_only: bool = False, + use_auth_token: Optional[Union[bool, str]] = None, + device_map: Optional[Union[Dict, str, torch.device]] = None, + **model_kwargs, + ): + """ + Loads the QuarkQuantizer and model. + + Args: + model_name_or_path (`Union[str, Path]`): + Can be either the model id of a model repo on the Hugging Face Hub, or a path to a local directory + containing a model. + subfolder (`str`, defaults to `""`): + In case the model files are located inside a subfolder of the model directory / repo on the Hugging + Face Hub, you can specify the subfolder name here. + revision (`Optional[str]`, *optional*, defaults to `None`): + Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id. + cache_dir (`Optional[str]`, *optional*): + Path to a directory in which a downloaded pretrained model weights have been cached if the standard cache should not be used. + trust_remote_code (`bool`, defaults to `False`): + Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories + you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the + model repository. + force_download (`bool`, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + local_files_only (`Optional[bool]`, defaults to `False`): + Whether or not to only look at local files (i.e., do not try to download the model). + use_auth_token (`Optional[str]`, defaults to `None`): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `transformers-cli login` (stored in `~/.huggingface`). + """ + + # TODO: fix + # task = TasksManager.infer_task_from_model(model_name_or_path) + task = "text-generation" + + device = None + if not isinstance(device_map, dict) and device_map not in ["auto", "balanced"]: + device = device_map + device_map = None + + model = TasksManager.get_model_from_task( + task, + model_name_or_path, + subfolder=subfolder, + revision=revision, + cache_dir=cache_dir, + use_auth_token=use_auth_token, + local_files_only=local_files_only, + force_download=force_download, + trust_remote_code=trust_remote_code, + device_map=device_map, + device=device, + framework="pt", + **model_kwargs, + ).eval() + + return cls(model, quantization_config, model_name_or_path) + + def quantize( + self, + dataloader: Optional[Dataset] = None, + ) -> torch.nn.Module: + """_summary_ + + Args: + dataloader (`Optional[Union[DataLoader[torch.Tensor], DataLoader[List[Dict[str, torch.Tensor]]], + DataLoader[Dict[str, torch.Tensor]]]]`, defaults to `None`): + The DataLoader providing data that the quantization process will use for calibration. This can be a simple DataLoader returning + tensors, or a more complex structure returning either a list of dictionaries or a dictionary of tensors. + + Returns: + torch.nn.Module: Quantized model + """ + if not self.quantizer.is_all_dynamic and dataloader is None: + raise ValueError("A calibration dataset is required for the quantization method.") + + self.model = self.quantizer.quantize_model(self.model, dataloader) + + return self.model + + def save_pretrained(self, save_directory: str, no_weight_matrix_merge=False): + """ + Save the quantized model to the specified directory. + + Args: + save_directory (`str`): + Directory to save the quantized model to. + no_weight_matrix_merge (`bool`, defaults to `False`): + Whether to merge weight matrix when dump quantized model + """ + if self.model_type != "llama": + raise ValueError("Only models with dtype `llama` can be saved.") + model = self.quantizer.freeze(self.model) + + with torch.inference_mode(): + export_config = EMPTY_EXPORTER_CONFIG if no_weight_matrix_merge else DEFAULT_EXPORTER_CONFIG + + exporter = ModelExporter(config=export_config, export_dir=save_directory) + exporter.export_model_info(model, self.model_type, self.model_dtype, export_type="vllm-adopt") + + model.config.save_pretrained(save_directory) + + def get_calibration_data( + self, + dataset_name: str, + num_samples: int = 100, + dataset_config_name: Optional[str] = None, + dataset_split: Optional[str] = None, + preprocess_function: Optional[Callable] = None, + preprocess_batch: bool = True, + seed: int = 2016, + token: Optional[Union[bool, str]] = None, + batch_size: int = 1, + ) -> Union[ + DataLoader[torch.Tensor], DataLoader[List[Dict[str, torch.Tensor]]], DataLoader[Dict[str, torch.Tensor]] + ]: + """ + Creates the calibration `datasets.Dataset` to use for the post-training static quantization calibration step. + + Args: + dataset_name (`str`): + The dataset repository name on the Hugging Face Hub or path to a local directory containing data files + to load to use for the calibration step. + num_samples (`int`, defaults to 100): + The maximum number of samples composing the calibration dataset. + dataset_config_name (`Optional[str]`, defaults to `None`): + The name of the dataset configuration. + dataset_split (`Optional[str]`, defaults to `None`): + Which split of the dataset to use to perform the calibration step. + preprocess_function (`Optional[Callable]`, defaults to `None`): + Processing function to apply to each example after loading dataset. + preprocess_batch (`bool`, defaults to `True`): + Whether the `preprocess_function` should be batched. + seed (`int`, defaults to 2016): + The random seed to use when shuffling the calibration dataset. + use_auth_token (`Optional[Union[bool,str]]`, defaults to `None`): + Deprecated. Please use the `token` argument instead. + token (`Optional[Union[bool,str]]`, defaults to `None`): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `huggingface_hub.constants.HF_TOKEN_PATH`). + + Returns: + The calibration `datasets.Dataset` to use for the post-training static quantization calibration + step. + """ + calib_dataset = load_dataset( + dataset_name, + name=dataset_config_name, + split=dataset_split, + token=token, + ) + + if num_samples is not None: + num_samples = min(num_samples, len(calib_dataset)) + calib_dataset = calib_dataset.shuffle(seed=seed).select(range(num_samples)) + + if preprocess_function is not None: + processed_calib_dataset = calib_dataset.map(preprocess_function, batched=preprocess_batch) + else: + processed_calib_dataset = calib_dataset + + dataloader = DataLoader(processed_calib_dataset, batch_size=batch_size, shuffle=False) + + return dataloader