InvokeAI/invokeai/backend/bnb.py

518 lines
22 KiB
Python

from typing import Any, Optional, Set, Type
import bitsandbytes as bnb
import torch
# The utils in this file take ideas from
# https://github.com/Lightning-AI/pytorch-lightning/blob/1551a16b94f5234a4a78801098f64d0732ef5cb5/src/lightning/fabric/plugins/precision/bitsandbytes.py
# Patterns:
# - Quantize:
# - Initialize model on meta device
# - Replace layers
# - Load state_dict to cpu
# - Load state_dict into model
# - Quantize on GPU
# - Extract state_dict
# - Save
# - Load:
# - Initialize model on meta device
# - Replace layers
# - Load state_dict to cpu
# - Load state_dict into model on cpu
# - Move to GPU
# class InvokeInt8Params(bnb.nn.Int8Params):
# """Overrides `bnb.nn.Int8Params` to add the following functionality:
# - Make it possible to load a quantized state dict without putting the weight on a "cuda" device.
# """
# def quantize(self, device: Optional[torch.device] = None):
# device = device or torch.device("cuda")
# if device.type != "cuda":
# raise RuntimeError(f"Int8Params quantization is only supported on CUDA devices ({device=}).")
# # https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L291-L302
# B = self.data.contiguous().half().cuda(device)
# if self.has_fp16_weights:
# self.data = B
# else:
# # we store the 8-bit rows-major weight
# # we convert this weight to the turning/ampere weight during the first inference pass
# CB, CBt, SCB, SCBt, coo_tensorB = bnb.functional.double_quant(B)
# del CBt
# del SCBt
# self.data = CB
# self.CB = CB
# self.SCB = SCB
class Invoke2Linear8bitLt(torch.nn.Linear):
"""This class is the base module for the [LLM.int8()](https://arxiv.org/abs/2208.07339) algorithm."""
def __init__(
self,
input_features: int,
output_features: int,
bias=True,
has_fp16_weights=True,
memory_efficient_backward=False,
threshold=0.0,
index=None,
device=None,
):
"""
Initialize Linear8bitLt class.
Args:
input_features (`int`):
Number of input features of the linear layer.
output_features (`int`):
Number of output features of the linear layer.
bias (`bool`, defaults to `True`):
Whether the linear class uses the bias term as well.
"""
super().__init__(input_features, output_features, bias, device)
assert not memory_efficient_backward, "memory_efficient_backward is no longer required and the argument is deprecated in 0.37.0 and will be removed in 0.39.0"
self.state = bnb.MatmulLtState()
self.index = index
self.state.threshold = threshold
self.state.has_fp16_weights = has_fp16_weights
self.state.memory_efficient_backward = memory_efficient_backward
if threshold > 0.0 and not has_fp16_weights:
self.state.use_pool = True
self.weight = Int8Params(self.weight.data, has_fp16_weights=has_fp16_weights, requires_grad=has_fp16_weights)
self._register_load_state_dict_pre_hook(maybe_rearrange_weight)
def _save_to_state_dict(self, destination, prefix, keep_vars):
super()._save_to_state_dict(destination, prefix, keep_vars)
# we only need to save SCB as extra data, because CB for quantized weights is already stored in weight.data
scb_name = "SCB"
# case 1: .cuda was called, SCB is in self.weight
param_from_weight = getattr(self.weight, scb_name)
# case 2: self.init_8bit_state was called, SCB is in self.state
param_from_state = getattr(self.state, scb_name)
# case 3: SCB is in self.state, weight layout reordered after first forward()
layout_reordered = self.state.CxB is not None
key_name = prefix + f"{scb_name}"
format_name = prefix + "weight_format"
if not self.state.has_fp16_weights:
if param_from_weight is not None:
destination[key_name] = param_from_weight if keep_vars else param_from_weight.detach()
destination[format_name] = torch.tensor(0, dtype=torch.uint8)
elif param_from_state is not None and not layout_reordered:
destination[key_name] = param_from_state if keep_vars else param_from_state.detach()
destination[format_name] = torch.tensor(0, dtype=torch.uint8)
elif param_from_state is not None:
destination[key_name] = param_from_state if keep_vars else param_from_state.detach()
weights_format = self.state.formatB
# At this point `weights_format` is an str
if weights_format not in LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING:
raise ValueError(f"Unrecognized weights format {weights_format}")
weights_format = LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING[weights_format]
destination[format_name] = torch.tensor(weights_format, dtype=torch.uint8)
def _load_from_state_dict(
self,
state_dict,
prefix,
local_metadata,
strict,
missing_keys,
unexpected_keys,
error_msgs,
):
super()._load_from_state_dict(
state_dict,
prefix,
local_metadata,
strict,
missing_keys,
unexpected_keys,
error_msgs,
)
unexpected_copy = list(unexpected_keys)
for key in unexpected_copy:
input_name = key[len(prefix) :]
if input_name == "SCB":
if self.weight.SCB is None:
# buffers not yet initialized, can't access them directly without quantizing first
raise RuntimeError(
"Loading a quantized checkpoint into non-quantized Linear8bitLt is "
"not supported. Please call module.cuda() before module.load_state_dict()",
)
input_param = state_dict[key]
self.weight.SCB.copy_(input_param)
if self.state.SCB is not None:
self.state.SCB = self.weight.SCB
unexpected_keys.remove(key)
def init_8bit_state(self):
self.state.CB = self.weight.CB
self.state.SCB = self.weight.SCB
self.weight.CB = None
self.weight.SCB = None
def forward(self, x: torch.Tensor):
self.state.is_training = self.training
if self.weight.CB is not None:
self.init_8bit_state()
# weights are cast automatically as Int8Params, but the bias has to be cast manually
if self.bias is not None and self.bias.dtype != x.dtype:
self.bias.data = self.bias.data.to(x.dtype)
out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)
if not self.state.has_fp16_weights:
if self.state.CB is not None and self.state.CxB is not None:
# we converted 8-bit row major to turing/ampere format in the first inference pass
# we no longer need the row-major weight
del self.state.CB
self.weight.data = self.state.CxB
return out
class InvokeLinear8bitLt(bnb.nn.Linear8bitLt):
"""Wraps `bnb.nn.Linear8bitLt` and adds the following functionality:
- enables instantiation directly on the device
- re-quantizaton when loading the state dict
"""
def __init__(
self, *args: Any, device: Optional[torch.device] = None, threshold: float = 6.0, **kwargs: Any
) -> None:
super().__init__(*args, device=device, threshold=threshold, **kwargs)
# If the device is CUDA or we are under a CUDA context manager, quantize the weight here, so we don't end up
# filling the device memory with float32 weights which could lead to OOM
# if torch.tensor(0, device=device).device.type == "cuda":
# self.quantize_()
# self._register_load_state_dict_pre_hook(partial(_quantize_on_load_hook, self.quantize_))
# self.register_load_state_dict_post_hook(_ignore_missing_weights_hook)
def _load_from_state_dict(
self,
state_dict,
prefix,
local_metadata,
strict,
missing_keys,
unexpected_keys,
error_msgs,
):
super()._load_from_state_dict(
state_dict,
prefix,
local_metadata,
strict,
missing_keys,
unexpected_keys,
error_msgs,
)
unexpected_copy = list(unexpected_keys)
for key in unexpected_copy:
input_name = key[len(prefix) :]
if input_name == "SCB":
if self.weight.SCB is None:
# buffers not yet initialized, can't access them directly without quantizing first
raise RuntimeError(
"Loading a quantized checkpoint into non-quantized Linear8bitLt is "
"not supported. Please call module.cuda() before module.load_state_dict()",
)
input_param = state_dict[key]
self.weight.SCB.copy_(input_param)
if self.state.SCB is not None:
self.state.SCB = self.weight.SCB
unexpected_keys.remove(key)
def quantize_(self, weight: Optional[torch.Tensor] = None, device: Optional[torch.device] = None) -> None:
"""Inplace quantize."""
if weight is None:
weight = self.weight.data
if weight.data.dtype == torch.int8:
# already quantized
return
assert isinstance(self.weight, bnb.nn.Int8Params)
self.weight = self.quantize(self.weight, weight, device)
@staticmethod
def quantize(
int8params: bnb.nn.Int8Params, weight: torch.Tensor, device: Optional[torch.device]
) -> bnb.nn.Int8Params:
device = device or torch.device("cuda")
if device.type != "cuda":
raise RuntimeError(f"Unexpected device type: {device.type}")
# https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L291-L302
B = weight.contiguous().to(device=device, dtype=torch.float16)
if int8params.has_fp16_weights:
int8params.data = B
else:
CB, CBt, SCB, SCBt, _ = bnb.functional.double_quant(B)
del CBt
del SCBt
int8params.data = CB
int8params.CB = CB
int8params.SCB = SCB
return int8params
# class _Linear4bit(bnb.nn.Linear4bit):
# """Wraps `bnb.nn.Linear4bit` to enable: instantiation directly on the device, re-quantizaton when loading the
# state dict, meta-device initialization, and materialization."""
# def __init__(self, *args: Any, device: Optional[torch.device] = None, **kwargs: Any) -> None:
# super().__init__(*args, device=device, **kwargs)
# self.weight = cast(bnb.nn.Params4bit, self.weight) # type: ignore[has-type]
# self.bias = cast(Optional[torch.nn.Parameter], self.bias) # type: ignore[has-type]
# # if the device is CUDA or we are under a CUDA context manager, quantize the weight here, so we don't end up
# # filling the device memory with float32 weights which could lead to OOM
# if torch.tensor(0, device=device).device.type == "cuda":
# self.quantize_()
# self._register_load_state_dict_pre_hook(partial(_quantize_on_load_hook, self.quantize_))
# self.register_load_state_dict_post_hook(_ignore_missing_weights_hook)
# def quantize_(self, weight: Optional[torch.Tensor] = None, device: Optional[torch.device] = None) -> None:
# """Inplace quantize."""
# if weight is None:
# weight = self.weight.data
# if weight.data.dtype == torch.uint8:
# # already quantized
# return
# assert isinstance(self.weight, bnb.nn.Params4bit)
# self.weight = self.quantize(self.weight, weight, device)
# @staticmethod
# def quantize(
# params4bit: bnb.nn.Params4bit, weight: torch.Tensor, device: Optional[torch.device]
# ) -> bnb.nn.Params4bit:
# device = device or torch.device("cuda")
# if device.type != "cuda":
# raise RuntimeError(f"Unexpected device type: {device.type}")
# # https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L156-L159
# w = weight.contiguous().to(device=device, dtype=torch.half)
# w_4bit, quant_state = bnb.functional.quantize_4bit(
# w,
# blocksize=params4bit.blocksize,
# compress_statistics=params4bit.compress_statistics,
# quant_type=params4bit.quant_type,
# )
# return _replace_param(params4bit, w_4bit, quant_state)
# def to_empty(self, *, device: _DEVICE, recurse: bool = True) -> Self:
# if self.weight.dtype == torch.uint8: # was quantized
# # cannot init the quantized params directly
# weight = torch.empty(self.weight.quant_state.shape, device=device, dtype=torch.half)
# else:
# weight = torch.empty_like(self.weight.data, device=device)
# device = torch.device(device)
# if device.type == "cuda": # re-quantize
# self.quantize_(weight, device)
# else:
# self.weight = _replace_param(self.weight, weight)
# if self.bias is not None:
# self.bias = _replace_param(self.bias, torch.empty_like(self.bias, device=device))
# return self
def convert_model_to_bnb_llm_int8(model: torch.nn.Module, ignore_modules: set[str]):
linear_cls = InvokeLinear8bitLt
_convert_linear_layers(model, linear_cls, ignore_modules)
# TODO(ryand): Is this necessary?
# set the compute dtype if necessary
# for m in model.modules():
# if isinstance(m, bnb.nn.Linear4bit):
# m.compute_dtype = self.dtype
# m.compute_type_is_set = False
# class BitsandbytesPrecision(Precision):
# """Plugin for quantizing weights with `bitsandbytes <https://github.com/TimDettmers/bitsandbytes>`__.
# .. warning:: This is an :ref:`experimental <versioning:Experimental API>` feature.
# .. note::
# The optimizer is not automatically replaced with ``bitsandbytes.optim.Adam8bit`` or equivalent 8-bit optimizers.
# Args:
# mode: The quantization mode to use.
# dtype: The compute dtype to use.
# ignore_modules: The submodules whose Linear layers should not be replaced, for example. ``{"lm_head"}``.
# This might be desirable for numerical stability. The string will be checked in as a prefix, so a value like
# "transformer.blocks" will ignore all linear layers in all of the transformer blocks.
# """
# def __init__(
# self,
# mode: Literal["nf4", "nf4-dq", "fp4", "fp4-dq", "int8", "int8-training"],
# dtype: Optional[torch.dtype] = None,
# ignore_modules: Optional[Set[str]] = None,
# ) -> None:
# if dtype is None:
# # try to be smart about the default selection
# if mode.startswith("int8"):
# dtype = torch.float16
# else:
# dtype = (
# torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
# )
# if mode.startswith("int8") and dtype is not torch.float16:
# # this limitation is mentioned in https://huggingface.co/blog/hf-bitsandbytes-integration#usage
# raise ValueError(f"{mode!r} only works with `dtype=torch.float16`, but you chose `{dtype}`")
# globals_ = globals()
# mode_to_cls = {
# "nf4": globals_["_NF4Linear"],
# "nf4-dq": globals_["_NF4DQLinear"],
# "fp4": globals_["_FP4Linear"],
# "fp4-dq": globals_["_FP4DQLinear"],
# "int8-training": globals_["_Linear8bitLt"],
# "int8": globals_["_Int8LinearInference"],
# }
# self._linear_cls = mode_to_cls[mode]
# self.dtype = dtype
# self.ignore_modules = ignore_modules or set()
# @override
# def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
# # avoid naive users thinking they quantized their model
# if not any(isinstance(m, torch.nn.Linear) for m in module.modules()):
# raise TypeError(
# "You are using the bitsandbytes precision plugin, but your model has no Linear layers. This plugin"
# " won't work for your model."
# )
# # convert modules if they haven't been converted already
# if not any(isinstance(m, (bnb.nn.Linear8bitLt, bnb.nn.Linear4bit)) for m in module.modules()):
# # this will not quantize the model but only replace the layer classes
# _convert_layers(module, self._linear_cls, self.ignore_modules)
# # set the compute dtype if necessary
# for m in module.modules():
# if isinstance(m, bnb.nn.Linear4bit):
# m.compute_dtype = self.dtype
# m.compute_type_is_set = False
# return module
# def _quantize_on_load_hook(quantize_fn: Callable[[torch.Tensor], None], state_dict: OrderedDict, *_: Any) -> None:
# # There is only one key that ends with `*.weight`, the other one is the bias
# weight_key = next((name for name in state_dict if name.endswith("weight")), None)
# if weight_key is None:
# return
# # Load the weight from the state dict and re-quantize it
# weight = state_dict.pop(weight_key)
# quantize_fn(weight)
# def _ignore_missing_weights_hook(module: torch.nn.Module, incompatible_keys: _IncompatibleKeys) -> None:
# # since we manually loaded the weight in the `_quantize_on_load_hook` hook, we need to avoid this missing key false
# # positive
# for key in reversed(incompatible_keys.missing_keys):
# if key.endswith("weight"):
# incompatible_keys.missing_keys.remove(key)
def _convert_linear_layers(
module: torch.nn.Module, linear_cls: Type, ignore_modules: Set[str], prefix: str = ""
) -> None:
for name, child in module.named_children():
fullname = f"{prefix}.{name}" if prefix else name
if isinstance(child, torch.nn.Linear) and not any(fullname.startswith(s) for s in ignore_modules):
has_bias = child.bias is not None
# since we are going to copy over the child's data, the device doesn't matter. I chose CPU
# to avoid spiking CUDA memory even though initialization is slower
# 4bit layers support quantizing from meta-device params so this is only relevant for 8-bit
_Linear4bit = globals()["_Linear4bit"]
device = torch.device("meta" if issubclass(linear_cls, _Linear4bit) else "cpu")
replacement = linear_cls(
child.in_features,
child.out_features,
bias=has_bias,
device=device,
)
if has_bias:
replacement.bias = _replace_param(replacement.bias, child.bias.data.clone())
state = {"quant_state": replacement.weight.quant_state if issubclass(linear_cls, _Linear4bit) else None}
replacement.weight = _replace_param(replacement.weight, child.weight.data.clone(), **state)
module.__setattr__(name, replacement)
else:
_convert_linear_layers(child, linear_cls, ignore_modules, prefix=fullname)
# def _replace_linear_layers(
# model: torch.nn.Module,
# linear_layer_type: Literal["Linear8bitLt", "Linear4bit"],
# modules_to_not_convert: set[str],
# current_key_name: str | None = None,
# ):
# has_been_replaced = False
# for name, module in model.named_children():
# if current_key_name is None:
# current_key_name = []
# current_key_name.append(name)
# if isinstance(module, torch.nn.Linear) 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)
# proceed = True
# for key in modules_to_not_convert:
# if (
# (key in current_key_name_str) and (key + "." in current_key_name_str)
# ) or key == current_key_name_str:
# proceed = False
# break
# if proceed:
# # Load bnb module with empty weight and replace ``nn.Linear` module
# if bnb_quantization_config.load_in_8bit:
# bnb_module = bnb.nn.Linear8bitLt(
# module.in_features,
# module.out_features,
# module.bias is not None,
# has_fp16_weights=False,
# threshold=bnb_quantization_config.llm_int8_threshold,
# )
# elif bnb_quantization_config.load_in_4bit:
# bnb_module = bnb.nn.Linear4bit(
# module.in_features,
# module.out_features,
# module.bias is not None,
# bnb_quantization_config.bnb_4bit_compute_dtype,
# compress_statistics=bnb_quantization_config.bnb_4bit_use_double_quant,
# quant_type=bnb_quantization_config.bnb_4bit_quant_type,
# )
# else:
# raise ValueError("load_in_8bit and load_in_4bit can't be both False")
# bnb_module.weight.data = module.weight.data
# if module.bias is not None:
# bnb_module.bias.data = module.bias.data
# bnb_module.requires_grad_(False)
# setattr(model, name, bnb_module)
# has_been_replaced = True
# if len(list(module.children())) > 0:
# _, _has_been_replaced = _replace_with_bnb_layers(
# module, bnb_quantization_config, modules_to_not_convert, current_key_name
# )
# has_been_replaced = has_been_replaced | _has_been_replaced
# # Remove the last key for recursion
# current_key_name.pop(-1)
# return model, has_been_replaced