modules_to_save module whose name contains the LoRA prefix ("lora_") is silently not restored on load

[zelinewang]

System Info

peft main @ cea8213 (0.19.2.dev0), torch 2.13, transformers 5.13, CPU-only.

Who can help?

This is a state-dict save/load correctness issue in modules_to_save handling.

Reproduction

A modules_to_save module whose name contains the adapter's parameter prefix (for LoRA, lora_) is silently not restored on load. The clearest example is a module literally named lora_head:

import copy, torch
from torch import nn
from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict

class ModelWithLoraNamedHead(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Linear(2, 2, bias=False)
        self.lora_head = nn.Linear(2, 2, bias=False)   # ordinary module, name contains "lora_"
    def forward(self, x):
        return self.lora_head(self.body(x))

config = LoraConfig(target_modules=["body"], modules_to_save=["lora_head"])
source = get_peft_model(ModelWithLoraNamedHead(), config)
with torch.no_grad():
    source.base_model.model.lora_head.modules_to_save["default"].weight.fill_(7.0)

state_dict = get_peft_model_state_dict(source)
target = get_peft_model(ModelWithLoraNamedHead(), copy.deepcopy(config))
result = set_peft_model_state_dict(target, state_dict)

print("unexpected keys:", [k for k in result.unexpected_keys if "lora_head" in k])
loaded = target.base_model.model.lora_head.modules_to_save["default"].weight
print("weight restored:", bool(torch.allclose(loaded, torch.full_like(loaded, 7.0))))

Output:

unexpected keys: [
  'base_model.model.lora_head.default.original_module.weight',
  'base_model.model.lora_head.default.modules_to_save.weight',
  'base_model.model.lora_head.default.modules_to_save.default.weight',
]
weight restored: False

The trained lora_head weight stays at its initialization, and because loading is non-strict the mismatch is silent — a caller that doesn't inspect result.unexpected_keys keeps running with the wrong head weights.

Expected behavior

The modules_to_save weight should round-trip regardless of the module's name, as it does for any name that doesn't contain lora_.

Root cause

In src/peft/utils/save_and_load.py, _insert_adapter_name_into_state_dict decides whether a key belongs to the adapter with a substring test:

if parameter_prefix in key:   # parameter_prefix == "lora_"

For a modules_to_save module named lora_head, "lora_" in "...lora_head..." is True, so the key is treated as a LoRA parameter and the adapter name is spliced in at the wrong position. A prefix/boundary-aware check (rather than an unanchored in) would classify these keys correctly.

Notes

I came across this while auditing save_and_load.py for edge cases; it needs a module name that contains the lora_ prefix, so it's narrow in practice, but the result is a silent wrong-weight load. I have a small fix along the lines above plus a regression test ready. Per .ai/AGENTS.md, I wanted to check in first rather than open a PR directly — would you welcome one? (Disclosure: the audit and this write-up were done with AI assistance; the reproduction above was run against the real editable checkout.)

[DaoyuanLi2816]

Confirmed — ran your repro against a fresh origin/main checkout and got the identical output (same 3 unexpected keys, weight restored: False). The root cause diagnosis looks right: _insert_adapter_name_into_state_dict's parameter_prefix in key is an unanchored substring test, so a modules_to_save module whose own name happens to contain lora_ gets misclassified. A boundary-aware check (matching on the actual key-path segment rather than a raw substring) would fix this without touching the common case.

This looks like exactly the kind of narrow-but-real correctness bug worth a PR — you've already got the repro and diagnosis nailed down, so I'd say go for it.

[Kaif10]

I'd like to work on this, if a maintainer is happy to assign it.

Root cause as I read it: in src/peft/utils/save_and_load.py, _insert_adapter_name_into_state_dict classifies keys with an unanchored substring test (if parameter_prefix in key, with parameter_prefix == "lora_"). A modules_to_save module whose name contains that substring — e.g. lora_head — matches, so its key is rewritten as if it were adapter weight, lands under the wrong key on load, ends up in unexpected_keys, and is silently dropped under the non-strict load. The trained head weights are lost with no error.

Proposed fix: make the classification boundary-aware — match lora_ only as a full .-delimited path segment (the actual adapter-parameter naming) rather than as an arbitrary substring of a module name — so modules_to_save modules that merely contain the string aren't misclassified. I'd add a regression test next to the existing modules_to_save round-trip tests asserting the weights of a lora_-named saved module are restored after save/load.

I can develop and test this CPU-only (the repro in the issue runs on CPU). Disclosure: I'll use AI assistance in preparing the change; I understand the diagnosis and fix and will be able to defend it in review. Happy to adjust the approach if you'd prefer a different anchoring.

[DaoyuanLi2816]

Opened #3451 with a fix and regression test. It also covers a companion save-side issue I found while tracing the root cause (the same substring check over-matches on the way to get_peft_model_state_dict, not just on load), and I verified the fix generalizes to other PEFT types (BOFT, VeRA, IA3, LoHa, VBLoRA) beyond the LoRA case in the reproduction script above.

@Kaif10 I saw your comment asking to be assigned — apologies for the overlap, there was no open PR yet when I went ahead and finished this. Happy to defer if a maintainer would rather go with your version instead.

[Kaif10]

No worries at all @DaoyuanLi2816 — and thanks for the heads-up. Your PR is the more complete fix: I'd deliberately scoped my local version to the load path only, and you've also handled the save-side over-match and verified it generalizes across BOFT/VeRA/IA3/LoHa/VBLoRA, which is clearly the better outcome for the issue. Happy to defer to #3451. I'll leave it to the maintainers, but from my side your version is the one to go with.