Incorrect weight shape for Linear4bit Lora

[SimonKitSangChu]

I tried to apply 4-bit training on Lora for EsmModel. However, there has been an error specifically for 4-bit training. The error disappears perfectly once load_in_4bit=True is commented out.

Code to reproduce error:

import torch
from transformers import EsmForMaskedLM
from peft import prepare_model_for_kbit_training, get_peft_model, LoraConfig

model_name_or_path = 'facebook/esm2_t30_150M_UR50D'
model = EsmForMaskedLM.from_pretrained(
        model_name_or_path,
        device_map='auto',
        load_in_4bit=True,
)
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)

config = LoraConfig(
    r=4,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    # task_type="CAUSAL_LM",  # no option on MASKED_LM yet; skipped
    target_modules=target_modules,
)
model = get_peft_model(model, config)
model = model.cuda()

layer = model.base_model.esm.encoder.layer[0]
outputs = layer.attention.self.query(torch.rand(2, 3, 640, device=model.device))  # 640 is the hidden_dim

Error received:

Traceback (most recent call last):
  File "test.py", line 105, in <module>
    outputs = layer.attention.self.query(torch.rand(2, 3, 640, device=model.device))
  File "python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "python3.10/site-packages/peft/tuners/lora.py", line 769, in forward
    result = super().forward(x)
  File "python3.10/site-packages/bitsandbytes/nn/modules.py", line 219, in forward
    out = bnb.matmul_4bit(x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state)
  File "python3.10/site-packages/bitsandbytes/autograd/_functions.py", line 564, in matmul_4bit
    return MatMul4Bit.apply(A, B, out, bias, quant_state)
  File "python3.10/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "python3.10/site-packages/bitsandbytes/autograd/_functions.py", line 512, in forward
    output = torch.nn.functional.linear(A, F.dequantize_fp4(B, state).to(A.dtype).t(), bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (6x640 and 1x204800)

Notice the weight matrix is not of the size (640, 640). Commenting out 4-bit loading results in a weight of (640, 640) shape.

Package versions:

transformers.__version__ = '4.30.0.dev0'
peft.__version__ = '0.4.0.dev0'

Let me know if you also need bitsandbytes version.

[SimonKitSangChu]

On a similar track, loading EsmModel in 8-bit experiences another error.

...
model = EsmForMaskedLM.from_pretrained(
        model_name_or_path,
        device_map='auto',
        load_in_8bit=True,
)
...
  File "python3.10/site-packages/bitsandbytes/nn/modules.py", line 344, in _save_to_state_dict
    param_from_weight = getattr(self.weight, weight_name)
AttributeError: 'Parameter' object has no attribute 'SCB'

[jzhang38]

Hi, @SimonKitSangChu I met exactly the same problem. Have you managed to resolve it?

[SimonKitSangChu]

Not yet. Interestingly I didn't experience this in the example checkpoint bigscience/mt0-large.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[debraj135]

Is this resolved?

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[mbhenaff]

I'm also encountering this error. Has anyone found a solution?

[jieWANGforwork]

Can anyone load 4bit LLM with lora to finetune? I met the same error, could please give solutions?

[DwanZhang-AI]

Same here!

[younesbelkada]

hi there, do you still face the same error with the latest updates of all packages?

pip install -U peft transformers bitsandbytes

Can you also share a small reproducible snippet?

[ragavsachdeva]

I was experiencing the same issue today and I believe the culprit it is:

model = model.cuda()

get rid of it and it works.

Note:get_peft_model already returns the model on cuda. I have absolutely no idea why calling .cuda() causes problems though.

[Yu-Fangxu]

hi there, do you still face the same error with the latest updates of all packages?

pip install -U peft transformers bitsandbytes

Can you also share a small reproducible snippet?

This helps me solve this problem, previous version: transformers: '4.38.2' peft: '0.9.0' bitsandbytes: '0.42.0' updated version: bitsandbytes-0.43.1 peft-0.10.0 tokenizers-0.19.1 transformers-4.40.1