DoRA RuntimeError: shape is invalid when Conv2d(groups=2)

[victolee0]

System Info

peft==0.15.2
torch==2.2.2

Who can help?

@BenjaminBossan @sayakpaul

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

from peft import LoraConfig, get_peft_model
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv = torch.nn.Conv2d(96, 192, (3, 3), stride=(2, 2), padding=(1, 1), bias=False, groups=2)
        self.fc = torch.nn.Linear(192, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.fc(x)
        return x


model = MyModel()

lora_config = LoraConfig(
    r=16,
    target_modules=["conv"],
    use_dora=True
)

model = get_peft_model(model, lora_config)

Expected behavior

RuntimeError: shape '[192, 48, 3, 3]' is invalid for input of size 165888

165888 / (192 * 48 * 3 * 3) = 2

when groups=2 got an error
if groups=1 works fine

[BenjaminBossan]

Thanks for the report. The snippet is not quite complete, could you please show the (shape of the) input you feed into the model? Also, which PEFT version are you using?

[SP1029]

@BenjaminBossan I think the snippet by @victolee0 is complete. I tried reproducing the error and was able to successfully reproduce on google colab. Following is my understanding of why the error is occurring in peft==0.15.2.

The PR #2403 fixes the groups argument which was being ignored in layers.py for LoRA. This is not present in peft 0.15.2, which likely is the source of this issue. Before this PR the lora_B had a incorrect size, but for lora it was working fine since lora_B @ lora_A is never directly added with the weight of the base layer. This is because during forward pass both of these terms are first multiplied with the input and then added. Also merge operation is not supported.

Before PR:
self.lora_B[adapter_name] = conv_layer(r, self.out_features, out_kernel, out_stride, bias=lora_bias)

After PR:

self.lora_B[adapter_name] = conv_layer(Add commentMore actions
        r, self.out_features // base_layer.groups, out_kernel, out_stride, bias=lora_bias
    )

In case of DoRA, when model = get_peft_model(model, lora_config) is run the dora layers are initialized. During this the weight norm is computed for which the lora_B @ lora_A has to be combined to the base layer's weight as seen below. But since these two terms are of different dimension this is not possible, hence we get the error. Also if we observe, in the last line of the error the number of parameters are differing by a factor of 2, denoting the missing groups=2.

peft/src/peft/tuners/lora/dora.py

Lines 52 to 53 in e3710e0

	lora_weight = torch.mm(lora_B.flatten(start_dim=1), lora_A.flatten(start_dim=1))
	lora_weight = lora_weight.reshape(weight.shape)

Code Snippet for error:
Same as above

Stack-Trace for error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-37-4f821cf5b80a>](https://localhost:8080/#) in <cell line: 0>()
     22 )
     23 
---> 24 model = get_peft_model(model, lora_config)

