543 lines
23 KiB
Python
543 lines
23 KiB
Python
|
import importlib.metadata
|
|||
|
import inspect
|
|||
|
import warnings
|
|||
|
from copy import deepcopy
|
|||
|
from inspect import signature
|
|||
|
|
|||
|
from packaging import version
|
|||
|
|
|||
|
from ..utils import (
|
|||
|
get_available_devices,
|
|||
|
is_accelerate_available,
|
|||
|
is_bitsandbytes_available,
|
|||
|
is_bitsandbytes_multi_backend_available,
|
|||
|
is_torch_available,
|
|||
|
logging,
|
|||
|
)
|
|||
|
|
|||
|
|
|||
|
if is_bitsandbytes_available():
|
|||
|
import bitsandbytes as bnb
|
|||
|
import torch
|
|||
|
import torch.nn as nn
|
|||
|
|
|||
|
from ..pytorch_utils import Conv1D
|
|||
|
|
|||
|
if is_accelerate_available():
|
|||
|
import accelerate
|
|||
|
from accelerate import init_empty_weights
|
|||
|
from accelerate.hooks import add_hook_to_module, remove_hook_from_module
|
|||
|
from accelerate.utils import find_tied_parameters
|
|||
|
|
|||
|
logger = logging.get_logger(__name__)
|
|||
|
|
|||
|
|
|||
|
def set_module_quantized_tensor_to_device(module, tensor_name, device, value=None, quantized_stats=None):
|
|||
|
"""
|
|||
|
A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing
|
|||
|
`param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function). The
|
|||
|
function is adapted from `set_module_tensor_to_device` function from accelerate that is adapted to support the
|
|||
|
class `Int8Params` from `bitsandbytes`.
|
|||
|
|
|||
|
Args:
|
|||
|
module (`torch.nn.Module`):
|
|||
|
The module in which the tensor we want to move lives.
|
|||
|
tensor_name (`str`):
|
|||
|
The full name of the parameter/buffer.
|
|||
|
device (`int`, `str` or `torch.device`):
|
|||
|
The device on which to set the tensor.
|
|||
|
value (`torch.Tensor`, *optional*):
|
|||
|
The value of the tensor (useful when going from the meta device to any other device).
|
|||
|
quantized_stats (`dict[str, Any]`, *optional*):
|
|||
|
Dict with items for either 4-bit or 8-bit serialization
|
|||
|
"""
|
|||
|
# Recurse if needed
|
|||
|
if "." in tensor_name:
|
|||
|
splits = tensor_name.split(".")
|
|||
|
for split in splits[:-1]:
|
|||
|
new_module = getattr(module, split)
|
|||
|
if new_module is None:
|
|||
|
raise ValueError(f"{module} has no attribute {split}.")
|
|||
|
module = new_module
|
|||
|
tensor_name = splits[-1]
|
|||
|
|
|||
|
if tensor_name not in module._parameters and tensor_name not in module._buffers:
|
|||
|
raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
|
|||
|
is_buffer = tensor_name in module._buffers
|
|||
|
old_value = getattr(module, tensor_name)
|
|||
|
|
|||
|
if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None:
|
|||
|
raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
|
|||
|
|
|||
|
prequantized_loading = quantized_stats is not None
|
|||
|
if is_buffer or not is_bitsandbytes_available():
|
|||
|
is_8bit = False
|
|||
|
is_4bit = False
|
|||
|
else:
|
|||
|
is_4bit = hasattr(bnb.nn, "Params4bit") and isinstance(module._parameters[tensor_name], bnb.nn.Params4bit)
|
|||
|
is_8bit = isinstance(module._parameters[tensor_name], bnb.nn.Int8Params)
|
|||
|
|
|||
|
if is_8bit or is_4bit:
|
|||
|
param = module._parameters[tensor_name]
|
|||
|
if param.device.type != "cuda":
|
|||
|
if value is None:
|
|||
|
new_value = old_value.to(device)
|
|||
|
elif isinstance(value, torch.Tensor):
|
|||
|
new_value = value.to("cpu")
|
|||
|
else:
|
|||
|
new_value = torch.tensor(value, device="cpu")
|
|||
|
|
|||
|
# Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
|
|||
|
# Since weights are saved in the correct "orientation", we skip transposing when loading.
|
|||
|
if issubclass(module.source_cls, Conv1D) and not prequantized_loading:
|
|||
|
new_value = new_value.T
|
|||
|
|
|||
|
kwargs = old_value.__dict__
|
|||
|
|
|||
|
if prequantized_loading != (new_value.dtype in (torch.int8, torch.uint8)):
|
|||
|
raise ValueError(
|
|||
|
f"Value dtype `{new_value.dtype}` is not compatible with parameter quantization status."
|
|||
|
)
|
|||
|
|
|||
|
if is_8bit:
|
|||
|
is_8bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) > version.parse(
|
|||
|
"0.37.2"
|
|||
|
)
|
|||
|
if new_value.dtype in (torch.int8, torch.uint8) and not is_8bit_serializable:
|
|||
|
raise ValueError(
|
|||
|
"Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
|
|||
|
"Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
|
|||
|
)
|
|||
|
new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(device)
|
|||
|
if prequantized_loading:
|
|||
|
setattr(new_value, "SCB", quantized_stats["SCB"].to(device))
|
|||
|
elif is_4bit:
|
|||
|
if prequantized_loading:
|
|||
|
is_4bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse(
|
|||
|
"0.41.3"
|
|||
|
)
|
|||
|
if new_value.dtype in (torch.int8, torch.uint8) and not is_4bit_serializable:
|
|||
|
raise ValueError(
|
|||
|
"Detected 4-bit weights but the version of bitsandbytes is not compatible with 4-bit serialization. "
|
|||
|
"Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
|
|||
|
)
|
|||
|
new_value = bnb.nn.Params4bit.from_prequantized(
|
|||
|
data=new_value,
|
|||
|
quantized_stats=quantized_stats,
|
|||
|
requires_grad=False,
|
|||
|
device=device,
|
|||
|
**kwargs,
|
|||
|
)
|
|||
|
else:
|
|||
|
new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(device)
|
|||
|
module._parameters[tensor_name] = new_value
|
|||
|
|
|||
|
else:
|
|||
|
if value is None:
|
|||
|
new_value = old_value.to(device)
|
|||
|
elif isinstance(value, torch.Tensor):
|
|||
|
new_value = value.to(device)
|
|||
|
else:
|
|||
|
new_value = torch.tensor(value, device=device)
|
|||
|
|
|||
|
if is_buffer:
|
|||
|
module._buffers[tensor_name] = new_value
|
|||
|
else:
|
|||
|
new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad)
|
|||
|
module._parameters[tensor_name] = new_value
|
|||
|
|
|||
|
|
|||
|
def _replace_with_bnb_linear(
|
|||
|
model,
|
|||
|
modules_to_not_convert=None,
|
|||
|
current_key_name=None,
|
|||
|
quantization_config=None,
|
|||
|
has_been_replaced=False,
|
|||
|
):
|
|||
|
"""
|
|||
|
Private method that wraps the recursion for module replacement.
|
|||
|
|
|||
|
Returns the converted model and a boolean that indicates if the conversion has been successful or not.
|
|||
|
"""
|
|||
|
for name, module in model.named_children():
|
|||
|
if current_key_name is None:
|
|||
|
current_key_name = []
|
|||
|
current_key_name.append(name)
|
|||
|
|
|||
|
if (isinstance(module, nn.Linear) or isinstance(module, Conv1D)) and name not in modules_to_not_convert:
|
|||
|
# Check if the current key is not in the `modules_to_not_convert`
|
|||
|
current_key_name_str = ".".join(current_key_name)
|
|||
|
if not any(
|
|||
|
(key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert
|
|||
|
):
|
|||
|
with init_empty_weights():
|
|||
|
if isinstance(module, Conv1D):
|
|||
|
in_features, out_features = module.weight.shape
|
|||
|
else:
|
|||
|
in_features = module.in_features
|
|||
|
out_features = module.out_features
|
|||
|
|
|||
|
if quantization_config.quantization_method() == "llm_int8":
|
|||
|
model._modules[name] = bnb.nn.Linear8bitLt(
|
|||
|
in_features,
|
|||
|
out_features,
|
|||
|
module.bias is not None,
|
|||
|
has_fp16_weights=quantization_config.llm_int8_has_fp16_weight,
|
|||
|
threshold=quantization_config.llm_int8_threshold,
|
|||
|
)
|
|||
|
has_been_replaced = True
|
|||
|
else:
|
|||
|
if (
|
|||
|
quantization_config.llm_int8_skip_modules is not None
|
|||
|
and name in quantization_config.llm_int8_skip_modules
|
|||
|
):
|
|||
|
pass
|
|||
|
else:
|
|||
|
extra_kwargs = (
|
|||
|
{"quant_storage": quantization_config.bnb_4bit_quant_storage}
|
|||
|
if "quant_storage" in list(signature(bnb.nn.Linear4bit).parameters)
|
|||
|
else {}
|
|||
|
)
|
|||
|
model._modules[name] = bnb.nn.Linear4bit(
|
|||
|
in_features,
|
|||
|
out_features,
|
|||
|
module.bias is not None,
|
|||
|
quantization_config.bnb_4bit_compute_dtype,
|
|||
|
compress_statistics=quantization_config.bnb_4bit_use_double_quant,
|
|||
|
quant_type=quantization_config.bnb_4bit_quant_type,
|
|||
|
**extra_kwargs,
|
|||
|
)
|
|||
|
has_been_replaced = True
|
|||
|
# Store the module class in case we need to transpose the weight later
|
|||
|
model._modules[name].source_cls = type(module)
|
|||
|
# Force requires grad to False to avoid unexpected errors
|
|||
|
model._modules[name].requires_grad_(False)
|
|||
|
if len(list(module.children())) > 0:
|
|||
|
_, has_been_replaced = _replace_with_bnb_linear(
|
|||
|
module,
|
|||
|
modules_to_not_convert,
|
|||
|
current_key_name,
|
|||
|
quantization_config,
|
|||
|
has_been_replaced=has_been_replaced,
|
|||
|
)
|
|||
|
# Remove the last key for recursion
|
|||
|
current_key_name.pop(-1)
|
|||
|
return model, has_been_replaced
|
|||
|
|
|||
|
|
|||
|
def replace_with_bnb_linear(model, modules_to_not_convert=None, current_key_name=None, quantization_config=None):
|
|||
|
"""
|
|||
|
A helper function to replace all `torch.nn.Linear` modules by `bnb.nn.Linear8bit` modules from the `bitsandbytes`
|
|||
|
library. This will enable running your models using mixed int8 precision as described by the paper `LLM.int8():
|
|||
|
8-bit Matrix Multiplication for Transformers at Scale`. Make sure `bitsandbytes` compiled with the correct CUDA
|
|||
|
version of your hardware is installed before running this function. `pip install -i https://test.pypi.org/simple/
|
|||
|
bitsandbytes`
|
|||
|
|
|||
|
The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should
|
|||
|
be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no
|
|||
|
CPU/GPU memory is required to run this function. Int8 mixed-precision matrix decomposition works by separating a
|
|||
|
matrix multiplication into two streams: (1) and systematic feature outlier stream matrix multiplied in fp16
|
|||
|
(0.01%), (2) a regular stream of int8 matrix multiplication (99.9%). With this method, int8 inference with no
|
|||
|
predictive degradation is possible for very large models (>=176B parameters).
|
|||
|
|
|||
|
Parameters:
|
|||
|
model (`torch.nn.Module`):
|
|||
|
Input model or `torch.nn.Module` as the function is run recursively.
|
|||
|
modules_to_not_convert (`list[`str`]`, *optional*, defaults to `["lm_head"]`):
|
|||
|
Names of the modules to not convert in `Linear8bitLt`. In practice we keep the `lm_head` in full precision
|
|||
|
for numerical stability reasons.
|
|||
|
current_key_name (`list[`str`]`, *optional*):
|
|||
|
An array to track the current key of the recursion. This is used to check whether the current key (part of
|
|||
|
it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or
|
|||
|
`disk`).
|
|||
|
quantization_config ('transformers.utils.quantization_config.BitsAndBytesConfig'):
|
|||
|
To configure and manage settings related to quantization, a technique used to compress neural network models
|
|||
|
by reducing the precision of the weights and activations, thus making models more efficient in terms of both
|
|||
|
storage and computation.
|
|||
|
"""
|
|||
|
modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
|
|||
|
model, has_been_replaced = _replace_with_bnb_linear(
|
|||
|
model, modules_to_not_convert, current_key_name, quantization_config
|
|||
|
)
|
|||
|
|
|||
|
if not has_been_replaced:
|
|||
|
logger.warning(
|
|||
|
"You are loading your model in 8bit or 4bit but no linear modules were found in your model."
|
|||
|
" Please double check your model architecture, or submit an issue on github if you think this is"
|
|||
|
" a bug."
|
|||
|
)
|
|||
|
|
|||
|
return model
|
|||
|
|
|||
|
|
|||
|
# For backward compatibility
|
|||
|
def replace_8bit_linear(*args, **kwargs):
|
|||
|
warnings.warn(
|
|||
|
"`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead",
|
|||
|
FutureWarning,
|
|||
|
)
|
|||
|
return replace_with_bnb_linear(*args, **kwargs)
|
|||
|
|
|||
|
|
|||
|
# For backward compatibility
|
|||
|
def set_module_8bit_tensor_to_device(*args, **kwargs):
|
|||
|
warnings.warn(
|
|||
|
"`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead",
|
|||
|
FutureWarning,
|
|||
|
)
|
|||
|
return set_module_quantized_tensor_to_device(*args, **kwargs)
|
|||
|
|
|||
|
|
|||
|
def get_keys_to_not_convert(model):
|
|||
|
r"""
|
|||
|
An utility function to get the key of the module to keep in full precision if any For example for CausalLM modules
|
|||
|
we may want to keep the lm_head in full precision for numerical stability reasons. For other architectures, we want
|
|||
|
to keep the tied weights of the model. The function will return a list of the keys of the modules to not convert in
|
|||
|
int8.
|
|||
|
|
|||
|
Parameters:
|
|||
|
model (`torch.nn.Module`):
|
|||
|
Input model
|
|||
|
"""
|
|||
|
# Create a copy of the model and tie the weights, then
|
|||
|
# check if it contains tied weights
|
|||
|
tied_model = deepcopy(model) # this has 0 cost since it is done inside `init_empty_weights` context manager`
|
|||
|
tied_model.tie_weights()
|
|||
|
|
|||
|
tied_params = find_tied_parameters(tied_model)
|
|||
|
# For compatibility with Accelerate < 0.18
|
|||
|
if isinstance(tied_params, dict):
|
|||
|
tied_keys = sum(list(tied_params.values()), []) + list(tied_params.keys())
|
|||
|
else:
|
|||
|
tied_keys = sum(tied_params, [])
|
|||
|
has_tied_params = len(tied_keys) > 0
|
|||
|
|
|||
|
# If there is not tied weights, we want to keep the lm_head(output_embedding) in full precision
|
|||
|
if not has_tied_params:
|
|||
|
output_emb = model.get_output_embeddings()
|
|||
|
if output_emb is not None:
|
|||
|
list_last_module = [name for name, module in model.named_modules() if id(module) == id(output_emb)]
|
|||
|
return list_last_module
|
|||
|
|
|||
|
# otherwise, no tied weights, no output embedding defined, simply keep the last module in full precision
|
|||
|
list_modules = list(model.named_parameters())
|
|||
|
list_last_module = [list_modules[-1][0]]
|
|||
|
# add last module together with tied weights
|
|||
|
intersection = set(list_last_module) - set(tied_keys)
|
|||
|
list_untouched = list(set(tied_keys)) + list(intersection)
|
|||
|
|
|||
|
# remove ".weight" from the keys
|
|||
|
names_to_remove = [".weight", ".bias"]
|
|||
|
filtered_module_names = []
|
|||
|
for name in list_untouched:
|
|||
|
for name_to_remove in names_to_remove:
|
|||
|
if name_to_remove in name:
|
|||
|
name = name.replace(name_to_remove, "")
|
|||
|
filtered_module_names.append(name)
|
|||
|
|
|||
|
return filtered_module_names
|
|||
|
|
|||
|
|
|||
|
# Copied from PEFT: https://github.com/huggingface/peft/blob/47b3712898539569c02ec5b3ed4a6c36811331a1/src/peft/utils/integrations.py#L41
|
|||
|
def dequantize_bnb_weight(weight: "torch.nn.Parameter", dtype: "torch.dtype", state=None):
|
|||
|
"""
|
|||
|
Helper function to dequantize 4bit or 8bit bnb weights.
|
|||
|
|
|||
|
If the weight is not a bnb quantized weight, it will be returned as is.
|
|||
|
"""
|
|||
|
if not isinstance(weight, torch.nn.Parameter):
|
|||
|
raise TypeError(f"Input weight should be of type nn.Parameter, got {type(weight)} instead")
|
|||
|
|
|||
|
cls_name = weight.__class__.__name__
|
|||
|
if cls_name not in ("Params4bit", "Int8Params"):
|
|||
|
return weight
|
|||
|
|
|||
|
if cls_name == "Params4bit":
|
|||
|
output_tensor = bnb.functional.dequantize_4bit(weight.data, weight.quant_state)
|
|||
|
logger.warning_once(
|
|||
|
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`"
|
|||
|
)
|
|||
|
return output_tensor.to(dtype)
|
|||
|
|
|||
|
if state.SCB is None:
|
|||
|
state.SCB = weight.SCB
|
|||
|
|
|||
|
if hasattr(bnb.functional, "int8_vectorwise_dequant"):
|
|||
|
# Use bitsandbytes API if available (requires v0.45.0+)
|
|||
|
dequantized = bnb.functional.int8_vectorwise_dequant(weight.data, state.SCB)
|
|||
|
else:
|
|||
|
# Multiply by (scale/127) to dequantize.
|
|||
|
dequantized = weight.data * state.SCB.view(-1, 1) * 7.874015718698502e-3
|
|||
|
|
|||
|
return dequantized.to(dtype)
|
|||
|
|
|||
|
|
|||
|
def _create_accelerate_new_hook(old_hook):
|
|||
|
r"""
|
|||
|
Creates a new hook based on the old hook. Use it only if you know what you are doing !
|
|||
|
This method is a copy of: https://github.com/huggingface/peft/blob/748f7968f3a31ec06a1c2b0328993319ad9a150a/src/peft/utils/other.py#L245
|
|||
|
with some changes
|
|||
|
"""
|
|||
|
old_hook_cls = getattr(accelerate.hooks, old_hook.__class__.__name__)
|
|||
|
old_hook_attr = old_hook.__dict__
|
|||
|
filtered_old_hook_attr = {}
|
|||
|
old_hook_init_signature = inspect.signature(old_hook_cls.__init__)
|
|||
|
for k in old_hook_attr.keys():
|
|||
|
if k in old_hook_init_signature.parameters:
|
|||
|
filtered_old_hook_attr[k] = old_hook_attr[k]
|
|||
|
new_hook = old_hook_cls(**filtered_old_hook_attr)
|
|||
|
return new_hook
|
|||
|
|
|||
|
|
|||
|
def _dequantize_and_replace(
|
|||
|
model,
|
|||
|
dtype,
|
|||
|
modules_to_not_convert=None,
|
|||
|
current_key_name=None,
|
|||
|
quantization_config=None,
|
|||
|
has_been_replaced=False,
|
|||
|
):
|
|||
|
"""
|
|||
|
Converts a quantized model into its dequantized original version. The newly converted model will have
|
|||
|
some performance drop compared to the original model before quantization - use it only for specific usecases
|
|||
|
such as QLoRA adapters merging.
|
|||
|
|
|||
|
Returns the converted model and a boolean that indicates if the conversion has been successful or not.
|
|||
|
"""
|
|||
|
quant_method = quantization_config.quantization_method()
|
|||
|
|
|||
|
target_cls = bnb.nn.Linear8bitLt if quant_method == "llm_int8" else bnb.nn.Linear4bit
|
|||
|
|
|||
|
for name, module in model.named_children():
|
|||
|
if current_key_name is None:
|
|||
|
current_key_name = []
|
|||
|
current_key_name.append(name)
|
|||
|
|
|||
|
if isinstance(module, target_cls) and name not in modules_to_not_convert:
|
|||
|
# Check if the current key is not in the `modules_to_not_convert`
|
|||
|
current_key_name_str = ".".join(current_key_name)
|
|||
|
|
|||
|
if not any(
|
|||
|
(key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert
|
|||
|
):
|
|||
|
bias = getattr(module, "bias", None)
|
|||
|
|
|||
|
device = module.weight.device
|
|||
|
with init_empty_weights():
|
|||
|
new_module = torch.nn.Linear(module.in_features, module.out_features, bias=bias is not None)
|
|||
|
|
|||
|
if quant_method == "llm_int8":
|
|||
|
state = module.state
|
|||
|
else:
|
|||
|
state = None
|
|||
|
|
|||
|
new_module.weight = torch.nn.Parameter(dequantize_bnb_weight(module.weight, dtype, state))
|
|||
|
|
|||
|
if bias is not None:
|
|||
|
new_module.bias = bias
|
|||
|
|
|||
|
# Create a new hook and attach it in case we use accelerate
|
|||
|
if hasattr(module, "_hf_hook"):
|
|||
|
old_hook = module._hf_hook
|
|||
|
new_hook = _create_accelerate_new_hook(old_hook)
|
|||
|
|
|||
|
remove_hook_from_module(module)
|
|||
|
add_hook_to_module(new_module, new_hook)
|
|||
|
|
|||
|
new_module.to(device)
|
|||
|
model._modules[name] = new_module
|
|||
|
has_been_replaced = True
|
|||
|
if len(list(module.children())) > 0:
|
|||
|
_, has_been_replaced = _dequantize_and_replace(
|
|||
|
module,
|
|||
|
dtype,
|
|||
|
modules_to_not_convert,
|
|||
|
current_key_name,
|
|||
|
quantization_config,
|
|||
|
has_been_replaced=has_been_replaced,
|
|||
|
)
|
|||
|
# Remove the last key for recursion
|
|||
|
current_key_name.pop(-1)
|
|||
|
return model, has_been_replaced
|
|||
|
|
|||
|
|
|||
|
def dequantize_and_replace(
|
|||
|
model,
|
|||
|
modules_to_not_convert=None,
|
|||
|
quantization_config=None,
|
|||
|
):
|
|||
|
model, has_been_replaced = _dequantize_and_replace(
|
|||
|
model,
|
|||
|
model.dtype,
|
|||
|
modules_to_not_convert=modules_to_not_convert,
|
|||
|
quantization_config=quantization_config,
|
|||
|
)
|
|||
|
|
|||
|
if not has_been_replaced:
|
|||
|
logger.warning(
|
|||
|
"For some reason the model has not been properly dequantized. You might see unexpected behavior."
|
|||
|
)
|
|||
|
|
|||
|
return model
|
|||
|
|
|||
|
|
|||
|
def _validate_bnb_multi_backend_availability(raise_exception):
|
|||
|
import bitsandbytes as bnb
|
|||
|
|
|||
|
bnb_supported_devices = getattr(bnb, "supported_torch_devices", set())
|
|||
|
available_devices = set(get_available_devices())
|
|||
|
|
|||
|
if not available_devices.intersection(bnb_supported_devices):
|
|||
|
if raise_exception:
|
|||
|
err_msg = (
|
|||
|
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}`. "
|
|||
|
"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"
|
|||
|
)
|
|||
|
|
|||
|
logger.error(err_msg)
|
|||
|
raise RuntimeError(err_msg)
|
|||
|
|
|||
|
logger.warning("No supported devices found for bitsandbytes multi-backend.")
|
|||
|
return False
|
|||
|
|
|||
|
logger.debug("Multi-backend validation successful.")
|
|||
|
return True
|
|||
|
|
|||
|
|
|||
|
def _validate_bnb_cuda_backend_availability(raise_exception):
|
|||
|
if not is_torch_available():
|
|||
|
return False
|
|||
|
|
|||
|
import torch
|
|||
|
|
|||
|
if not torch.cuda.is_available():
|
|||
|
log_msg = (
|
|||
|
"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. "
|
|||
|
"Please check currently supported platforms and installation instructions at https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend"
|
|||
|
)
|
|||
|
if raise_exception:
|
|||
|
logger.error(log_msg)
|
|||
|
raise RuntimeError(log_msg)
|
|||
|
|
|||
|
logger.warning(log_msg)
|
|||
|
return False
|
|||
|
|
|||
|
logger.debug("CUDA backend validation successful.")
|
|||
|
return True
|
|||
|
|
|||
|
|
|||
|
def validate_bnb_backend_availability(raise_exception=False):
|
|||
|
"""
|
|||
|
Validates if the available devices are supported by bitsandbytes, optionally raising an exception if not.
|
|||
|
"""
|
|||
|
if not is_bitsandbytes_available():
|
|||
|
if importlib.util.find_spec("bitsandbytes") and version.parse(
|
|||
|
importlib.metadata.version("bitsandbytes")
|
|||
|
) < version.parse("0.43.1"):
|
|||
|
return _validate_bnb_cuda_backend_availability(raise_exception)
|
|||
|
return False
|
|||
|
|
|||
|
if is_bitsandbytes_multi_backend_available():
|
|||
|
return _validate_bnb_multi_backend_availability(raise_exception)
|
|||
|
return _validate_bnb_cuda_backend_availability(raise_exception)
|