vera-pissa method added

llm/tools/merge_vera_params.py

@@ -76,8 +76,7 @@ def merge():
     ) and args.device == "cpu":
         raise ValueError("We can not apply bfloat16 or nf4/fp4 vera merge on cpu.")
 
-    vera_config.merge_weights = False
-    # with device_guard(): 会导致svd无法进行分解
+    # with device_guard() will cause SVD decomposition to fail
     model = AutoModelForCausalLM.from_pretrained(
         vera_config.base_model_name_or_path,
         config=config,

paddlenlp/peft/vera/vera_layers.py

@@ -48,7 +48,6 @@
         self.merged = False
 
         if pissa_init:
-            assert self.vera_alpha == self.r, "pissa method requires vera_alpha=r, scaling=1"
             self.scaling = 1.0
             self.vera_A = self.create_parameter(
                 shape=[in_features, r],
@@ -102,7 +101,7 @@
         dtype = weight.dtype
 
         if dtype != paddle.float32:
             weight = weight.astype(paddle.float32)
 
         U, S, Vh = paddle.linalg.svd(weight.data, full_matrices=False)
 
@@ -120,20 +119,20 @@
         self.weight.set_value(weight)
 
     def merge(self):
         if not self.merged:
             diag_b = paddle.diag(self.vera_b)
             diag_d = paddle.diag(self.vera_d)
             new_weight = self.weight + self.vera_A @ diag_d @ self.vera_B @ diag_b * self.scaling
             self.weight.set_value(new_weight)
             self.merged = True
 
     def unmerge(self):
         if self.merged:
             diag_b = paddle.diag(self.vera_b)
             diag_d = paddle.diag(self.vera_d)
             new_weight = self.weight - self.vera_A @ diag_d @ self.vera_B @ diag_b * self.scaling
             self.weight.set_value(new_weight)
             self.merged = False
 
     def forward(self, input: paddle.Tensor, *args, **kwargs):
         result = F.linear(x=input, weight=self.weight, bias=self.bias, name=self.name)
@@ -145,5 +144,5 @@
         return result
 
     def extra_repr(self):
         name = f", name={self.name}" if self.name else ""
         return f"in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, rank={self.r}{name}"

paddlenlp/peft/vera/vera_model.py

@@ -24,7 +24,6 @@
 from paddle.distributed.fleet.meta_parallel import PipelineLayer
 
 from ...transformers.model_utils import PretrainedModel, _add_variant, dtype_guard
-from ...transformers.utils import weight_name_suffix
 from ...utils.env import VERA_WEIGHTS_NAME
 from ...utils.log import logger
 from .vera_config import VeRAConfig
@@ -46,9 +45,9 @@
             self.model = self.get_vera_model(model, vera_config)
         self.is_pipelinemodel = False
         if issubclass(type(self.model), PipelineLayer):
             raise NotImplementedError("vera don't support pipeline parallel now")
         if vera_config.tensor_parallel_degree > 1:
             raise NotImplementedError("vera don't support tensor parallel now")
         self.forward = self.model.forward
 
     @classmethod
@@ -77,14 +76,14 @@
                 vera_config_tensor_parallel_degree > 1
                 and vera_config_tensor_parallel_degree != model.config.tensor_parallel_degree
             ):
                 raise NotImplementedError(
                     f"{vera_config_tensor_parallel_degree} is not equal to {model.config.tensor_parallel_degree}. Please merge VeRA weights first."
                 )
 
             # set vera state dict
             vera_model.set_state_dict(vera_state_dict)
         else:
             logger.error(f"VeRA weights not found under {vera_path}, creating VeRA weights from scratch")
 
         return vera_model
 
@@ -103,15 +102,15 @@
         save_model_config = kwargs.get("save_model_config", True)
 
         if self.is_pipelinemodel:
             self.model._single_to_pp_mapping = None
         if self.quantized and merge_tensor_parallel and self.vera_config.tensor_parallel_degree > 1:
             merge_tensor_parallel = False
             logger.warning(
                 "Quantized strategy does not support merge_tensor_parallel. Set merge_tensor_parallel to False."
             )
         if self.is_pipelinemodel and merge_tensor_parallel and self.vera_config.tensor_parallel_degree > 1:
             merge_tensor_parallel = False
             logger.warning(
                 "Pipeline parallism does not support merge_tensor_parallel. Set merge_tensor_parallel to False."
             )
 
@@ -128,9 +127,6 @@
         logger.info(f"vera config to save is {vera_config_to_save}")
 
         trainable_state_dict = self.get_trainable_state_dict()
-        if vera_config_to_save.tensor_parallel_degree > 1:
-            if variant is None:
-                variant = weight_name_suffix()
 
         # save vera weight
         vera_weight_name = _add_variant(VERA_WEIGHTS_NAME, variant)
@@ -143,7 +139,7 @@
             if save_model_config:
                 model_config_to_save = copy.deepcopy(self.model.config)
                 if merge_tensor_parallel:
                     model_config_to_save.tensor_parallel_degree = -1
                 model_config_to_save.save_pretrained(save_directory)
 
     def _find_and_replace_module(self, model, module_name, vera_config, enable_vera):
@@ -178,17 +174,17 @@
         setattr(parent_module, attribute_chain[-1], vera_module)
 
     def _find_and_restore_module(self, module_name):
         parent_module = self.model
         attribute_chain = module_name.split(".")
         for name in attribute_chain[:-1]:
             parent_module = getattr(parent_module, name)
         module = getattr(parent_module, attribute_chain[-1])
         original_model_class = self.restore_layer_map[module.__class__]
         original_module = original_model_class(in_features=module.weight.shape[0], out_features=module.weight.shape[1])
         original_module.weight = module.weight
         if module.bias is not None:
             original_module.bias = module.bias
         setattr(parent_module, attribute_chain[-1], original_module)
 
     def get_trainable_state_dict(self):
         trainable_state_dict = OrderedDict()
@@ -199,14 +195,14 @@
         return trainable_state_dict
 
     def print_trainable_parameters(self) -> None:
         freeze_numel = 0
         trainable_numel = 0
         for _, weight in self.model.state_dict().items():
             if weight.stop_gradient:
                 freeze_numel += np.prod(weight.shape)
             else:
                 trainable_numel += np.prod(weight.shape)
         logger.debug(
             f"Frozen parameters: {freeze_numel:.2e} || Trainable parameters:{trainable_numel:.2e} || Total parameters:{freeze_numel+trainable_numel:.2e}|| Trainable:{trainable_numel / (freeze_numel+trainable_numel):.2%}"
         )
 
@@ -215,14 +211,14 @@
             if isinstance(layer, VeRALinear):
                 for name, weight in layer.state_dict().items():
                     if self.vera_config.trainable_bias in ["vera", "all"] and "bias" in name:
                         weight.stop_gradient = False
                     elif "vera" in name:
                         # notfreezeB=True, vera_b, vera_d, vera_B is trainable
                         # notfreezeB=False, vera_b, vera_d is trainable
                         if "vera_b" in name or "vera_d" in name:
                             weight.stop_gradient = False
                         elif "vera_B" in name and notfreezeB:
                             weight.stop_gradient = False
                         else:
                             weight.stop_gradient = True
                     else:
@@ -230,26 +226,26 @@
             else:
                 for name, weight in layer.state_dict().items():
                     if self.vera_config.trainable_bias == "all" and "bias" in name:
                         weight.stop_gradient = False
                     else:
                         weight.stop_gradient = True
         if self.vera_config.trainable_modules is not None:
             for name, weight in self.model.state_dict().items():
                 if any(
                     re.fullmatch(trainable_module, name) for trainable_module in self.vera_config.trainable_modules
                 ):
                     weight.stop_gradient = False
 
     def get_vera_model(self, model: Union[PretrainedModel, nn.Layer], vera_config: VeRAConfig):
 
         if vera_config.target_modules is None:
             return model
         elif isinstance(vera_config.target_modules, str):
             target_modules = [vera_config.target_modules]
             if vera_config.enable_vera_list is None:
                 enable_vera_list = [vera_config.enable_vera_list]
             else:
                 raise TypeError(
                     f"Invalid `enable_vera_list` value: {vera_config.enable_vera_list}. Since `target_modules` is `str`, `enable_vera_list` must be `None` or `List[bool]`"
                 )
         else:
@@ -257,7 +253,7 @@
             if vera_config.enable_vera_list is None:
                 enable_vera_list = [None for _ in range(len(target_modules))]
             else:
                 raise TypeError(
                     f"Invalid `enable_vera_list` value: {vera_config.enable_vera_list}. Since `target_modules` is `List[str]`, `enable_vera_list` must be `None` or `List[Optional[List[bool]]]`"
                 )
 
@@ -269,23 +265,19 @@
         return model
 
     def restore_original_model(self):
-        # make sure W and vera weights are not merged before we restore the original model
-        if self.vera_config.merge_weights:
-            self.train()
-
         for layer_name, layer in self.model.named_sublayers():
             if isinstance(layer, VeRALinear):
                 self._find_and_restore_module(layer_name)
             else:
                 raise NotImplementedError(f"{layer} restoration is not supported yet.")
         return self.model
 
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
         try:
             return super().__getattr__(name)  # defer to nn.Layer's logic
         except AttributeError:
             return getattr(self.model, name)
 
     def train(self):
         self.training = True

test_file.txt/vera_config.json

@@ -0,0 +1,15 @@
+{
+  "base_model_name_or_path": null,
+  "do_qat": false,
+  "dtype": null,
+  "enable_vera_list": null,
+  "head_dim": null,
+  "pissa_init": false,
+  "r": 8,
+  "target_modules": null,
+  "tensor_parallel_degree": -1,
+  "trainable_bias": null,
+  "trainable_modules": null,
+  "vera_alpha": 8,
+  "vera_dropout": 0.0
+}

tests/peft/test_vera.py

@@ -16,7 +16,7 @@
 import os
 import re
 import unittest
-from tempfile import TemporaryDirectory
+from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 import numpy as np
 import paddle
@@ -41,7 +41,13 @@ def test_r_raise_exception(self):
 
     def test_forward(self):
         vera_layer = VeRALinear(
-            in_features=16, out_features=16, r=4, vera_dropout=0.1, vera_alpha=4, base_linear_module=nn.Linear(16, 16)
+            in_features=16,
+            out_features=16,
+            r=4,
+            vera_dropout=0.1,
+            vera_alpha=4,
+            base_linear_module=nn.Linear(16, 16),
+            pissa_init=True,
         )
         x = paddle.randn([2, 4, 16], "float32")
         output = vera_layer(x)
@@ -104,10 +110,7 @@ class TestVeraModel(unittest.TestCase):
     @parameterized.expand([(None,), ("all",), ("vera",)])
     def test_vera_model_constructor(self, bias):
         vera_config = VeRAConfig(
-            target_modules=[".*q_proj.*", ".*v_proj.*"],
-            r=4,
-            vera_alpha=4,
-            head_dim=2,
+            target_modules=[".*q_proj.*", ".*v_proj.*"], r=4, vera_alpha=4, head_dim=2, pissa_init=True
         )
         # turn off plm dropout for to test train vs test
         model = AutoModel.from_pretrained(
@@ -156,6 +159,17 @@ def test_vera_model_save_load(self):
             config_loaded_results = config_loaded_vera_model(input_ids)
             self.assertTrue(paddle.allclose(original_results[0], config_loaded_results[0]))
 
+    def test_restore_original_model(self):
+        vera_config = VeRAConfig(
+            target_modules=[".*q_proj.*", ".*v_proj.*"],
+            r=4,
+            vera_alpha=4,
+        )
+        model = AutoModel.from_pretrained("__internal_testing__/tiny-random-bert")
+        vera_model = VeRAModel(model, vera_config)
+        with self.assertRaises(NotImplementedError):
+            vera_model.restore_original_model()
+
     def test_vera_module_raise_exception(self):
         vera_config = VeRAConfig(
             target_modules=[".*norm1.*"],
@@ -176,6 +190,17 @@ def test_save_load(self):
             loaded_vera_config = VeRAConfig.from_pretrained(tempdir)
             self.assertEqual(vera_config, loaded_vera_config)
 
+    def test_save_load_err(self):
+        with NamedTemporaryFile("w+t") as f:
+            with self.assertRaises(ValueError):
+                VeRAConfig.from_pretrained(f.name)
+
+    def test_save_pretrained_file_error(self):
+        with NamedTemporaryFile("w+t") as f:
+            vera_config = VeRAConfig()
+            with self.assertRaises(AssertionError):
+                vera_config.save_pretrained(f.name)
+
 
 if __name__ == "__main__":
     unittest.main()