diff --git a/src/peft/tuners/shira/layer.py b/src/peft/tuners/shira/layer.py index db80e73569..0437eef548 100644 --- a/src/peft/tuners/shira/layer.py +++ b/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) diff --git a/src/peft/tuners/shira/mask_functions.py b/src/peft/tuners/shira/mask_functions.py index 3304989b3f..80bb733639 100644 --- a/src/peft/tuners/shira/mask_functions.py +++ b/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 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..0d7a2e7ca9 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, "random_seed": 42}, + ), ( VeraConfig, {"r": 8, "target_modules": ["q_proj", "v_proj"]}, diff --git a/tests/test_shira.py b/tests/test_shira.py index 9845ee426e..1f55b8a659 100644 --- a/tests/test_shira.py +++ b/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