Refactor to better support LoRA variants (#2443)

Provide a framework to add LoRA variants like DoRA. This way, it will be
easier in the future to add variants to LoRA without the need to either
copy all the LoRA code and make small changes, or clutter the existing
LoRA code with countless if statements. Adding more LoRA variants in the
future will not balloon the size of the proper LoRA implementation.

The new approach is to add LoraVariant subclass to
peft/tuners/lora/variants.py. Typically, this will require one subclass
per supported layer type. The subclass should basically be stateless by
only implementing static methods (which will facilitate composition,
e.g. if a new variant can be combined with DoRA). The subclass needs to
provide a set of methods (init, merge_safe, merge_unsafe, unmerge,
forward). In the LoRA code itself, these methods will be called if the
corresponding adapter uses the subclass.

The choice which variant to dispatch to is determined by the
resolve_lora_variant method. It is called during update_layer and can be
overridden by each LoRA layer (so e.g. lora.Linear dispatches to another
class than lora.Embedding, or bnb LoRA layers could theoretically
dispatch to a different class than normal lora.Linear).

For now, the only LoRA variant is DoRA. This has been refactored to use
the new approach.

Author: csqaiub

src/peft/tuners/lora/bnb.py

@@ -24,7 +24,7 @@
 from peft.utils.integrations import dequantize_bnb_weight
 from peft.utils.other import transpose
 
-from .layer import LoraLayer
+from .layer import LoraLayer, LoraVariant
 
 
 if is_bnb_available():
@@ -60,6 +60,14 @@ def __init__(
                 lora_bias=lora_bias,
             )
 
+        def resolve_lora_variant(self, *, use_dora: bool, **kwargs) -> Optional[LoraVariant]:
+            if not use_dora:
+                return None
+
+            from .variants import DoraLinearVariant
+
+            return DoraLinearVariant()
+
         def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
             """
             Merge the active adapter weights into the base weights
@@ -85,7 +93,6 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                 warnings.warn(
                     "Merge lora module to 8-bit linear may get different generations due to rounding errors."
                 )
-                lora_data = self.get_delta_weight(active_adapter)
 
                 weight = self.get_base_layer().weight
                 state = self.get_base_layer().state
@@ -95,22 +102,11 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                 # Dequantize the result of identity matrix and int8 weight because bitsandbytes does not support int8
                 # dequantization directly
                 output = dequantize_bnb_weight(weight, state=state)
-                if not self.use_dora[active_adapter]:
+                if active_adapter not in self.lora_variant:  # vanilla LoRA
+                    lora_data = self.get_delta_weight(active_adapter)
                     w_data = output.to(lora_data.dtype).to(lora_data.device) + lora_data
                 else:
-                    # handle dora
-                    # since output already includes scaling, set it to 1 here
-                    weight_norm = (
-                        self.lora_magnitude_vector[active_adapter]
-                        .get_weight_norm(output, lora_data, scaling=1)
-                        .detach()
-                    )
-                    # We need to cache weight_norm because it has to be based on the original weights. We
-                    # cannot calculate it on the fly based on the merged weights when unmerging because its a
-                    # different value
-                    self._cache_store(f"{active_adapter}-weight_norm", weight_norm)
-                    dora_factor = self.lora_magnitude_vector[active_adapter].weight / weight_norm
-                    w_data = dora_factor.view(-1, 1) * (output + lora_data)
+                    w_data = self.lora_variant[active_adapter].merge_safe(self, active_adapter, output)
 
                 if safe_merge and not torch.isfinite(w_data).all():
                     raise ValueError(
@@ -120,6 +116,7 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                 self.get_base_layer().weight = bnb.nn.Int8Params(
                     w_data.to("cpu"), requires_grad=False, has_fp16_weights=weight.has_fp16_weights
                 ).to(weight.device)
+
                 if self.lora_bias[active_adapter]:
                     bias_data = self.get_base_layer().bias.data + self.lora_B[active_adapter].bias
                     if safe_merge and not torch.isfinite(bias_data):
@@ -146,20 +143,18 @@ def unmerge(self) -> None:
                 warnings.warn(
                     "Unmerge lora module to 8-bit linear may get different generations due to rounding errors."
                 )
-                lora_data = self.get_delta_weight(active_adapter)
 
                 weight = self.get_base_layer().weight
                 state = self.get_base_layer().state
                 if state.SCB is None:
                     state.SCB = weight.SCB
                 output = dequantize_bnb_weight(weight, state=state)
 
-                if not self.use_dora[active_adapter]:
+                if active_adapter not in self.lora_variant:  # vanilla LoRA
+                    lora_data = self.get_delta_weight(active_adapter)
                     w_data = output.to(lora_data.dtype).to(lora_data.device) - lora_data
                 else:
-                    weight_norm = self._cache_pop(f"{active_adapter}-weight_norm")
-                    dora_factor = self.lora_magnitude_vector[active_adapter].weight / weight_norm
-                    w_data = output.data / dora_factor.view(-1, 1) - lora_data
+                    w_data = self.lora_variant[active_adapter].unmerge(self, active_adapter, output)
 
                 self.get_base_layer().weight = bnb.nn.Int8Params(
                     w_data.to("cpu"), requires_grad=False, has_fp16_weights=weight.has_fp16_weights
@@ -243,26 +238,20 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
                         expected_dtype = result.dtype
                         x = self._cast_input_dtype(x, lora_A.weight.dtype)
 
-                    if not self.use_dora[active_adapter]:
+                    if active_adapter not in self.lora_variant:  # vanilla LoRA
                         output = lora_B(lora_A(dropout(x))) * scaling
+                        if requires_conversion:
+                            output = output.to(expected_dtype)
+                        result = result + output
                     else:
-                        if isinstance(dropout, torch.nn.Identity) or not self.training:
-                            base_result = result
-                        else:
-                            x = dropout(x)
-                            base_result = None
-
-                        output = self.lora_magnitude_vector[active_adapter](
-                            x,
-                            lora_A=lora_A,
-                            lora_B=lora_B,
-                            scaling=scaling,
-                            base_layer=self.get_base_layer(),
-                            base_result=base_result,
+                        result = self.lora_variant[active_adapter].forward(
+                            self,
+                            active_adapter=active_adapter,
+                            x=x,
+                            result=result,
                         )
-                    if requires_conversion:
-                        output = output.to(expected_dtype)
-                    result = result + output
+                        if requires_conversion:
+                            result = result.to(expected_dtype)
 
             return result
 
@@ -326,6 +315,14 @@ def __init__(
                 lora_bias=lora_bias,
             )
 
+        def resolve_lora_variant(self, *, use_dora: bool) -> Optional[LoraVariant]:
+            if not use_dora:
+                return None
+
+            from .variants import DoraLinearVariant
+
+            return DoraLinearVariant()
+
         def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
             """
             Merge the active adapter weights into the base weights
@@ -354,37 +351,27 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                 # Refer to https://gist.github.com/ChrisHayduk/1a53463331f52dca205e55982baf9930
                 weight = self.get_base_layer().weight
                 kwargs = weight.__dict__
-                lora_data = self.get_delta_weight(active_adapter)
 
                 output = dequantize_bnb_weight(weight, state=weight.quant_state)
-                if not self.use_dora[active_adapter]:
+                if active_adapter not in self.lora_variant:  # vanilla LoRA
+                    lora_data = self.get_delta_weight(active_adapter)
                     w_data = output + lora_data
                 else:
-                    # handle dora
-                    # since output already includes scaling, set it to 1 here
-                    weight_norm = (
-                        self.lora_magnitude_vector[active_adapter]
-                        .get_weight_norm(output, lora_data, scaling=1)
-                        .detach()
-                    )
-                    # We need to cache weight_norm because it has to be based on the original weights. We
-                    # cannot calculate it on the fly based on the merged weights when unmerging because its a
-                    # different value
-                    self._cache_store(f"{active_adapter}-weight_norm", weight_norm)
-                    dora_factor = self.lora_magnitude_vector[active_adapter].weight / weight_norm
-                    w_data = dora_factor.view(-1, 1) * (output + lora_data)
+                    w_data = self.lora_variant[active_adapter].merge_safe(self, active_adapter, output)
 
                 if safe_merge and not torch.isfinite(w_data).all():
                     raise ValueError(
                         f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken"
                     )
+
                 if "bnb_quantized" in kwargs:
                     kwargs["bnb_quantized"] = False
                 kwargs["requires_grad"] = False
                 kwargs.pop("data", None)
                 # torch.compile can introduce attributes preceded by '_', remove them
                 kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
                 self.get_base_layer().weight = bnb.nn.Params4bit(w_data.to("cpu"), **kwargs).to(weight.device)
+
                 if self.lora_bias[active_adapter]:
                     bias_data = self.get_base_layer().bias.data + self.lora_B[active_adapter].bias
                     if safe_merge and not torch.isfinite(bias_data):
@@ -411,23 +398,22 @@ def unmerge(self) -> None:
                     "Unmerge lora module to 4-bit linear may get different generations due to rounding errors."
                 )
 
