Repository: huggingface/peft

Issue:
Faster init with LoRA
Purpose: accelerate peft model initialisation with LoRA config.
see issue description in https://github.com/huggingface/peft/issues/871

The idea is to avoid costly calls of weight initialisation functions in CPU, when those weights are not necessary and will be replaced by the original weights anyways.

This PR initially changes only init for Linear layers. 

potential further improvements:
- apply to Embedding and Conv2d layers
- apply to compressed layer types, like Linear4bit, Linear8bit
- init LoRA modules (A, B) right on target device, not CPU


Discussion:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/peft/pr_872). All of your documentation changes will be reflected on that endpoint.
I wonder if we can use [`skip_init`](https://pytorch.org/docs/stable/generated/torch.nn.utils.skip_init.html) or something else from PyTorch instead of monkey patching a hard-coded list of initialization methods.
> I wonder if we can use [`skip_init`](https://pytorch.org/docs/stable/generated/torch.nn.utils.skip_init.html) or something else from PyTorch instead of monkey patching a hard-coded list of initialization methods.

well, I borrowed this idea [from GPTQ code](https://github.com/IST-DASLab/gptq/blob/2d65066eeb06a5c9ff5184d8cebdf33662c67faf/llama.py#L15). But I support your idea to avoid monkey patching.
Let me try this improvement - I will write back soon.
@BenjaminBossan !
`skip_init()` turned out not very convenient to apply from inside of a subclass. However, looking at its source code, I got an idea to use `meta` device. It works much faster than my original code, since meta tensors are faster than empty cpu tensors. Thanks for the suggestion!
To accomodate this change I also:
	- made `self.to(self.weight.device)` in LoraLayer.update_layer() conditional, to avoid moving loras to 'meta' device
	- removed  self.weight.requires_grad = False in `Linear.__init__()` for it is not relevant.   # lora.py:850

Now with Llama-2-7B application of Guanaco LoRA adapter takes 1.5 seconds vs 103 seconds originally. 

@poedator Great, thanks for experimenting. This looks like a very nice speedup indeed. I'm not 100% sure if we can remove all of that code, but let's see. Could you please run `make style` on your code? Then the code quality checks should pass and the unit tests should run.
I added couple more commits to pass device param to LoRA layer constructors to initialize LoRA adapters right at the target device. Apparently I need to do some more local testing ...
As a matter of introduction: I am a junior researcher @ Yandex, 2d name in this SpQR paper on LLM quantization: https://arxiv.org/abs/2306.03078. We use HF packages daily and want them to get even better ).
Thanks for applying the fixes and providing more context. 

Unfortunately, the tests are failing right now. The main reason seems to be that the code path when `init_lora_weights=False` is not properly implemented yet, resulting in the weights staying on meta device. This needs to be fixed and it might require a bit of fiddling, as the logic is a bit all over the place. It would probably help if we cleaned up the initialization logic, but that requires some non-trivial refactoring.

To proceed, would you be up for fixing the failing tests for the code path I mentioned?

**Edit**: We merged #807, so there is a merge conflict now. But it is easy to fix. Please apply your changes to this file: https://github.com/huggingface/peft/blob/main/src/peft/tuners/lora/layer.py instead of `lora.py` and you should be good.
@BenjaminBossan, thank you for your guidance on this PR.

I updated the PR to catch up with #807 and also limited my changes to just the Linear layer init to prove the concept. 
Some of the tests with `init_lora_weights=False` are also failing locally (test_merge_layers_09 and 10). This is puzzling because they do not seem to involve LoRA for Linear. Nevertheless I want to see this PR coming thru, so let me try to fix it with your guidance.

Pls be more specific on the failing tests. do you want me to update the tests code or to change some LoRA code?
By the way, I re-ran the failed tests in `main` branch and they failed too. What could that be? test names:
- test_lora_bnb_4bit_quantization_from_pretrained_safetensors
- test_lora_causal_lm_mutli_gpu_inference
- test_lora_seq2seq_lm_mutli_gpu_inference
- test_merge_layers_09_Conv2d_1
- test_merge_layers_10_Conv2d_2
(Python 3.10.9, Ubuntu 18.04, A100, mucho RAM, accelerate 0.21.0, bitsandbytes  0.41.1, peft 0.6.0.dev0, torch 2.0.1, transformers 4.30.2)

On PR scope - let me know if you want me to include code for all of the other layer types. I wanted to do it initially but got confused by failed tests and scaled back.
Also I may try to pass device argument to LoRAs inits so they are created on Cuda where they belong (still faster). But that may be better left to a separate PR.
Thanks for the updates and the further investigation. Regarding the failing tests, the idea is that all existing tests should pass after the change. There could be tests that need adjustment, but I'm not aware of any. Anyway, for me, all the tests passed. As you can see, the CI passes too. So not sure why they failed for you.

I think the strategy you suggested, which is to start with one layer type first, is valid. We can do that and once that has proven to work well, we can tackle the others.

I did some further digging into the code and IIUC, I think the weight initialization happening in these two lines is completely unnecessary:

https://github.com/huggingface/peft/blob/0b2f950cc212dd8acda983760c0aba2bf4872a00/src/peft/tuners/lora/layer.py#L142
https://github.com/huggingface/peft/blob/0b2f950cc212dd8acda983760c0aba2bf4872a00/src/peft/tuners/lora/layer.py#L151

It creates a randomly initialized `self.weight` (and possibly `self.bias`) on the lora layer, which is then initialized with kaiming init twice! And later, it is replaced by the original weight of the target layer, so the whole initialization was unnecessary:

https://github.com/huggingface/peft/blob/0b2f950cc212dd8acda983760c0aba2bf4872a00/src/peft/tuners/lora/model.py#L213

Therefore, I think this step could be completely skipped. Since we're talking about the full sized weight here, not the smaller LoRA weights, this could take quite some time. The fix you proposed would create that weight on the meta device, but I wonder if we can rewrite the code to completely skip the initialization instead. If you want to explore that, feel free to do so, otherwise I'll take a look at it later. If you have some benchmark script to test the initialization speed, please share it.

> Also I may try to pass device argument to LoRAs inits so they are created on Cuda where they belong (still faster). But that may be better left to a separate PR.

Yes, that's indeed best left for a separate PR.
@poedator I attempted to remove the redundant initializations completely in #887. If you have time to test potential speedups of that branch, it would be great.
> @poedator I attempted to remove the redundant initializations completely in #887. If you have time to test potential speedups of that branch, it would be great.

@BenjaminBossan, 
I just tested #887 speed, it is comparable to my PR and both are much faster than `main` branch. 
see script and results here: https://gist.github.com/poedator/45331cd9c7837cbfea0ad578b3ce98ed

let me know if you prefer to proceed with #887 or with this PR.  I admit that None is even better than tensor on `meta` device.

#872 passed the tests as of today. Apparently there is something in my local setup that causes errors when merging back the adapters for conv2d. I can investigate it further, but could you share versions of libs that you use for testing. My env is [described above](https://github.com/huggingface/peft/pull/872#issuecomment-1697849147).
Thanks for providing the script @poedator. Could you please paste your measurements for reference? Unfortunately, I still haven't gotten access to Llama 2 (really not sure why it's still pending), so I can't test that model. For the others, I got:

```
# main, x3
test 1 with model bert-base took  0.186 sec.
test 2 with model bloomz-1b7 took 3.740 sec.
test 1 with model bert-base took  0.208 sec.
test 2 with model bloomz-1b7 took 3.889 sec.
test 1 with model bert-base took  0.187 sec.
test 2 with model bloomz-1b7 took 3.762 sec.

# PR 887, x3
test 1 with model bert-base took  0.019 sec.
test 2 with model bloomz-1b7 took 0.057 sec.
test 1 with model bert-base took  0.019 sec.
test 2 with model bloomz-1b7 took 0.030 sec.
test 1 with model bert-base took  0.019 sec.
test 2 with model bloomz-1b7 took 0.029 sec.
```

So this looks like a very decent speed up, ~10x, although in practice a lot of time is spent on loading the model in the first place, so overall clock time is only 11 vs 7 sec.

> let me know if you prefer to proceed with #887 or with this PR.

My impression is that it would be even better not to initialize the weights at all instead of initializing them on meta device. Having useless inits is just causing confusion. The main advantage with initializing with meta is that the code changes are smaller, so they're less likely to accidentally break something. But from what I can tell, removing the inits should also not break anything.

@younesbelkada @pacman100 Do you have an opinion on that? How should we proceed?
> Could you please paste your measurements for reference? 
my results were at the top of the [gist with code](https://gist.github.com/poedator/45331cd9c7837cbfea0ad578b3ce98ed)
Here I added 2 more tests with Llama 1 - 7B.  the benefit is more notieceable with LLMs. Guanaco tests are slower because they involve Embedding layers, not covered by either our PRs.
```
Switched to branch 'pr887'       # PR #887 by Benjamin Bossan
test 1 with model bert-base took  0.048 sec.
test 2 with model bloomz-1b7 took 0.607 sec.
test 3  Llama-2-7B + Guanaco took 5.863 sec.
test 4 with model Llama-1-7B took 0.094 sec.
test 5  Llama-1-7B + Guanaco took 6.165 sec.

Switched to branch 'fast_init'   # PR #872 by poedator@
test 1 with model bert-base took  0.063 sec.      ~3x
test 2 with model bloomz-1b7 took 0.528 sec.      ~7x
test 3  Llama-2-7B + Guanaco took 5.618 sec.      ~15x
test 4 with model Llama-1-7B took 0.116 sec.      ~100x
test 5  Llama-1-7B + Guanaco took 6.273 sec.      ~12x

Switched to branch 'main'         # main version, commit 85013987aa82aa1af3da1236b6902556ce3e483e
test 1 with model bert-base took   0.224 sec.
test 2 with model bloomz-1b7 took  3.842 sec.
test 3  Llama-2-7B + Guanaco took 76.174 sec.  <---- terribly slow (vs 7 sec for loading whole 7B model)
test 4 with model Llama-1-7B took 11.607 sec.
test 5  Llama-1-7B + Guanaco took 77.542 sec.  <---- terribly slow
```

The current way of proceeding looks great IMO as the code change is quite minimal for a nice gain, also I think `meta` device has been introduced in torch>=1.13 which is the current min requirement for PEFT: https://github.com/huggingface/peft/blob/main/setup.py#L45
Thinking a bit and discussing offline with @BenjaminBossan , there might be one problem with this approach if users uses `LoraLayer` as a standalone module, they might end up with weights being initialized on the meta device. This is fine as tests are passing but the module is public so there might be some users using it already, so let's keep that in mind
#887 will also make `LoraLayer` unusable as a standalone module as it will remove the weight attribute per my understanding from my discussion with Benjamin
Thank you @poedator and @BenjaminBossan for making the LoRA init super fast  🚀. I really like both approaches, favouring Benjamin's as it avoids the inits of weights altogether.
Anything else missing for merging this PR?
Since #887 is merged, I think we can close this one. Progress for other layers/adapters is tracked in #896.

Resolving patch:
diff --git a/src/peft/tuners/lora/layer.py b/src/peft/tuners/lora/layer.py
--- a/src/peft/tuners/lora/layer.py
+++ b/src/peft/tuners/lora/layer.py
@@ -43,6 +43,18 @@ def __init__(self, in_features: int, out_features: int, **kwargs):
         self.out_features = out_features
         self.kwargs = kwargs
 
+    def _init_empty_weights(self, cls, *args, **kwargs) -> None:
+        # A helper method that allows to initialize the layer of the given class without spending time to initialize the
+        # model weights. The implementation is inspired by
+        # https://pytorch.org/docs/stable/generated/torch.nn.utils.skip_init.html but this function cannot be used
+        # directly.
+        # Instead of this approach, it would be possible to bypass the __init__ of the class but that runs the risk of
+        # omitting important logic inside that __init__.
+        kwargs = kwargs.copy()
+        final_device = kwargs.pop("device", "cpu")
+        cls.__init__(self, *args, device="meta", **kwargs)
+        self.to_empty(device=final_device)
+
     def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
         self.r[adapter_name] = r
         self.lora_alpha[adapter_name] = lora_alpha
@@ -63,7 +75,7 @@ def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weig
         weight = getattr(self, "weight", None)
         if weight is not None:
             # the layer is already completely initialized, this is an update
-            self.to(weight.device)
+            self.to(weight.device, dtype=weight.dtype)
 
     def update_layer_conv2d(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
         self.r[adapter_name] = r
@@ -84,7 +96,11 @@ def update_layer_conv2d(self, adapter_name, r, lora_alpha, lora_dropout, init_lo
             self.scaling[adapter_name] = lora_alpha / r
         if init_lora_weights:
             self.reset_lora_parameters(adapter_name)
-        self.to(self.weight.device)
+
+        weight = getattr(self, "weight", None)
+        if weight is not None:
+            # the layer is already completely initialized, this is an update
+            self.to(self.weight.device, dtype=weight.dtype)
 
     def update_layer_embedding(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
         self.r[adapter_name] = r
@@ -97,14 +113,18 @@ def update_layer_embedding(self, adapter_name, r, lora_alpha, lora_dropout, init
         self.lora_dropout[adapter_name] = lora_dropout_layer
         # Actual trainable parameters
         if r > 0:
-            weight_A = torch.randn((r, self.in_features), dtype=self.weight.dtype, device=self.weight.device)
-            weight_B = torch.randn((self.out_features, r), dtype=self.weight.dtype, device=self.weight.device)
+            weight_A = torch.randn((r, self.in_features))
+            weight_B = torch.randn((self.out_features, r))
             self.lora_embedding_A[adapter_name] = nn.Parameter(weight_A)
             self.lora_embedding_B[adapter_name] = nn.Parameter(weight_B)
             self.scaling[adapter_name] = lora_alpha / r
         if init_lora_weights:
             self.reset_lora_parameters(adapter_name)
-        self.to(self.weight.device)
+
+        weight = getattr(self, "weight", None)
+        if weight is not None:
+            # the layer is already completely initialized, this is an update
+            self.to(self.weight.device, dtype=weight.dtype)
 
     def reset_lora_parameters(self, adapter_name):
         if adapter_name in self.lora_A.keys():
@@ -145,6 +165,8 @@ def __init__(
         # this gets the init from nn.Linear's super perspective, i.e.
         # nn.Module.__init__, which should always be called
         super(nn.Linear, self).__init__()
+        # Note that we don't use self._init_empty_weights() for Linear because it is a bit slower and the benefit of
+        # added robustness is not big enough for Linear.
 
         LoraLayer.__init__(self, in_features=in_features, out_features=out_features)
         # Freezing the pre-trained weight matrix
@@ -226,13 +248,8 @@ def __init__(
         **kwargs,
     ) -> None:
         init_lora_weights = kwargs.pop("init_lora_weights", True)
-
-        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
+        self._init_empty_weights(nn.Embedding, num_embeddings, embedding_dim, **kwargs)
         LoraLayer.__init__(self, in_features=num_embeddings, out_features=embedding_dim)
-
-        self.weight.requires_grad = False
-
-        nn.Embedding.reset_parameters(self)
         self.update_layer_embedding(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
         self.active_adapter = adapter_name
 
@@ -306,8 +323,8 @@ def __init__(
         **kwargs,
     ) -> None:
         init_lora_weights = kwargs.pop("init_lora_weights", True)
+        self._init_empty_weights(nn.Conv2d, in_channels, out_channels, kernel_size, stride=stride, padding=padding)
 
-        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride, padding)
         LoraLayer.__init__(
             self,
             in_features=in_channels,
@@ -316,10 +333,7 @@ def __init__(
             stride=stride,
             padding=padding,
         )
-        # Freezing the pre-trained weight matrix
-        self.weight.requires_grad = False
 
-        nn.Conv2d.reset_parameters(self)
         self.update_layer_conv2d(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
         self.active_adapter = adapter_name
 


Test patch:
diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py
--- a/tests/test_custom_models.py
+++ b/tests/test_custom_models.py
@@ -368,7 +368,7 @@ def test_adding_multiple_adapters_with_bias_raises(self, test_name, model_id, co
 class TestRepr(unittest.TestCase):
     """Tests related to the repr of adapted models"""
 
-    def test_repr_lora(self):
+    def test_repr_lora_linear(self):
         config = LoraConfig(target_modules=["lin0"])
         model = get_peft_model(MLP(), config)
         print_output = repr(model.model.lin0)
@@ -377,3 +377,35 @@ def test_repr_lora(self):
         self.assertTrue("lora_A" in print_output)
         self.assertTrue("lora_B" in print_output)
         self.assertTrue("default" in print_output)
+
+    def test_repr_lora_embedding(self):
+        config = LoraConfig(target_modules=["emb"])
+        model = get_peft_model(ModelEmbConv1D(), config)
+        print_output = repr(model.model.emb)
+        self.assertTrue(print_output.startswith("Embedding"))
+        self.assertTrue("100, 5" in print_output)
+        self.assertTrue("lora_embedding_A" in print_output)
+        self.assertTrue("lora_embedding_B" in print_output)
+        self.assertTrue("default" in print_output)
+
+    def test_repr_lora_conv1d(self):
+        config = LoraConfig(target_modules=["conv1d"])
+        model = get_peft_model(ModelEmbConv1D(), config)
+        print_output = repr(model.model.conv1d)
+        self.assertTrue(print_output.startswith("Linear"))
+        self.assertTrue("in_features=5, out_features=1" in print_output)
+        self.assertTrue("lora_A" in print_output)
+        self.assertTrue("lora_B" in print_output)
+        self.assertTrue("default" in print_output)
+
+    def test_repr_lora_conv2d(self):
+        config = LoraConfig(target_modules=["conv2d"])
+        model = get_peft_model(ModelConv2D(), config)
+        print_output = repr(model.model.conv2d)
+        self.assertTrue(print_output.startswith("Conv2d"))
+        self.assertTrue("5, 10" in print_output)
+        self.assertTrue("kernel_size=(3, 3)" in print_output)
+        self.assertTrue("stride=(1, 1)" in print_output)
+        self.assertTrue("lora_A" in print_output)
+        self.assertTrue("lora_B" in print_output)
+        self.assertTrue("default" in print_output)