From 36c0d28c59c73f3c516d86e848562e5d72802c43 Mon Sep 17 00:00:00 2001
From: sterlind <55418321+sterlind@users.noreply.github.com>
Date: Sat, 25 Mar 2023 18:25:42 -0700
Subject: [PATCH] Copy lora.py from alpaca_lora_4bit

See https://github.com/johnsmith0031/alpaca_lora_4bit/blob/main/peft/tuners/lora.py
---
 src/peft/tuners/lora.py | 99 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 90 insertions(+), 9 deletions(-)

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 0f65cbf..588c39b 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -32,10 +32,18 @@ def is_bnb_available():
     return importlib.util.find_spec("bitsandbytes") is not None
 
 
+def is_gptq_available():
+    return importlib.util.find_spec("quant") is not None
+
+
 if is_bnb_available():
     import bitsandbytes as bnb
 
 
+if is_gptq_available():
+    import quant
+
+
 @dataclass
 class LoraConfig(PeftConfig):
     """
@@ -163,6 +171,8 @@ class LoraModel(torch.nn.Module):
                         new_module = MergedLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
                 elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None:
                     new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs)
+                elif isinstance(target, Autograd4bitQuantLinear) and self.peft_config.enable_lora is None:
+                    new_module = Linear4bitLt(target.in_features, target.out_features, bias=bias, **kwargs)
                 elif self.peft_config.enable_lora is not None:
                     kwargs.update({"enable_lora": self.peft_config.enable_lora})
                     if isinstance(target, Conv1D):
@@ -193,17 +203,31 @@ class LoraModel(torch.nn.Module):
 
     def _replace_module(self, parent_module, child_name, new_module, old_module):
         setattr(parent_module, child_name, new_module)
-        new_module.weight = old_module.weight
-        if old_module.bias is not None:
+        if isinstance(old_module, Autograd4bitQuantLinear) and isinstance(new_module, Linear4bitLt):
+            new_module.qweight = old_module.qweight
+            new_module.scales = old_module.scales
+            new_module.zeros = old_module.zeros
             new_module.bias = old_module.bias
-        if getattr(old_module, "state", None) is not None:
-            new_module.state = old_module.state
-            new_module.to(old_module.weight.device)
+            if getattr(old_module, "state", None) is not None:
+                new_module.state = old_module.state
+                new_module.to(old_module.qweight.device)
 
-        # dispatch to correct device
-        for name, module in new_module.named_modules():
-            if "lora_" in name:
-                module.to(old_module.weight.device)
+            # dispatch to correct device
+            for name, module in new_module.named_modules():
+                if "lora_" in name:
+                    module.to(old_module.qweight.device)
+        else:
+            new_module.weight = old_module.weight
+            if old_module.bias is not None:
+                new_module.bias = old_module.bias
+            if getattr(old_module, "state", None) is not None:
+                new_module.state = old_module.state
+                new_module.to(old_module.weight.device)
+
+            # dispatch to correct device
+            for name, module in new_module.named_modules():
+                if "lora_" in name:
+                    module.to(old_module.weight.device)
 
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
@@ -614,3 +638,60 @@ if is_bnb_available():
                     output = self.zero_pad(after_B) * self.scaling
                     result += output
             return result
+
+if is_gptq_available():
+
+    from autograd_4bit import Autograd4bitQuantLinear
+
+    class Linear4bitLt(Autograd4bitQuantLinear, 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,
+        ):
+            Autograd4bitQuantLinear.__init__(
+                self,
+                in_features,
+                out_features
+            )
+            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.qweight.requires_grad = False
+                self.scales.requires_grad = False
+                self.zeros.requires_grad = False
+                self.bias.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 forward(self, x: torch.Tensor):
+            result = super().forward(x)
+
+            if self.disable_adapters:
+                return result
+            elif self.r > 0:
+                if not torch.is_autocast_enabled():
+                    expected_dtype = result.dtype
+
+                    if x.dtype != torch.float32:
+                        x = x.float()
+                    output = self.lora_B(self.lora_A(self.lora_dropout(x))).to(expected_dtype) * self.scaling
+                    result += output
+                else:
+                    output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
+                    result += output
+            return result