[utils] add merge_lora utility function (#227)

* add merge_lora utility function * forward contrib credits from original script * some changes * make style * fix tets * finally fix tests * Update tests/test_peft_model.py * adapt from suggestions * adapt * Update src/peft/tuners/lora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> * fix 8bit * Update src/peft/tuners/lora.py Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> --------- Co-authored-by: edbeeching <edbeeching@users.noreply.github.com> Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
2026-09-09 11:28:32 +08:00 · 2023-03-30 13:45:37 +02:00
parent 542f2470e7
commit 8f63f565c6
3 changed files with 151 additions and 42 deletions
@@ -82,6 +82,10 @@ class LoraConfig(PeftConfig):
            "the final layer `classifier/score` are randomly initialized and as such need to be trainable and saved."
        },
    )
+    init_lora_weights: bool = field(
+        default=True,
+        metadata={"help": "Whether to initialize the weights of the Lora layers."},
+    )

    def __post_init__(self):
        self.peft_type = PeftType.LORA
@@ -135,6 +139,7 @@ class LoraModel(torch.nn.Module):
            "fan_in_fan_out": self.peft_config.fan_in_fan_out,
            "merge_weights": (self.peft_config.merge_weights or self.peft_config.inference_mode)
            and not is_hf_device_map_available,
+            "init_lora_weights": self.peft_config.init_lora_weights,
        }
        key_list = [key for key, _ in self.model.named_modules()]
        for key in key_list:
@@ -233,6 +238,37 @@ class LoraModel(torch.nn.Module):
    def disable_adapter_layers(self):
        self._set_adapter_layers(enabled=False)

+    def merge_and_unload(self):
+        r"""
+        This method merges the LoRa layers into the base model. This is needed if someone wants to use the base model
+        as a standalone model.
+        """
+        if self.config.model_type == "gpt2":
+            raise ValueError("GPT2 models are not supported for merging LORA layers")
+
+        if getattr(self.model, "is_loaded_in_8bit", False):
+            raise ValueError("Cannot merge LORA layers when the model is loaded in 8-bit mode")
+
+        key_list = [key for key, _ in self.model.named_modules() if "lora" not in key]
+        for key in key_list:
+            parent, target, target_name = self._get_submodules(key)
+            if isinstance(target, LoraLayer):
+                bias = target.bias is not None
+                new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias)
+
+                # manually merge if not merged
+                if not target.merged:
+                    # merge weights per: https://arxiv.org/pdf/2106.09685.pdf / page 4
+                    if target.r > 0:
+                        target.weight.data += (
+                            transpose(target.lora_B.weight @ target.lora_A.weight, target.fan_in_fan_out)
+                            * target.scaling
+                        ).to(target.weight.dtype)
+                    target.merged = True
+
+                self._replace_module(parent, target_name, new_module, target)
+        return self.model
+

 # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
 # and modified to work with PyTorch FSDP
@@ -297,6 +333,8 @@ class Linear(nn.Linear, LoraLayer):
        merge_weights: bool = True,
        **kwargs,
    ):
+        init_lora_weights = kwargs.pop("init_lora_weights", True)
+
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)

@@ -308,7 +346,8 @@ class Linear(nn.Linear, LoraLayer):
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
-        self.reset_parameters()
+        if init_lora_weights:
+            self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.T

@@ -375,6 +414,8 @@ class MergedLinear(nn.Linear, LoraLayer):
        merge_weights: bool = True,
        **kwargs,
    ):
+        init_lora_weights = kwargs.pop("init_lora_weights", True)
+
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
        if out_features % len(enable_lora) != 0:
@@ -398,7 +439,9 @@ class MergedLinear(nn.Linear, LoraLayer):
            self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1)
            self.lora_ind[enable_lora, :] = True
            self.lora_ind = self.lora_ind.view(-1)
-        self.reset_parameters()
+
+        if init_lora_weights:
+            self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.T

