test_compressed_tensors.py 2.36 KB
Newer Older
1
2
3
4
5
6
7
8
"""Test model set-up and weight loading for sparseml-quantized models.

Run `pytest tests/quantization/test_compressed_tensors.py`.
"""

import torch

from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
9
10
    CompressedTensorsLinearMethod, CompressedTensorsW8A8DynamicToken,
    CompressedTensorsW8A8StaticTensor)
11
12
13
14


def test_compressed_tensors_w8a8_static_setup(vllm_runner):
    model_path = "nm-testing/tinyllama-one-shot-static-quant-test-compressed"
15
16
17
18
    with vllm_runner(model_path, quantization="sparseml",
                     enforce_eager=True) as llm:
        model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model  # noqa: E501
        layer = model.model.layers[0]
19

20
21
22
23
        qkv_proj = layer.self_attn.qkv_proj
        o_proj = layer.self_attn.o_proj
        gate_up_proj = layer.mlp.gate_up_proj
        down_proj = layer.mlp.down_proj
24

25
26
27
28
29
30
        assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
        assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
        assert isinstance(gate_up_proj.quant_method,
                          CompressedTensorsLinearMethod)
        assert isinstance(down_proj.quant_method,
                          CompressedTensorsLinearMethod)
31

32
        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8StaticTensor)
33

34
35
36
        assert qkv_proj.weight.dtype is torch.int8
        assert o_proj.weight.dtype is torch.int8
        assert gate_up_proj.weight.dtype is torch.int8
37

38
39
40
        assert qkv_proj.weight_scale.shard_splitter is not None
        assert qkv_proj.weight_scale.logical_widths is not None
        assert qkv_proj.input_scale.dtype is torch.float32
41
42
43
44


def test_compressed_tensors_w8a8_dynanmic_per_token(vllm_runner):
    model_path = "nm-testing/tinyllama-one-shot-dynamic-test"
45
46
47
48
49
50
51
52
53
54
55
56
    with vllm_runner(model_path,
                     quantization="sparseml",
                     enforce_eager=True,
                     dtype=torch.float16) as llm:
        model = llm.model.llm_engine.model_executor.driver_worker.model_runner.model  # noqa: E501
        layer = model.model.layers[0]

        qkv_proj = layer.self_attn.qkv_proj

        assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
        assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8DynamicToken)
        assert qkv_proj.weight.dtype is torch.int8