@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 156 lines 5.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm.lora.lora_model import LoRAModel
7 from vllm.lora.peft_helper import PEFTHelper
8 from vllm.lora.utils import parse_fine_tuned_lora_name
9 from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
10 from vllm.model_executor.models.utils import WeightsMapper
11
12 lora_lst = ["baichuan7B", "baichuan7B-zero", "baichuan7B-zero-regex", "chatglm3-6b"]
13 BAICHUAN_LORA_MODULES = [
14 "W_pack",
15 "o_proj",
16 "gate_up_proj",
17 "down_proj",
18 ]
19
20 MOCK_PACKED_MAPPING = {
21 "W_pack": ["W_pack"],
22 "gate_up_proj": [
23 "gate_proj",
24 "up_proj",
25 ],
26 }
27
28
29 @pytest.mark.parametrize("lora_name", lora_lst)
30 def test_load_checkpoints(
31 lora_name,
32 baichuan_lora_files,
33 baichuan_zero_lora_files,
34 baichuan_regex_lora_files,
35 chatglm3_lora_files,
36 ):
37 expected_lora_lst: list[str] = []
38 for module in BAICHUAN_LORA_MODULES:
39 if module in MOCK_PACKED_MAPPING:
40 expected_lora_lst.extend(MOCK_PACKED_MAPPING[module])
41 else:
42 expected_lora_lst.append(module)
43 expected_lora_modules = set(expected_lora_lst)
44 if lora_name == "baichuan7B":
45 peft_helper = PEFTHelper.from_local_dir(
46 baichuan_lora_files, max_position_embeddings=4096
47 )
48 # For the baichuan7B model, load it's LoRA,
49 # and the test should pass.
50 LoRAModel.from_local_checkpoint(
51 baichuan_lora_files,
52 expected_lora_modules,
53 peft_helper=peft_helper,
54 lora_model_id=1,
55 device="cpu",
56 model_vocab_size=64000,
57 )
58 elif lora_name == "baichuan7B-zero":
59 # Test that the target_modules contain prefix
60 # such as "model.layers.0.self_atten.W_pack", and
61 # the test should pass.
62 peft_helper = PEFTHelper.from_local_dir(
63 baichuan_zero_lora_files, max_position_embeddings=4096
64 )
65 LoRAModel.from_local_checkpoint(
66 baichuan_zero_lora_files,
67 expected_lora_modules,
68 peft_helper=peft_helper,
69 lora_model_id=1,
70 device="cpu",
71 model_vocab_size=64000,
72 )
73 elif lora_name == "baichuan7B-zero-regex":
74 # Test that the `target_modules` in the form of regular expressions,
75 # such as `model\\..*(W_pack|o_proj)`, and the test should pass.
76 peft_helper = PEFTHelper.from_local_dir(
77 baichuan_regex_lora_files, max_position_embeddings=4096
78 )
79 LoRAModel.from_local_checkpoint(
80 baichuan_regex_lora_files,
81 expected_lora_modules,
82 peft_helper=peft_helper,
83 lora_model_id=1,
84 device="cpu",
85 model_vocab_size=64000,
86 )
87 else:
88 # For the baichuan7B model, load chatglm3-6b's LoRA,
89 # and the test should raise the following error.
90 expected_error = "Please verify that the loaded LoRA module is correct" # noqa: E501
91 peft_helper = PEFTHelper.from_local_dir(
92 chatglm3_lora_files, max_position_embeddings=4096
93 )
94 with pytest.raises(ValueError, match=expected_error):
95 LoRAModel.from_local_checkpoint(
96 chatglm3_lora_files,
97 expected_lora_modules,
98 peft_helper=peft_helper,
99 lora_model_id=1,
100 device="cpu",
101 model_vocab_size=64000,
102 )
103
104
105 def test_lora_weights_mapping(baichuan_lora_files):
106 expected_lora_lst: list[str] = []
107 for module in BAICHUAN_LORA_MODULES:
108 if module in MOCK_PACKED_MAPPING:
109 expected_lora_lst.extend(MOCK_PACKED_MAPPING[module])
110 else:
111 expected_lora_lst.append(module)
112 expected_lora_modules = set(expected_lora_lst)
113 hf_to_vllm_mapper = WeightsMapper(
114 orig_to_new_prefix={
115 "model.": "language_model.model.",
116 },
117 orig_to_new_substr={
118 ".layers.": ".baichuan_layers.",
119 },
120 )
121 peft_helper = PEFTHelper.from_local_dir(
122 baichuan_lora_files, max_position_embeddings=4096
123 )
124 lora_model = LoRAModel.from_local_checkpoint(
125 baichuan_lora_files,
126 expected_lora_modules,
127 peft_helper=peft_helper,
128 lora_model_id=1,
129 device="cpu",
130 model_vocab_size=64000,
131 weights_mapper=hf_to_vllm_mapper,
132 )
133 for name in lora_model.loras:
134 assert name.startswith(hf_to_vllm_mapper.orig_to_new_prefix["model."])
135 assert ".baichuan_layers." in name
136
137
138 def test_gemma4_lora_weights_mapping():
139 mapper = Gemma4ForCausalLM.hf_to_vllm_mapper
140 name = "base_model.model.model.language_model.layers.9.mlp.down_proj.lora_A.weight"
141 assert parse_fine_tuned_lora_name(name, mapper) == (
142 "model.layers.9.mlp.down_proj",
143 True,
144 )
145
146
147 def test_gemma4_moe_lora_weights_mapping():
148 mapper = Gemma4ForCausalLM.hf_to_vllm_mapper
149 name = (
150 "base_model.model.model.language_model.layers.9.experts."
151 "gate_up_proj.lora_B.weight"
152 )
153 assert parse_fine_tuned_lora_name(name, mapper) == (
154 "model.layers.9.experts.gate_up_proj",
155 False,
156 )