@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 201 lines 6.51 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections import OrderedDict
5 from typing import NamedTuple
6 from unittest.mock import MagicMock, patch
7
8 import pytest
9 from huggingface_hub.utils import HfHubHTTPError
10 from torch import nn
11
12 from vllm.lora.utils import (
13 get_adapter_absolute_path,
14 parse_fine_tuned_lora_name,
15 replace_submodule,
16 )
17 from vllm.model_executor.models.utils import WeightsMapper
18
19
20 class LoRANameParserTestConfig(NamedTuple):
21 name: str
22 module_name: str
23 is_lora_a: bool
24 weights_mapper: WeightsMapper | None = None
25
26
27 def test_parse_fine_tuned_lora_name_valid():
28 fixture = [
29 LoRANameParserTestConfig(
30 "base_model.model.lm_head.lora_A.weight", "lm_head", True, False
31 ),
32 LoRANameParserTestConfig(
33 "base_model.model.lm_head.lora_B.weight", "lm_head", False, False
34 ),
35 LoRANameParserTestConfig(
36 "base_model.model.model.embed_tokens.lora_embedding_A",
37 "model.embed_tokens",
38 True,
39 ),
40 LoRANameParserTestConfig(
41 "base_model.model.model.embed_tokens.lora_embedding_B",
42 "model.embed_tokens",
43 False,
44 ),
45 LoRANameParserTestConfig(
46 "base_model.model.model.layers.9.mlp.down_proj.lora_A.weight",
47 "model.layers.9.mlp.down_proj",
48 True,
49 ),
50 LoRANameParserTestConfig(
51 "base_model.model.model.layers.9.mlp.down_proj.lora_B.weight",
52 "model.layers.9.mlp.down_proj",
53 False,
54 ),
55 LoRANameParserTestConfig(
56 "language_model.layers.9.mlp.down_proj.lora_A.weight",
57 "language_model.layers.9.mlp.down_proj",
58 True,
59 ),
60 LoRANameParserTestConfig(
61 "language_model.layers.9.mlp.down_proj.lora_B.weight",
62 "language_model.layers.9.mlp.down_proj",
63 False,
64 ),
65 # Test with WeightsMapper
66 LoRANameParserTestConfig(
67 "base_model.model.model.layers.9.mlp.down_proj.lora_A.weight",
68 "language_model.model.layers.9.mlp.down_proj",
69 True,
70 weights_mapper=WeightsMapper(
71 orig_to_new_prefix={"model.": "language_model.model."}
72 ),
73 ),
74 LoRANameParserTestConfig(
75 "base_model.model.model.layers.9.mlp.down_proj.lora_B.weight",
76 "language_model.model.layers.9.mlp.down_proj",
77 False,
78 weights_mapper=WeightsMapper(
79 orig_to_new_prefix={"model.": "language_model.model."}
80 ),
81 ),
82 LoRANameParserTestConfig(
83 "model.layers.9.mlp.down_proj.lora_A.weight",
84 "language_model.model.layers.9.mlp.down_proj",
85 True,
86 weights_mapper=WeightsMapper(
87 orig_to_new_prefix={"model.": "language_model.model."}
88 ),
89 ),
90 LoRANameParserTestConfig(
91 "model.layers.9.mlp.down_proj.lora_B.weight",
92 "language_model.model.layers.9.mlp.down_proj",
93 False,
94 weights_mapper=WeightsMapper(
95 orig_to_new_prefix={"model.": "language_model.model."}
96 ),
97 ),
98 ]
99 for name, module_name, is_lora_a, weights_mapper in fixture:
100 assert (module_name, is_lora_a) == parse_fine_tuned_lora_name(
101 name, weights_mapper
102 )
103
104
105 def test_parse_fine_tuned_lora_name_invalid():
106 fixture = {
107 "base_model.weight",
108 "base_model.model.weight",
109 }
110 for name in fixture:
111 with pytest.raises(ValueError, match="unsupported LoRA weight"):
112 parse_fine_tuned_lora_name(name)
113
114
115 def test_replace_submodule():
116 model = nn.Sequential(
117 OrderedDict(
118 [
119 ("dense1", nn.Linear(764, 100)),
120 ("act1", nn.ReLU()),
121 ("dense2", nn.Linear(100, 50)),
122 (
123 "seq1",
124 nn.Sequential(
125 OrderedDict(
126 [
127 ("dense1", nn.Linear(100, 10)),
128 ("dense2", nn.Linear(10, 50)),
129 ]
130 )
131 ),
132 ),
133 ("act2", nn.ReLU()),
134 ("output", nn.Linear(50, 10)),
135 ("outact", nn.Sigmoid()),
136 ]
137 )
138 )
139
140 sigmoid = nn.Sigmoid()
141
142 replace_submodule(model, "act1", sigmoid)
143 assert dict(model.named_modules())["act1"] == sigmoid
144
145 dense2 = nn.Linear(1, 5)
146 replace_submodule(model, "seq1.dense2", dense2)
147 assert dict(model.named_modules())["seq1.dense2"] == dense2
148
149
150 # Unit tests for get_adapter_absolute_path
151 @patch("os.path.isabs")
152 def test_get_adapter_absolute_path_absolute(mock_isabs):
153 path = "/absolute/path/to/lora"
154 mock_isabs.return_value = True
155 assert get_adapter_absolute_path(path) == path
156
157
158 @patch("os.path.expanduser")
159 def test_get_adapter_absolute_path_expanduser(mock_expanduser):
160 # Path with ~ that needs to be expanded
161 path = "~/relative/path/to/lora"
162 absolute_path = "/home/user/relative/path/to/lora"
163 mock_expanduser.return_value = absolute_path
164 assert get_adapter_absolute_path(path) == absolute_path
165
166
167 @patch("os.path.exists")
168 @patch("os.path.abspath")
169 def test_get_adapter_absolute_path_local_existing(mock_abspath, mock_exist):
170 # Relative path that exists locally
171 path = "relative/path/to/lora"
172 absolute_path = "/absolute/path/to/lora"
173 mock_exist.return_value = True
174 mock_abspath.return_value = absolute_path
175 assert get_adapter_absolute_path(path) == absolute_path
176
177
178 @patch("huggingface_hub.HfApi.snapshot_download")
179 @patch("os.path.exists")
180 def test_get_adapter_absolute_path_huggingface(mock_exist, mock_snapshot_download):
181 # Hugging Face model identifier
182 path = "org/repo"
183 absolute_path = "/mock/snapshot/path"
184 mock_exist.return_value = False
185 mock_snapshot_download.return_value = absolute_path
186 assert get_adapter_absolute_path(path) == absolute_path
187
188
189 @patch("huggingface_hub.HfApi.snapshot_download")
190 @patch("os.path.exists")
191 def test_get_adapter_absolute_path_huggingface_error(
192 mock_exist, mock_snapshot_download
193 ):
194 # Hugging Face model identifier with download error
195 path = "org/repo"
196 mock_exist.return_value = False
197 mock_snapshot_download.side_effect = HfHubHTTPError(
198 "failed to query model info",
199 response=MagicMock(),
200 )
201 assert get_adapter_absolute_path(path) == path