Fix LoRA tests (#696)

This PR updates `test/lora/utils.py` based on latest rebase.
HabanaAI · Jan 20, 2025 · b10992b · b10992b
1 parent 018ce62
commit b10992b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git 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
@@ -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)