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 importlib15import re16import types17from collections import defaultdict18from typing import TYPE_CHECKING, Optional, Union19 20from packaging import version21 22from .base import HfQuantizer23from .quantizers_utils import get_module_from_name24 25 26if TYPE_CHECKING:27 from ..modeling_utils import PreTrainedModel28 29from safetensors import safe_open30 31from ..utils import is_torch_available, is_torchao_available, logging32 33 34if is_torch_available():35 import torch36 import torch.nn as nn37 38if is_torchao_available():39 import torchao40 41 if version.parse(importlib.metadata.version("torchao")) >= version.parse("0.14.0"):42 from torchao.prototype.safetensors.safetensors_support import (43 flatten_tensor_state_dict,44 unflatten_tensor_state_dict,45 )46 from torchao.prototype.safetensors.safetensors_utils import is_metadata_torchao47 48 49logger = logging.get_logger(__name__)50 51 52def fuzzy_match_size(config_name: str) -> Optional[str]:53 """54 Extract the size digit from strings like "4weight", "8weight".55 Returns the digit as an integer if found, otherwise None.56 """57 config_name = config_name.lower()58 59 str_match = re.search(r"(\d)weight", config_name)60 61 if str_match:62 return str_match.group(1)63 64 return None65 66 67def _quantization_type(weight):68 from torchao.dtypes import AffineQuantizedTensor69 from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor70 71 if isinstance(weight, AffineQuantizedTensor):72 return f"{weight.__class__.__name__}({weight._quantization_type()})"73 74 if isinstance(weight, LinearActivationQuantizedTensor):75 return f"{weight.__class__.__name__}(activation={weight.input_quant_func}, weight={_quantization_type(weight.original_weight_tensor)})"76 77 78def _linear_extra_repr(self):79 weight = _quantization_type(self.weight)80 if weight is None:81 return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight=None"82 else:83 return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={weight}"84 85 86if is_torchao_available():87 SUPPORTED_SAFE_SERIALIZATION_CONFIGS = [88 torchao.quantization.Float8WeightOnlyConfig,89 torchao.quantization.Float8DynamicActivationFloat8WeightConfig,90 ]91 92 TORCHAO_VERSION = version.parse(importlib.metadata.version("torchao"))93 94 95class TorchAoHfQuantizer(HfQuantizer):96 """97 Quantizer for torchao: https://github.com/pytorch/ao/98 """99 100 requires_parameters_quantization = True101 requires_calibration = False102 required_packages = ["torchao"]103 104 def __init__(self, quantization_config, **kwargs):105 super().__init__(quantization_config, **kwargs)106 107 if isinstance(self.quantization_config.quant_type, str):108 is_int_4 = "int4" in self.quantization_config.quant_type109 else:110 config_name = self.quantization_config.quant_type.__class__.__name__111 is_int_4 = fuzzy_match_size(config_name) == "4"112 113 # TODO: better way to get the serialized key names? Hard to read from torchao codebase114 if is_int_4:115 self.weight_ao_keys = ["qdata", "scale", "zero_point"]116 else:117 self.weight_ao_keys = ["qdata", "scale"]118 # Instead of serializing the simple torch.Tensor like usual, torchao adds a `:_data` suffix so we need this119 self.full_ao_keys = self.weight_ao_keys + ["_data"]120 121 def validate_environment(self, *args, **kwargs):122 if not is_torchao_available():123 raise ImportError("Loading an torchao quantized model requires torchao library (`pip install torchao`)")124 125 self.offload = False126 device_map = kwargs.get("device_map")127 if isinstance(device_map, dict):128 if ("disk" in device_map.values() or "cpu" in device_map.values()) and len(device_map) > 1:129 self.offload = True130 if self.pre_quantized and "disk" in device_map.values():131 raise ValueError(132 "You are attempting to perform disk offload with a pre-quantized torchao model "133 "This is not supported yet . Please remove the disk device from the device_map."134 )135 if self.pre_quantized:136 weights_only = kwargs.get("weights_only")137 if weights_only:138 torch_version = version.parse(importlib.metadata.version("torch"))139 if torch_version < version.parse("2.5.0"):140 raise RuntimeError(141 f"In order to use torchao pre-quantized model, you need to have torch>=2.5.0. However, the current version is {torch_version}."142 f" You can also set with `weights_only=False` in `from_pretrained` if you don't want to update torch"143 )144 145 def update_dtype(self, dtype):146 if self.quantization_config.quant_type == "int4_weight_only":147 if dtype is not None and dtype != torch.bfloat16:148 logger.warning_once(149 f"Setting dtype to {dtype} for int4_weight_only quantization, but only bfloat16 is supported right now. Please set the dtype to bfloat16."150 )151 if dtype is None:152 logger.warning_once(153 "Setting dtype to torch.bfloat16 for int4_weight_only quantization since only bfloat16 is supported right now. Please set dtype=torch.bfloat16 to remove this warning."154 )155 dtype = torch.bfloat16156 if self.quantization_config.quant_type == "int8_dynamic_activation_int8_weight":157 if dtype is None:158 logger.info(159 "Setting dtype to torch.float32 for int8_dynamic_activation_int8_weight quantization as no dtype was specified in from_pretrained"160 )161 # we need to set the dtype, otherwise we have dtype mismatch when performing the quantized linear op162 dtype = torch.float32163 return dtype164 165 def get_state_dict_and_metadata(self, model, safe_serialization: Optional[bool] = False):166 """167 If the model is safe serializable, we flatten the state dict of tensor subclasses so that it is compatible with168 the safetensors format.169 """170 if type(self.quantization_config.quant_type) in SUPPORTED_SAFE_SERIALIZATION_CONFIGS and safe_serialization:171 if TORCHAO_VERSION >= version.parse("0.14.0"):172 return flatten_tensor_state_dict(model.state_dict())173 else:174 raise RuntimeError(175 f"In order to use safetensors with torchao, please use torchao version >= 0.14.0. Current version: {TORCHAO_VERSION}"176 )177 else:178 return None, {}179 180 def adjust_target_dtype(self, dtype: "torch.dtype") -> "torch.dtype":181 if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"):182 from accelerate.utils import CustomDtype183 184 # Import AOBaseConfig directly since we know we have the right version185 if self.quantization_config._get_ao_version() > version.Version("0.9.0"):186 from torchao.core.config import AOBaseConfig187 188 quant_type = self.quantization_config.quant_type189 if isinstance(quant_type, AOBaseConfig):190 # Extract size digit using fuzzy match on the class name191 config_name = quant_type.__class__.__name__192 size_digit = fuzzy_match_size(config_name)193 194 # Map the extracted digit to appropriate dtype195 if size_digit == "4":196 return CustomDtype.INT4197 else:198 # Default to int8199 return torch.int8200 201 # Original mapping for non-AOBaseConfig types202 map_to_target_dtype = {203 "int4_weight_only": CustomDtype.INT4,204 "int8_weight_only": torch.int8,205 "int8_dynamic_activation_int8_weight": torch.int8,206 "autoquant": None,207 }208 return map_to_target_dtype[self.quantization_config.quant_type]209 else:210 raise ValueError(211 "You are using `device_map='auto'` on a torchao quantized model. To automatically compute"212 " the appropriate device map, you should upgrade your `accelerate` library with "213 "`pip install --upgrade accelerate`"214 )215 216 def adjust_max_memory(self, max_memory: dict[str, Union[int, str]]) -> dict[str, Union[int, str]]:217 # need more space for the quantization parameters (e.g. scale). Tested with int4 wo and group size = 128218 max_memory = {key: val * 0.9 for key, val in max_memory.items()}219 return max_memory220 221 def _process_model_before_weight_loading(222 self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[list[str]] = None, **kwargs223 ):224 self.modules_to_not_convert = self.get_modules_to_not_convert(225 model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules226 )227 if self.quantization_config.include_input_output_embeddings:228 input_emb = model.get_input_embeddings()229 input_emb_names = [name for name, module in model.named_modules() if id(module) == id(input_emb)]230 output_emb = model.get_output_embeddings()231 output_emb_names = [name for name, module in model.named_modules() if id(module) == id(output_emb)]232 self.modules_to_not_convert = [233 x for x in self.modules_to_not_convert if x not in input_emb_names + output_emb_names234 ]235 return236 237 def update_unexpected_keys(self, model, unexpected_keys: list[str]) -> list[str]:238 return [k for k in unexpected_keys if not any(k.endswith(x) for x in self.full_ao_keys)]239 240 def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:241 if self.quantization_config.quant_type == "autoquant":242 return False243 244 # check if the param_name is not in self.modules_to_not_convert245 if any(key + "." in param_name or key == param_name for key in self.modules_to_not_convert):246 return False247 elif any(param_name.endswith(f":{x}") for x in self.full_ao_keys):248 return True249 else:250 # we only quantize the weight of nn.Linear and nn.Embedding251 module, tensor_name = get_module_from_name(model, param_name)252 _QUANTIZABLE = [torch.nn.Linear]253 if self.quantization_config.include_input_output_embeddings:254 _QUANTIZABLE.append(torch.nn.Embedding)255 return isinstance(module, tuple(_QUANTIZABLE)) and tensor_name == "weight"256 257 def create_quantized_param(258 self,259 model: "PreTrainedModel",260 param_value: "torch.Tensor",261 param_name: str,262 target_device: "torch.device",263 **kwargs,264 ):265 """266 Each nn.Linear layer that needs to be quantized is processed here.267 First, we set the value the weight tensor, then we move it to the target device. Finally, we quantize the module.268 """269 from torchao.quantization import quantize_270 271 full_name = param_name272 # Those are the pre quantized weights273 if ":" in param_name:274 param_name = param_name.rsplit(":", 1)[0]275 module, tensor_name = get_module_from_name(model, param_name)276 277 if self.pre_quantized:278 # If it's a bias, no need to do anything special (except removing the ":_data" part of the key, but was279 # already done) - if it's unsafe-serialized (i.e. not safetensors), not need for anything either280 is_unsafe_serialization = ":" not in full_name281 if tensor_name == "bias" or is_unsafe_serialization:282 module._parameters[tensor_name] = torch.nn.Parameter(283 param_value.to(target_device), requires_grad=param_value.requires_grad284 )285 return286 # Sanity check for the new serialization format287 elif not (TORCHAO_VERSION >= version.parse("0.14.0") and is_metadata_torchao(self.metadata)):288 raise ValueError("To use `safetensors` serialization, you should have `torchao>=0.14.0` installed")289 290 # Save the states for later quantization when they are all gathered291 if not hasattr(self, "ao_params"):292 self.ao_params = defaultdict(dict)293 self.ao_params[param_name].update({full_name: param_value})294 295 # We are ready for quantization in this case (we retrieved all the needed keys)296 if len(self.ao_params[param_name]) == len(self.weight_ao_keys):297 new_param = unflatten_tensor_state_dict(self.ao_params[param_name], self.metadata)[param_name]298 # Set it299 module._parameters[tensor_name] = torch.nn.Parameter(300 new_param.to(target_device), requires_grad=new_param.requires_grad301 )302 303 # Free memory304 del self.ao_params[param_name]305 306 # Add repr to the module307 if isinstance(module, nn.Linear):308 module.extra_repr = types.MethodType(_linear_extra_repr, module)309 else:310 module._parameters[tensor_name] = torch.nn.Parameter(311 param_value, requires_grad=param_value.requires_grad312 ).to(target_device)313 # if we are quantizing tied parameters, to avoid tying the quantized weights314 # the correct order to do it is315 # 1. load the weight to model316 # 2. run tie_weights to populate the weights317 # 3. quantize318 input_embed = model.get_input_embeddings()319 if self.quantization_config.untie_embedding_weights and id(module) == id(input_embed):320 model.tie_weights()321 setattr(model.config.get_text_config(decoder=True), "tie_word_embeddings", False)322 323 # handle ModuleFqnToConfig, introduced in torchao 0.12.0+324 if self.quantization_config._get_ao_version() >= version.Version("0.12.0"):325 from torchao.quantization import ModuleFqnToConfig326 327 config = self.quantization_config.get_apply_tensor_subclass()328 if isinstance(config, ModuleFqnToConfig):329 module_fqn, _ = param_name.rsplit(".", 1)330 c = None331 if module_fqn in config.module_fqn_to_config:332 c = config.module_fqn_to_config[module_fqn]333 else:334 c = config.module_fqn_to_config.get("_default", None)335 if c is not None:336 # filter_fn: not filtering out any modules337 quantize_(module, c, filter_fn=lambda x, fqn: True)338 return339 340 quantize_(module, self.quantization_config.get_apply_tensor_subclass())341 342 def _process_model_after_weight_loading(self, model, **kwargs):343 """No process required for torchao quantized model"""344 if self.quantization_config.quant_type == "autoquant":345 from torchao import autoquant346 from torchao.quantization import ALL_AUTOQUANT_CLASS_LIST347 348 model = torch.compile(model, mode="max-autotune")349 model = autoquant(350 model,351 qtensor_class_list=ALL_AUTOQUANT_CLASS_LIST,352 set_inductor_config=False,353 **self.quantization_config.quant_type_kwargs,354 )355 return model356 return357 358 def is_serializable(self, safe_serialization=None) -> bool:359 if safe_serialization:360 _is_torchao_serializable = type(361 self.quantization_config.quant_type362 ) in SUPPORTED_SAFE_SERIALIZATION_CONFIGS and TORCHAO_VERSION >= version.parse("0.14.0")363 if not _is_torchao_serializable:364 logger.warning(365 f"torchao quantized model only supports safe serialization for {SUPPORTED_SAFE_SERIALIZATION_CONFIGS}, \366 and torchao version >= 0.14.0, please set `safe_serialization` to False for \367 {type(self.quantization_config.quant_type)} and {TORCHAO_VERSION}."368 )369 return _is_torchao_serializable370 371 _is_torchao_serializable = version.parse(importlib.metadata.version("huggingface_hub")) >= version.parse(372 "0.25.0"373 )374 if not _is_torchao_serializable:375 logger.warning("torchao quantized model is only serializable after huggingface_hub >= 0.25.0 ")376 if self.offload and self.quantization_config.modules_to_not_convert is None:377 logger.warning(378 "The model contains offloaded modules and these modules are not quantized. We don't recommend saving the model as we won't be able to reload them."379 "If you want to specify modules to not quantize, please specify modules_to_not_convert in the quantization_config."380 )381 return False382 return _is_torchao_serializable383 384 def get_accelerator_warm_up_factor(self):385 """386 This factor is used in caching_allocator_warmup to determine how many bytes to pre-allocate for accelerator warmup.387 - A factor of 2 means we pre-allocate the full memory footprint of the model.388 - A factor of 4 means we pre-allocate half of that, and so on389 390 However, when using TorchAO, calculating memory usage with param.numel() * param.element_size() doesn't give the correct size for quantized weights (like int4 or int8)391 That's because TorchAO internally represents quantized tensors using subtensors and metadata, and the reported element_size() still corresponds to the dtype392 not the actual bit-width of the quantized data.393 394 To correct for this:395 - Use a division factor of 8 for int4 weights396 - Use a division factor of 4 for int8 weights397 """398 if self.quantization_config._get_ao_version() > version.Version("0.9.0"):399 from torchao.core.config import AOBaseConfig400 401 quant_type = self.quantization_config.quant_type402 # For autoquant case, it will be treated in the string implementation below in map_to_target_dtype403 if isinstance(quant_type, AOBaseConfig):404 # Extract size digit using fuzzy match on the class name405 config_name = quant_type.__class__.__name__406 size_digit = fuzzy_match_size(config_name)407 408 if size_digit == "4":409 return 8410 else:411 return 4412 413 # Original mapping for non-AOBaseConfig types414 map_to_target_dtype = {415 "int4_weight_only": 8,416 "int8_weight_only": 4,417 "int8_dynamic_activation_int8_weight": 4,418 "autoquant": 4,419 }420 421 return map_to_target_dtype[self.quantization_config.quant_type]422 423 @property424 def is_trainable(self) -> bool:425 supported_quant_types_for_training = [426 "int8_weight_only",427 "int8_dynamic_activation_int8_weight",428 ]429 return self.quantization_config.quant_type in supported_quant_types_for_training430 431 @property432 def is_compileable(self) -> bool:433 return True434 435 def set_metadata(self, checkpoint_files: list[str]):436 if checkpoint_files[0].endswith(".safetensors"):437 metadata = {}438 for checkpoint in checkpoint_files:439 with safe_open(checkpoint, framework="pt") as f:440 metadata_ = f.metadata() or {}441 metadata.update(metadata_)442 # Save it443 self.metadata = metadata444