peft

1
# Copyright 2023-present the HuggingFace Inc. team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import pytest
16
import torch
17
from torch import nn
18

19
from peft import LoraConfig, get_peft_model
20

21

22
class ModelWithModuleDict(nn.Module):
23
    def __init__(self):
24
        super().__init__()
25
        self.other_layer = nn.Linear(10, 10)
26
        self.module = nn.ModuleDict({"foo": nn.Linear(10, 10)})
27

28
    def forward(self):
29
        return self.module["foo"](torch.rand(1, 10))
30

31

32
class ModelWithModuleList(nn.Module):
33
    def __init__(self):
34
        super().__init__()
35
        self.other_layer = nn.Linear(10, 10)
36
        self.module = nn.ModuleList([nn.Linear(10, 10)])
37

38
    def forward(self):
39
        return self.module[0](torch.rand(1, 10))
40

41

42
class ModelWithParameterDict(nn.Module):
43
    def __init__(self):
44
        super().__init__()
45
        self.other_layer = nn.Linear(10, 10)
46
        self.module = nn.ParameterDict({"foo": nn.Parameter(torch.rand(10, 10))})
47

48
    def forward(self):
49
        return self.module["foo"]
50

51

52
class ModelWithParameterList(nn.Module):
53
    def __init__(self):
54
        super().__init__()
55
        self.other_layer = nn.Linear(10, 10)
56
        self.module = nn.ParameterList([nn.Parameter(torch.rand(10, 10))])
57

58
    def forward(self):
59
        return self.module[0]
60

61

62
@pytest.mark.parametrize(
63
    "cls", [ModelWithModuleDict, ModelWithModuleList, ModelWithParameterDict, ModelWithParameterList]
64
)
65
def test_modules_to_save_targets_module_dict_raises(cls):
66
    model = cls()
67
    peft_config = LoraConfig(
68
        target_modules=["other_layer"],
69
        modules_to_save=["module"],
70
    )
71
    model()  # sanity check that the model would normally work
72

73
    msg = "modules_to_save cannot be applied to modules of type"
74
    with pytest.raises(TypeError, match=msg):
75
        get_peft_model(model=model, peft_config=peft_config)
76