Aluode/PerceptionLabPortable
0
1import importlib.metadata2import inspect3import warnings4from copy import deepcopy5from inspect import signature6 7from packaging import version8 9from ..utils import (10 get_available_devices,11 is_accelerate_available,12 is_bitsandbytes_available,13 is_bitsandbytes_multi_backend_available,14 is_torch_available,15 logging,16)17 18 19if is_bitsandbytes_available():20 import bitsandbytes as bnb21 import torch22 import torch.nn as nn23 24 from ..pytorch_utils import Conv1D25 26if is_accelerate_available():27 import accelerate28 from accelerate import init_empty_weights29 from accelerate.hooks import add_hook_to_module, remove_hook_from_module30 from accelerate.utils import find_tied_parameters31 32logger = logging.get_logger(__name__)33 34 35def set_module_quantized_tensor_to_device(module, tensor_name, device, value=None, quantized_stats=None):36 """37 A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing38 `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function). The39 function is adapted from `set_module_tensor_to_device` function from accelerate that is adapted to support the40 class `Int8Params` from `bitsandbytes`.41 42 Args:43 module (`torch.nn.Module`):44 The module in which the tensor we want to move lives.45 tensor_name (`str`):46 The full name of the parameter/buffer.47 device (`int`, `str` or `torch.device`):48 The device on which to set the tensor.49 value (`torch.Tensor`, *optional*):50 The value of the tensor (useful when going from the meta device to any other device).51 quantized_stats (`dict[str, Any]`, *optional*):52 Dict with items for either 4-bit or 8-bit serialization53 """54 # Recurse if needed55 if "." in tensor_name:56 splits = tensor_name.split(".")57 for split in splits[:-1]:58 new_module = getattr(module, split)59 if new_module is None:60 raise ValueError(f"{module} has no attribute {split}.")61 module = new_module62 tensor_name = splits[-1]63 64 if tensor_name not in module._parameters and tensor_name not in module._buffers:65 raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")66 is_buffer = tensor_name in module._buffers67 old_value = getattr(module, tensor_name)68 69 if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None:70 raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")71 72 prequantized_loading = quantized_stats is not None73 if is_buffer or not is_bitsandbytes_available():74 is_8bit = False75 is_4bit = False76 else:77 is_4bit = hasattr(bnb.nn, "Params4bit") and isinstance(module._parameters[tensor_name], bnb.nn.Params4bit)78 is_8bit = isinstance(module._parameters[tensor_name], bnb.nn.Int8Params)79 80 if is_8bit or is_4bit:81 param = module._parameters[tensor_name]82 if param.device.type != "cuda":83 if value is None:84 new_value = old_value.to(device)85 elif isinstance(value, torch.Tensor):86 new_value = value.to("cpu")87 else:88 new_value = torch.tensor(value, device="cpu")89 90 # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.91 # Since weights are saved in the correct "orientation", we skip transposing when loading.92 if issubclass(module.source_cls, Conv1D) and not prequantized_loading:93 new_value = new_value.T94 95 kwargs = old_value.__dict__96 97 if prequantized_loading != (new_value.dtype in (torch.int8, torch.uint8)):98 raise ValueError(99 f"Value dtype `{new_value.dtype}` is not compatible with parameter quantization status."100 )101 102 if is_8bit:103 is_8bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) > version.parse(104 "0.37.2"105 )106 if new_value.dtype in (torch.int8, torch.uint8) and not is_8bit_serializable:107 raise ValueError(108 "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "109 "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."110 )111 new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(device)112 if prequantized_loading:113 setattr(new_value, "SCB", quantized_stats["SCB"].to(device))114 elif is_4bit:115 if prequantized_loading:116 is_4bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse(117 "0.41.3"118 )119 if new_value.dtype in (torch.int8, torch.uint8) and not is_4bit_serializable:120 raise ValueError(121 "Detected 4-bit weights but the version of bitsandbytes is not compatible with 4-bit serialization. "122 "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."123 )124 new_value = bnb.nn.Params4bit.from_prequantized(125 data=new_value,126 quantized_stats=quantized_stats,127 requires_grad=False,128 device=device,129 **kwargs,130 )131 else:132 new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(device)133 module._parameters[tensor_name] = new_value134 135 else:136 if value is None:137 new_value = old_value.to(device)138 elif isinstance(value, torch.Tensor):139 new_value = value.to(device)140 else:141 new_value = torch.tensor(value, device=device)142 143 if is_buffer:144 module._buffers[tensor_name] = new_value145 else:146 new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)147 module._parameters[tensor_name] = new_value148 149 150def _replace_with_bnb_linear(151 model,152 modules_to_not_convert=None,153 current_key_name=None,154 quantization_config=None,155 has_been_replaced=False,156):157 """158 Private method that wraps the recursion for module replacement.159 160 Returns the converted model and a boolean that indicates if the conversion has been successful or not.161 """162 for name, module in model.named_children():163 if current_key_name is None:164 current_key_name = []165 current_key_name.append(name)166 167 if (isinstance(module, (nn.Linear, Conv1D))) and name not in modules_to_not_convert:168 # Check if the current key is not in the `modules_to_not_convert`169 current_key_name_str = ".".join(current_key_name)170 if not any(171 (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert172 ):173 with init_empty_weights():174 if isinstance(module, Conv1D):175 in_features, out_features = module.weight.shape176 else:177 in_features = module.in_features178 out_features = module.out_features179 180 if quantization_config.quantization_method() == "llm_int8":181 model._modules[name] = bnb.nn.Linear8bitLt(182 in_features,183 out_features,184 module.bias is not None,185 has_fp16_weights=quantization_config.llm_int8_has_fp16_weight,186 threshold=quantization_config.llm_int8_threshold,187 )188 has_been_replaced = True189 else:190 if (191 quantization_config.llm_int8_skip_modules is not None192 and name in quantization_config.llm_int8_skip_modules193 ):194 pass195 else:196 extra_kwargs = (197 {"quant_storage": quantization_config.bnb_4bit_quant_storage}198 if "quant_storage" in list(signature(bnb.nn.Linear4bit).parameters)199 else {}200 )201 model._modules[name] = bnb.nn.Linear4bit(202 in_features,203 out_features,204 module.bias is not None,205 quantization_config.bnb_4bit_compute_dtype,206 compress_statistics=quantization_config.bnb_4bit_use_double_quant,207 quant_type=quantization_config.bnb_4bit_quant_type,208 **extra_kwargs,209 )210 has_been_replaced = True211 # Store the module class in case we need to transpose the weight later212 model._modules[name].source_cls = type(module)213 # Force requires grad to False to avoid unexpected errors214 model._modules[name].requires_grad_(False)215 if len(list(module.children())) > 0:216 _, has_been_replaced = _replace_with_bnb_linear(217 module,218 modules_to_not_convert,219 current_key_name,220 quantization_config,221 has_been_replaced=has_been_replaced,222 )223 # Remove the last key for recursion224 current_key_name.pop(-1)225 return model, has_been_replaced226 227 228def replace_with_bnb_linear(model, modules_to_not_convert=None, current_key_name=None, quantization_config=None):229 """230 A helper function to replace all `torch.nn.Linear` modules by `bnb.nn.Linear8bit` modules from the `bitsandbytes`231 library. This will enable running your models using mixed int8 precision as described by the paper `LLM.int8():232 8-bit Matrix Multiplication for Transformers at Scale`. Make sure `bitsandbytes` compiled with the correct CUDA233 version of your hardware is installed before running this function. `pip install -i https://test.pypi.org/simple/234 bitsandbytes`235 236 The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should237 be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no238 CPU/GPU memory is required to run this function. Int8 mixed-precision matrix decomposition works by separating a239 matrix multiplication into two streams: (1) and systematic feature outlier stream matrix multiplied in fp16240 (0.01%), (2) a regular stream of int8 matrix multiplication (99.9%). With this method, int8 inference with no241 predictive degradation is possible for very large models (>=176B parameters).242 243 Parameters:244 model (`torch.nn.Module`):245 Input model or `torch.nn.Module` as the function is run recursively.246 modules_to_not_convert (`list[`str`]`, *optional*, defaults to `["lm_head"]`):247 Names of the modules to not convert in `Linear8bitLt`. In practice we keep the `lm_head` in full precision248 for numerical stability reasons.249 current_key_name (`list[`str`]`, *optional*):250 An array to track the current key of the recursion. This is used to check whether the current key (part of251 it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or252 `disk`).253 quantization_config ('transformers.utils.quantization_config.BitsAndBytesConfig'):254 To configure and manage settings related to quantization, a technique used to compress neural network models255 by reducing the precision of the weights and activations, thus making models more efficient in terms of both256 storage and computation.257 """258 modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert259 model, has_been_replaced = _replace_with_bnb_linear(260 model, modules_to_not_convert, current_key_name, quantization_config261 )262 263 if not has_been_replaced:264 logger.warning(265 "You are loading your model in 8bit or 4bit but no linear modules were found in your model."266 " Please double check your model architecture, or submit an issue on github if you think this is"267 " a bug."268 )269 270 return model271 272 273# For backward compatibility274def replace_8bit_linear(*args, **kwargs):275 warnings.warn(276 "`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead",277 FutureWarning,278 )279 return replace_with_bnb_linear(*args, **kwargs)280 281 282# For backward compatibility283def set_module_8bit_tensor_to_device(*args, **kwargs):284 warnings.warn(285 "`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead",286 FutureWarning,287 )288 return set_module_quantized_tensor_to_device(*args, **kwargs)289 290 291def get_keys_to_not_convert(model):292 r"""293 An utility function to get the key of the module to keep in full precision if any For example for CausalLM modules294 we may want to keep the lm_head in full precision for numerical stability reasons. For other architectures, we want295 to keep the tied weights of the model. The function will return a list of the keys of the modules to not convert in296 int8.297 298 Parameters:299 model (`torch.nn.Module`):300 Input model301 """302 # Create a copy of the model and tie the weights, then303 # check if it contains tied weights304 tied_model = deepcopy(model) # this has 0 cost since it is done inside `init_empty_weights` context manager`305 tied_model.tie_weights()306 307 tied_params = find_tied_parameters(tied_model)308 # For compatibility with Accelerate < 0.18309 if isinstance(tied_params, dict):310 tied_keys = sum(list(tied_params.values()), []) + list(tied_params.keys())311 else:312 tied_keys = sum(tied_params, [])313 has_tied_params = len(tied_keys) > 0314 315 # If there is not tied weights, we want to keep the lm_head（output_embedding) in full precision316 if not has_tied_params:317 output_emb = model.get_output_embeddings()318 if output_emb is not None:319 list_last_module = [name for name, module in model.named_modules() if id(module) == id(output_emb)]320 return list_last_module321 322 # otherwise, no tied weights, no output embedding defined, simply keep the last module in full precision323 list_modules = list(model.named_parameters())324 list_last_module = [list_modules[-1][0]]325 # add last module together with tied weights326 intersection = set(list_last_module) - set(tied_keys)327 list_untouched = list(set(tied_keys)) + list(intersection)328 329 # remove ".weight" from the keys330 names_to_remove = [".weight", ".bias"]331 filtered_module_names = []332 for name in list_untouched:333 for name_to_remove in names_to_remove:334 if name_to_remove in name:335 name = name.replace(name_to_remove, "")336 filtered_module_names.append(name)337 338 return filtered_module_names339 340 341# Copied from PEFT: https://github.com/huggingface/peft/blob/47b3712898539569c02ec5b3ed4a6c36811331a1/src/peft/utils/integrations.py#L41342def dequantize_bnb_weight(weight: "torch.nn.Parameter", dtype: "torch.dtype", state=None):343 """344 Helper function to dequantize 4bit or 8bit bnb weights.345 346 If the weight is not a bnb quantized weight, it will be returned as is.347 """348 if not isinstance(weight, torch.nn.Parameter):349 raise TypeError(f"Input weight should be of type nn.Parameter, got {type(weight)} instead")350 351 cls_name = weight.__class__.__name__352 if cls_name not in ("Params4bit", "Int8Params"):353 return weight354 355 if cls_name == "Params4bit":356 output_tensor = bnb.functional.dequantize_4bit(weight.data, weight.quant_state)357 logger.warning_once(358 f"The model is going to be dequantized in {output_tensor.dtype} - if you want to upcast it to another dtype, make sure to pass the desired dtype when quantizing the model through `bnb_4bit_quant_type` argument of `BitsAndBytesConfig`"359 )360 return output_tensor.to(dtype)361 362 if state.SCB is None:363 state.SCB = weight.SCB364 365 if hasattr(bnb.functional, "int8_vectorwise_dequant"):366 # Use bitsandbytes API if available (requires v0.45.0+)367 dequantized = bnb.functional.int8_vectorwise_dequant(weight.data, state.SCB)368 else:369 # Multiply by (scale/127) to dequantize.370 dequantized = weight.data * state.SCB.view(-1, 1) * 7.874015718698502e-3371 372 return dequantized.to(dtype)373 374 375def _create_accelerate_new_hook(old_hook):376 r"""377 Creates a new hook based on the old hook. Use it only if you know what you are doing !378 This method is a copy of: https://github.com/huggingface/peft/blob/748f7968f3a31ec06a1c2b0328993319ad9a150a/src/peft/utils/other.py#L245379 with some changes380 """381 old_hook_cls = getattr(accelerate.hooks, old_hook.__class__.__name__)382 old_hook_attr = old_hook.__dict__383 filtered_old_hook_attr = {}384 old_hook_init_signature = inspect.signature(old_hook_cls.__init__)385 for k in old_hook_attr:386 if k in old_hook_init_signature.parameters:387 filtered_old_hook_attr[k] = old_hook_attr[k]388 new_hook = old_hook_cls(**filtered_old_hook_attr)389 return new_hook390 391 392def _dequantize_and_replace(393 model,394 dtype,395 modules_to_not_convert=None,396 current_key_name=None,397 quantization_config=None,398 has_been_replaced=False,399):400 """401 Converts a quantized model into its dequantized original version. The newly converted model will have402 some performance drop compared to the original model before quantization - use it only for specific usecases403 such as QLoRA adapters merging.404 405 Returns the converted model and a boolean that indicates if the conversion has been successful or not.406 """407 quant_method = quantization_config.quantization_method()408 409 target_cls = bnb.nn.Linear8bitLt if quant_method == "llm_int8" else bnb.nn.Linear4bit410 411 for name, module in model.named_children():412 if current_key_name is None:413 current_key_name = []414 current_key_name.append(name)415 416 if isinstance(module, target_cls) and name not in modules_to_not_convert:417 # Check if the current key is not in the `modules_to_not_convert`418 current_key_name_str = ".".join(current_key_name)419 420 if not any(421 (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert422 ):423 bias = getattr(module, "bias", None)424 425 device = module.weight.device426 with init_empty_weights():427 new_module = torch.nn.Linear(module.in_features, module.out_features, bias=bias is not None)428 429 if quant_method == "llm_int8":430 state = module.state431 else:432 state = None433 434 new_module.weight = torch.nn.Parameter(dequantize_bnb_weight(module.weight, dtype, state))435 436 if bias is not None:437 new_module.bias = bias438 439 # Create a new hook and attach it in case we use accelerate440 if hasattr(module, "_hf_hook"):441 old_hook = module._hf_hook442 new_hook = _create_accelerate_new_hook(old_hook)443 444 remove_hook_from_module(module)445 add_hook_to_module(new_module, new_hook)446 447 new_module.to(device)448 model._modules[name] = new_module449 has_been_replaced = True450 if len(list(module.children())) > 0:451 _, has_been_replaced = _dequantize_and_replace(452 module,453 dtype,454 modules_to_not_convert,455 current_key_name,456 quantization_config,457 has_been_replaced=has_been_replaced,458 )459 # Remove the last key for recursion460 current_key_name.pop(-1)461 return model, has_been_replaced462 463 464def dequantize_and_replace(465 model,466 modules_to_not_convert=None,467 quantization_config=None,468):469 model, has_been_replaced = _dequantize_and_replace(470 model,471 model.dtype,472 modules_to_not_convert=modules_to_not_convert,473 quantization_config=quantization_config,474 )475 476 if not has_been_replaced:477 logger.warning(478 "For some reason the model has not been properly dequantized. You might see unexpected behavior."479 )480 481 return model482 483 484def _validate_bnb_multi_backend_availability(raise_exception):485 import bitsandbytes as bnb486 487 bnb_supported_devices = getattr(bnb, "supported_torch_devices", set())488 available_devices = set(get_available_devices())489 490 if not available_devices.intersection(bnb_supported_devices):491 if raise_exception:492 err_msg = (493 f"None of the available devices `available_devices = {available_devices or None}` are supported by the bitsandbytes version you have installed: `bnb_supported_devices = {bnb_supported_devices}`. "494 "Please check the docs to see if the backend you intend to use is available and how to install it: https://huggingface.co/docs/bitsandbytes/main/en/installation"495 )496 497 logger.error(err_msg)498 raise RuntimeError(err_msg)499 500 logger.warning("No supported devices found for bitsandbytes multi-backend.")501 return False502 503 logger.debug("Multi-backend validation successful.")504 return True505 506 507def _validate_bnb_cuda_backend_availability(raise_exception):508 if not is_torch_available():509 return False510 511 import torch512 513 if not torch.cuda.is_available():514 log_msg = (515 "CUDA is required but not available for bitsandbytes. Please consider installing the multi-platform enabled version of bitsandbytes, which is currently a work in progress. "516 "Please check currently supported platforms and installation instructions at https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend"517 )518 if raise_exception:519 logger.error(log_msg)520 raise RuntimeError(log_msg)521 522 logger.warning(log_msg)523 return False524 525 logger.debug("CUDA backend validation successful.")526 return True527 528 529def validate_bnb_backend_availability(raise_exception=False):530 """531 Validates if the available devices are supported by bitsandbytes, optionally raising an exception if not.532 """533 if not is_bitsandbytes_available():534 if importlib.util.find_spec("bitsandbytes") and version.parse(535 importlib.metadata.version("bitsandbytes")536 ) < version.parse("0.43.1"):537 return _validate_bnb_cuda_backend_availability(raise_exception)538 return False539 540 if is_bitsandbytes_multi_backend_available():541 return _validate_bnb_multi_backend_availability(raise_exception)542 return _validate_bnb_cuda_backend_availability(raise_exception)543