src/peft/tuners/waveft/model.py

# Copyright 2025-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import warnings

import torch
from transformers.pytorch_utils import Conv1D

from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists
from peft.utils import (
    TRANSFORMERS_MODELS_TO_WAVEFT_TARGET_MODULES_MAPPING,
)
from peft.utils.other import get_pattern_key

from .layer import WaveFTLayer, WaveFTLinear


class WaveFTModel(BaseTuner):
    prefix: str = "waveft_"
    tuner_layer_cls: type[BaseTunerLayer] = WaveFTLayer
    target_module_mapping = TRANSFORMERS_MODELS_TO_WAVEFT_TARGET_MODULES_MAPPING

    def _calculate_proportional_parameters(self, model: torch.nn.Module, waveft_config):
        """Calculate proportional parameter allocation for all target modules."""
        target_modules_info = []
        for name, module in model.named_modules():
            if check_target_module_exists(waveft_config, name):
                # Handle case where module is already wrapped with WaveFT
                if isinstance(module, WaveFTLayer):
                    # Use the base layer for dimension calculations
                    base_module = module.base_layer
                    if isinstance(base_module, torch.nn.Linear):
                        input_dim, output_dim = base_module.in_features, base_module.out_features
                    elif isinstance(base_module, Conv1D):
                        input_dim, output_dim = base_module.weight.shape[1], base_module.weight.shape[0]
                    else:
                        continue
                elif isinstance(module, torch.nn.Linear):
                    input_dim, output_dim = module.in_features, module.out_features
                elif isinstance(module, Conv1D):
                    input_dim, output_dim = module.weight.shape[1], module.weight.shape[0]
                else:
                    continue
                target_modules_info.append((name, input_dim, output_dim))

        if not target_modules_info:
            raise ValueError("No target modules found for proportional parameter allocation.")

        total_sum = sum(input_dim * output_dim for (_, input_dim, output_dim) in target_modules_info)
        num_layers = len(target_modules_info)
        total_budget = waveft_config.n_frequency * num_layers

        n_frequency_dict = {}
        for name, input_dim, output_dim in target_modules_info:
            layer_ratio = (input_dim * output_dim) / total_sum
            n_freq = round(layer_ratio * total_budget)
            n_frequency_dict[name] = n_freq

        return n_frequency_dict

    def _create_and_replace(
        self,
        waveft_config,
        adapter_name,
        target,
        target_name,
        parent,
        current_key,
        **optional_kwargs,
    ):
        if current_key is None:
            raise ValueError("Current Key shouldn't be `None`")

        # Calculate proportional parameters if needed (only once per adapter)
        if waveft_config.proportional_parameters:
            if not hasattr(self, "_proportional_params_cache"):
                self._proportional_params_cache = {}
            if adapter_name not in self._proportional_params_cache:
                n_frequency_dict = self._calculate_proportional_parameters(self.model, waveft_config)
                self._proportional_params_cache[adapter_name] = n_frequency_dict

        # Determine n_frequency: Priority order:
        # 1. From proportional parameter cache (if proportional_parameters=True)
        # 2. From optional_kwargs (if passed directly)
        # 3. From n_frequency_pattern in config
        # 4. From default n_frequency in config
        n_frequency = None
        if (
            waveft_config.proportional_parameters
            and hasattr(self, "_proportional_params_cache")
            and adapter_name in self._proportional_params_cache
        ):
            n_frequency = self._proportional_params_cache[adapter_name].get(current_key)

        if n_frequency is None and "n_frequency" in optional_kwargs:
            n_frequency = optional_kwargs["n_frequency"]

        if n_frequency is None:
            pattern_keys = list(waveft_config.n_frequency_pattern.keys())
            target_name_key = get_pattern_key(pattern_keys, current_key)
            n_frequency = waveft_config.n_frequency_pattern.get(target_name_key, waveft_config.n_frequency)

        # Determine wavelet_family
        wavelet_family = None
        if "wavelet_family" in optional_kwargs:
            wavelet_family = optional_kwargs["wavelet_family"]
        if wavelet_family is None:
            wavelet_family = waveft_config.wavelet_family

        scaling = waveft_config.scaling
        random_loc_seed = waveft_config.random_loc_seed
        bias = hasattr(target, "bias") and target.bias is not None
        # Prepare kwargs for module creation/update
        kwargs = {
            "n_frequency": n_frequency,
            "scaling": scaling,
            "fan_in_fan_out": waveft_config.fan_in_fan_out,
            "init_weights": waveft_config.init_weights,
            "random_loc_seed": waveft_config.random_loc_seed,
            "wavelet_family": wavelet_family,  # Use determined wavelet family
        }
        kwargs["bias"] = bias

        if isinstance(target, WaveFTLayer):
            target.update_layer(
                adapter_name,
                n_frequency,
                scaling,
                waveft_config.init_weights,
                random_loc_seed,
                wavelet_family=wavelet_family,  # Pass determined wavelet family
                use_idwt=waveft_config.use_idwt,
            )
        else:
            new_module = self._create_new_module(waveft_config, adapter_name, target, **kwargs)
            if adapter_name != self.active_adapter:
                new_module.requires_grad_(False)
            self._replace_module(parent, target_name, new_module, target)

    @staticmethod
    def _create_new_module(waveft_config, adapter_name, target, **kwargs):
        if isinstance(target, BaseTunerLayer):
            target_base_layer = target.get_base_layer()
        else:
            target_base_layer = target

        if isinstance(target_base_layer, torch.nn.Linear):
            if kwargs["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."
                )
                kwargs["fan_in_fan_out"] = waveft_config.fan_in_fan_out = False
        elif isinstance(target_base_layer, Conv1D):
            kwargs["is_target_conv_1d_layer"] = True
            if not kwargs["fan_in_fan_out"]:
                warnings.warn(
                    "fan_in_fan_out is set to False but the target module is `Conv1D`. Setting fan_in_fan_out to True."
                )
                kwargs["fan_in_fan_out"] = waveft_config.fan_in_fan_out = True
        else:
            raise ValueError(
                f"Target module {target} is not supported. Currently, only the following modules are supported: "
                "`torch.nn.Linear`."
            )

        kwargs["wavelet_family"] = waveft_config.wavelet_family
        kwargs["use_idwt"] = waveft_config.use_idwt
        new_module = WaveFTLinear(target, adapter_name, **kwargs)

        return new_module

    def delete_adapter(self, adapter_name: str) -> None:
        """
        Deletes an existing adapter.

        Args:
            adapter_name (str): Name of the adapter to be deleted.
        """
        super().delete_adapter(adapter_name)
        # Clean up proportional parameters cache
        if hasattr(self, "_proportional_params_cache") and adapter_name in self._proportional_params_cache:
            del self._proportional_params_cache[adapter_name]