FIX Multiple issues with low precision base models

As a user, it should be possible to manually cast the base model to a
lower precision dtype, float16 or bfloat16, and still have the different
PEFT methods work correctly. Currently, this is not the case for many
PEFT methods, as can be replicated by the added tests.

To understand the problem, it helps to take a step back. By default,
PEFT will treat the adapter weights with high precision, i.e. with
float32. When the base model is lower precision, the user needs to pass
inputs in lower precision too, as otherwise self.base_layer(x) would
fail. However, this low precision input clashes with the high precision
adapter weights.

The solution implemented in this PR is to cast the input to a higher
dtype [1]. That way, the whole adapter operation is conducted in high
precision. Only once that has finished will the final result be cast to
the original dtype. This should lead to better results, but it may
require more memory. Note that this is how LoRA is implemented, so the
changes in this PR bring the other methods more in line with what LoRA
does.

If the user does not want the adapter to be in float32, they can always
pass autocast_adapter_dtype=False when calling get_peft_model or
PeftModel.from_pretrained. This is also tested.

Besides adjusting the forward method to account for these changes, the
merge and unmerge methods also often had to be adjusted, as they did not
correctly account for the base model dtype. Now, those methods should
always conserve the original dtype of the base model.

Note that if, for whatever reason, the input casting in [1] is not
desired, users can use the disable_input_dtype_casting context manager
to disable it (more context information on this feature can be found in
PR huggingface#2353). I updated the corresponding code to be agnostic to the
specific PEFT method (beforehand, it was only for LoRA).

Author: BenjaminBossan

src/peft/helpers.py

@@ -22,6 +22,7 @@
 
 from .peft_model import PeftConfig, PeftModel
 from .tuners.lora import LoraLayer
+from .tuners.tuners_utils import BaseTunerLayer
 
 
 def update_forward_signature(model: PeftModel) -> None:
@@ -218,8 +219,6 @@ def disable_input_dtype_casting(model: nn.Module, active: bool = True):
     """
     Context manager disables input dtype casting to the dtype of the weight.
 
-    Currently specifically works for LoRA.
-
     Parameters:
         model (nn.Module):
             The model containing PEFT modules whose input dtype casting is to be adjusted.
@@ -237,7 +236,7 @@ def disable_input_dtype_casting(model: nn.Module, active: bool = True):
 
     original_values = {}
     for name, module in model.named_modules():
-        if not isinstance(module, LoraLayer):
+        if not isinstance(module, BaseTunerLayer):
             continue
         original_values[name] = module.cast_input_dtype_enabled
         module.cast_input_dtype_enabled = False
@@ -246,7 +245,7 @@ def disable_input_dtype_casting(model: nn.Module, active: bool = True):
         yield
     finally:
         for name, module in model.named_modules():
-            if not isinstance(module, LoraLayer):
+            if not isinstance(module, BaseTunerLayer):
                 continue
             if name in original_values:
                 module.cast_input_dtype_enabled = original_values[name]

src/peft/tuners/boft/layer.py

@@ -220,6 +220,8 @@ def __init__(self, base_layer: nn.Module, **kwargs) -> None:
         # Mark the weight as unmerged
         self._disable_adapters = False
         self.merged_adapters = []
+        # flag to enable/disable casting of input to weight dtype during forward call
+        self.cast_input_dtype_enabled = True
         self.kwargs = kwargs
 
         base_layer = self.get_base_layer()
@@ -503,13 +505,14 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
         for active_adapter in adapter_names:
             if active_adapter in self.boft_R.keys():
                 base_layer = self.get_base_layer()
+                orig_dtype = base_layer.weight.dtype
                 if safe_merge:
                     # Note that safe_merge will be slower than the normal merge
                     # because of the copy operation.
                     orig_weight = base_layer.weight.data.clone()
                     butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter)
                     orig_weight = torch.transpose(orig_weight, 0, 1)
-                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight)
+                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight.to(butterfly_oft_mat.dtype))
                     orig_weight = torch.transpose(orig_weight, 0, 1)
                     orig_weight = orig_weight * boft_s
 
@@ -518,16 +521,16 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                             f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken"
                         )
 
-                    self.base_layer.weight.data = orig_weight.contiguous()
+                    self.base_layer.weight.data = orig_weight.contiguous().to(orig_dtype)
                 else:
                     butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter)
                     orig_weight = base_layer.weight.data.clone()
                     orig_weight = torch.transpose(orig_weight, 0, 1)
-                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight)
+                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight.to(butterfly_oft_mat.dtype))
                     orig_weight = torch.transpose(orig_weight, 0, 1)
                     orig_weight = orig_weight * boft_s
 
-                    self.base_layer.weight.data = orig_weight.contiguous()
+                    self.base_layer.weight.data = orig_weight.contiguous().to(orig_dtype)
 
                 self.merged_adapters.append(active_adapter)
 
@@ -538,17 +541,20 @@ def unmerge(self) -> None:
         if not self.merged:
             warnings.warn("Already unmerged. Nothing to do.")
             return
+
         while len(self.merged_adapters) > 0:
             active_adapter = self.merged_adapters.pop()
+            base_layer = self.get_base_layer()
+            orig_dtype = base_layer.weight.dtype
             if active_adapter in self.boft_R.keys():
                 butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter)
 
-                orig_weight = self.get_base_layer().weight.data.clone()
+                orig_weight = base_layer.weight.data.clone()
                 orig_weight = torch.transpose(orig_weight, 0, 1)
-                orig_weight = torch.mm(butterfly_oft_mat.t(), orig_weight)
+                orig_weight = torch.mm(butterfly_oft_mat.t(), orig_weight.to(butterfly_oft_mat.dtype))
                 orig_weight = torch.transpose(orig_weight, 0, 1)
 
-                self.get_base_layer().weight.data = orig_weight * (1 / boft_s)
+                base_layer.weight.data = (orig_weight * (1 / boft_s)).to(orig_dtype)
 
     def get_delta_weight(self, adapter) -> tuple[torch.Tensor, torch.Tensor]:
         """
@@ -804,6 +810,7 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
         for active_adapter in adapter_names:
             if active_adapter in self.boft_R.keys():
                 base_layer = self.get_base_layer()
+                orig_dtype = base_layer.weight.dtype
                 if safe_merge:
                     # Note that safe_merge will be slower than the normal merge
                     # because of the copy operation.
@@ -814,14 +821,14 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                         self.out_features, self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0]
                     )
                     orig_weight = torch.transpose(orig_weight, 0, 1)
