Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from vllm.model_executor.layers.quantization.gguf_transcode import (
|
|
reconstruction_error,
|
|
transcode_affine,
|
|
)
|
|
from vllm.transformers_utils.gguf_tensor_reader import dequantize, quant_size
|
|
|
|
|
|
def source_blocks(weight_type, rows=32, k=768, scale=0.00390625):
|
|
if weight_type in (41, 42):
|
|
k = 640
|
|
block, size = quant_size(weight_type)
|
|
data = np.random.default_rng(20261003 + weight_type).integers(
|
|
0,
|
|
256,
|
|
(rows * k // block, size),
|
|
dtype=np.uint8,
|
|
)
|
|
offset = 80 if weight_type == 10 else (0 if weight_type in (41, 42) else size - 2)
|
|
data[:, offset : offset + 2] = np.frombuffer(np.float16(scale).tobytes(), np.uint8)
|
|
if weight_type == 10:
|
|
data[:, 82:84] = np.frombuffer(np.float16(-0.0078125).tobytes(), np.uint8)
|
|
return data.reshape(rows, -1)
|
|
|
|
|
|
def oracle(data, weight_type):
|
|
if weight_type != 41:
|
|
return dequantize(data, weight_type)
|
|
# The installed gguf-py has the Q1 enum but no dequantizer. This follows
|
|
# llama.cpp dequantize_row_q1_0: least-significant bit first in each byte.
|
|
blocks = data.reshape(-1, 18)
|
|
d = blocks[:, :2].copy().view("<f2").astype(np.float32)
|
|
values = (
|
|
np.unpackbits(blocks[:, 2:], axis=1, bitorder="little").astype(np.float32) * 2
|
|
- 1
|
|
)
|
|
return (values * d).reshape(data.shape[0], -1)
|
|
|
|
|
|
@pytest.mark.parametrize("weight_type", [10, 34, 35, 41, 42])
|
|
def test_u2_affine_and_ternary_match_official_reconstruction(weight_type):
|
|
data = source_blocks(weight_type)
|
|
canonical = transcode_affine(data, weight_type)
|
|
assert canonical.bits == 2
|
|
np.testing.assert_array_equal(canonical.dequantize(), oracle(data, weight_type))
|
|
assert reconstruction_error(canonical, oracle(data, weight_type))["max_abs"] == 0
|
|
for rank in range(4):
|
|
shard = canonical.tp_slice(rank, 4, axis=1)
|
|
width = canonical.codes.shape[1] // 4
|
|
np.testing.assert_array_equal(
|
|
shard.dequantize(),
|
|
canonical.dequantize()[:, rank * width : (rank + 1) * width],
|
|
)
|
|
|
|
|
|
def test_q1_scale_at_fp16_limit_does_not_overflow_canonical_coefficients():
|
|
data = source_blocks(41, scale=65504)
|
|
canonical = transcode_affine(data, 41)
|
|
assert np.all(np.isfinite(canonical.scales))
|
|
np.testing.assert_array_equal(canonical.dequantize(), oracle(data, 41))
|
|
|
|
|
|
def test_q2_k_expansion_reports_fp16_rounding():
|
|
data = source_blocks(10, scale=0.001337)
|
|
result = reconstruction_error(transcode_affine(data, 10), oracle(data, 10))
|
|
assert 0 < result["relative_l2"] < 0.001
|