vllm/tests/lora/test_layers_utils.py at 152459c4d38ae848d15540a8d54401b4f576eec9 · AlejandroParedesLT/vllm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import pytest
import torch
from torch import nn

from vllm.lora.layers.utils import _get_lora_device

pytestmark = pytest.mark.skip_global_cleanup


def _param() -> nn.Parameter:
    return nn.Parameter(torch.empty(1), requires_grad=False)


def test_get_lora_device_unquantized():
    base_layer = nn.Module()
    base_layer.weight = _param()
    assert _get_lora_device(base_layer) == base_layer.weight.device


def test_get_lora_device_gptq_awq():
    base_layer = nn.Module()
    base_layer.qweight = _param()
    assert _get_lora_device(base_layer) == base_layer.qweight.device


def test_get_lora_device_ark_linear():
    base_layer = nn.Module()
    base_layer.ark_linear = nn.Module()
    base_layer.ark_linear.qweight = _param()
    assert _get_lora_device(base_layer) == base_layer.ark_linear.qweight.device


def test_get_lora_device_unsupported_raises():
    base_layer = nn.Module()
    with pytest.raises(ValueError, match="Unsupported base layer"):
        _get_lora_device(base_layer)