-                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight)
+                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight.to(butterfly_oft_mat.dtype))
                     orig_weight = torch.transpose(orig_weight, 0, 1)
                     orig_weight = orig_weight * boft_s
                     orig_weight = orig_weight.view(
                         self.out_features, self.in_features, base_layer.kernel_size[0], base_layer.kernel_size[0]
                     )
 
-                    self.base_layer.weight.data = orig_weight.contiguous()
+                    self.base_layer.weight.data = orig_weight.contiguous().to(orig_dtype)
                 else:
                     butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter)
 
@@ -830,14 +837,14 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                         self.out_features, self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0]
                     )
                     orig_weight = torch.transpose(orig_weight, 0, 1)
-                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight)
+                    orig_weight = torch.mm(butterfly_oft_mat, orig_weight.to(butterfly_oft_mat.dtype))
                     orig_weight = torch.transpose(orig_weight, 0, 1)
                     orig_weight = orig_weight * boft_s
                     orig_weight = orig_weight.view(
                         self.out_features, self.in_features, base_layer.kernel_size[0], base_layer.kernel_size[0]
                     )
 
-                    self.base_layer.weight.data = orig_weight.contiguous()
+                    self.base_layer.weight.data = orig_weight.contiguous().to(orig_dtype)
 
                 self.merged_adapters.append(active_adapter)
 
