diff --git a/setup.py b/setup.py index 575d265..b51aa93 100644 --- a/setup.py +++ b/setup.py @@ -43,7 +43,6 @@ setup( "torch>=1.13.0", "transformers", "accelerate", - "bitsandbytes", ], extras_require=extras, classifiers=[ diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index f7fa9d4..5ff7bbc 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import importlib import math import warnings from dataclasses import asdict, dataclass, field @@ -23,11 +24,17 @@ import torch.nn as nn import torch.nn.functional as F from transformers.pytorch_utils import Conv1D -import bitsandbytes as bnb - from ..utils import PeftConfig, PeftType, transpose +def is_bnb_available(): + return importlib.util.find_spec("bitsandbytes") is not None + + +if is_bnb_available(): + import bitsandbytes as bnb + + @dataclass class LoraConfig(PeftConfig): """ @@ -106,6 +113,12 @@ class LoraModel(torch.nn.Module): self.forward = self.model.forward def _find_and_replace(self): + loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) + if loaded_in_8bit and not is_bnb_available(): + raise ImportError( + "To use Lora with 8-bit quantization, please install the `bitsandbytes` package. " + "You can install it with `pip install bitsandbytes`." + ) is_target_modules_in_base_model = False kwargs = { "r": self.peft_config.r, @@ -121,7 +134,7 @@ class LoraModel(torch.nn.Module): is_target_modules_in_base_model = True parent, target, target_name = self._get_submodules(key) bias = target.bias is not None - if isinstance(target, bnb.nn.Linear8bitLt) and self.peft_config.enable_lora is None: + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt) and self.peft_config.enable_lora is None: kwargs.update( { "has_fp16_weights": target.state.has_fp16_weights, @@ -194,6 +207,27 @@ class LoraModel(torch.nn.Module): # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License (MIT). See LICENSE in the repo root for license information. # ------------------------------------------------------------------------------------------ + + +# had to adapt it for `lora_only` to work +def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: + for n, p in model.named_parameters(): + if "lora_" not in n: + p.requires_grad = False + if bias == "none": + return + elif bias == "all": + for n, p in model.named_parameters(): + if "bias" in n: + p.requires_grad = True + elif bias == "lora_only": + for m in model.modules(): + if isinstance(m, LoraLayer) and hasattr(m, "bias") and m.bias is not None: + m.bias.requires_grad = True + else: + raise NotImplementedError + + class LoraLayer: def __init__( self, @@ -379,64 +413,47 @@ class MergedLinear(nn.Linear, LoraLayer): return result -class Linear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): - # Lora implemented in a dense layer - def __init__( - self, - in_features, - out_features, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0.0, - **kwargs, - ): - bnb.nn.Linear8bitLt.__init__( +if is_bnb_available(): + + class Linear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): + # Lora implemented in a dense layer + def __init__( self, in_features, out_features, - bias=kwargs.get("bias", True), - has_fp16_weights=kwargs.get("has_fp16_weights", True), - memory_efficient_backward=kwargs.get("memory_efficient_backward", False), - threshold=kwargs.get("threshold", 0.0), - index=kwargs.get("index", None), - ) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False) - # Actual trainable parameters - if r > 0: - self.lora_A = nn.Linear(in_features, r, bias=False) - self.lora_B = nn.Linear(r, out_features, bias=False) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - self.reset_parameters() + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + **kwargs, + ): + bnb.nn.Linear8bitLt.__init__( + self, + in_features, + out_features, + bias=kwargs.get("bias", True), + has_fp16_weights=kwargs.get("has_fp16_weights", True), + memory_efficient_backward=kwargs.get("memory_efficient_backward", False), + threshold=kwargs.get("threshold", 0.0), + index=kwargs.get("index", None), + ) + LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False) + # Actual trainable parameters + if r > 0: + self.lora_A = nn.Linear(in_features, r, bias=False) + self.lora_B = nn.Linear(r, out_features, bias=False) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + self.reset_parameters() - def reset_parameters(self): - if hasattr(self, "lora_A"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + def reset_parameters(self): + if hasattr(self, "lora_A"): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B.weight) - def forward(self, x: torch.Tensor): - result = super().forward(x) - if self.r > 0: - result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling - return result - - -# had to adapt it for `lora_only` to work -def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: - for n, p in model.named_parameters(): - if "lora_" not in n: - p.requires_grad = False - if bias == "none": - return - elif bias == "all": - for n, p in model.named_parameters(): - if "bias" in n: - p.requires_grad = True - elif bias == "lora_only": - for m in model.modules(): - if isinstance(m, LoraLayer) and hasattr(m, "bias") and m.bias is not None: - m.bias.requires_grad = True - else: - raise NotImplementedError + def forward(self, x: torch.Tensor): + result = super().forward(x) + if self.r > 0: + result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling + return result