A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 22:35 UTC
Add lora test for tp>1 case for TPU. (#21970)
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
XiongfeiWei committed
Aug 1, 2025 at 11:56 UTC
d84b97a3e33ed79aaba7552bfe5889d363875562
1 file changed
+16
-7
tests/tpu/lora/test_lora.py
+16
-7
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
+from torch_xla._internal import tpu
5
6
import vllm
7
from vllm.lora.request import LoRARequest
@@ -27,25 +28,31 @@ def use_v1_only(monkeypatch: pytest.MonkeyPatch):
28
yield
29
30
30
-def setup_vllm(num_loras: int) -> vllm.LLM:
31
+def setup_vllm(num_loras: int, tp: int) -> vllm.LLM:
32
return vllm.LLM(model="Qwen/Qwen2.5-3B-Instruct",
33
num_scheduler_steps=1,
34
max_model_len=256,
35
max_seq_len_to_capture=256,
36
max_num_seqs=8,
37
+ tensor_parallel_size=tp,
38
enable_lora=True,
39
max_loras=num_loras,
40
max_lora_rank=8)
41
42
41
-def test_single_lora():
43
+TPU_TENSOR_PARALLEL_SIZES = [1, tpu.num_available_chips()
44
+ ] if tpu.num_available_chips() > 1 else [1]
45
+
46
+
47
+@pytest.mark.parametrize("tp", TPU_TENSOR_PARALLEL_SIZES)
48
+def test_single_lora(tp: int):
49
"""
50
This test ensures we can run a single LoRA adapter on the TPU backend.
51
We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_1_adapter" which
52
will force Qwen2.5-3B-Instruct to claim 1+1=1.
53
"""
54
48
- llm = setup_vllm(1)
55
+ llm = setup_vllm(1, tp)
56
57
prompt = "What is 1+1? \n"
58
@@ -63,7 +70,8 @@ def test_single_lora():
70
assert int(answer) == 1
71
72
66
-def test_lora_hotswapping():
73
+@pytest.mark.parametrize("tp", TPU_TENSOR_PARALLEL_SIZES)
74
+def test_lora_hotswapping(tp: int):
75
"""
76
This test ensures we can run multiple LoRA adapters on the TPU backend, even
77
if we only have space to store 1.
@@ -79,7 +87,7 @@ def test_lora_hotswapping():
87
for i in range(1, 5)
88
]
89
82
- llm = setup_vllm(1)
90
+ llm = setup_vllm(1, tp)
91
92
prompt = "What is 1+1? \n"
93
@@ -94,7 +102,8 @@ def test_lora_hotswapping():
102
assert int(answer) == i + 1
103
104
97
-def test_multi_lora():
105
+@pytest.mark.parametrize("tp", TPU_TENSOR_PARALLEL_SIZES)
106
+def test_multi_lora(tp: int):
107
"""
108
This test ensures we can run multiple LoRA adapters on the TPU backend, when
109
we have enough space to store all of them.
@@ -109,7 +118,7 @@ def test_multi_lora():
118
for i in range(1, 5)
119
]
120
112
- llm = setup_vllm(4)
121
+ llm = setup_vllm(4, tp)
122
123
prompt = "What is 1+1? \n"
124