@@ -850,26 +857,28 @@ def unmerge(self) -> None:
             return
         while len(self.merged_adapters) > 0:
             active_adapter = self.merged_adapters.pop()
+            base_layer = self.get_base_layer()
+            orig_dtype = base_layer.weight.dtype
             if active_adapter in self.boft_R.keys():
                 butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter)
 
-                orig_weight = self.get_base_layer().weight.data.clone()
+                orig_weight = base_layer.weight.data.clone()
                 orig_weight = orig_weight.view(
                     self.out_features,
-                    self.in_features * self.get_base_layer().kernel_size[0] * self.get_base_layer().kernel_size[0],
+                    self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0],
                 )
                 orig_weight = torch.transpose(orig_weight, 0, 1)
-                orig_weight = torch.mm(butterfly_oft_mat.t(), orig_weight)
+                orig_weight = torch.mm(butterfly_oft_mat.t(), orig_weight.to(butterfly_oft_mat.dtype))
                 orig_weight = torch.transpose(orig_weight, 0, 1)
                 orig_weight = orig_weight * (1 / boft_s)
                 orig_weight = orig_weight.view(
                     self.out_features,
                     self.in_features,
-                    self.get_base_layer().kernel_size[0],
-                    self.get_base_layer().kernel_size[0],
+                    base_layer.kernel_size[0],
+                    base_layer.kernel_size[0],
                 )
 
-                self.get_base_layer().weight.data = orig_weight
+                self.get_base_layer().weight.data = orig_weight.to(orig_dtype)
 
     def get_delta_weight(self, adapter) -> tuple[torch.Tensor, torch.Tensor]:
         """
@@ -968,10 +977,12 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
             scaled_rotated_weight = scaled_rotated_weight.view(
                 self.out_features, self.in_features, self.base_layer.kernel_size[0], self.base_layer.kernel_size[0]
             )
+            x = self._cast_input_dtype(x, scaled_rotated_weight.dtype)
+            bias = self._cast_input_dtype(self.base_layer.bias, scaled_rotated_weight.dtype)
             result = F.conv2d(
                 input=x,
                 weight=scaled_rotated_weight,
-                bias=self.base_layer.bias,
+                bias=bias,
                 padding=self.base_layer.padding[0],
                 stride=self.base_layer.stride[0],
             )

src/peft/tuners/bone/layer.py

@@ -35,6 +35,8 @@ def __init__(self, base_layer: nn.Module, **kwargs) -> None:
         # Mark the weight as unmerged
         self._disable_adapters = False
         self.merged_adapters = []
+        # flag to enable/disable casting of input to weight dtype during forward call
+        self.cast_input_dtype_enabled = True
         self.kwargs = kwargs
 
         base_layer = self.get_base_layer()
@@ -150,6 +152,7 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = N
         for active_adapter in adapter_names:
             if active_adapter in self.bone_block.keys():
                 base_layer = self.get_base_layer()
+                orig_dtype = base_layer.weight.dtype
                 if safe_merge:
                     # Note that safe_merge will be slower than the normal merge
                     # because of the copy operation.
@@ -166,14 +169,14 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = N
                             f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken"
                         )
 
-                    self.base_layer.weight.data = orig_weight
+                    self.base_layer.weight.data = orig_weight.to(orig_dtype)
                 else:
                     if self.bone_fn == "bat":
                         delta_weight = self.get_delta_weight(active_adapter, self.base_layer.weight.data)
-                        self.base_layer.weight.data += delta_weight
+                        self.base_layer.weight.data += delta_weight.to(orig_dtype)
                     else:
                         delta_weight = self.get_delta_weight_bone(active_adapter, self.base_layer.weight.data)
-                        self.base_layer.weight.data = delta_weight
+                        self.base_layer.weight.data = delta_weight.to(orig_dtype)
                 self.merged_adapters.append(active_adapter)
 
     def unmerge(self) -> None:
@@ -183,16 +186,19 @@ def unmerge(self) -> None:
         if not self.merged:
             warnings.warn("Already unmerged. Nothing to do.")
             return
+
         while len(self.merged_adapters) > 0:
             active_adapter = self.merged_adapters.pop()
+            base_layer = self.get_base_layer()
+            orig_dtype = base_layer.weight.dtype
             if active_adapter in self.bone_block.keys():
                 orig_weight = self.get_base_layer().weight.data.clone()
                 if self.bone_fn == "bat":
                     delta_weight = self.get_delta_weight(active_adapter, orig_weight, re=True)
                 else:
                     delta_weight = self.get_delta_weight_bone(active_adapter, orig_weight, re=True)
 
-                self.get_base_layer().weight.data = delta_weight
+                base_layer.weight.data = delta_weight.to(orig_dtype)
 
     def get_delta_weight(self, adapter, orig_weight, re: bool = False) -> torch.Tensor:
         """
