wehub-resource-sync
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests that the auto_gptq quantization method works correctly.
|
|
|
|
Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.quantization.utils import is_quant_method_supported
|
|
from vllm.model_executor.layers.quantization.auto_gptq import (
|
|
AutoGPTQConfig,
|
|
AutoGPTQLinearMethod,
|
|
)
|
|
|
|
PROMPT = "On the surface of Mars, we found"
|
|
|
|
MODELS = [
|
|
"TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",
|
|
]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not is_quant_method_supported("auto_gptq"),
|
|
reason="auto_gptq is not supported on this GPU type.",
|
|
)
|
|
@pytest.mark.parametrize("model_id", MODELS)
|
|
def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
|
|
"""Test that quantization='auto_gptq' loads and runs correctly."""
|
|
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
|
|
|
with vllm_runner(
|
|
model_id,
|
|
dtype=torch.float16,
|
|
quantization="auto_gptq",
|
|
max_model_len=2048,
|
|
enforce_eager=True,
|
|
) as llm:
|
|
|
|
def check_model(model):
|
|
for name, submodule in model.named_modules():
|
|
if name == "model.layers.0.self_attn.qkv_proj":
|
|
assert isinstance(submodule.quant_method, AutoGPTQLinearMethod)
|
|
break
|
|
|
|
llm.apply_model(check_model)
|
|
|
|
outputs = llm.generate_greedy([PROMPT], max_tokens=8)
|
|
assert outputs
|
|
assert len(outputs[0][1]) > 0
|
|
|
|
|
|
def test_auto_gptq_config_get_name():
|
|
"""Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
|
|
assert AutoGPTQConfig.get_name() == "auto_gptq"
|