| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | from torch import nn |
| 7 | |
| 8 | from vllm.lora.layers.utils import _get_lora_device |
| 9 | |
| 10 | pytestmark = pytest.mark.skip_global_cleanup |
| 11 | |
| 12 | |
| 13 | def _param() -> nn.Parameter: |
| 14 | return nn.Parameter(torch.empty(1), requires_grad=False) |
| 15 | |
| 16 | |
| 17 | def test_get_lora_device_unquantized(): |
| 18 | base_layer = nn.Module() |
| 19 | base_layer.weight = _param() |
| 20 | assert _get_lora_device(base_layer) == base_layer.weight.device |
| 21 | |
| 22 | |
| 23 | def test_get_lora_device_gptq_awq(): |
| 24 | base_layer = nn.Module() |
| 25 | base_layer.qweight = _param() |
| 26 | assert _get_lora_device(base_layer) == base_layer.qweight.device |
| 27 | |
| 28 | |
| 29 | def test_get_lora_device_ark_linear(): |
| 30 | base_layer = nn.Module() |
| 31 | base_layer.ark_linear = nn.Module() |
| 32 | base_layer.ark_linear.qweight = _param() |
| 33 | assert _get_lora_device(base_layer) == base_layer.ark_linear.qweight.device |
| 34 | |
| 35 | |
| 36 | def test_get_lora_device_unsupported_raises(): |
| 37 | base_layer = nn.Module() |
| 38 | with pytest.raises(ValueError, match="Unsupported base layer"): |
| 39 | _get_lora_device(base_layer) |