Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_cpu_offload.py 2.13 KB
Newer Older
Simon Mo's avatar
[Misc] Add SPDX-FileCopyrightText (#19100)  
Simon Mo committed
1 
# SPDX-License-Identifier: Apache-2.0
Harry Mellor's avatar
Convert formatting to use `ruff` instead of `yapf` + `isort` (#26247)  
Harry Mellor committed
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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73 
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Expanded quantized model tests for CPU offloading
# Base tests: tests/basic_correctness/test_cpu_offload.py

import pytest

from tests.quantization.utils import is_quant_method_supported

from ..utils import compare_two_settings


@pytest.mark.skipif(
    not is_quant_method_supported("fp8"),
    reason="fp8 is not supported on this GPU type.",
)
def test_cpu_offload_fp8():
    # Test loading a quantized checkpoint
    compare_two_settings(
        "neuralmagic/Qwen2-1.5B-Instruct-FP8",
        [],
        ["--cpu-offload-gb", "1"],
        max_wait_seconds=480,
    )


@pytest.mark.skipif(
    not is_quant_method_supported("gptq_marlin"),
    reason="gptq_marlin is not supported on this GPU type.",
)
def test_cpu_offload_gptq(monkeypatch):
    # This quant method is sensitive to dummy weights, so we force real weights
    monkeypatch.setenv("VLLM_TEST_FORCE_LOAD_FORMAT", "auto")
    # Test GPTQ Marlin
    compare_two_settings(
        "Qwen/Qwen2-1.5B-Instruct-GPTQ-Int4",
        [],
        ["--cpu-offload-gb", "1"],
        max_wait_seconds=480,
    )


@pytest.mark.skipif(
    not is_quant_method_supported("awq_marlin"),
    reason="awq_marlin is not supported on this GPU type.",
)
def test_cpu_offload_awq(monkeypatch):
    # This quant method is sensitive to dummy weights, so we force real weights
    monkeypatch.setenv("VLLM_TEST_FORCE_LOAD_FORMAT", "auto")
    # Test AWQ Marlin
    compare_two_settings(
        "Qwen/Qwen2-1.5B-Instruct-AWQ",
        [],
        ["--cpu-offload-gb", "1"],
        max_wait_seconds=480,
    )


@pytest.mark.skipif(
    not is_quant_method_supported("gptq_marlin"),
    reason="gptq_marlin is not supported on this GPU type.",
)
def test_cpu_offload_compressed_tensors(monkeypatch):
    # This quant method is sensitive to dummy weights, so we force real weights
    monkeypatch.setenv("VLLM_TEST_FORCE_LOAD_FORMAT", "auto")
    # Test wNa16
    compare_two_settings(
        "nm-testing/tinyllama-oneshot-w4a16-channel-v2",
        [],
        ["--cpu-offload-gb", "1"],
        max_wait_seconds=480,
    )