From 3128dc8dda12c1f15bbcd0462e0f041579b20a6e Mon Sep 17 00:00:00 2001 From: nemo Date: Tue, 19 May 2026 16:11:22 +0200 Subject: [PATCH 1/6] Add quantization support for SHiRA This builds upon PR #3117 and was used to review the effectiveness of the changes. Apparently the changes are good! On the way, some things were cleaned up like determining the in/out features using the appropriate utility function during init and masking. --- src/peft/tuners/shira/layer.py | 53 +++++++++++++++---------- src/peft/tuners/shira/mask_functions.py | 8 ++-- src/peft/tuners/shira/model.py | 23 +++++++++-- tests/test_quantization.py | 6 ++- 4 files changed, 61 insertions(+), 29 deletions(-) diff --git a/src/peft/tuners/shira/layer.py b/src/peft/tuners/shira/layer.py index db80e73569..f6d70a1d48 100644 --- a/src/peft/tuners/shira/layer.py +++ b/src/peft/tuners/shira/layer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import warnings from typing import Optional @@ -20,7 +19,8 @@ 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 @@ -37,20 +37,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 +82,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 +92,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 +157,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_weights += self.get_delta_weight(active_adapter) + orig_weight = self.get_base_weight().clone() + orig_weight += 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 +181,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 +196,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: @@ -205,13 +211,15 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: elif self.merged: result = self.base_layer(x, *args, **kwargs) else: - new_weight = copy.deepcopy(self.base_layer.weight.data) + base_result = self.base_layer(x) + 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 = F.linear(x, new_weight, bias=self.base_layer.bias) + result = base_result + F.linear(x, new_weight) return result @@ -222,3 +230,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) diff --git a/src/peft/tuners/shira/mask_functions.py b/src/peft/tuners/shira/mask_functions.py index 3304989b3f..3f96b962ba 100644 --- a/src/peft/tuners/shira/mask_functions.py +++ b/src/peft/tuners/shira/mask_functions.py @@ -55,9 +55,11 @@ 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_shira_weights = r * (shape[0] + shape[1]) random_generator = torch.Generator() if random_seed is not None: @@ -65,8 +67,8 @@ def random_mask(base_layer: nn.Module, r: int, random_seed: Optional[int] = None 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)) + val = torch.ones(*idx.shape, dtype=bool) + mask = torch.zeros(*shape).to(val).view(1, -1) mask = mask.scatter_(1, idx.unsqueeze(0), val.unsqueeze(0)).view(shape) return mask diff --git a/src/peft/tuners/shira/model.py b/src/peft/tuners/shira/model.py index d07b9a2d49..78be90d7ab 100644 --- a/src/peft/tuners/shira/model.py +++ b/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, diff --git a/tests/test_quantization.py b/tests/test_quantization.py index 30980d593f..e015075e7f 100644 --- a/tests/test_quantization.py +++ b/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}, + ), ( VeraConfig, {"r": 8, "target_modules": ["q_proj", "v_proj"]}, From 7ff4b5d7d25b85a827d7479544f4ebdd7a738b71 Mon Sep 17 00:00:00 2001 From: nemo Date: Tue, 30 Jun 2026 20:35:59 +0200 Subject: [PATCH 2/6] Move mask to device As suggested by Benjamin. --- src/peft/tuners/shira/mask_functions.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/peft/tuners/shira/mask_functions.py b/src/peft/tuners/shira/mask_functions.py index 3f96b962ba..e5b4a3d6f9 100644 --- a/src/peft/tuners/shira/mask_functions.py +++ b/src/peft/tuners/shira/mask_functions.py @@ -64,11 +64,10 @@ def random_mask(base_layer: nn.Module, r: int, random_seed: Optional[int] = None 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(*idx.shape, dtype=bool) - mask = torch.zeros(*shape).to(val).view(1, -1) + device = base_layer.weight.device + idx = (torch.randperm(base_layer.weight.numel(), 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 From 102ddff9b898972b2dd8cc36ae803772410cf674 Mon Sep 17 00:00:00 2001 From: nemo Date: Tue, 30 Jun 2026 20:36:35 +0200 Subject: [PATCH 3/6] forward: Use original way of gathering delta weights Apparently avoiding the copy is costing quite a lot of runtime performance. So, I reverted the change back to the iterative addition with final nn.linear call. The test and mask changes are necessary to make the random mask generation consistent across (and among) test calls. This was especially important for bnb 4bit quantization since `weight.numel()` for 4bit weights returns `numel // 2` due to packing. --- src/peft/tuners/shira/layer.py | 5 ++--- src/peft/tuners/shira/mask_functions.py | 3 ++- tests/test_quantization.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/peft/tuners/shira/layer.py b/src/peft/tuners/shira/layer.py index f6d70a1d48..d0e657aade 100644 --- a/src/peft/tuners/shira/layer.py +++ b/src/peft/tuners/shira/layer.py @@ -211,15 +211,14 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: elif self.merged: result = self.base_layer(x, *args, **kwargs) else: - base_result = self.base_layer(x) - new_weight = torch.zeros(self.out_features, self.in_features).to(base_result) + new_weight = self.get_base_weight().clone() 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) + result = F.linear(x, new_weight, bias=self.base_layer.bias) return result diff --git a/src/peft/tuners/shira/mask_functions.py b/src/peft/tuners/shira/mask_functions.py index e5b4a3d6f9..80bb733639 100644 --- a/src/peft/tuners/shira/mask_functions.py +++ b/src/peft/tuners/shira/mask_functions.py @@ -60,12 +60,13 @@ def mask_fn(base_layer, r): def random_mask(base_layer: nn.Module, r: int, random_seed: Optional[int] = None, **kwargs) -> torch.tensor: 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) device = base_layer.weight.device - idx = (torch.randperm(base_layer.weight.numel(), generator=random_generator)[:num_shira_weights]).to(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) diff --git a/tests/test_quantization.py b/tests/test_quantization.py index e015075e7f..0d7a2e7ca9 100644 --- a/tests/test_quantization.py +++ b/tests/test_quantization.py @@ -178,7 +178,7 @@ def _quant_id(backend): ), ( ShiraConfig, - {"r": 8}, + {"r": 8, "random_seed": 42}, ), ( VeraConfig, From 09bebe1e1ace8c4b06c958697b501af452c03c5e Mon Sep 17 00:00:00 2001 From: nemo Date: Mon, 6 Jul 2026 17:09:46 +0200 Subject: [PATCH 4/6] Warn when using more efficient forward There are now two forward implementations: a slower fallback implementation using the base layer's forward and a faster implementation that uses the base weights instead, ignoring potential forward/backward hooks. Since the latter is significantly faster in training and inference it is a worthy implementation but we need at least warn the user that their forward/backward hooks are not being called. --- src/peft/tuners/shira/layer.py | 35 ++++++++++++++++++++++++++++++++++ tests/test_shira.py | 11 +++++++++++ 2 files changed, 46 insertions(+) diff --git a/src/peft/tuners/shira/layer.py b/src/peft/tuners/shira/layer.py index d0e657aade..72a0e514c6 100644 --- a/src/peft/tuners/shira/layer.py +++ b/src/peft/tuners/shira/layer.py @@ -13,6 +13,7 @@ # limitations under the License. import warnings +from functools import lru_cache from typing import Optional import torch @@ -111,6 +112,26 @@ def set_scale(self, adapter, scale): return self.scaling[adapter] = scale + @lru_cache(None) # noqa: B019 + def _warn_once_about_module_hooks(self): + # 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 = self.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 Linear(nn.Module, ShiraLayer): # SHiRA implemented in a dense layer @@ -210,7 +231,21 @@ 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: + self._warn_once_about_module_hooks() new_weight = self.get_base_weight().clone() for active_adapter in self.active_adapters: diff --git a/tests/test_shira.py b/tests/test_shira.py index 9845ee426e..a933d60732 100644 --- a/tests/test_shira.py +++ b/tests/test_shira.py @@ -15,6 +15,7 @@ # This test file is for tests specific to SHiRA. import os +from unittest.mock import patch import pytest import torch @@ -276,3 +277,13 @@ def test_shira_dtypes(self, dtype): inputs = torch.randn(5, 10).to(dtype) output = peft_model(inputs) # should not raise assert output.dtype == dtype + + def test_shira_warns_about_hooks(self): + model = MLP() + + with patch('peft.tuners.shira.layer.ShiraLayer._warn_once_about_module_hooks') as m: + config = ShiraConfig(r=2, target_modules=["lin1", "lin2"]) + peft_model = get_peft_model(model, config) + inputs = torch.randn(5, 10) + _ = peft_model(inputs) + assert m.call_count > 0 From b54bbf46be369a3d20f9edd71eacb0103a80aedf Mon Sep 17 00:00:00 2001 From: nemo Date: Mon, 6 Jul 2026 17:36:47 +0200 Subject: [PATCH 5/6] Make style --- tests/test_shira.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_shira.py b/tests/test_shira.py index a933d60732..36f2a21d60 100644 --- a/tests/test_shira.py +++ b/tests/test_shira.py @@ -281,7 +281,7 @@ def test_shira_dtypes(self, dtype): def test_shira_warns_about_hooks(self): model = MLP() - with patch('peft.tuners.shira.layer.ShiraLayer._warn_once_about_module_hooks') as m: + with patch("peft.tuners.shira.layer.ShiraLayer._warn_once_about_module_hooks") as m: config = ShiraConfig(r=2, target_modules=["lin1", "lin2"]) peft_model = get_peft_model(model, config) inputs = torch.randn(5, 10) From c0020e5ccf315980a53d11222ecb316866f8b707 Mon Sep 17 00:00:00 2001 From: nemo Date: Tue, 7 Jul 2026 13:27:48 +0200 Subject: [PATCH 6/6] Improve warning test and reviewer feedback The tests now make sure that the warning is only ever issued once and the code reflects that test. Previously, each instance of ShiraLayer had its own warning cache and therefore was able to emit the warning more than once. Now the state is global. --- src/peft/tuners/shira/layer.py | 44 ++++++++++++++++++---------------- tests/test_shira.py | 39 ++++++++++++++++++++++++------ 2 files changed, 55 insertions(+), 28 deletions(-) diff --git a/src/peft/tuners/shira/layer.py b/src/peft/tuners/shira/layer.py index 72a0e514c6..0437eef548 100644 --- a/src/peft/tuners/shira/layer.py +++ b/src/peft/tuners/shira/layer.py @@ -26,6 +26,28 @@ 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",) @@ -112,26 +134,6 @@ def set_scale(self, adapter, scale): return self.scaling[adapter] = scale - @lru_cache(None) # noqa: B019 - def _warn_once_about_module_hooks(self): - # 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 = self.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 Linear(nn.Module, ShiraLayer): # SHiRA implemented in a dense layer @@ -245,7 +247,7 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: result = base_result + F.linear(x, new_weight) else: - self._warn_once_about_module_hooks() + _warn_once_about_module_hooks(self) new_weight = self.get_base_weight().clone() for active_adapter in self.active_adapters: diff --git a/tests/test_shira.py b/tests/test_shira.py index 36f2a21d60..1f55b8a659 100644 --- a/tests/test_shira.py +++ b/tests/test_shira.py @@ -15,7 +15,6 @@ # This test file is for tests specific to SHiRA. import os -from unittest.mock import patch import pytest import torch @@ -278,12 +277,38 @@ def test_shira_dtypes(self, dtype): output = peft_model(inputs) # should not raise assert output.dtype == dtype - def test_shira_warns_about_hooks(self): + @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() - with patch("peft.tuners.shira.layer.ShiraLayer._warn_once_about_module_hooks") as m: - config = ShiraConfig(r=2, target_modules=["lin1", "lin2"]) - peft_model = get_peft_model(model, config) - inputs = torch.randn(5, 10) + 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) - assert m.call_count > 0 + + 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