@@ -30,17 +30,19 @@ from peft import (
 from .testing_common import PeftTestConfigManager


-# This has to be in the order: model_id, lora_kwargs, prefix_tuning_kwargs, prompt_encoder_kwargs, prompt_tuning_kwargs
 PEFT_DECODER_MODELS_TO_TEST = [
-    # ("HuggingFaceM4/tiny-random-LlamaForCausalLM", {}, {}, {}, {}), wait until the next `transformers` release
-    ("hf-internal-testing/tiny-random-OPTForCausalLM", {}, {}, {}, {}),
-    ("hf-internal-testing/tiny-random-GPTNeoXForCausalLM", {}, {}, {}, {}),
-    ("hf-internal-testing/tiny-random-GPT2LMHeadModel", {}, {}, {}, {}),
-    ("hf-internal-testing/tiny-random-BloomForCausalLM", {}, {}, {}, {}),
-    ("hf-internal-testing/tiny-random-gpt_neo", {}, {}, {}, {}),
-    ("hf-internal-testing/tiny-random-GPTJForCausalLM", {}, {}, {}, {}),
+    "hf-internal-testing/tiny-random-OPTForCausalLM",
+    "hf-internal-testing/tiny-random-GPTNeoXForCausalLM",
+    "hf-internal-testing/tiny-random-GPT2LMHeadModel",
+    "hf-internal-testing/tiny-random-BloomForCausalLM",
+    "hf-internal-testing/tiny-random-gpt_neo",
+    "hf-internal-testing/tiny-random-GPTJForCausalLM",
 ]

+FULL_GRID = {
+    "model_ids": PEFT_DECODER_MODELS_TO_TEST,
+}
+

 class PeftTestMixin:
    torch_device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -54,10 +56,6 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
    We use parametrized.expand for debugging purposes to test each model individually.
    """

-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
-    def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs):
-        self._test_model_attr(model_id, config_cls, config_kwargs)
-
    def _test_model_attr(self, model_id, config_cls, config_kwargs):
        model = AutoModelForCausalLM.from_pretrained(model_id)
        config = config_cls(
@@ -70,6 +68,10 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
        self.assertTrue(hasattr(model, "from_pretrained"))
        self.assertTrue(hasattr(model, "push_to_hub"))

+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
+    def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs):
+        self._test_model_attr(model_id, config_cls, config_kwargs)
+
    def _test_prepare_for_training(self, model_id, config_cls, config_kwargs):
        model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
        config = config_cls(
@@ -111,7 +113,7 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):

        self.assertTrue(dummy_output.requires_grad)

-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
    def test_prepare_for_training_parametrized(self, test_name, model_id, config_cls, config_kwargs):
        self._test_prepare_for_training(model_id, config_cls, config_kwargs)

@@ -157,10 +159,61 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
            # check if `config.json` is not present
            self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))

-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
    def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs):
        self._test_save_pretrained(model_id, config_cls, config_kwargs)

+    def _test_merge_layers(self, model_id, config_cls, config_kwargs):
+        model = AutoModelForCausalLM.from_pretrained(model_id)
+        config = config_cls(
+            base_model_name_or_path=model_id,
+            **config_kwargs,
+        )
+        model = get_peft_model(model, config)
+        model = model.to(self.torch_device)
+
+        if config.peft_type != "LORA":
+            with self.assertRaises(AttributeError):
+                model = model.merge_and_unload()
+        elif model.config.model_type == "gpt2":
+            with self.assertRaises(ValueError):
+                model = model.merge_and_unload()
+        else:
+            dummy_input = torch.LongTensor([[1, 2, 3, 2, 1]]).to(self.torch_device)
+            model.eval()
+            logits_lora = model(dummy_input)[0]
+
+            model = model.merge_and_unload()
+
+            logits_merged = model(dummy_input)[0]
+
+            transformers_model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
+
+            logits_transformers = transformers_model(dummy_input)[0]
+
+            self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-3, rtol=1e-3))
+            self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-3, rtol=1e-3))
+
+            with tempfile.TemporaryDirectory() as tmp_dirname:
+                model.save_pretrained(tmp_dirname)
+
+                model_from_pretrained = AutoModelForCausalLM.from_pretrained(tmp_dirname).to(self.torch_device)
+
+                logits_merged_from_pretrained = model_from_pretrained(dummy_input)[0]
+
+                self.assertTrue(torch.allclose(logits_merged, logits_merged_from_pretrained, atol=1e-3, rtol=1e-3))
+
+    @parameterized.expand(
+        PeftTestConfigManager.get_grid_parameters(
+            {
+                "model_ids": PEFT_DECODER_MODELS_TO_TEST,
+                "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]},
+            },
+        )
+    )
+    def test_merge_layers(self, test_name, model_id, config_cls, config_kwargs):
+        self._test_merge_layers(model_id, config_cls, config_kwargs)
+
    def _test_generate(self, model_id, config_cls, config_kwargs):
        model = AutoModelForCausalLM.from_pretrained(model_id)
        config = config_cls(
@@ -180,6 +233,6 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
            # check if `generate` raises an error if no positional arguments are passed
            _ = model.generate(input_ids, attention_mask=attention_mask)

-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID))
    def test_generate(self, test_name, model_id, config_cls, config_kwargs):
        self._test_generate(model_id, config_cls, config_kwargs)
@@ -71,34 +71,47 @@ class ClassInstantier(OrderedDict):

        return super().__getitem__(key, *args, **kwargs)

-    def get_grid_parameters(self, model_list):
+    def get_grid_parameters(self, grid_parameters, filter_params_func=None):
        r"""
        Returns a list of all possible combinations of the parameters in the config classes.
-        """
-        grid_parameters = []
-        for model_tuple in model_list:
-            model_id, lora_kwargs, prefix_tuning_kwargs, prompt_encoder_kwargs, prompt_tuning_kwargs = model_tuple
-            for key, value in self.items():
-                peft_method = value[1].copy()
-                if key == "lora":
-                    # update value[1] if necessary
-                    if lora_kwargs is not None:
-                        peft_method.update(lora_kwargs)
-                elif key == "prefix_tuning":
-                    # update value[1] if necessary
-                    if prefix_tuning_kwargs is not None:
-                        peft_method.update(prefix_tuning_kwargs)
-                elif key == "prompt_encoder":
-                    # update value[1] if necessary
-                    if prompt_encoder_kwargs is not None:
-                        peft_method.update(prompt_encoder_kwargs)
-                else:
-                    # update value[1] if necessary
-                    if prompt_tuning_kwargs is not None:
-                        peft_method.update(prompt_tuning_kwargs)
-                grid_parameters.append((f"test_{model_id}_{key}", model_id, value[0], peft_method))

-        return grid_parameters
+        Args:
+            grid_parameters (`dict`):
+                A dictionary containing the parameters to be tested. There should be at least the key "model_ids" which
+                contains a list of model ids to be tested. The other keys should be the name of the config class
+                post-fixed with "_kwargs" and the value should be a dictionary containing the parameters to be tested
+                for that config class.
+            filter_params_func (`callable`, `optional`):
+                A function that takes a list of tuples and returns a list of tuples. This function is used to filter
+                out the tests that needs for example to be skipped.
+
+        Returns:
+            generated_tests (`list`):
+                A list of tuples containing the name of the test, the model id, the config class and the config class
+                kwargs.
+        """
+        generated_tests = []
+        model_list = grid_parameters["model_ids"]
+
+        for model_id in model_list:
+            for key, value in self.items():
+                if "{}_kwargs".format(key) in grid_parameters:
+                    peft_configs = []
+                    current_peft_config = value[1].copy()
+                    for current_key, current_value in grid_parameters[f"{key}_kwargs"].items():
+                        for kwarg in current_value:
+                            current_peft_config.update({current_key: kwarg})
+                            peft_configs.append(current_peft_config)
+                else:
+                    peft_configs = [value[1].copy()]
+
+                for peft_config in peft_configs:
+                    generated_tests.append((f"test_{model_id}_{key}", model_id, value[0], peft_config))
+
+        if filter_params_func is not None:
+            generated_tests = filter_params_func(generated_tests)
+
+        return generated_tests


 PeftTestConfigManager = ClassInstantier(CLASSES_MAPPING)