73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
from typing import Any, Dict, List
|
|
|
|
import torch
|
|
|
|
from vllm.model_executor.quantization_utils.base import QuantizationConfig
|
|
|
|
|
|
class AWQConfig(QuantizationConfig):
|
|
"""Config class for AWQ.
|
|
|
|
Reference: https://arxiv.org/abs/2306.00978
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
weight_bits: int,
|
|
group_size: int,
|
|
zero_point: bool,
|
|
) -> None:
|
|
self.weight_bits = weight_bits
|
|
self.group_size = group_size
|
|
self.zero_point = zero_point
|
|
|
|
if self.weight_bits != 4:
|
|
raise ValueError(
|
|
"Currently, only 4-bit weight quantization is supported for "
|
|
f"AWQ, but got {self.weight_bits} bits.")
|
|
self.pack_factor = 32 // self.weight_bits
|
|
|
|
def __repr__(self) -> str:
|
|
return (f"AWQConfig(weight_bits={self.weight_bits}, "
|
|
f"group_size={self.group_size}, "
|
|
f"zero_point={self.zero_point})")
|
|
|
|
@classmethod
|
|
def get_name(cls) -> str:
|
|
return "awq"
|
|
|
|
@classmethod
|
|
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
|
|
return [torch.half]
|
|
|
|
@classmethod
|
|
def get_min_capability(cls) -> int:
|
|
# The AWQ kernel only supports Ampere or newer GPUs.
|
|
return 80
|
|
|
|
@classmethod
|
|
def get_config_filenames(cls) -> List[str]:
|
|
return [
|
|
"quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq
|
|
"quantize_config.json", # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq # pylint: disable=line-too-long
|
|
]
|
|
|
|
@classmethod
|
|
def from_config(cls, config: Dict[str, Any]) -> "AWQConfig":
|
|
weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
|
|
group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
|
|
zero_point = cls.get_from_keys(config, ["zero_point"])
|
|
return cls(weight_bits, group_size, zero_point)
|
|
|
|
@classmethod
|
|
def get_packed_tensor_names(cls) -> List[str]:
|
|
return ["qweight", "qzeros"]
|
|
|
|
@classmethod
|
|
def get_transposed_tensor_names(cls) -> List[str]:
|
|
return ["qweight", "qzeros", "scales"]
|
|
|
|
@classmethod
|
|
def get_tp_tensor_names(cls) -> List[str]:
|
|
return ["qweight", "qzeros", "scales"]
|