QLoRA (4bit base model) not working with OLoRA or PISSA initialisation
Author: Haakooto
Pull request: no
Labels: name: bug

name: PEFT

### System Info

- `transformers` version: 4.43.2
- Platform: Linux-5.15.0-117-generic-x86_64-with-glibc2.35
- Python version: 3.12.4
- Huggingface_hub version: 0.24.2
- Safetensors version: 0.4.3
- Accelerate version: 0.33.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: distributed
- Using GPU in script?: yes
- GPU type: NVIDIA H100 PCIe


### Who can help?

@SunMarc

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Loading model with BitsAndBytesConfig(load_in_4bit=True), and apply adapter with LoraConfig(init_lora_weights="olora"):

```python
import torch
import torch.distributed as dist
from torch.distributed.elastic.multiprocessing.errors import record
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    DataCollatorForLanguageModeling,
    BitsAndBytesConfig,
)
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig, TaskType
from datasets import Dataset

@record
def main(local_rank: int, qlora: bool, lora_init: str):
    device = torch.device("cuda", local_rank)
    model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

    tokenizer = AutoTokenizer.from_pretrained(model_path, legacy=False)
    tokenizer.pad_token = tokenizer.eos_token

    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        target_modules="all-linear",
        init_lora_weights=lora_init,
    )

    training_args = SFTConfig(
        do_train=True,
        output_dir="./test",
        ddp_backend="nccl",
        bf16=True,
        local_rank=local_rank,
        disable_tqdm=local_rank != 0,
        ddp_find_unused_parameters=False,
        max_seq_length=tokenizer.model_max_length,
        dataset_text_field="text",
    )
    
    # dummy data
    dataset = Dataset.from_dict({"text": ["Hello, how are you?", "I am fine, thank you!"]})
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
    )
    
    q4bit_conf = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_storage=torch.bfloat16,
    )

    model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
            quantization_config=q4bit_conf if qlora else None,
            device_map=device
        )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
        data_collator=data_collator,
        peft_config=peft_config,
    )
    trainer.train()

if __name__ == "__main__":
    qlora = True            # When False, training finishes successfully
    lora_init = "olora"    # When not olora or pissa(_niter_16), training finishes successfully
    dist.init_process_group(backend="nccl")
    local_rank = dist.get_rank()
    main(local_rank, qlora, lora_init)
    dist.destroy_process_group()
```

Running script with `torchrun --nproc_per_node=1 train.py`. Get same result when using 2 GPUs.

### Expected behavior

`RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 1x1)`

Full traceback is attatched here: [error.json](https://github.com/user-attachments/files/16547637/error.json) The final traces are 
```
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 438, in forward
      query_states = self.q_proj(hidden_states)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/peft/tuners/lora/bnb.py", line 489, in forward
      output = lora_B(lora_A(dropout(x))) * scaling
                      ^^^^^^^^^^^^^^^^^^
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/home/local/NTU/hator/miniconda3/envs/py12/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 117, in forward
      return F.linear(input, self.weight, self.bias)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 1x1)

============================================================
```

This only happens when both loading model in 4bit and initialising LoRA with "olora" or "pissa".

Comments:
I have narrowed down the issue, and it seems to be all in the peft-repository. I have created a new issue there: https://github.com/huggingface/peft/issues/1999

Thanks for updating, I will close this one then!