diff --git a/tests/lora/test_lora_hpu.py b/tests/lora/test_lora_hpu.py index bfdf1228b1b55..f3a3d42dffb7c 100644 --- a/tests/lora/test_lora_hpu.py +++ b/tests/lora/test_lora_hpu.py @@ -41,7 +41,7 @@ def createLoraMask(indices, batch_size, seq_len, max_loras, max_lora_rank, @pytest.mark.parametrize("rank", RANKS) @pytest.mark.parametrize("dtype", DTYPES) def test_apply_lora(m, n, k, rank, dtype) -> None: - manager = DummyLoRAManager() + manager = DummyLoRAManager(device="hpu") module_name = "module" weight = torch.rand([m, n], device="hpu", dtype=dtype) @@ -111,7 +111,7 @@ def test_apply_lora_packed_2slice(m, n, k, rank, dtype) -> None: if m // 2 not in TENSOR_SIZES: pytest.skip("m//2 must be in TENSOR_SIZES") - manager = DummyLoRAManager() + manager = DummyLoRAManager(device="hpu") module_name = "module" weight = torch.rand([m // 2, n], device="hpu", dtype=dtype) @@ -183,7 +183,7 @@ def test_apply_lora_packed_2slice(m, n, k, rank, dtype) -> None: @pytest.mark.parametrize("rank", RANKS) @pytest.mark.parametrize("dtype", DTYPES) def test_apply_lora_packed_3slice(qkv, n, k, rank, dtype) -> None: - manager = DummyLoRAManager() + manager = DummyLoRAManager(device="hpu") module_name = "module" weight_q = torch.empty(qkv[0], n, device="hpu", dtype=dtype) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 8096288777271..a35aa47f1b4b8 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -3,7 +3,6 @@ import torch from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights -from vllm.utils import get_device class DummyLoRAManager: @@ -32,10 +31,10 @@ def init_random_lora( lora_alpha=1, lora_a=torch.rand([weight.shape[1], rank], dtype=weight.dtype, - device=get_device()), + device=self._device), lora_b=torch.rand([rank, weight.shape[0]], dtype=weight.dtype, - device=get_device()), + device=self._device), ) if generate_embeddings_tensor: lora.embeddings_tensor = torch.rand( @@ -61,8 +60,8 @@ def init_lora( module_name, rank=rank, lora_alpha=1, - lora_a=torch.rand([input_dim, rank], device=get_device()), - lora_b=torch.rand([rank, output_dim], device=get_device()), + lora_a=torch.rand([input_dim, rank], device=self._device), + lora_b=torch.rand([rank, output_dim], device=self._device), embeddings_tensor=embeddings_tensor, ) self.set_module_lora(module_name, lora)