main
py 48 lines 1.61 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 get_adapter_absolute_path
9 from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
10
11 # Provide absolute path and huggingface lora ids
12 lora_fixture_name = ["llama32_lora_files", "llama32_lora_huggingface_id"]
13 LLAMA_LORA_MODULES = [
14 "qkv_proj",
15 "o_proj",
16 "gate_up_proj",
17 "down_proj",
18 "embed_tokens",
19 "lm_head",
20 ]
21
22
23 @pytest.mark.parametrize("lora_fixture_name", lora_fixture_name)
24 def test_load_checkpoints_from_huggingface(lora_fixture_name, request):
25 lora_name = request.getfixturevalue(lora_fixture_name)
26 packed_modules_mapping = Qwen3ForCausalLM.packed_modules_mapping
27
28 expected_lora_lst: list[str] = []
29 for module in LLAMA_LORA_MODULES:
30 if module in packed_modules_mapping:
31 expected_lora_lst.extend(packed_modules_mapping[module])
32 else:
33 expected_lora_lst.append(module)
34 expected_lora_modules = set(expected_lora_lst)
35 lora_path = get_adapter_absolute_path(lora_name)
36
37 # lora loading should work for either absolute path and huggingface id.
38 peft_helper = PEFTHelper.from_local_dir(lora_path, 4096)
39 lora_model = LoRAModel.from_local_checkpoint(
40 lora_path,
41 expected_lora_modules,
42 peft_helper=peft_helper,
43 lora_model_id=1,
44 device="cpu",
45 )
46
47 # Assertions to ensure the model is loaded correctly
48 assert lora_model is not None, "LoRAModel is not loaded correctly"