12 frames
[/usr/local/lib/python3.11/dist-packages/peft/mapping_func.py](https://localhost:8080/#) in get_peft_model(model, peft_config, adapter_name, mixed, autocast_adapter_dtype, revision, low_cpu_mem_usage)
    111 
    112     if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not peft_config.is_prompt_learning:
--> 113         return PeftModel(
    114             model,
    115             peft_config,

[/usr/local/lib/python3.11/dist-packages/peft/peft_model.py](https://localhost:8080/#) in __init__(self, model, peft_config, adapter_name, autocast_adapter_dtype, low_cpu_mem_usage)
    130             ctx = init_empty_weights if low_cpu_mem_usage else nullcontext
    131             with ctx():
--> 132                 self.base_model = cls(model, {adapter_name: peft_config}, adapter_name)
    133 
    134         self.set_additional_trainable_modules(peft_config, adapter_name)

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/model.py](https://localhost:8080/#) in __init__(self, model, config, adapter_name, low_cpu_mem_usage)
    140 
    141     def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False) -> None:
--> 142         super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
    143 
    144     def _check_new_adapter_config(self, config: LoraConfig) -> None:

[/usr/local/lib/python3.11/dist-packages/peft/tuners/tuners_utils.py](https://localhost:8080/#) in __init__(self, model, peft_config, adapter_name, low_cpu_mem_usage)
    178         self._pre_injection_hook(self.model, self.peft_config[adapter_name], adapter_name)
    179         if peft_config != PeftType.XLORA or peft_config[adapter_name] != PeftType.XLORA:
--> 180             self.inject_adapter(self.model, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
    181 
    182         # Copy the peft_config in the injected model.

[/usr/local/lib/python3.11/dist-packages/peft/tuners/tuners_utils.py](https://localhost:8080/#) in inject_adapter(self, model, adapter_name, autocast_adapter_dtype, low_cpu_mem_usage)
    506                 ctx = init_empty_weights if low_cpu_mem_usage else nullcontext
    507                 with ctx():
--> 508                     self._create_and_replace(peft_config, adapter_name, target, target_name, parent, current_key=key)
    509 
    510         if not self.targeted_module_names and not uses_dummy_target_modules:

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/model.py](https://localhost:8080/#) in _create_and_replace(self, lora_config, adapter_name, target, target_name, parent, current_key)
    235         else:
    236             device_map = self.model.hf_device_map if hasattr(self.model, "hf_device_map") else None
--> 237             new_module = self._create_new_module(lora_config, adapter_name, target, device_map=device_map, **kwargs)
    238             if adapter_name not in self.active_adapters:
    239                 # adding an additional adapter: it is not automatically trainable

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/model.py](https://localhost:8080/#) in _create_new_module(lora_config, adapter_name, target, **kwargs)
    340         new_module = None
    341         for dispatcher in dispatchers:
--> 342             new_module = dispatcher(target, adapter_name, lora_config=lora_config, **kwargs)
    343             if new_module is not None:  # first match wins
    344                 break

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/layer.py](https://localhost:8080/#) in dispatch_default(target, adapter_name, lora_config, **kwargs)
   1745     elif isinstance(target_base_layer, torch.nn.Conv2d):
   1746         kwargs.update(lora_config.loftq_config)
-> 1747         new_module = Conv2d(target, adapter_name, **kwargs)
   1748     elif isinstance(target_base_layer, torch.nn.Conv3d):
   1749         kwargs.update(lora_config.loftq_config)

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/layer.py](https://localhost:8080/#) in __init__(self, *args, **kwargs)
   1319     # Lora implemented in a conv2d layer
   1320     def __init__(self, *args, **kwargs):
-> 1321         super().__init__(*args, **kwargs)
   1322         if not self._kernel_dim == 4:
   1323             raise ValueError(f"Conv2d layer kernel must have 4 dimensions, not {self._kernel_dim}")

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/layer.py](https://localhost:8080/#) in __init__(self, base_layer, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights, use_rslora, use_dora, lora_bias, **kwargs)
   1032         self._kernel_dim = base_layer.weight.dim()
   1033 
-> 1034         self.update_layer(
   1035             adapter_name,
   1036             r,

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/layer.py](https://localhost:8080/#) in update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights, use_rslora, use_dora, lora_bias)
   1082 
   1083         if use_dora:
-> 1084             self.dora_init(adapter_name)
   1085             self.use_dora[adapter_name] = True
   1086         else:

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/layer.py](https://localhost:8080/#) in dora_init(self, adapter_name)
   1102         lora_B = self.lora_B[adapter_name].weight
   1103         scaling = self.scaling[adapter_name]
-> 1104         dora_layer.update_layer(base_layer=self.get_base_layer(), lora_A=lora_A, lora_B=lora_B, scaling=scaling)
   1105         self.lora_magnitude_vector[adapter_name] = dora_layer
   1106 

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/dora.py](https://localhost:8080/#) in update_layer(self, base_layer, lora_A, lora_B, scaling, place_on_cpu)
     51             if weight.data.ndim >= 4:  # For handling LoRAs applied to Conv layers.
     52                 lora_weight = torch.mm(lora_B.flatten(start_dim=1), lora_A.flatten(start_dim=1))
---> 53                 lora_weight = lora_weight.reshape(weight.shape)
     54             else:
     55                 lora_weight = lora_B @ lora_A

RuntimeError: shape '[192, 48, 3, 3]' is invalid for input of size 165888

[BenjaminBossan]

Thanks for the investigation @SP1029. You are right, with PEFT v0.15.2, the error can be reproduced. I falsely assumed that the error would occur during the foward call, which is why I asked for a complete example.

Just to confirm, with PEFT installed from main, there is no error? At least I don't get anything when trying.

[SP1029]

Yes @BenjaminBossan on main with the above code snippet there is no error. But just to validate, I tried a forward pass on the model, it gave rise to a new error.

New Snippet

from peft import LoraConfig, get_peft_model
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv = torch.nn.Conv2d(96, 192, (3, 3), stride=(2, 2), padding=(1, 1), bias=False, groups=2)
        self.fc = torch.nn.Linear(192, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.fc(x)
        return x


model = MyModel()

lora_config = LoraConfig(
    r=16,
    target_modules=["conv"]
    use_dora=True
)

model = get_peft_model(model, lora_config)

x = torch.zeros([1,96,192,384])
model(x).shape

Normal Model:
with model = get_peft_model(model, lora_config) commented it works

LoRA groups=1:
works

DoRA groups=1:
works

LoRA groups=2:

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/layer.py](https://localhost:8080/#) in forward(self, x, *args, **kwargs)
   1285 
   1286                 if active_adapter not in self.lora_variant:  # vanilla LoRA
-> 1287                     result = result + lora_B(lora_A(dropout(x))) * scaling
   1288                 else:
   1289                     result = self.lora_variant[active_adapter].forward(

RuntimeError: The size of tensor a (192) must match the size of tensor b (96) at non-singleton dimension 1

DoRA groups=2:

[/usr/local/lib/python3.11/dist-packages/peft/tuners/lora/dora.py](https://localhost:8080/#) in forward(self, x, lora_A, lora_B, scaling, base_layer, base_result)
    176                 base_result = base_result - bias.view(*bias_shape)
    177 
--> 178         result_dora = (mag_norm_scale - 1) * base_result + mag_norm_scale * lora_B(lora_A(x)) * scaling
    179         return result_dora
    180 

RuntimeError: The size of tensor a (192) must match the size of tensor b (96) at non-singleton dimension 1

[SP1029]

The main places conv layers under lora seemed to be in layers.py and dora.py. At these locations they are used as lora_B(lora_A(input)) form. By the following change its dimensions would be compatible with base_layer(input).

@BenjaminBossan How would the following overall idea be? With this all the above cases work.

src/peft/tuners/lora/layer.py
@1099
     self.lora_A[adapter_name] = conv_layer(self.in_features, r, kernel_size, stride, padding, bias=False)
-    self.lora_B[adapter_name] = conv_layer(
-        r, self.out_features // base_layer.groups, out_kernel, out_stride, bias=lora_bias
-    )
+    self.lora_B[adapter_name] = conv_layer(
+        r, self.out_features, out_kernel, out_stride, base_layer.groups, lora_bias
+    )

src/peft/tuners/lora/dora.py
@52
-    lora_weight = torch.mm(lora_B.flatten(start_dim=1), lora_A.flatten(start_dim=1))
+    r = lora_A.shape[0]
+    lora_weight = torch.mm(lora_B.view([-1,r]), lora_A.view([r,-1]))
@148
-    lora_weight = torch.mm(lora_B.weight.flatten(start_dim=1), lora_A.weight.flatten(start_dim=1))
+    r = lora_A.shape[0]
+    lora_weight = torch.mm(lora_B.weight.view([-1,r]), lora_A.weight.view([r,-1]))

A point is that since in_features and out_features should be divisible by groups, so this definition of lora_B introduces a requirement of rank r to be divisible by groups.
Do suggest any alternate suggestions which seem to better or anything which might have been overlooked.

[BenjaminBossan]

Thanks for further testing @SP1029. I can confirm the issue you posted, it appears that our test coverage is not good enough when it comes to conv + groups. As to the proposed solution, it might have been something I tested at some point, I'm not quite sure as I don't have a record. As you correctly say, having this restriction reduces the usefulness of supporting the feature, but if the current implementation does not work for all shapes, it is also restrictive. Maybe we can give your proposal a try in a draft PR and see how it fares with the unit tests (+ add some more tests for better coverage).

As discussed in #2403, I don't think there is a "right" solution for this problem, but we need to consider the pros and cons of each proposal.

Also pinging @gslama12 who worked on this.

[SP1029]

Added a draft PR.
On running pytest tests/test_custom_models.py -k groups -v, there were existing tests which failed due to the constrain of divisibility. E.g. the in_features,out_features multiples of 5 not divisible by 8, and rank=8 causing ValueError: in_channels must be divisible by groups.
Would it be fine if I get back to this (adding more tests) by Thursday end, since slight busy till then?

[BenjaminBossan]

Thanks a lot for the draft PR @SP1029. Regarding the tests, it is important that we have a few to understand what does and does not work. But don't worry about the time, take as much as you need. If I make some progress before then, I'll post it.

[gslama12]

Given your example I just realised that the fix in #2403 only targets depthwise convolutions where groups=in_features. Thanks for pointing this out! I will take a look at this over the weekend.

As for the suggested fix:
I think that the constraint, that rank must be divisible by groups makes the groups argument somewhat less useful. To my knowledge one of the most common way to use groups is groups=in_features and in this use case, the groups argument will mostly be larger than the rank.

[BenjaminBossan]

#2549 is merged.