-                lora_data = self.get_delta_weight(active_adapter)
                 weight = self.get_base_layer().weight
                 kwargs = weight.__dict__
                 output = dequantize_bnb_weight(weight, state=weight.quant_state)
 
-                if not self.use_dora[active_adapter]:
+                if active_adapter not in self.lora_variant:  # vanilla LoRA
+                    lora_data = self.get_delta_weight(active_adapter)
                     w_data = output - lora_data
                 else:
-                    weight_norm = self._cache_pop(f"{active_adapter}-weight_norm")
-                    dora_factor = self.lora_magnitude_vector[active_adapter].weight / weight_norm
-                    w_data = output.data / dora_factor.view(-1, 1) - lora_data
+                    w_data = self.lora_variant[active_adapter].unmerge(self, active_adapter, output)
 
                 if "bnb_quantized" in kwargs:
                     kwargs["bnb_quantized"] = False
                 kwargs["requires_grad"] = False
                 kwargs.pop("data", None)
                 self.get_base_layer().weight = bnb.nn.Params4bit(w_data.to("cpu"), **kwargs).to(weight.device)
+
                 if self.lora_bias[active_adapter]:
                     self.get_base_layer().bias.data -= self.lora_B[active_adapter].bias
 
@@ -512,26 +498,20 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
                         expected_dtype = result.dtype
                         x = self._cast_input_dtype(x, lora_A.weight.dtype)
 
