huggingface / huggingface/peft

add_weighted_adapter with rank_pattern produces adapters that cannot be reloaded (linear, ties, dare, magnitude_prune)

Open
#3,737 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
21.7k
Forks
2.5k
Avg merge
4d 12h
Merged PRs (30d)
59

Description

### System Info

- peft: main (`0e8d0ae8`)
- transformers: 5.16.1
- torch: 2.13.0
- Python 3.11
- CPU only (not device specific)

### Who can help?

@benjaminbossan @githubnemo

### Reproduction

When LoRA adapters that use `rank_pattern` are combined with `add_weighted_adapter` and one of the `linear`, `ties`, `dare_linear`, `dare_ties` or `magnitude_prune` combination types, the resulting adapter cannot be loaded again after saving. If the adapters use different ranks for the same module (but the same maximum rank), `add_weighted_adapter` fails right away. `svd` and `cat` are not affected.

```python
import tempfile

import torch
from torch import nn

from peft import LoraConfig, PeftModel, get_peft_model

class MLP(nn.Module):
def __init__(self):
super().__init__()
self.lin0 = nn.Linear(20, 20)
self.lin1 = nn.Linear(20, 20)

def forward(self, x):
return self.lin1(self.lin0(x))

def make_config(rank_pattern):
return LoraConfig(target_modules=["lin0", "lin1"], r=8, rank_pattern=rank_pattern, init_lora_weights=False)

# 1. Both adapters use the same rank_pattern: combining works, but the result cannot be reloaded
for combination_type in ["linear", "ties", "dare_linear", "dare_ties", "magnitude_prune", "svd", "cat"]:
torch.manual_seed(0)
model = get_peft_model(MLP(), make_config({"lin1": 4}), adapter_name="a")
model.add_adapter("b", make_config({"lin1": 4}))
kwargs = {} if combination_type in ("linear", "svd", "cat") else {"density": 0.5}
model.add_weighted_adapter(["a", "b"], [0.5, 0.5], "merged", combination_type=combination_type, **kwargs)
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir, selected_adapters=["merged"])
try:
PeftModel.from_pretrained(MLP(), f"{tmp_dir}/merged")
print(f"{combination_type}: reload OK")
except RuntimeError as e:
print(f"{combination_type}: reload FAILED: {str(e).splitlines()[1].strip()}")

# 2. Same maximum rank, but different ranks for the same module: combining fails immediately
torch.manual_seed(0)
model = get_peft_model(MLP(), make_config({"lin1": 4}), adapter_name="a")
model.add_adapter("b", make_config({}))
try:
model.add_weighted_adapter(["a", "b"], [0.5, 0.5], "merged", combination_type="linear")
except RuntimeError as e:
print(f"linear with different per-module ranks: FAILED: {e}")
```

Output:

```
linear: reload FAILED: size mismatch for base_model.model.lin1.lora_A.default.weight: copying a param with shape torch.Size([4, 20]) from checkpoint, the shape in current model is torch.Size([8, 20]).
ties: reload FAILED: size mismatch for base_model.model.lin1.lora_A.default.weight: copying a param with shape torch.Size([4, 20]) from checkpoint, the shape in current model is torch.Size([8, 20]).
dare_linear: reload FAILED: size mismatch for base_model.model.lin1.lora_A.default.weight: copying a param with shape torch.Size([4, 20]) from checkpoint, the shape in current model is torch.Size([8, 20]).
dare_ties: reload FAILED: size mismatch for base_model.model.lin1.lora_A.default.weight: copying a param with shape torch.Size([4, 20]) from checkpoint, the shape in current model is torch.Size([8, 20]).
magnitude_prune: reload FAILED: size mismatch for base_model.model.lin1.lora_A.default.weight: copying a param with shape torch.Size([4, 20]) from checkpoint, the shape in current model is torch.Size([8, 20]).
svd: reload OK
cat: reload OK
linear with different per-module ranks: FAILED: stack expects each tensor to be equal size, but got [4, 20] at entry 0 and [8, 20] at entry 1
```

The reload failure also happens when combining a single adapter with weight 1.0, even though the forward pass of the combined adapter matches the original adapter.

#### Cause

`_check_add_weighted_adapter` only compares the maximum rank of each adapter (`max(config.r, *config.rank_pattern.values())`) and uses it as the rank of the new adapter:

https://github.com/huggingface/peft/blob/0e8d0ae8ab94f189f28b845e293d7452d7892d91/src/peft/tuners/lora/model.py#L637-L650

The config of the new adapter then gets this single rank and an empty `rank_pattern` (introduced in #2550 as a follow-up to #2512, both of which targeted `cat`):

https://github.com/huggingface/peft/blob/0e8d0ae8ab94f189f28b845e293d7452d7892d91/src/peft/tuners/lora/model.py#L758-L765

`svd` and `cat` write into the allocated tensors, so their shapes match the config. The `linear` family instead replaces `.data` with the weighted sum of the source `lora_A`/`lora_B` weights:

https://github.com/huggingface/peft/blob/0e8d0ae8ab94f189f28b845e293d7452d7892d91/src/peft/tuners/lora/model.py#L901-L951

So for a module whose rank in `rank_pattern` is lower than the maximum, the new adapter holds tensors of the lower rank while the saved config says `r=8` with no `rank_pattern`. When loading, the adapter is created with rank 8 for every module and the state dict no longer fits. If the source adapters have different ranks for the same module, `torch.stack` in `task_arithmetic` fails instead.

The existing test `test_add_weighted_adapter_cat_with_rank_pattern` only covers `cat`.

### Expected behavior

- Adapters with `rank_pattern` combined with the `linear` family should produce an adapter whose config matches its weights, so that it can be saved and loaded again. For example, the new adapter's `rank_pattern` could record the rank of each module.
- If the source adapters use different ranks for the same module, which the `linear` family cannot combine, `add_weighted_adapter` should raise a clear `ValueError` rather than a `torch.stack` error. The check in `_check_add_weighted_adapter` currently only compares the maximum rank.

I plan to open a PR with a fix and tests for all affected combination types once a maintainer approves.

Contributor guide

Open the contributing guide

Research direction

Start in src/peft/tuners/lora/model.py at _check_add_weighted_adapter and add_weighted_adapter, then review the existing test test_add_weighted_adapter_cat_with_rank_pattern. Reproduce the rank_pattern cases for the linear family and inspect how the resulting configuration and weights are saved. Done means affected combinations reload successfully and mismatched per-module ranks produce a clear ValueError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.