Aluode/PerceptionLabPortable
0
1# Copyright 2024 The HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import importlib15from collections import defaultdict16from functools import cached_property17from typing import TYPE_CHECKING, Optional, Union18 19from packaging import version20 21from .base import HfQuantizer22from .quantizers_utils import get_module_from_name23 24 25if TYPE_CHECKING:26 from ..modeling_utils import PreTrainedModel27 28from ..utils import (29 ACCELERATE_MIN_VERSION,30 is_accelerate_available,31 is_bitsandbytes_available,32 is_torch_available,33 is_torch_hpu_available,34 is_torch_npu_available,35 is_torch_xpu_available,36 logging,37)38 39 40if is_torch_available():41 import torch42 43 from ..pytorch_utils import Conv1D44 45logger = logging.get_logger(__name__)46 47 48class Bnb4BitHfQuantizer(HfQuantizer):49 """50 4-bit quantization from bitsandbytes.py quantization method:51 before loading: converts transformer layers into Linear4bit during loading: load 16bit weight and pass to the52 layer object after: quantizes individual weights in Linear4bit into 4bit at the first .cuda() call53 saving:54 from state dict, as usual; saves weights and `quant_state` components55 loading:56 need to locate `quant_state` components and pass to Param4bit constructor57 """58 59 use_keep_in_fp32_modules = True60 requires_parameters_quantization = True61 requires_calibration = False62 63 required_packages = ["bitsandbytes", "accelerate"]64 65 def __init__(self, quantization_config, **kwargs):66 super().__init__(quantization_config, **kwargs)67 68 if self.quantization_config.llm_int8_skip_modules is not None:69 self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules70 71 # This describes the additional items that are saved on the state dict (on the params themselves)72 self.bnb_keys = [73 f"quant_state.bitsandbytes__{self.quantization_config.bnb_4bit_quant_type}",74 "absmax",75 "quant_map",76 ]77 if self.quantization_config.bnb_4bit_use_double_quant:78 self.bnb_keys.extend(["nested_absmax", "nested_quant_map"])79 80 def validate_environment(self, *args, **kwargs):81 if not is_accelerate_available():82 raise ImportError(83 f"Using `bitsandbytes` 4-bit quantization requires Accelerate: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`"84 )85 if not is_bitsandbytes_available(check_library_only=True):86 raise ImportError(87 "Using `bitsandbytes` 4-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`"88 )89 if not is_torch_available():90 raise ImportError(91 "The bitsandbytes library requires PyTorch but it was not found in your environment. "92 "You can install it with `pip install torch`."93 )94 # `bitsandbytes` versions older than 0.43.1 eagerly require CUDA at import time,95 # so those versions of the library are practically only available when CUDA is too.96 if version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.43.1"):97 if not torch.cuda.is_available():98 raise ImportError(99 "The installed version of bitsandbytes (<0.43.1) requires CUDA, but CUDA is not available. "100 "You may need to install PyTorch with CUDA support or upgrade bitsandbytes to >=0.43.1."101 )102 103 from ..integrations import validate_bnb_backend_availability104 from ..utils import is_bitsandbytes_multi_backend_available105 106 bnb_multibackend_is_enabled = is_bitsandbytes_multi_backend_available()107 validate_bnb_backend_availability(raise_exception=True)108 109 if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):110 raise ValueError(111 "Converting into 4-bit or 8-bit weights from tf/flax weights is currently not supported, please make"112 " sure the weights are in PyTorch format."113 )114 115 device_map = kwargs.get("device_map")116 if (117 device_map is not None118 and isinstance(device_map, dict)119 and not self.quantization_config.llm_int8_enable_fp32_cpu_offload120 ):121 device_map_without_lm_head = {122 key: device_map[key] for key in device_map if key not in self.modules_to_not_convert123 }124 if set(device_map.values()) == {"cpu"} and bnb_multibackend_is_enabled:125 pass126 elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values():127 raise ValueError(128 "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the "129 "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules "130 "in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to "131 "`from_pretrained`. Check "132 "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu "133 "for more details. "134 )135 136 def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":137 if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"):138 from accelerate.utils import CustomDtype139 140 if target_dtype != torch.int8:141 logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization")142 return CustomDtype.INT4143 else:144 raise ValueError(145 "You are using `device_map='auto'` on a 4bit loaded version of the model. To automatically compute"146 " the appropriate device map, you should upgrade your `accelerate` library,"147 "`pip install --upgrade accelerate` or install it from source to support fp4 auto device map"148 "calculation. You may encounter unexpected behavior, or pass your own device map"149 )150 151 def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:152 return [k for k in unexpected_keys if not any(k.endswith(x) for x in self.bnb_keys)]153 154 def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:155 import bitsandbytes as bnb156 157 # They are on the params themselves, so we cannot easily extract the module from the name158 if any(param_name.endswith(x) for x in self.bnb_keys):159 return True160 module, name = get_module_from_name(model, param_name)161 return isinstance(module, bnb.nn.Linear4bit) and name != "bias"162 163 def get_param_name(self, param_name: str) -> str:164 """165 Get the right param_name in order to get the module associated with the param.166 This is useful for quantized stats lile absmax or quant_map as we need to update the param_name to get the module as they are stored in ...weight.absmax.167 """168 if self.pre_quantized:169 # We need to get the param name of quantized weights and not its components. Otherwise, we won't be able to get the nn.Module associated.170 if any(param_name.endswith(x) for x in self.bnb_keys):171 param_name = (172 param_name.rsplit(".", 1)[0] if "quant_state." not in param_name else param_name.rsplit(".", 2)[0]173 )174 return param_name175 176 def create_quantized_param(177 self,178 model: "PreTrainedModel",179 param_value: "torch.Tensor",180 param_name: str,181 target_device: "torch.device",182 **kwargs,183 ):184 import bitsandbytes as bnb185 186 full_name = param_name187 188 # update param name to get the weights instead of the quantized stats189 param_name = self.get_param_name(param_name)190 module, tensor_name = get_module_from_name(model, param_name)191 192 # `torch.Tensor.to(<int num>)` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)).193 if isinstance(target_device, int) and is_torch_npu_available():194 target_device = f"npu:{target_device}"195 196 # construct `new_value` for the module._parameters[tensor_name]197 if self.pre_quantized:198 module_name = param_name.rsplit(".", 1)[0]199 # Save the states for later quantization when they are all gathered200 if not hasattr(self, "param_quant_stats"):201 self.param_quant_stats = defaultdict(dict)202 self.param_quant_stats[module_name].update({full_name: param_value})203 204 # We are ready for quantization in this case (note, the +1 is for the weight itself)205 if len(self.param_quant_stats[module_name]) == len(self.bnb_keys) + 1:206 param_kwargs = {}207 if self.is_bnb_supports_quant_storage_module:208 param_kwargs["module"] = module209 210 weight = self.param_quant_stats[module_name].pop(f"{module_name}.weight")211 new_value = bnb.nn.Params4bit.from_prequantized(212 data=weight,213 quantized_stats=self.param_quant_stats[module_name],214 requires_grad=False,215 device=target_device,216 **param_kwargs,217 )218 # Set it219 module._parameters[tensor_name] = new_value220 # Delete the states221 del self.param_quant_stats[module_name]222 else:223 new_value = param_value.to("cpu")224 old_value = getattr(module, tensor_name)225 226 # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.227 # Since weights are saved in the correct "orientation", we skip transposing when loading.228 if issubclass(module.source_cls, Conv1D):229 new_value = new_value.T230 231 kwargs = old_value.__dict__232 kwargs.pop("_is_hf_initialized", None)233 new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device)234 235 module._parameters[tensor_name] = new_value236 237 # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.adjust_max_memory238 def adjust_max_memory(self, max_memory: dict[str, Union[int, str]]) -> dict[str, Union[int, str]]:239 # need more space for buffers that are created during quantization240 max_memory = {key: val * 0.90 for key, val in max_memory.items()}241 return max_memory242 243 # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.update_dtype244 def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":245 if dtype is None:246 # We force the `dtype` to be float16, this is a requirement from `bitsandbytes`247 logger.info(248 "Overriding dtype=%s with `dtype=torch.float16` due to "249 "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. "250 "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"251 " dtype=torch.float16 to remove this warning.",252 dtype,253 )254 dtype = torch.float16255 return dtype256 257 def update_device_map(self, device_map):258 if device_map is None:259 if torch.cuda.is_available():260 device_map = {"": torch.cuda.current_device()}261 elif is_torch_npu_available():262 device_map = {"": f"npu:{torch.npu.current_device()}"}263 elif is_torch_hpu_available():264 device_map = {"": f"hpu:{torch.hpu.current_device()}"}265 elif is_torch_xpu_available():266 device_map = {"": torch.xpu.current_device()}267 else:268 device_map = {"": "cpu"}269 logger.info(270 "The device_map was not initialized. "271 f"Setting device_map to {device_map}. "272 "If you want to use the model for inference, please set device_map ='auto' "273 )274 return device_map275 276 # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_before_weight_loading277 def _process_model_before_weight_loading(278 self,279 model: "PreTrainedModel",280 device_map,281 keep_in_fp32_modules: Optional[list[str]] = None,282 **kwargs,283 ):284 from ..integrations import replace_with_bnb_linear285 286 llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload287 288 self.modules_to_not_convert = self.get_modules_to_not_convert(289 model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules290 )291 292 # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`293 if isinstance(device_map, dict) and len(device_map.keys()) > 1:294 keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]295 296 if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload:297 raise ValueError(298 "If you want to offload some keys to `cpu` or `disk`, you need to set "299 "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be "300 " converted to 8-bit but kept in 32-bit."301 )302 self.modules_to_not_convert.extend(keys_on_cpu)303 304 model = replace_with_bnb_linear(305 model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config306 )307 308 model.config.quantization_config = self.quantization_config309 310 # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_after_weight_loading with 8bit->4bit311 def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):312 model.is_loaded_in_4bit = True313 model.is_4bit_serializable = self.is_serializable()314 return model315 316 def is_serializable(self, safe_serialization=None):317 _is_4bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.41.3")318 319 if not _is_4bit_serializable:320 logger.warning(321 "You are calling `save_pretrained` to a 4-bit converted model, but your `bitsandbytes` version doesn't support it. "322 "If you want to save 4-bit models, make sure to have `bitsandbytes>=0.41.3` installed."323 )324 return False325 326 return True327 328 @cached_property329 def is_bnb_supports_quant_storage_module(self) -> bool:330 """331 determines if the current version of bitsandbytes supports332 the `module` parameter in `Params4bit.from_prequantized`333 :return:334 """335 return version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.43.3")336 337 @property338 def is_trainable(self) -> bool:339 return True340 341 def _dequantize(self, model):342 from ..integrations import dequantize_and_replace343 344 model = dequantize_and_replace(345 model, self.modules_to_not_convert, quantization_config=self.quantization_config346 )347 return model348