| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| from typing import TYPE_CHECKING, Optional |
|
|
| from .base import HfQuantizer |
|
|
|
|
| if TYPE_CHECKING: |
| from ..modeling_utils import PreTrainedModel |
|
|
| from ..utils import is_accelerate_available, is_eetq_available, is_torch_available, logging |
| from .quantizers_utils import get_module_from_name |
|
|
|
|
| if is_torch_available(): |
| import torch |
|
|
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| class EetqHfQuantizer(HfQuantizer): |
| """ |
| 8-bit quantization from EETQ quantization method: |
| before loading: converts transformer layers into W8A16Linear during loading: load 16bit weight and pass to the |
| layer object after: quantizes individual weights in Linear8bitLt into 8bit at first .cuda() call |
| """ |
|
|
| requires_parameters_quantization = True |
| requires_calibration = False |
|
|
| required_packages = ["eetq", "accelerate"] |
|
|
| def __init__(self, quantization_config, **kwargs): |
| super().__init__(quantization_config, **kwargs) |
| self.quantization_config = quantization_config |
|
|
| def validate_environment(self, *args, **kwargs): |
| if not is_eetq_available(): |
| raise ImportError( |
| "Using `eetq` 8-bit quantization requires eetq." |
| "Please install the latest version of eetq from : https://github.com/NetEase-FuXi/EETQ" |
| ) |
|
|
| try: |
| import eetq |
| except ImportError as exc: |
| if "shard_checkpoint" in str(exc): |
| |
| |
| |
| raise ImportError( |
| "You are using a version of EETQ that is incompatible with the current transformers version. " |
| "Either downgrade transformers to <= v4.46.3 or, if available, upgrade EETQ to > v1.0.0." |
| ) from exc |
| else: |
| raise |
|
|
| if not is_accelerate_available(): |
| raise ImportError("Loading an EETQ quantized model requires accelerate (`pip install accelerate`)") |
|
|
| if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): |
| raise ValueError( |
| "Converting into 8-bit weights from tf/flax weights is currently not supported, please make" |
| " sure the weights are in PyTorch format." |
| ) |
|
|
| if not torch.cuda.is_available(): |
| raise RuntimeError("No GPU found. A GPU is needed for quantization.") |
|
|
| device_map = kwargs.get("device_map") |
| if device_map is None: |
| logger.warning_once( |
| "You have loaded an EETQ model on CPU and have a CUDA device available, make sure to set " |
| "your model on a GPU device in order to run your model." |
| ) |
| elif device_map is not None: |
| if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()): |
| raise ValueError( |
| "You are attempting to load an EETQ model with a device_map that contains a CPU or disk device." |
| " This is not supported. Please remove the CPU or disk device from the device_map." |
| ) |
|
|
| def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype": |
| if dtype is None: |
| dtype = torch.float16 |
| logger.info( |
| "Overriding dtype=%s with `dtype=torch.float16` due to " |
| "requirements of `eetq` to enable model loading in 8-bit. " |
| "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass" |
| " dtype=torch.float16 to remove this warning.", |
| dtype, |
| ) |
| elif dtype != torch.float16: |
| logger.info("We suggest you to set `dtype=torch.float16` for better efficiency with EETQ.") |
| return dtype |
|
|
| def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: |
| from eetq import EetqLinear |
|
|
| module, tensor_name = get_module_from_name(model, param_name) |
|
|
| if isinstance(module, EetqLinear): |
| if self.pre_quantized or tensor_name == "bias": |
| return False |
| else: |
| return True |
| return False |
|
|
| def create_quantized_param( |
| self, |
| model: "PreTrainedModel", |
| param_value: "torch.Tensor", |
| param_name: str, |
| target_device: "torch.device", |
| **kwargs, |
| ): |
| from eetq import EetqLinear, quantize_and_preprocess_weights |
|
|
| module, tensor_name = get_module_from_name(model, param_name) |
| new_value, weight_scale = quantize_and_preprocess_weights(param_value) |
|
|
| |
| if isinstance(module, EetqLinear): |
| if self.pre_quantized or tensor_name == "bias": |
| if tensor_name == "weight" and param_value.dtype != torch.int8: |
| raise ValueError("Expect quantized weights but got an unquantized weight") |
| else: |
| if tensor_name == "weight_scale": |
| raise ValueError("Expect unquantized weights but got a quantized weight_scale") |
|
|
| module._buffers[tensor_name] = new_value.to(target_device) |
| module.register("weight_scales", weight_scale.to(target_device)) |
|
|
| def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): |
| return model |
|
|
| def _process_model_before_weight_loading( |
| self, |
| model: "PreTrainedModel", |
| keep_in_fp32_modules: Optional[list[str]] = None, |
| **kwargs, |
| ): |
| from ..integrations import replace_with_eetq_linear |
|
|
| self.modules_to_not_convert = self.get_modules_to_not_convert( |
| model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules |
| ) |
|
|
| model = replace_with_eetq_linear( |
| model, |
| modules_to_not_convert=self.modules_to_not_convert, |
| quantization_config=self.quantization_config, |
| pre_quantized=self.pre_quantized, |
| ) |
|
|
| model.config.quantization_config = self.quantization_config |
|
|
| def is_serializable(self, safe_serialization=None): |
| return True |
|
|
| @property |
| def is_trainable(self) -> bool: |
| return True |
|
|
