test_lora_checkpoints.py 4.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
import pytest

from vllm.lora.models import LoRAModel
7
from vllm.lora.peft_helper import PEFTHelper
8
from vllm.model_executor.models.baichuan import BaiChuanBaseForCausalLM
9
from vllm.model_executor.models.utils import WeightsMapper
10

11
lora_lst = ["baichuan7B", "baichuan7B-zero", "baichuan7B-zero-regex", "chatglm3-6b"]
12
13
14
15
16
17
BAICHUAN_LORA_MODULES = [
    "W_pack",
    "o_proj",
    "gate_up_proj",
    "down_proj",
]
18

19
20
21
22
23
24

@pytest.mark.parametrize("lora_name", lora_lst)
def test_load_checkpoints(
    lora_name,
    baichuan_lora_files,
    baichuan_zero_lora_files,
25
    baichuan_regex_lora_files,
26
27
    chatglm3_lora_files,
):
28
    packed_modules_mapping = BaiChuanBaseForCausalLM.packed_modules_mapping
29

30
    expected_lora_lst: list[str] = []
31
    for module in BAICHUAN_LORA_MODULES:
32
        if module in packed_modules_mapping:
33
            expected_lora_lst.extend(packed_modules_mapping[module])
34
        else:
35
36
            expected_lora_lst.append(module)
    expected_lora_modules = set(expected_lora_lst)
37
    if lora_name == "baichuan7B":
38
39
40
        peft_helper = PEFTHelper.from_local_dir(
            baichuan_lora_files, max_position_embeddings=4096
        )
41
42
43
44
45
        # For the baichuan7B model, load it's LoRA,
        # and the test should pass.
        LoRAModel.from_local_checkpoint(
            baichuan_lora_files,
            expected_lora_modules,
46
            peft_helper=peft_helper,
47
48
            lora_model_id=1,
            device="cpu",
49
            model_vocab_size=64000,
50
        )
51
    elif lora_name == "baichuan7B-zero":
52
        # Test that the target_modules contain prefix
53
54
        # such as "model.layers.0.self_atten.W_pack", and
        # the test should pass.
55
56
57
        peft_helper = PEFTHelper.from_local_dir(
            baichuan_zero_lora_files, max_position_embeddings=4096
        )
58
59
60
        LoRAModel.from_local_checkpoint(
            baichuan_zero_lora_files,
            expected_lora_modules,
61
            peft_helper=peft_helper,
62
63
            lora_model_id=1,
            device="cpu",
64
            model_vocab_size=64000,
65
        )
66
67
68
    elif lora_name == "baichuan7B-zero-regex":
        # Test that the `target_modules` in the form of regular expressions,
        # such as `model\\..*(W_pack|o_proj)`, and the test should pass.
69
70
71
        peft_helper = PEFTHelper.from_local_dir(
            baichuan_regex_lora_files, max_position_embeddings=4096
        )
72
73
74
        LoRAModel.from_local_checkpoint(
            baichuan_regex_lora_files,
            expected_lora_modules,
75
            peft_helper=peft_helper,
76
77
            lora_model_id=1,
            device="cpu",
78
            model_vocab_size=64000,
79
        )
80
81
82
83
    else:
        # For the baichuan7B model, load chatglm3-6b's LoRA,
        # and the test should raise the following error.
        expected_error = "Please verify that the loaded LoRA module is correct"  # noqa: E501
84
85
86
        peft_helper = PEFTHelper.from_local_dir(
            chatglm3_lora_files, max_position_embeddings=4096
        )
87
88
89
90
        with pytest.raises(ValueError, match=expected_error):
            LoRAModel.from_local_checkpoint(
                chatglm3_lora_files,
                expected_lora_modules,
91
                peft_helper=peft_helper,
92
93
                lora_model_id=1,
                device="cpu",
94
                model_vocab_size=64000,
95
            )
96
97


98
def test_lora_weights_mapping(baichuan_lora_files):
99
    packed_modules_mapping = BaiChuanBaseForCausalLM.packed_modules_mapping
100

101
    expected_lora_lst: list[str] = []
102
    for module in BAICHUAN_LORA_MODULES:
103
        if module in packed_modules_mapping:
104
            expected_lora_lst.extend(packed_modules_mapping[module])
105
        else:
106
107
            expected_lora_lst.append(module)
    expected_lora_modules = set(expected_lora_lst)
108
109
110
111
112
113
114
115
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "model.": "language_model.model.",
        },
        orig_to_new_substr={
            ".layers.": ".baichuan_layers.",
        },
    )
116
117
118
    peft_helper = PEFTHelper.from_local_dir(
        baichuan_lora_files, max_position_embeddings=4096
    )
119
120
121
    lora_model = LoRAModel.from_local_checkpoint(
        baichuan_lora_files,
        expected_lora_modules,
122
        peft_helper=peft_helper,
123
124
        lora_model_id=1,
        device="cpu",
125
        model_vocab_size=64000,
126
127
128
129
        weights_mapper=hf_to_vllm_mapper,
    )
    for name in lora_model.loras:
        assert name.startswith(hf_to_vllm_mapper.orig_to_new_prefix["model."])
130
        assert ".baichuan_layers." in name