ENH: Refactor LoRA bnb layers for faster initialization

[BenjaminBossan]

Partly addresses #896

Description

After speeding up normal LoRA layer initialization, this PR improves initialization speed of bnb LoRA layers.

The method to achieve this is different from the one used before, namely this time the base layer is stored as a reference on the LoRA layer. This allows us to avoid calling __init__ on the bnb layer, which is what is slow.

Notes

We cannot use the same method as for the normal LoRA layers, (i.e. calling the super class's __init__ with meta device) because the bnb layers have extra logic that still creates unnecessary weights (this is something that could potentially be fixed by bnb).

However, the way used here could also be a solution to the normal layers, so if we want to have consistency, the normal layers could be refactored to use the same approach.

Another small advantage of this approach is that we can now correctly handle any *args, **kwargs that are passed to forward. If we had this consistently in all layers, it would allow us to handle extra arguments like lora_scale in the future.

Interestingly, even though we now save the base layer as a reference, which results in a different state_dict, the existing models can still be loaded successfully. This is because the adapter state_dict is not affected by the change, so users can still load their existing adapters.

The only problem would occur if users dump the whole model, i.e. base model and adapter, using torch.save and then trying to load with torch.load. For those users, we could theoretically provide a script to convert the state_dict (i.e. renaming some keys).

To ensure that the old adapters can still be loaded successfully, I'm working at the same time on adding regression tests. I'll create a separate PR for those to avoid blowing up this one.

Tests

I ran a test on bloomz-1b1 for how long it takes to create the PeftModel, the results are:

8bit: 1108.34 ms > 26.82 ms
4bit: 1101.96 ms > 23.69 ms

Script was:

import time
from contextlib import contextmanager

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model

@contextmanager
def timed(msg):
    start = time.perf_counter()
    yield
    end = time.perf_counter()
    print(f"{msg} took {1000 * (end - start):.2f} ms")

torch.manual_seed(1000)
model_id = "bigscience/bloomz-1b1"
test_8bit = True

if test_8bit:
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        load_in_8bit=True,
    )
    config = LoraConfig(
        r=8,
        init_lora_weights=False,
    )
else:
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=False,
        bnb_4bit_compute_type=torch.float32,
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=bnb_config,
        torch_dtype=torch.float32,
    )
    config = LoraConfig(
        r=8,
        init_lora_weights=False,
    )

with timed("creating peft model"):
    get_peft_model(model, config)

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[pacman100]

Thank you @BenjaminBossan for speeding up the creation of quantized peft models 🚀. I like the approach and it simplifies and makes it more extendable. Left a couple of comments.

[pacman100]

Thank you @BenjaminBossan, good to merge after resolving the conflicts.