Add quantization support for SHiRA

src/peft/tuners/shira/layer.py

@@ -12,19 +12,42 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 import warnings
+from functools import lru_cache
 from typing import Optional
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
+from peft.tuners.tuners_utils import BaseTunerLayer, _get_in_out_features, check_adapters_to_merge
+from peft.utils import quantization_extra_repr, resolve_quantization_backend
 
 from .config import ShiraConfig
 
 
+# use a LRU_cache so that the warning is only ever called once and not repeated for every layer/step/epoch
+@lru_cache(None)
+def _warn_once_about_module_hooks(shira_layer):
+    # ShiRA has a forward method that uses the base weights instead of the .forward call of the base layer.
+    # This ignores any hook set on the base layer. Inform the user about this so that they can register the
+    # hooks on the PEFT module instead.
+    base_layer = shira_layer.get_base_layer()
+    if any(
+        [
+            base_layer._forward_hooks,
+            base_layer._forward_pre_hooks,
+            base_layer._backward_hooks,
+            base_layer._backward_pre_hooks,
+        ]
+    ):
+        warnings.warn(
+            "One of the base layers adapted with ShiRA has backward/forward (pre) hooks set which will be ignored "
+            "by the adapter's forward implementation. Please set the hooks on the adapted layer instead (i.e., "
+            "apply the hooks on the same path but after applying the PEFT config)."
+        )
+
+
 class ShiraLayer(BaseTunerLayer):
     # List all names of layers that may contain trainable adapter weights
     adapter_layer_names = ("shira_weight",)
@@ -37,20 +60,19 @@ def __init__(self, base_layer: nn.Module, **kwargs):
         self.scaling = {}
         self.shira_weight = nn.ParameterDict({})
         self.shira_indices = {}
-        self.weight_shape = base_layer.weight.shape  # Assumes SHiRA is on some layer with "weight" parameter
+        self.quantization_backend = resolve_quantization_backend(
+            self.get_base_layer(), get_apply_tensor_subclass=kwargs.get("get_apply_tensor_subclass")
+        )
 
         # Mark the weight as unmerged
         self._disable_adapters = False
         self.merged_adapters = []
 
         base_layer = self.get_base_layer()
-        if isinstance(base_layer, nn.Linear):
-            in_features, out_features = base_layer.in_features, base_layer.out_features
-        else:
-            raise NotImplementedError("Only nn.Linear layers supported currently")
+        self.in_features, self.out_features = _get_in_out_features(base_layer)
+        if None in (self.in_features, self.out_features):
+            raise TypeError("Only nn.Linear layers supported currently")
 
