from typing import Optional import torch import torch.nn as nn from enum import Enum from dataclasses import asdict from tqdm import tqdm from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists, onload_layer from peft.utils import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, ModulesToSaveWrapper, _get_submodules from .layer import RotationLayer, Linear TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING.copy() class RotationTuner(BaseTuner): prefix: str = "rotation_" tuner_layer_class = RotationLayer target_module_mapping = TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING @staticmethod def _check_target_module_exists(rotation_config, key: str) -> bool: return check_target_module_exists(rotation_config, key) def _create_and_replace( self, rotation_config, adapter_name: str, target: nn.Module, target_name: str, parent: nn.Module, current_key: str, **optional_kwargs, ) -> None: """ Create and replace a target module with a rotation-augmented version. This method is called when an existing module is already a RotationLayer and needs to have a new adapter added to it. Args: rotation_config: Configuration for the rotation adapter adapter_name: Name of the adapter to add target: The target module to augment target_name: Name of the target module parent: Parent module containing the target current_key: Full key path to the current module **optional_kwargs: Additional optional arguments Raises: ValueError: If current_key is not provided """ if current_key is None: raise ValueError("current_key must be provided to create Rotation layer") # Check if target is already a RotationLayer if isinstance(target, RotationLayer): target.update_layer( adapter_name=adapter_name, r=rotation_config.r, T=rotation_config.T, num_rotations=rotation_config.num_rotations, ) else: # Create new rotation layer new_module = self._create_new_module( rotation_config=rotation_config, adapter_name=adapter_name, target=target, **optional_kwargs, ) if new_module is not None: self._replace_module(parent, target_name, new_module, target) def _replace_module(self, parent, child_name, new_module, child): setattr(parent, child_name, new_module) # child layer wraps the original module, unpack it if hasattr(child, "base_layer"): child = child.base_layer meta = torch.device("meta") # dispatch to correct device for name, module in new_module.named_modules(): if (self.prefix in name) or ("ranknum" in name): if hasattr(child, "qweight"): weight = child.qweight elif hasattr(child, "W_q"): weight = child.W_q elif hasattr(child, "weight"): weight = child.weight elif getattr(child, "in_proj_weight", None) is not None: # MHA weight = child.in_proj_weight else: weight = next(child.parameters()) if not any(p.device == meta for p in module.parameters()): module.to(weight.device) def _mark_only_adapters_as_trainable(self, model): # First, freeze all parameters for n, p in model.named_parameters(): if self.prefix not in n: p.requires_grad = False else: p.requires_grad = True # Handle bias parameters based on config for active_adapter in self.active_adapters: bias_config = self.peft_config[active_adapter].bias if bias_config == "none": continue elif bias_config == "all": # Enable all bias parameters for n, p in model.named_parameters(): if "bias" in n: p.requires_grad = True elif bias_config == "rotation_only": # Enable only bias in rotation layers for name, m in model.named_modules(): if isinstance(m, RotationLayer): if hasattr(m, "bias") and m.bias is not None: m.bias.requires_grad = True else: raise NotImplementedError( f"Requested bias configuration '{bias_config}' is not implemented. " f"Supported values: 'none', 'all', 'rotation_only'" ) @staticmethod def _create_new_module( rotation_config, adapter_name: str, target: nn.Module, **kwargs, ) -> Optional[nn.Module]: """ Create a new rotation-augmented module. Args: rotation_config: Configuration for the rotation adapter adapter_name: Name of the adapter target: Base module to augment **kwargs: Additional arguments Returns: New RotationLayer module wrapping the target, or None if unsupported """ if isinstance(target, nn.Linear): return Linear( base_layer=target, adapter_name=adapter_name, r=rotation_config.r, T=rotation_config.T, num_rotations=rotation_config.num_rotations, **kwargs, ) else: # Unsupported layer type print( f"Rotation layer does not support {type(target).__name__} yet. " f"Skipping this module." ) return None def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" try: return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: if name == "model": # see #1892: prevent infinite recursion if class is not initialized raise return getattr(self.model, name) def get_peft_config_as_dict(self, inference: bool = False): config_dict = {} for key, value in self.peft_config.items(): config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} if inference: config["inference_mode"] = True config_dict[key] = config return config def _set_adapter_layers(self, enabled=True): for module in self.model.modules(): if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): module.enable_adapters(enabled) def enable_adapter_layers(self) -> None: """Enable all adapters. Call this if you have previously disabled all adapters and want to re-enable them. """ self._set_adapter_layers(enabled=True) def disable_adapter_layers(self): for active_adapter in self.active_adapters: val = self.peft_config[active_adapter].bias if val != "none": msg = ( f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same " "output as the base model would without adaption." ) print(msg) self._set_adapter_layers(enabled=False) def set_adapter(self, adapter_name): """Set the active adapter(s). Additionally, this function will set the specified adapters to trainable (i.e., requires_grad=True). If this is not desired, use the following code. ```py >>> for name, param in model_peft.named_parameters(): ... if ...: # some check on name (ex. if 'lora' in name) ... param.requires_grad = False ``` Args: adapter_name (`str` or `list[str]`): Name of the adapter(s) to be activated. """ for module in self.model.modules(): if isinstance(module, RotationLayer): if module.merged: print("Adapter cannot be set when the model is merged. Unmerging the model first.") module.unmerge() module.set_adapter(adapter_name) self.active_adapter = adapter_name def merge_adapter(self, adapter_names: Optional[list[str]] = None) -> None: """ Merge adapter weights into the base model weights. This can speed up inference by eliminating the need for runtime rotation computations. Args: adapter_names: List of adapter names to merge. If None, merges all active adapters. """ for module in self.model.modules(): if isinstance(module, RotationLayer): module.merge(safe_merge=False, adapter_names=adapter_names) def unmerge_adapter(self) -> None: """ Unmerge adapter weights from the base model weights. This reverses the merge operation, restoring dynamic adapter behavior. """ for module in self.model.modules(): if isinstance(module, RotationLayer): module.unmerge() @staticmethod def _prepare_adapter_config(peft_config, model_config): if peft_config.target_modules is None: if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING: raise ValueError("Please specify `target_modules` in `peft_config`") peft_config.target_modules = set( TRANSFORMERS_MODELS_TO_ROTATION_TARGET_MODULES_MAPPING[model_config["model_type"]] ) return peft_config def _check_new_adapter_config(self, config) -> None: """ Check the validity of a new adapter configuration. Args: config: Configuration to validate Raises: ValueError: If configuration is invalid """ # Validate rank if config.r <= 0: raise ValueError(f"r must be positive, got {config.r}") # Validate num_rotations if config.num_rotations <= 0: raise ValueError( f"num_rotations must be positive, got {config.num_rotations}" ) # Validate bias configuration valid_bias_configs = ["none", "all", "rotation_only"] if hasattr(config, "bias") and config.bias not in valid_bias_configs: raise ValueError( f"Invalid bias configuration '{config.bias}'. " f"Must be one of {valid_bias_configs}" ) def _unload_and_optionally_merge( self, merge=True, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[list[str]] = None, ): if merge: self._check_merge_allowed() key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] desc = "Unloading " + ("and merging " if merge else "") + "model" for key in tqdm(key_list, disable=not progressbar, desc=desc): try: parent, target, target_name = _get_submodules(self.model, key) except AttributeError: continue with onload_layer(target): if hasattr(target, "unload_and_optionally_merge_module"): # if layers have special unloading method, like MultiheadAttention, use that unloaded_module = target.unload_and_optionally_merge_module( merge=merge, safe_merge=safe_merge, adapter_names=adapter_names ) self._replace_module(parent, target_name, unloaded_module, target) elif hasattr(target, "base_layer"): if merge: target.merge(safe_merge=safe_merge, adapter_names=adapter_names) self._replace_module(parent, target_name, target.get_base_layer(), target) return self.model def delete_adapter(self, adapter_name: str) -> None: """ Deletes an existing adapter. Args: adapter_name (str): Name of the adapter to be deleted. """ if adapter_name not in list(self.peft_config.keys()): raise ValueError(f"Adapter {adapter_name} does not exist") del self.peft_config[adapter_name] key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] new_adapter = None for key in key_list: _, target, _ = _get_submodules(self.model, key) if isinstance(target, RotationLayer): target.delete_adapter(adapter_name) if new_adapter is None: new_adapter = target.active_adapters[:] self.active_adapter = new_adapter or [] self._delete_auxiliary_adapter(adapter_name, new_active_adapters=new_adapter) def merge_and_unload( self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[list[str]] = None ) -> torch.nn.Module: r""" This method merges the OFT layers into the base model. This is needed if someone wants to use the base model as a standalone model. Args: progressbar (`bool`): whether to show a progressbar indicating the unload and merge process safe_merge (`bool`): whether to activate the safe merging check to check if there is any potential Nan in the adapter weights adapter_names (`List[str]`, *optional*): The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults to `None`. """ return self._unload_and_optionally_merge( progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names ) def unload(self) -> torch.nn.Module: """ Gets back the base model by removing all the oft modules without merging. This gives back the original base model. """ return self._unload_and_optionally_merge(merge=False)