@@ -213,12 +219,15 @@ def get_delta_weight(self, adapter, orig_weight, re: bool = False) -> torch.Tens
 
         if cast_to_fp32:
             weight_bone = weight_bone.float()
+        orig_weight = orig_weight.to(weight_bone.dtype)
 
         r = weight_bone.size(-1)
         if re:
             o = orig_weight.reshape(orig_weight.size(0) // r, r, orig_weight.size(1) // r, r).permute(2, 0, 1, 3)
             one = torch.eye(weight_bone.size(-1)).to(weight_bone.device)
+            # inverse must be in float32, after that the dtype can be adjusted if needed
             inv_I_plus_b = torch.inverse(one + weight_bone)
+            inv_I_plus_b = inv_I_plus_b.to(weight_bone.dtype)
             w = (o - weight_bone) @ inv_I_plus_b
             output_tensor = w.permute(1, 2, 0, 3).reshape(*orig_weight.shape)
         else:
@@ -318,7 +327,9 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
                     delta_weight = self.get_delta_weight(active_adapter, orig_weight)
                     orig_weight = orig_weight + delta_weight
 
-                result = F.linear(input=x, weight=orig_weight, bias=self.base_layer.bias)
+                x = self._cast_input_dtype(x, orig_weight.dtype)
+                bias = self.base_layer.bias.to(orig_weight.dtype)
+                result = F.linear(input=x, weight=orig_weight, bias=bias)
             else:
                 result = self.base_layer(x, *args, **kwargs)
                 for active_adapter in self.active_adapters:
@@ -329,6 +340,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
                     if x.size(-1) % r != 0:
                         padding_size = (r - x.size(-1) % r) % r
                         x = F.pad(x, (0, padding_size))
+                    x = self._cast_input_dtype(x, bone.dtype)
                     result = result + torch.sum(x.reshape(*x.shape[:-1], x.size(-1) // r, r), dim=-2) @ bone
 
         result = result.to(previous_dtype)

src/peft/tuners/fourierft/layer.py

@@ -84,12 +84,13 @@ def reset_fourier_parameters(self, adapter_name):
             nn.init.zeros_(self.fourierft_spectrum[adapter_name])
 
     def get_delta_weight(self, adapter) -> torch.Tensor:
+        # careful: ifft2 does not work with float16 or bfloat16
         spectrum = self.fourierft_spectrum[adapter]
         indices = self.indices[adapter].to(spectrum.device)
-        dense_spectrum = torch.zeros(self.out_features, self.in_features, device=spectrum.device, dtype=spectrum.dtype)
-        dense_spectrum[indices[0, :], indices[1, :]] = spectrum
+        dense_spectrum = torch.zeros(self.out_features, self.in_features, device=spectrum.device)
+        dense_spectrum[indices[0, :], indices[1, :]] = spectrum.float()
         delta_weight = torch.fft.ifft2(dense_spectrum).real * self.fourierft_scaling[adapter]
-        return delta_weight
+        return delta_weight.to(spectrum.dtype)
 
 
 class FourierFTLinear(nn.Module, FourierFTLayer):