-        self.in_features = in_features
-        self.out_features = out_features
         self.kwargs = kwargs
 
     def update_layer(
@@ -83,7 +105,7 @@ def update_layer(
         # https://github.com/pytorch/pytorch/issues/79542.
         shira_init_weight = torch.zeros(num_shira_weight) if init_weights else torch.randn(num_shira_weight)
         self.shira_weight[adapter_name] = nn.Parameter(
-            shira_init_weight.to(self.base_layer.weight.dtype).to(self.base_layer.weight.device),
+            shira_init_weight,
             requires_grad=True,
         )
 
@@ -93,7 +115,7 @@ def update_layer(
             self.shira_indices[adapter_name] = torch.cat(
                 [mask_indices[0].unsqueeze(0), mask_indices[1].unsqueeze(0)], 0
             ).to(torch.int)
-            self.shira_indices[adapter_name] = self.shira_indices[adapter_name].to(self.base_layer.weight.device)
+            self.shira_indices[adapter_name] = self.shira_indices[adapter_name]
 
             if self.shira_indices[adapter_name].shape[1] != self.shira_weight[adapter_name].shape[0]:
                 raise ValueError(
@@ -158,18 +180,20 @@ def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = N
                 if safe_merge:
                     # Note that safe_merge will be slower than the normal merge
                     # because of the copy operation.
-                    orig_weights = base_layer.weight.data.clone()
+                    orig_weight = self.get_base_weight().clone()
+                    orig_weight += self.get_delta_weight(active_adapter)
 
-                    orig_weights += self.get_delta_weight(active_adapter)
-
-                    if not torch.isfinite(orig_weights).all():
+                    if not torch.isfinite(orig_weight).all():
                         raise ValueError(
                             f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken"
                         )
 
-                    base_layer.weight.data = orig_weights
+                    self.set_base_weight(orig_weight)
                 else:
-                    base_layer.weight.data += self.get_delta_weight(active_adapter)
+                    orig_weight = self.get_base_weight()
+                    delta_weight = self.get_delta_weight(active_adapter)
+                    orig_weight += delta_weight
+                    self.set_base_weight(orig_weight)
                 self.merged_adapters.append(active_adapter)
 
     def unmerge(self) -> None:
@@ -180,7 +204,9 @@ def unmerge(self) -> None:
         while len(self.merged_adapters) > 0:
             active_adapter = self.merged_adapters.pop()
             if active_adapter in self.shira_weight.keys():
-                self.get_base_layer().weight.data -= self.get_delta_weight(active_adapter)
+                orig_weight = self.get_base_weight()
+                orig_weight -= self.get_delta_weight(active_adapter)
+                self.set_base_weight(orig_weight)
 
     def get_delta_weight(self, adapter) -> torch.Tensor:
         """
@@ -193,8 +219,11 @@ def get_delta_weight(self, adapter) -> torch.Tensor:
 
         # In multi-gpu environment, the indices are at the wrong gpu.  This is needed to correct this.
         self.shira_indices[adapter] = self.shira_indices[adapter].to(self.shira_weight[adapter].device)
+
         return torch.sparse_coo_tensor(
-            self.shira_indices[adapter], self.shira_weight[adapter] * self.scaling[adapter], self.weight_shape
+            self.shira_indices[adapter],
+            self.shira_weight[adapter] * self.scaling[adapter],
+            (self.out_features, self.in_features),
         )
 
     def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
@@ -204,8 +233,23 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
             result = self.base_layer(x, *args, **kwargs)
         elif self.merged:
             result = self.base_layer(x, *args, **kwargs)
+        elif self.quantization_backend and not self.quantization_backend.supports_merge:
+            # For the normal forward path, self.get_base_weight() needs to be called, but if the quantization backend
+            # doesn't support it, we use a pattern that avoid dequantizing the base weight. The disadvantage is that
+            # this is slower.
+            base_result = self.base_layer(x, *args, **kwargs)
+            new_weight = torch.zeros(self.out_features, self.in_features).to(base_result)
+
+            for active_adapter in self.active_adapters:
+                if active_adapter not in self.shira_weight.keys():
+                    continue
+                new_weight += self.get_delta_weight(active_adapter)
+
+            result = base_result + F.linear(x, new_weight)
         else:
-            new_weight = copy.deepcopy(self.base_layer.weight.data)
+            _warn_once_about_module_hooks(self)
+            new_weight = self.get_base_weight().clone()
+
             for active_adapter in self.active_adapters:
                 if active_adapter not in self.shira_weight.keys():
                     continue
@@ -222,3 +266,6 @@ def supports_lora_conversion(self, adapter_name: str = "default") -> bool:
     def __repr__(self) -> str:
         rep = super().__repr__()
         return "shira." + rep
+
+    def extra_repr(self) -> str:
+        return quantization_extra_repr(self)

src/peft/tuners/shira/mask_functions.py

@@ -55,18 +55,20 @@ def mask_fn(base_layer, r):
 import torch
 from torch import nn
 
+from peft.tuners.tuners_utils import _get_in_out_features
+
 
 def random_mask(base_layer: nn.Module, r: int, random_seed: Optional[int] = None, **kwargs) -> torch.tensor:
-    shape = base_layer.weight.shape
+    shape = _get_in_out_features(base_layer)[::-1]
+    num_base_weights = shape[0] * shape[1]
     num_shira_weights = r * (shape[0] + shape[1])
     random_generator = torch.Generator()
     if random_seed is not None:
         random_generator.manual_seed(random_seed)
-    idx = (torch.randperm(base_layer.weight.numel(), generator=random_generator)[:num_shira_weights]).to(
-        base_layer.weight.device
-    )
-    val = torch.ones_like(idx.type(base_layer.weight.dtype))
-    mask = torch.zeros_like(base_layer.weight.view(1, -1))
+    device = base_layer.weight.device
+    idx = (torch.randperm(num_base_weights, generator=random_generator)[:num_shira_weights]).to(device)
+    val = torch.ones(*idx.shape, dtype=bool).to(device)
+    mask = torch.zeros(*shape).to(val).view(1, -1).to(device)
     mask = mask.scatter_(1, idx.unsqueeze(0), val.unsqueeze(0)).view(shape)
 
     return mask

src/peft/tuners/shira/model.py

@@ -21,11 +21,23 @@
 from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer
 from peft.utils import (
     TRANSFORMERS_MODELS_TO_SHIRA_TARGET_MODULES_MAPPING,
+    get_quantization_kwargs,
+    resolve_quantization_backend,
 )
 
 from .layer import Linear, ShiraLayer
 
 
+def _get_tuner_layer_class(target_base_layer: torch.nn.Module) -> type[ShiraLayer] | None:
+    layer_cls: type[ShiraLayer] | None = None
+    if isinstance(target_base_layer, torch.nn.Linear):
+        layer_cls = Linear
+    elif (quant_backend := resolve_quantization_backend(target_base_layer)) is not None:
+        layer_cls = {"linear": Linear}.get(quant_backend.layer_type)
+
+    return layer_cls
+
+
 class ShiraModel(BaseTuner):
     """
     Creates a Sparse High Rank Adapter (SHiRA) Model from a pretrained model.
@@ -72,7 +84,7 @@ def _create_and_replace(
             raise ValueError("Current Key shouldn't be `None`")
 
         bias = hasattr(target, "bias") and target.bias is not None
-        kwargs = {}
+        kwargs = get_quantization_kwargs(self)
         kwargs["bias"] = bias
         if shira_config.mask_type == "random":
             kwargs["random_seed"] = shira_config.random_seed
@@ -91,6 +103,7 @@ def _create_and_replace(
                 mask,
                 shira_config.r,
                 config=shira_config,
+                **kwargs,
             )
         else:
             new_module = self._create_new_module(shira_config, adapter_name, target, **kwargs)
@@ -108,14 +121,16 @@ def _create_new_module(shira_config, adapter_name, target, **kwargs):
         else:
             target_base_layer = target
 
-        if isinstance(target_base_layer, torch.nn.Linear):
+        layer_cls = _get_tuner_layer_class(target_base_layer)
+
+        if layer_cls is Linear:
             if shira_config.fan_in_fan_out:
                 warnings.warn(
                     "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. "
                     "Setting fan_in_fan_out to False."
                 )
                 shira_config.fan_in_fan_out = False
-        else:
+        elif layer_cls is None:
             raise TypeError(
                 f"Target module {target} is not supported. Currently, only the following modules are supported: "
                 "`torch.nn.Linear`."
@@ -127,7 +142,7 @@ def _create_new_module(shira_config, adapter_name, target, **kwargs):
             else None
         )
 
-        new_module = Linear(
+        new_module = layer_cls(
             target,
             mask,
             adapter_name,

tests/test_quantization.py

@@ -23,7 +23,7 @@
 from accelerate.utils.memory import clear_device_cache
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig, TorchAoConfig
 
-from peft import BOFTConfig, MissConfig, VeraConfig, get_peft_model
+from peft import BOFTConfig, MissConfig, ShiraConfig, VeraConfig, get_peft_model
 from peft.import_utils import (
     is_bnb_4bit_available,
     is_bnb_available,
@@ -176,6 +176,10 @@ def _quant_id(backend):
         MissConfig,
         {"r": 2, "init_weights": "bat"},
     ),
+    (
+        ShiraConfig,
+        {"r": 8, "random_seed": 42},
+    ),
     (
         VeraConfig,
         {"r": 8, "target_modules": ["q_proj", "v_proj"]},

tests/test_shira.py

@@ -276,3 +276,39 @@ def test_shira_dtypes(self, dtype):
         inputs = torch.randn(5, 10).to(dtype)
         output = peft_model(inputs)  # should not raise
         assert output.dtype == dtype
+
+    @pytest.mark.parametrize(
+        "expected_warnings, hook_setter",
+        [
+            (0, lambda m: None),
+            (1, lambda m: m.register_forward_hook(lambda *x: None)),
+            (1, lambda m: m.register_backward_hook(lambda *x: None)),
+            (1, lambda m: m.register_forward_pre_hook(lambda *x: None)),
+        ],
+    )
+    def test_shira_warns_about_hooks(self, expected_warnings, hook_setter):
+        # ShiRA by default uses an efficient forward pass that only uses the base layer's weights,
+        # not its forward. This means that forward/backward hooks on the base layer are not called.
+        # We test that a warning is issued to the user to highlight this fact.
+        import peft
+
+        model = MLP()
+        hook_setter(model.lin1)
+
+        # Reset the 'warn only once' mechanic to make it possible to test this when shira
+        # was instantiated already by this or other tests. This is necessary because it is a global state.
+        peft.tuners.shira.layer._warn_once_about_module_hooks.cache_clear()
+
+        config = ShiraConfig(r=2, target_modules=["lin1", "lin2"])
+        peft_model = get_peft_model(model, config)
+        inputs = torch.randn(5, 10)
+
+        # Test multiple invocations of the layers to make sure the warning is only issued once.
+        # This violates PT031, so we disable it.
+        with pytest.warns() as record:  # noqa: PT031
+            _ = peft_model(inputs)
+            _ = peft_model(inputs)
+
+        warning_match = "One of the base layers adapted with ShiRA"
+        relevant_warnings = [w for w in record if warning_match in str(w.message)]
+        assert len(relevant_warnings) == expected_warnings