-                    if not self.use_dora[active_adapter]:
+                    if active_adapter not in self.lora_variant:  # vanilla LoRA
                         output = lora_B(lora_A(dropout(x))) * scaling
+                        if requires_conversion:
+                            output = output.to(expected_dtype)
+                        result = result + output
                     else:
-                        if isinstance(dropout, torch.nn.Identity) or not self.training:
-                            base_result = result
-                        else:
-                            x = dropout(x)
-                            base_result = None
-
-                        output = self.lora_magnitude_vector[active_adapter](
-                            x,
-                            lora_A=lora_A,
-                            lora_B=lora_B,
-                            scaling=scaling,
-                            base_layer=self.get_base_layer(),
-                            base_result=base_result,
+                        result = self.lora_variant[active_adapter].forward(
+                            self,
+                            active_adapter=active_adapter,
+                            x=x,
+                            result=result,
                         )
-                    if requires_conversion:
-                        output = output.to(expected_dtype)
-                    result = result + output
+                        if requires_conversion:
+                            result = result.to(expected_dtype)
 
             return result
 

src/peft/tuners/lora/hqq.py

@@ -23,7 +23,7 @@
 from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
 from peft.utils.other import transpose
 
-from .layer import LoraLayer
+from .layer import LoraLayer, LoraVariant
 
 
 if is_hqq_available():
@@ -63,6 +63,14 @@ def __init__(
                 lora_bias=lora_bias,
             )
 
+        def resolve_lora_variant(self, *, use_dora: bool, **kwargs) -> Optional[LoraVariant]:
+            if not use_dora:
+                return None
+
+            from .variants import DoraLinearVariant
+
+            return DoraLinearVariant()
+
         def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
             """
             Merge the active adapter weights into the base weights
@@ -87,26 +95,19 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
 
                 layer = self.get_base_layer()
                 quant_config = {**copy.deepcopy(layer.quant_config), "offload_meta": layer.offload_meta}
-                lora_data = self.get_delta_weight(active_adapter)
 
                 output = layer.dequantize()
-                if not self.use_dora[active_adapter]:
+                if active_adapter not in self.lora_variant:  # vanilla LoRA
+                    lora_data = self.get_delta_weight(active_adapter)
                     w_data = output + lora_data
                 else:
-                    # handle dora
-                    # since output already includes scaling, set it to 1 here
-                    weight_norm = self._get_weight_norm(output, lora_data, scaling=1).detach()
-                    # We need to cache weight_norm because it has to be based on the original weights. We
-                    # cannot calculate it on the fly based on the merged weights when unmerging because its a
-                    # different value
-                    self._cache_store(f"{active_adapter}-weight_norm", weight_norm)
-                    dora_factor = self.lora_magnitude_vector[active_adapter] / weight_norm
-                    w_data = dora_factor.view(-1, 1) * (output + lora_data)
+                    w_data = self.lora_variant[active_adapter].merge_safe(self, active_adapter, output)
 
                 if safe_merge and not torch.isfinite(w_data).all():
                     raise ValueError(
                         f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken"
                     )
+
                 new_hqq_layer = HQQLinear(None, quant_config, compute_dtype=layer.compute_dtype, device=layer.device)
                 quant_config.pop("offload_meta", None)
                 new_hqq_layer.quantize(w_data, **quant_config)
@@ -126,17 +127,15 @@ def unmerge(self) -> None:
                 if active_adapter not in self.lora_A.keys():
                     continue
 
-                lora_data = self.get_delta_weight(active_adapter)
                 layer = self.get_base_layer()
                 quant_config = {**copy.deepcopy(layer.quant_config), "offload_meta": layer.offload_meta}
                 output = layer.dequantize()
 
-                if not self.use_dora[active_adapter]:
-                    w_data = output - lora_data
+                if active_adapter not in self.lora_variant:  # vanilla LoRA
+                    lora_data = self.get_delta_weight(active_adapter)
+                    w_data = output.to(lora_data.dtype).to(lora_data.device) - lora_data
                 else:
-                    weight_norm = self._cache_pop(f"{active_adapter}-weight_norm")
-                    dora_factor = self.lora_magnitude_vector[active_adapter] / weight_norm
-                    w_data = output.data / dora_factor.view(-1, 1) - lora_data
+                    w_data = self.lora_variant[active_adapter].unmerge(self, active_adapter, output)
 
                 new_hqq_layer = HQQLinear(None, quant_config, compute_dtype=layer.compute_dtype, device=layer.device)
                 quant_config.pop("offload_meta", None)
@@ -218,23 +217,16 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
                         expected_dtype = result.dtype
                         x = self._cast_input_dtype(x, lora_A.weight.dtype)
 
-                    if not self.use_dora[active_adapter]:
+                    if active_adapter not in self.lora_variant:  # vanilla LoRA
                         result = result + lora_B(lora_A(dropout(x))) * scaling
                     else:
-                        if isinstance(dropout, torch.nn.Identity) or not self.training:
-                            base_result = result
-                        else:
-                            x = dropout(x)
-                            base_result = None
-
-                        result = result + self.lora_magnitude_vector[active_adapter](
-                            x,
-                            lora_A=lora_A,
-                            lora_B=lora_B,
-                            scaling=scaling,
-                            base_layer=self.get_base_layer(),
-                            base_result=base_result,
+                        result = self.lora_variant[active_adapter].forward(
+                            self,
+                            active_adapter=active_adapter,
+                            x=x,
+                            result=result,
                         )
+
                     if requires_conversion:
                         result = result.to(expected_dtype)