huggingface / huggingface/peft
add_weighted_adapter with rank_pattern produces adapters that cannot be reloaded (linear, ties, dare, magnitude_prune)
- 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
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