`modules_to_save` raises TypeError when a name is a suffix of a LoRA sub-module (e.g. "out", "A")

System Info

• peft: main (fd570a5f), 0.21.0
• transformers: 5.17.0
• torch: 2.14.0
• accelerate: 1.15.0
• Python: 3.11
• Platform: Linux, CPU only

Who can help?

@BenjaminBossan

Reproduction

import torch
from torch import nn
from peft import LoraConfig, get_peft_model


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(8, 8)
        self.out = nn.Linear(8, 2)

    def forward(self, x):
        return self.out(self.lin(x))


config = LoraConfig(r=4, target_modules=["lin"], modules_to_save=["out"])
get_peft_model(Net(), config)

Traceback (most recent call last):
  ...
  File ".../peft/utils/other.py", line 1118, in _set_trainable
    new_module = wrapper_cls(target, adapter_name, **wrapper_kwargs)
  File ".../peft/utils/other.py", line 605, in __init__
    super().__init__(module_to_save, adapter_name, tied_module=tied_module)
  File ".../peft/utils/other.py", line 327, in __init__
    self.check_module()
  File ".../peft/utils/other.py", line 352, in check_module
    raise TypeError(f"{self._error_message_name()} cannot be applied to modules of type {cls_name}")
TypeError: modules_to_save cannot be applied to modules of type <class 'torch.nn.modules.container.ModuleDict'>

The same happens with modules_to_save=["A"] (matches lin.lora_A) or ["B"] (matches lin.lora_B).

Cause

_set_trainable (src/peft/utils/other.py) matches modules_to_save entries with key.endswith(target_key). The wrapping runs after the LoRA layers are injected, so at that point the model also contains lin.lora_dropout, lin.lora_A, lin.lora_B, and "lin.lora_dropout".endswith("out") is true. The LoRA layer's own ModuleDict is then passed to ModulesToSaveWrapper, whose type check raises.

There is already a guard for the adapter-name collision case (lin.lora_A.default when the adapter is called default), which checks whether the grandparent is a BaseTunerLayer. The case where the parent is the tuner layer isn't covered.

Expected behavior

Only the real out module is wrapped in ModulesToSaveWrapper; the sub-modules that a PEFT layer creates for itself (lora_A, lora_B, lora_dropout, ...) should never be modules_to_save candidates.

Proposed fix

In _set_trainable, skip a match whose parent is a BaseTunerLayer (continue). This keeps the existing suffix matching untouched (I saw #1917, where it was explained that this must stay for backwards compatibility) and only stops PEFT from matching its own internals.

I have the fix (5 lines) and a parametrized regression test in tests/test_initialization.py next to test_modules_to_save_targets_lora_layer_raises ready. The test fails on main and passes with the fix; tests/test_initialization.py, tests/test_other.py, tests/test_trainable_tokens.py and the modules_to_save/trainable_token tests in tests/test_custom_models.py are unchanged otherwise, and make quality passes. I intend to provide the PR once approved.

AI assistance (Claude) was used to investigate and draft this; I have reviewed the change and run the tests myself.