FourierFT: add dynamic_scaling parameter

Author: frutiemax92

src/peft/tuners/fourierft/config.py

@@ -118,6 +118,11 @@ class FourierFTConfig(PeftConfig):
         default=False,
         metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"},
     )
+
+    dynamic_scaling: bool = field(
+        default=False,
+        metadata={"help": "Put scaling_factor as a learnable parameter"},
+    )
     target_modules: Optional[Union[list[str], str]] = field(
         default=None,
         metadata={

src/peft/tuners/fourierft/layer.py

@@ -25,14 +25,14 @@
 
 class FourierFTLayer(BaseTunerLayer):
     # All names of layers that may contain (trainable) adapter weights
-    adapter_layer_names = ("fourierft_spectrum",)
+    adapter_layer_names = ("fourierft_spectrum", "fourierft_scaling")
     # All names of other parameters that may contain adapter-related parameters
-    other_param_names = ("fourierft_n_frequency", "fourierft_scaling", "fourierft_random_loc_seed")
+    other_param_names = ("fourierft_n_frequency", "fourierft_random_loc_seed")
 
-    def __init__(self, base_layer: nn.Module, alpha, **kwargs) -> None:
+    def __init__(self, base_layer: nn.Module, **kwargs) -> None:
         self.base_layer = base_layer
         self.fourierft_n_frequency = {}
-        self.fourierft_scaling = {}
+        self.fourierft_scaling = nn.ParameterDict({})
         self.fourierft_spectrum = nn.ParameterDict({})
         self.indices = {}
         self.fourierft_random_loc_seed = {}
@@ -55,7 +55,7 @@ def __init__(self, base_layer: nn.Module, alpha, **kwargs) -> None:
             raise ValueError(f"Unsupported layer type {type(base_layer)}")
 
     def update_layer(
-        self, adapter_name, n_frequency, scaling, init_weights, random_loc_seed, inference_mode: bool = False, **kwargs
+        self, adapter_name, n_frequency, scaling, init_weights, random_loc_seed, dynamic_scaling, inference_mode: bool = False, **kwargs
     ):
         if n_frequency <= 0:
             raise ValueError(f"`n_frequency` should be a positive integer value but the value passed is {n_frequency}")
@@ -73,7 +73,8 @@ def update_layer(
         self.indices[adapter_name] = torch.stack(
             [self.indices[adapter_name] // self.in_features, self.indices[adapter_name] % self.in_features], dim=0
         )
-        self.fourierft_scaling[adapter_name] = scaling
+        self.fourierft_scaling[adapter_name] = nn.Parameter(torch.tensor(scaling, dtype=torch.float32))
+        self.fourierft_scaling[adapter_name].requires_grad = dynamic_scaling
         # Actual trainable parameters
         self.fourierft_spectrum[adapter_name] = nn.Parameter(torch.randn(n_frequency), requires_grad=True)
 
@@ -107,21 +108,22 @@ def __init__(
         n_frequency: int = 1000,
         alpha: float = None,
         scaling: float = 150.0,
+        dynamic_scaling: bool = False,
         fan_in_fan_out: bool = False,  # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
         init_weights: Union[bool, str] = False,
         random_loc_seed: int = 777,
         **kwargs,
     ) -> None:
         super().__init__()
-        FourierFTLayer.__init__(self, base_layer, alpha, **kwargs)
+        FourierFTLayer.__init__(self, base_layer, **kwargs)
 
         # apply alpha patch
         if alpha:
             n_frequency = int(alpha * self.in_features * self.out_features)
 
         self.fan_in_fan_out = fan_in_fan_out
         self._active_adapter = adapter_name
-        self.update_layer(adapter_name, n_frequency, scaling, init_weights, random_loc_seed)
+        self.update_layer(adapter_name, n_frequency, scaling, init_weights, random_loc_seed, dynamic_scaling)
 
     def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
         """
@@ -210,29 +212,30 @@ def __init__(
         n_frequency: int = 1000,
         alpha: float = None,
         scaling: float = 150.0,
+        dynamic_scaling: bool = False,
         fan_in_fan_out: bool = False,  # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
         init_weights: Union[bool, str] = False,
         random_loc_seed: int = 777,
         **kwargs,
     ) -> None:
         super().__init__()
-        FourierFTLayer.__init__(self, base_layer, alpha, **kwargs)
+        FourierFTLayer.__init__(self, base_layer, **kwargs)
 
                 # apply alpha patch
         if alpha:
             n_frequency = int(alpha * self.in_features * self.out_features)
-            
+
         self.fan_in_fan_out = fan_in_fan_out
         self._active_adapter = adapter_name
         self.kW = base_layer.kernel_size[0]
         self.kH = base_layer.kernel_size[1]
         self.stride = base_layer.stride
         self.padding = base_layer.padding
-        self.update_layer(adapter_name, n_frequency, scaling, init_weights, random_loc_seed)
+        self.update_layer(adapter_name, n_frequency, scaling, init_weights, random_loc_seed, dynamic_scaling)
 
 
     def update_layer(
-        self, adapter_name, n_frequency, scaling, init_weights, random_loc_seed, inference_mode: bool = False, **kwargs
+        self, adapter_name, n_frequency, scaling, init_weights, random_loc_seed, dynamic_scaling, inference_mode: bool = False, **kwargs
     ):
         if n_frequency <= 0:
             raise ValueError(f"`n_frequency` should be a positive integer value but the value passed is {n_frequency}")
@@ -241,6 +244,7 @@ def update_layer(
                 f"`n_frequency` should be less than or equal to the product of the input and output dimensions "
                 f"but the value passed is {n_frequency} and the product is {self.in_features * self.out_features}"
             )
+        
         self.fourierft_n_frequency[adapter_name] = n_frequency
         self.fourierft_random_loc_seed[adapter_name] = random_loc_seed
         self.indices[adapter_name] = torch.randperm(
@@ -250,7 +254,8 @@ def update_layer(
         self.indices[adapter_name] = torch.stack(
             [self.indices[adapter_name] // self.in_features, self.indices[adapter_name] % self.in_features], dim=0
         )
-        self.fourierft_scaling[adapter_name] = scaling
+        self.fourierft_scaling[adapter_name] = nn.Parameter(torch.tensor(scaling, dtype=torch.float32))
+        self.fourierft_scaling[adapter_name].requires_grad = dynamic_scaling
         # Actual trainable parameters
         self.fourierft_spectrum[adapter_name] = nn.Parameter(torch.randn(n_frequency, self.kW, self.kH), requires_grad=True)
 

src/peft/tuners/fourierft/model.py

@@ -96,13 +96,15 @@ def _create_and_replace(
 
         n_frequency = fourierft_config.n_frequency_pattern.get(target_name_key, fourierft_config.n_frequency)
         scaling = fourierft_config.scaling
+        dynamic_scaling = fourierft_config.dynamic_scaling
         alpha = fourierft_config.alpha
         random_loc_seed = fourierft_config.random_loc_seed
         bias = hasattr(target, "bias") and target.bias is not None
         kwargs = {
             "n_frequency": n_frequency,
             "alpha": alpha,
             "scaling": scaling,
+            "dynamic_scaling" : dynamic_scaling,
             "fan_in_fan_out": fourierft_config.fan_in_fan_out,
             "init_weights": fourierft_config.init_weights,
             "random_loc_seed": fourierft_config.random_loc_seed,