| from __future__ import annotations
|
|
|
| import argparse
|
| from math import prod
|
| from pathlib import Path
|
| import ctypes
|
| import logging
|
| import numpy as np
|
|
|
|
|
| import gguf
|
| from gguf.constants import GGMLQuantizationType
|
|
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
|
|
| 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), winmode=0)
|
|
|
| self.libggml.ggml_quantize_chunk.restype = ctypes.c_size_t
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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),
|
| )
|
|
|
| self.libggml.ggml_quantize_requires_imatrix.restype = ctypes.c_bool
|
| self.libggml.ggml_quantize_requires_imatrix.argtypes = (ctypes.c_int,)
|
|
|
| for t in (
|
| "q4_0", "q4_1", "q5_0", "q5_1", "q8_0",
|
| "q2_K", "q3_K", "q4_K", "q5_K", "q6_K",
|
| "tq1_0", "tq2_0",
|
| "iq2_xxs", "iq2_xs", "iq2_s", "iq3_xxs", "iq3_s", "iq1_s", "iq1_m",
|
| "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:
|
|
|
| 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")
|
| if self.libggml.ggml_quantize_requires_imatrix(qtype.value):
|
|
|
| qw = np.sum((data * data).reshape((-1, data.shape[-1])), axis=0).ctypes.data_as(c_float_p)
|
| else:
|
| qw = ctypes.cast(0, c_float_p)
|
| 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], qw)
|
| assert result.size == result_size
|
| return result
|
|
|
|
|
| def create_sample(ggml_quants: GGMLQuants, hidden_size, qtype: GGMLQuantizationType) -> np.ndarray:
|
| gguf_writer = gguf.GGUFWriter(f"Quant_{qtype.name}_{hidden_size}.gguf", "llama")
|
|
|
|
|
| for size in [768, 1024, 2048, 5120, 18944]:
|
| tensor = np.random.randn(size, hidden_size).astype(np.float32)
|
| shape_str = "x".join(map(str, tensor.shape))
|
|
|
| gguf_writer.add_tensor(f"tensor_{qtype.name}_{shape_str}", ggml_quants.quantize(tensor, qtype), raw_dtype=qtype)
|
|
|
| gguf_writer.write_header_to_file()
|
| gguf_writer.write_kv_data_to_file()
|
| gguf_writer.write_tensors_to_file()
|
|
|
| gguf_writer.close()
|
|
|
|
|
| if __name__ == "__main__":
|
| parser = argparse.ArgumentParser(description="Test Python (de)quantization against the reference C implementation")
|
| parser.add_argument("--libggml", type=Path, default="libggml.so", help="The path to libggml.so")
|
| parser.add_argument("--hidden_size", type=int, default=256, help="The hidden size of the sample tensor")
|
| parser.add_argument("--seed", type=int, default=0, help="The hidden size of the sample tensor")
|
|
|
| np.random.seed(0)
|
|
|
| args = parser.parse_args()
|
|
|
| logging.basicConfig(level=logging.DEBUG)
|
|
|
| ggml_quants = GGMLQuants(args.libggml)
|
|
|
| qtypes = [
|
| GGMLQuantizationType.IQ1_M,
|
| GGMLQuantizationType.IQ1_S,
|
| GGMLQuantizationType.IQ2_S,
|
| GGMLQuantizationType.IQ2_XS,
|
| GGMLQuantizationType.IQ2_XXS,
|
| GGMLQuantizationType.IQ3_S,
|
| GGMLQuantizationType.IQ3_XXS,
|
| GGMLQuantizationType.IQ4_NL,
|
| GGMLQuantizationType.IQ4_XS,
|
| GGMLQuantizationType.Q2_K,
|
| GGMLQuantizationType.Q3_K,
|
| GGMLQuantizationType.Q4_K,
|
| GGMLQuantizationType.Q5_K,
|
| GGMLQuantizationType.Q6_K,
|
| GGMLQuantizationType.Q4_0,
|
| GGMLQuantizationType.Q5_0,
|
| GGMLQuantizationType.Q8_0,
|
| ]
|
|
|
| for qtype in qtypes:
|
| create_sample(ggml_quants, args.hidden_size, qtype)
|
|
|
