mirror of
https://github.com/ggerganov/llama.cpp.git
synced 2024-10-29 22:20:15 +01:00
gguf-py : Numpy dequantization for most types
This commit is contained in:
Francis Couture-Harpin
parent
3a14e00366
commit
5a9edda7ca
@ -4,7 +4,7 @@ from typing import Any, Callable, Sequence
|
||||
|
||||
from numpy.typing import DTypeLike
|
||||
|
||||
from .constants import GGML_QUANT_SIZES, GGMLQuantizationType
|
||||
from .constants import GGML_QUANT_SIZES, GGMLQuantizationType, QK_K
|
||||
from .lazy import LazyNumpyTensor
|
||||
|
||||
import numpy as np
|
||||
@ -64,8 +64,10 @@ def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
|
||||
|
||||
|
||||
def dequantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
|
||||
if qtype == GGMLQuantizationType.F32 or qtype == GGMLQuantizationType.F16:
|
||||
return data.astype(np.float32, copy=False)
|
||||
if qtype == GGMLQuantizationType.F32:
|
||||
return data.view(np.float32)
|
||||
elif qtype == GGMLQuantizationType.F16:
|
||||
return data.view(np.float16).astype(np.float32)
|
||||
elif (q := _type_traits.get(qtype)) is not None:
|
||||
return q.dequantize(data)
|
||||
else:
|
||||
@ -187,6 +189,166 @@ class BF16(__Quant, qtype=GGMLQuantizationType.BF16):
|
||||
return (blocks.view(np.int16).astype(np.int32) << 16).view(np.float32)
|
||||
|
||||
|
||||
class Q4_0(__Quant, qtype=GGMLQuantizationType.Q4_0):
|
||||
@classmethod
|
||||
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
imax = abs(blocks).argmax(axis=-1, keepdims=True)
|
||||
max = np.take_along_axis(blocks, imax, axis=-1)
|
||||
|
||||
d = max / -8
|
||||
with np.errstate(divide="ignore"):
|
||||
id = np.where(d == 0, 0, 1 / d)
|
||||
# FIXME: Q4_0's reference rounding is cursed and depends on FMA
|
||||
qs = np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(8.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
|
||||
|
||||
qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
|
||||
qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
|
||||
|
||||
d = d.astype(np.float16).view(np.uint8)
|
||||
|
||||
return np.concatenate([d, qs], axis=-1)
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, qs = np.hsplit(blocks, [2])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
|
||||
qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.int8) - np.int8(8)
|
||||
|
||||
return (d * qs.astype(np.float32))
|
||||
|
||||
|
||||
class Q4_1(__Quant, qtype=GGMLQuantizationType.Q4_1):
|
||||
@classmethod
|
||||
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
max = blocks.max(axis=-1, keepdims=True)
|
||||
min = blocks.min(axis=-1, keepdims=True)
|
||||
|
||||
d = (max - min) / 15
|
||||
with np.errstate(divide="ignore"):
|
||||
id = np.where(d == 0, 0, 1 / d)
|
||||
qs = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
|
||||
|
||||
qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
|
||||
qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
|
||||
|
||||
d = d.astype(np.float16).view(np.uint8)
|
||||
m = min.astype(np.float16).view(np.uint8)
|
||||
|
||||
return np.concatenate([d, m, qs], axis=-1)
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, rest = np.hsplit(blocks, [2])
|
||||
m, qs = np.hsplit(rest, [2])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
m = m.view(np.float16).astype(np.float32)
|
||||
|
||||
qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.float32)
|
||||
|
||||
return (d * qs) + m
|
||||
|
||||
|
||||
class Q5_0(__Quant, qtype=GGMLQuantizationType.Q5_0):
|
||||
@classmethod
|
||||
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
imax = abs(blocks).argmax(axis=-1, keepdims=True)
|
||||
max = np.take_along_axis(blocks, imax, axis=-1)
|
||||
|
||||
d = max / -16
|
||||
with np.errstate(divide="ignore"):
|
||||
id = np.where(d == 0, 0, 1 / d)
|
||||
# FIXME: Q5_0's reference rounding is cursed and depends on FMA
|
||||
q = np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(16.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
|
||||
|
||||
qs = q.reshape((n_blocks, 2, cls.block_size // 2))
|
||||
qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
|
||||
|
||||
qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
|
||||
|
||||
d = d.astype(np.float16).view(np.uint8)
|
||||
|
||||
return np.concatenate([d, qh, qs], axis=-1)
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, rest = np.hsplit(blocks, [2])
|
||||
qh, qs = np.hsplit(rest, [4])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
qh = qh.view(np.uint32)
|
||||
|
||||
qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
|
||||
ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qh = (qh & np.uint32(0x01)).astype(np.uint8)
|
||||
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
|
||||
|
||||
qs = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(16)
|
||||
|
||||
return (d * qs.astype(np.float32))
|
||||
|
||||
|
||||
class Q5_1(__Quant, qtype=GGMLQuantizationType.Q5_1):
|
||||
@classmethod
|
||||
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
max = blocks.max(axis=-1, keepdims=True)
|
||||
min = blocks.min(axis=-1, keepdims=True)
|
||||
|
||||
d = (max - min) / 31
|
||||
with np.errstate(divide="ignore"):
|
||||
id = np.where(d == 0, 0, 1 / d)
|
||||
q = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
|
||||
|
||||
qs = q.reshape((n_blocks, 2, cls.block_size // 2))
|
||||
qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
|
||||
|
||||
qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
|
||||
|
||||
d = d.astype(np.float16).view(np.uint8)
|
||||
m = min.astype(np.float16).view(np.uint8)
|
||||
|
||||
return np.concatenate([d, m, qh, qs], axis=-1)
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, rest = np.hsplit(blocks, [2])
|
||||
m, rest = np.hsplit(rest, [2])
|
||||
qh, qs = np.hsplit(rest, [4])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
m = m.view(np.float16).astype(np.float32)
|
||||
qh = qh.view(np.uint32)
|
||||
|
||||
qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
|
||||
ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qh = (qh & np.uint32(0x01)).astype(np.uint8)
|
||||
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
|
||||
|
||||
qs = (ql | (qh << np.uint8(4))).astype(np.float32)
|
||||
|
||||
return (d * qs) + m
|
||||
|
||||
|
||||
class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0):
|
||||
@classmethod
|
||||
# Implementation of Q8_0 with bit-exact same results as reference implementation in ggml-quants.c
|
||||
@ -211,3 +373,227 @@ class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0):
|
||||
x = x.view(np.int8).astype(np.float32)
|
||||
|
||||
return (x * d)
|
||||
|
||||
|
||||
class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K):
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
scales, rest = np.hsplit(blocks, [QK_K // 16])
|
||||
qs, rest = np.hsplit(rest, [QK_K // 4])
|
||||
d, dmin = np.hsplit(rest, [2])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
dmin = dmin.view(np.float16).astype(np.float32)
|
||||
|
||||
# (n_blocks, 16, 1)
|
||||
dl = (d * (scales & 0xF).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
|
||||
ml = (dmin * (scales >> 4).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
|
||||
|
||||
shift = np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
||||
|
||||
qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & np.uint8(3)
|
||||
|
||||
qs = qs.reshape((n_blocks, QK_K // 16, 16)).astype(np.float32)
|
||||
|
||||
qs = dl * qs - ml
|
||||
|
||||
return qs.reshape((n_blocks, -1))
|
||||
|
||||
|
||||
class Q3_K(__Quant, qtype=GGMLQuantizationType.Q3_K):
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
hmask, rest = np.hsplit(blocks, [QK_K // 8])
|
||||
qs, rest = np.hsplit(rest, [QK_K // 4])
|
||||
scales, d = np.hsplit(rest, [12])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
|
||||
# The scales are packed at 6-bit each in this pattern:
|
||||
# 0: IIIIAAAA
|
||||
# 1: JJJJBBBB
|
||||
# 2: KKKKCCCC
|
||||
# 3: LLLLDDDD
|
||||
# 4: MMMMEEEE
|
||||
# 5: NNNNFFFF
|
||||
# 6: OOOOGGGG
|
||||
# 7: PPPPHHHH
|
||||
# 8: MMIIEEAA
|
||||
# 9: NNJJFFBB
|
||||
# 10: OOKKGGCC
|
||||
# 11: PPLLHHDD
|
||||
lscales, hscales = np.hsplit(scales, [8])
|
||||
lscales = lscales.reshape((n_blocks, 1, 8)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
|
||||
lscales = lscales.reshape((n_blocks, 16))
|
||||
hscales = hscales.reshape((n_blocks, 1, 4)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 4, 1))
|
||||
hscales = hscales.reshape((n_blocks, 16))
|
||||
scales = (lscales & np.uint8(0x0F)) | ((hscales & np.uint8(0x03)) << np.uint8(4))
|
||||
scales = (scales.astype(np.int8) - np.int8(32)).astype(np.float32)
|
||||
|
||||
dl = (d * scales).reshape((n_blocks, 16, 1))
|
||||
|
||||
ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
||||
qh = hmask.reshape(n_blocks, -1, 1, 32) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1))
|
||||
ql = ql.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(3)
|
||||
qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(1))
|
||||
qh = qh ^ np.uint8(1) # strangely, the offset is zero when the bitmask is 1
|
||||
q = (ql.astype(np.int8) - (qh << np.uint8(2)).astype(np.int8)).astype(np.float32)
|
||||
|
||||
return (dl * q).reshape((n_blocks, QK_K))
|
||||
|
||||
|
||||
class Q4_K(__Quant, qtype=GGMLQuantizationType.Q4_K):
|
||||
K_SCALE_SIZE = 12
|
||||
|
||||
@staticmethod
|
||||
def get_scale_min(scales: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||
n_blocks = scales.shape[0]
|
||||
scales = scales.view(np.uint8)
|
||||
### Unpacking the following: ###
|
||||
# 0 EEAAAAAA
|
||||
# 1 FFBBBBBB
|
||||
# 2 GGCCCCCC
|
||||
# 3 HHDDDDDD
|
||||
# 4 eeaaaaaa
|
||||
# 5 ffbbbbbb
|
||||
# 6 ggcccccc
|
||||
# 7 hhdddddd
|
||||
# 8 eeeeEEEE
|
||||
# 9 ffffFFFF
|
||||
# 10 ggggGGGG
|
||||
# 11 hhhhHHHH
|
||||
scales = scales.reshape((n_blocks, 3, 4))
|
||||
d, m, m_d = np.split(scales, 3, axis=-2)
|
||||
|
||||
sc = np.concatenate([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], axis=-1)
|
||||
min = np.concatenate([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], axis=-1)
|
||||
|
||||
return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8)))
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, rest = np.hsplit(blocks, [2])
|
||||
dmin, rest = np.hsplit(rest, [2])
|
||||
scales, qs = np.hsplit(rest, [cls.K_SCALE_SIZE])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
dmin = dmin.view(np.float16).astype(np.float32)
|
||||
|
||||
sc, m = Q4_K.get_scale_min(scales)
|
||||
|
||||
d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
|
||||
dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
|
||||
|
||||
qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 32)).astype(np.float32)
|
||||
|
||||
return (d * qs - dm).reshape((n_blocks, QK_K))
|
||||
|
||||
|
||||
class Q5_K(__Quant, qtype=GGMLQuantizationType.Q5_K):
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, rest = np.hsplit(blocks, [2])
|
||||
dmin, rest = np.hsplit(rest, [2])
|
||||
scales, rest = np.hsplit(rest, [Q4_K.K_SCALE_SIZE])
|
||||
qh, qs = np.hsplit(rest, [QK_K // 8])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
dmin = dmin.view(np.float16).astype(np.float32)
|
||||
|
||||
sc, m = Q4_K.get_scale_min(scales)
|
||||
|
||||
d = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
|
||||
dm = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
|
||||
|
||||
ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8, 1))
|
||||
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
|
||||
qh = (qh & np.uint8(0x01)).reshape((n_blocks, -1, 32))
|
||||
q = (ql | (qh << np.uint8(4))).astype(np.float32)
|
||||
|
||||
return (d * q - dm).reshape((n_blocks, QK_K))
|
||||
|
||||
|
||||
class Q6_K(__Quant, qtype=GGMLQuantizationType.Q6_K):
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
ql, rest = np.hsplit(blocks, [QK_K // 2])
|
||||
qh, rest = np.hsplit(rest, [QK_K // 4])
|
||||
scales, d = np.hsplit(rest, [QK_K // 16])
|
||||
|
||||
scales = scales.view(np.int8).astype(np.float32)
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
d = (d * scales).reshape((n_blocks, QK_K // 16, 1))
|
||||
|
||||
ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
|
||||
qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
||||
qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32))
|
||||
q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32)
|
||||
q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32)
|
||||
|
||||
return (d * q).reshape((n_blocks, QK_K))
|
||||
|
||||
|
||||
class IQ4_NL(__Quant, qtype=GGMLQuantizationType.IQ4_NL):
|
||||
QK4_NL = 32
|
||||
|
||||
kvalues = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, qs = np.hsplit(blocks, [2])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
|
||||
qs = qs.reshape((n_blocks, -1, 1, cls.QK4_NL // 2)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
|
||||
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 1))
|
||||
|
||||
kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
|
||||
qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1))
|
||||
|
||||
return (d * qs)
|
||||
|
||||
|
||||
class IQ4_XS(__Quant, qtype=GGMLQuantizationType.IQ4_XS):
|
||||
|
||||
@classmethod
|
||||
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
||||
n_blocks = blocks.shape[0]
|
||||
|
||||
d, rest = np.hsplit(blocks, [2])
|
||||
scales_h, rest = np.hsplit(rest, [2])
|
||||
scales_l, qs = np.hsplit(rest, [QK_K // 64])
|
||||
|
||||
d = d.view(np.float16).astype(np.float32)
|
||||
scales_h = scales_h.view(np.uint16)
|
||||
|
||||
scales_l = scales_l.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
|
||||
scales_h = scales_h.reshape((n_blocks, 1, -1)) >> np.array([2 * i for i in range(QK_K // 32)], dtype=np.uint16).reshape((1, -1, 1))
|
||||
scales_l = scales_l.reshape((n_blocks, -1)) & np.uint8(0x0F)
|
||||
scales_h = scales_h.reshape((n_blocks, -1)).astype(np.uint8) & np.uint8(0x03)
|
||||
|
||||
scales = (scales_l | (scales_h << np.uint8(4))).astype(np.int8) - np.int8(32)
|
||||
dl = (d * scales.astype(np.float32)).reshape((n_blocks, -1, 1))
|
||||
|
||||
qs = qs.reshape((n_blocks, -1, 1, 16)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
||||
qs = qs.reshape((n_blocks, -1, 32, 1)) & np.uint8(0x0F)
|
||||
|
||||
kvalues = np.array(IQ4_NL.kvalues, dtype=np.int8).reshape((1, 1, 1, -1))
|
||||
qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1, 32))
|
||||
|
||||
return (dl * qs).reshape((n_blocks, -1))
|
||||
|
||||
206
gguf-py/tests/test_quants.py
Executable file
206
gguf-py/tests/test_quants.py
Executable file
@ -0,0 +1,206 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
# Test gguf.quants so that it exactly matches the C implementation of the (de)quantization
|
||||
|
||||
# NOTE: this is kind of a mess, but at least it worked for initially testing the Python implementations.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from math import prod
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import ctypes
|
||||
import logging
|
||||
import numpy as np
|
||||
|
||||
# Necessary to load the local gguf package
|
||||
if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / 'gguf-py').exists():
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
import gguf
|
||||
from gguf.constants import GGMLQuantizationType
|
||||
|
||||
|
||||
logger = logging.getLogger("test-quants")
|
||||
|
||||
|
||||
c_float_p = ctypes.POINTER(ctypes.c_float)
|
||||
|
||||
|
||||
class ggml_init_params(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("mem_size", ctypes.c_size_t),
|
||||
("mem_buffer", ctypes.c_void_p),
|
||||
("no_alloc", ctypes.c_bool),
|
||||
]
|
||||
|
||||
|
||||
class GGMLQuants:
|
||||
libggml: ctypes.CDLL
|
||||
|
||||
def __init__(self, libggml: Path):
|
||||
self.libggml = ctypes.CDLL(str(libggml))
|
||||
self.libggml.ggml_quantize_chunk.restype = ctypes.c_size_t
|
||||
# enum ggml_type type,
|
||||
# const float * src,
|
||||
# void * dst,
|
||||
# int64_t start,
|
||||
# int64_t nrows,
|
||||
# int64_t n_per_row,
|
||||
# const float * imatrix) {
|
||||
self.libggml.ggml_quantize_chunk.argtypes = (
|
||||
ctypes.c_int,
|
||||
ctypes.POINTER(ctypes.c_float),
|
||||
ctypes.c_void_p,
|
||||
ctypes.c_int64,
|
||||
ctypes.c_int64,
|
||||
ctypes.c_int64,
|
||||
ctypes.POINTER(ctypes.c_float),
|
||||
)
|
||||
|
||||
for t in (
|
||||
"q4_0", "q4_1", "q5_0", "q5_1", "q8_0",
|
||||
"q2_K", "q3_K", "q4_K", "q5_K", "q6_K",
|
||||
"iq4_nl", "iq4_xs",
|
||||
):
|
||||
dequant_func: ctypes._NamedFuncPointer = getattr(self.libggml, "dequantize_row_" + t)
|
||||
dequant_func.restype = None
|
||||
dequant_func.argtypes = (ctypes.c_void_p, ctypes.POINTER(ctypes.c_float), ctypes.c_int64)
|
||||
|
||||
self.libggml.ggml_fp16_to_fp32_row.restype = None
|
||||
self.libggml.ggml_fp16_to_fp32_row.argtypes = (ctypes.POINTER(ctypes.c_uint16), ctypes.POINTER(ctypes.c_float), ctypes.c_int64)
|
||||
self.libggml.ggml_bf16_to_fp32_row.restype = None
|
||||
self.libggml.ggml_bf16_to_fp32_row.argtypes = (ctypes.POINTER(ctypes.c_uint16), ctypes.POINTER(ctypes.c_float), ctypes.c_int64)
|
||||
|
||||
self.libggml.ggml_init.argtypes = (ggml_init_params,)
|
||||
|
||||
self.libggml.ggml_init(ggml_init_params(1 * 1024 * 1024, 0, False))
|
||||
|
||||
def dequantize(self, tensor: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
|
||||
result = np.zeros(gguf.quant_shape_from_byte_shape(tensor.shape, qtype), dtype=np.float32, order="C")
|
||||
if qtype == GGMLQuantizationType.F32:
|
||||
# no-op
|
||||
result = tensor.view(np.float32)
|
||||
elif qtype == GGMLQuantizationType.F16:
|
||||
self.libggml.ggml_fp16_to_fp32_row(tensor.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), result.ctypes.data_as(c_float_p), result.size)
|
||||
elif qtype == GGMLQuantizationType.BF16:
|
||||
self.libggml.ggml_bf16_to_fp32_row(tensor.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), result.ctypes.data_as(c_float_p), result.size)
|
||||
else:
|
||||
lw_qname = qtype.name.lower()
|
||||
if lw_qname[-1] == "k":
|
||||
lw_qname = lw_qname[:-1] + "K"
|
||||
dequant_func: ctypes._NamedFuncPointer = getattr(self.libggml, "dequantize_row_" + lw_qname)
|
||||
dequant_func(tensor.ctypes.data_as(ctypes.c_void_p), result.ctypes.data_as(c_float_p), result.size)
|
||||
return result
|
||||
|
||||
def quantize(self, data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
|
||||
result = np.zeros(gguf.quant_shape_to_byte_shape(data.shape, qtype), dtype=np.uint8, order="C")
|
||||
result_size = self.libggml.ggml_quantize_chunk(qtype.value, data.ctypes.data_as(c_float_p), result.ctypes.data_as(ctypes.c_void_p), 0, prod(data.shape[:-1]), data.shape[-1], ctypes.cast(0, c_float_p))
|
||||
assert result.size == result_size
|
||||
return result
|
||||
|
||||
|
||||
def compare_tensors(t1: np.ndarray, t2: np.ndarray, qtype: GGMLQuantizationType) -> bool:
|
||||
same = np.array_equal(t1, t2)
|
||||
if same:
|
||||
return True
|
||||
else:
|
||||
block_size, type_size = gguf.GGML_QUANT_SIZES[qtype]
|
||||
if t1.dtype == np.float32:
|
||||
t1 = t1.reshape((-1, block_size))
|
||||
t2 = t2.reshape((-1, block_size))
|
||||
else:
|
||||
t1 = t1.reshape((-1, type_size))
|
||||
t2 = t2.reshape((-1, type_size))
|
||||
x = t1.view(np.uint8) ^ t2.view(np.uint8)
|
||||
diff_bits = np.count_nonzero(np.unpackbits(x, axis=-1), axis=-1)
|
||||
logger.debug(f"{diff_bits.shape=}")
|
||||
num_bad_blocks = np.count_nonzero(diff_bits, axis=0)
|
||||
logger.debug(f"{num_bad_blocks} bad blocks ({100 * num_bad_blocks / x.shape[0]:.6f}%)")
|
||||
bad_block_id = np.argmax(diff_bits, axis=0)
|
||||
logger.debug(f"Worst block id: {bad_block_id}")
|
||||
logger.debug(f"Sample bad block ({diff_bits[bad_block_id]} differing bits):\n{t1[bad_block_id]}\nReference:\n{t2[bad_block_id]}")
|
||||
|
||||
sum_diff_bits = np.sum(diff_bits)
|
||||
logger.debug(f"{sum_diff_bits} bits differ ({100 * sum_diff_bits/(x.size * 8):.6f}%)")
|
||||
return False
|
||||
|
||||
|
||||
def do_test(libggml_path: Path):
|
||||
ggml_quants = GGMLQuants(libggml_path)
|
||||
|
||||
np.set_printoptions(precision=None, threshold=(4 * 256) + 1, formatter={"int": lambda n: "0x%02X" % n})
|
||||
|
||||
r = np.random.randn(8, 1024, 1024).astype(np.float32, copy=False)
|
||||
|
||||
for qtype in (GGMLQuantizationType.F16, *gguf.quants._type_traits.keys()):
|
||||
has_dequantize = False
|
||||
has_quantize = False
|
||||
|
||||
try:
|
||||
gguf.dequantize(np.zeros((gguf.GGML_QUANT_SIZES[qtype][1]), dtype=np.uint8), qtype)
|
||||
has_dequantize = True
|
||||
except (NotImplementedError, AssertionError) as e:
|
||||
if isinstance(e, AssertionError):
|
||||
logger.error(f"Error with {qtype.name}: {e}")
|
||||
raise e
|
||||
try:
|
||||
gguf.quantize(np.zeros((gguf.GGML_QUANT_SIZES[qtype][0]), dtype=np.float32), qtype)
|
||||
has_quantize = True
|
||||
except (NotImplementedError, AssertionError) as e:
|
||||
if isinstance(e, AssertionError):
|
||||
logger.error(f"Error with {qtype.name}: {e}")
|
||||
raise e
|
||||
|
||||
if not has_dequantize and not has_quantize:
|
||||
continue
|
||||
|
||||
logger.info(f"Testing {qtype.name}")
|
||||
|
||||
rc = r.copy(order="C")
|
||||
|
||||
pyq = None
|
||||
|
||||
if has_quantize:
|
||||
logger.debug(f"Quantizing to {qtype.name} with Python")
|
||||
pyq = gguf.quants.quantize(rc, qtype)
|
||||
|
||||
logger.debug(f"Quantizing to {qtype.name} with C")
|
||||
ggq = ggml_quants.quantize(rc, qtype)
|
||||
|
||||
if has_quantize:
|
||||
assert pyq is not None
|
||||
if qtype == GGMLQuantizationType.F16:
|
||||
pyq = pyq.view(np.uint8)
|
||||
quant_equal = compare_tensors(pyq, ggq, qtype)
|
||||
|
||||
if not quant_equal:
|
||||
logger.error(f"Quantization to {qtype.name} does not match ❌")
|
||||
else:
|
||||
logger.info(f"Quantization to {qtype.name} matches exactly ✅")
|
||||
|
||||
if has_dequantize:
|
||||
logger.debug(f"Dequantizing from {qtype.name} with Python")
|
||||
pydq = gguf.quants.dequantize(ggq, qtype)
|
||||
logger.debug(f"Dequantizing from {qtype.name} with C")
|
||||
ggdq = ggml_quants.dequantize(ggq, qtype)
|
||||
|
||||
dequant_equal = compare_tensors(pydq, ggdq, qtype)
|
||||
|
||||
if not dequant_equal:
|
||||
logger.error(f"Dequantization from {qtype.name} does not match ❌")
|
||||
else:
|
||||
logger.info(f"Dequantization from {qtype.name} matches exactly ✅")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Test Python (de)quantization against the reference C implementation")
|
||||
parser.add_argument("--libggml", type=Path, default=Path(__file__).parent.parent.parent / "build" / "ggml" / "src" / "libggml.so", help="The path to libggml.so")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
|
||||
do_test(args.libggml)
|
||||
Loading…
Reference in New Issue
Block a user