char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Richard Davisonggml : add NVFP4 quantization type support (#19769)5eae9cb1d

char/aliceai-t5-moe
10.3 KiB248 linesraw
1#!/usr/bin/env python3
2
3# Test gguf.quants so that it exactly matches the C implementation of the (de)quantization
4
5# NOTE: this is kind of a mess, but at least it worked for initially testing the Python implementations.
6
7from __future__ import annotations
8
9import argparse
10from math import prod
11import os
12import sys
13from pathlib import Path
14import ctypes
15import logging
16import numpy as np
17
18# Necessary to load the local gguf package
19if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / 'gguf-py').exists():
20    sys.path.insert(0, str(Path(__file__).parent.parent))
21
22import gguf
23from gguf.constants import GGMLQuantizationType
24
25
26logger = logging.getLogger("test-quants")
27
28
29c_float_p = ctypes.POINTER(ctypes.c_float)
30
31
32class ggml_init_params(ctypes.Structure):
33    _fields_ = [
34        ("mem_size", ctypes.c_size_t),
35        ("mem_buffer", ctypes.c_void_p),
36        ("no_alloc", ctypes.c_bool),
37    ]
38
39
40class GGMLQuants:
41    libggml: ctypes.CDLL
42
43    def __init__(self, libggml: Path):
44        self.libggml = ctypes.CDLL(str(libggml))
45        self.libggml.ggml_quantize_chunk.restype = ctypes.c_size_t
46        # enum ggml_type   type,
47        #    const float * src,
48        #           void * dst,
49        #        int64_t   start,
50        #        int64_t   nrows,
51        #        int64_t   n_per_row,
52        #    const float * imatrix) {
53        self.libggml.ggml_quantize_chunk.argtypes = (
54            ctypes.c_int,
55            ctypes.POINTER(ctypes.c_float),
56            ctypes.c_void_p,
57            ctypes.c_int64,
58            ctypes.c_int64,
59            ctypes.c_int64,
60            ctypes.POINTER(ctypes.c_float),
61        )
62
63        self.libggml.ggml_quantize_requires_imatrix.restype = ctypes.c_bool
64        self.libggml.ggml_quantize_requires_imatrix.argtypes = (ctypes.c_int,)
65
66        for t in (
67            "q4_0", "q4_1", "q5_0", "q5_1", "q8_0",
68            "q2_K", "q3_K", "q4_K", "q5_K", "q6_K",
69            "tq1_0", "tq2_0",
70            "mxfp4",
71            "nvfp4",
72            "iq2_xxs", "iq2_xs", "iq2_s", "iq3_xxs", "iq3_s", "iq1_s", "iq1_m",
73            "iq4_nl", "iq4_xs",
74        ):
75            dequant_func: ctypes._NamedFuncPointer = getattr(self.libggml, "dequantize_row_" + t)
76            dequant_func.restype = None
77            dequant_func.argtypes = (ctypes.c_void_p, ctypes.POINTER(ctypes.c_float), ctypes.c_int64)
78
79        self.libggml.ggml_fp16_to_fp32_row.restype = None
80        self.libggml.ggml_fp16_to_fp32_row.argtypes = (ctypes.POINTER(ctypes.c_uint16), ctypes.POINTER(ctypes.c_float), ctypes.c_int64)
81        self.libggml.ggml_bf16_to_fp32_row.restype = None
82        self.libggml.ggml_bf16_to_fp32_row.argtypes = (ctypes.POINTER(ctypes.c_uint16), ctypes.POINTER(ctypes.c_float), ctypes.c_int64)
83
84        self.libggml.ggml_init.argtypes = (ggml_init_params,)
85
86        self.libggml.ggml_init(ggml_init_params(1 * 1024 * 1024, 0, False))
87
88    def dequantize(self, tensor: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
89        result = np.zeros(gguf.quant_shape_from_byte_shape(tensor.shape, qtype), dtype=np.float32, order="C")
90        if qtype == GGMLQuantizationType.F32:
91            # no-op
92            result = tensor.view(np.float32)
93        elif qtype == GGMLQuantizationType.F16:
94            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)
95        elif qtype == GGMLQuantizationType.BF16:
96            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)
97        else:
98            lw_qname = qtype.name.lower()
99            if lw_qname[-1] == "k":
100                lw_qname = lw_qname[:-1] + "K"
101            dequant_func: ctypes._NamedFuncPointer = getattr(self.libggml, "dequantize_row_" + lw_qname)
102            dequant_func(tensor.ctypes.data_as(ctypes.c_void_p), result.ctypes.data_as(c_float_p), result.size)
103        return result
104
105    def quantize(self, data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
106        result = np.zeros(gguf.quant_shape_to_byte_shape(data.shape, qtype), dtype=np.uint8, order="C")
107        if self.libggml.ggml_quantize_requires_imatrix(qtype.value):
108            # TODO: is a column-wise sum of squares appropriate?
109            qw = np.sum((data * data).reshape((-1, data.shape[-1])), axis=0).ctypes.data_as(c_float_p)
110        else:
111            qw = ctypes.cast(0, c_float_p)
112        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)
113        assert result.size == result_size
114        return result
115
116
117def compare_tensors(t1: np.ndarray, t2: np.ndarray, qtype: GGMLQuantizationType) -> bool:
118    same = np.array_equal(t1, t2)
119    if same:
120        return True
121    else:
122        block_size, type_size = gguf.GGML_QUANT_SIZES[qtype]
123        if t1.dtype == np.float32:
124            t1 = t1.reshape((-1, block_size))
125            t2 = t2.reshape((-1, block_size))
126        else:
127            t1 = t1.reshape((-1, type_size))
128            t2 = t2.reshape((-1, type_size))
129        x = t1.view(np.uint8) ^ t2.view(np.uint8)
130        diff_bits = np.count_nonzero(np.unpackbits(x, axis=-1), axis=-1)
131        num_bad_blocks = np.count_nonzero(diff_bits, axis=0)
132        if num_bad_blocks == 0 and t1.shape == t2.shape:
133            logger.debug("Bits are equal, but arrays don't match, likely contains NANs")
134            return True
135        logger.debug(f"{num_bad_blocks} bad blocks ({100 * num_bad_blocks / x.shape[0]:.6f}%)")
136        bad_block_id = np.argmax(diff_bits, axis=0)
137        logger.debug(f"Worst block id: {bad_block_id}")
138        logger.debug(f"Sample bad block ({diff_bits[bad_block_id]} differing bits):\n{t1[bad_block_id]}\nReference:\n{t2[bad_block_id]}")
139
140        sum_diff_bits = np.sum(diff_bits)
141        logger.debug(f"{sum_diff_bits} bits differ ({100 * sum_diff_bits / (x.size * 8):.6f}%)")
142        return False
143
144
145def do_test(libggml_path: Path, quick: bool = False, user_type: GGMLQuantizationType | None = None):
146    ggml_quants = GGMLQuants(libggml_path)
147
148    np.set_printoptions(precision=None, threshold=(4 * 256) + 1, formatter={"int": lambda n: "0x%02X" % n})
149
150    r = np.random.randn(8, 1024, 1024).astype(np.float32, copy=False)
151    # test zero blocks
152    r[0, 0, :] = 0
153    ## Maybe test infinities? (can make NANs, not really useful in practice)
154    # r[0, 1, 0] = np.inf
155    # r[0, 2, 0] = -np.inf
156    # r[0, 3, 0] = np.inf
157    # r[0, 3, 1] = -np.inf
158
159    for qtype in ((GGMLQuantizationType.F16, *gguf.quants._type_traits.keys()) if user_type is None else (user_type,)):
160        has_dequantize = False
161        has_quantize = False
162
163        try:
164            gguf.dequantize(np.zeros((gguf.GGML_QUANT_SIZES[qtype][1]), dtype=np.uint8), qtype)
165            has_dequantize = True
166        except (NotImplementedError, AssertionError) as e:
167            if isinstance(e, AssertionError):
168                logger.error(f"Error with {qtype.name}: {e}")
169                raise e
170        try:
171            gguf.quantize(np.zeros((gguf.GGML_QUANT_SIZES[qtype][0]), dtype=np.float32), qtype)
172            has_quantize = True
173        except (NotImplementedError, AssertionError) as e:
174            if isinstance(e, AssertionError):
175                logger.error(f"Error with {qtype.name}: {e}")
176                raise e
177
178        if not has_dequantize and not has_quantize:
179            continue
180
181        logger.info(f"Testing {qtype.name}")
182
183        rc = r.copy(order="C")
184
185        pyq = None
186        ggq = None
187
188        if has_quantize:
189            logger.debug(f"Quantizing to {qtype.name} with Python")
190            pyq = gguf.quants.quantize(rc, qtype)
191
192            logger.debug(f"Quantizing to {qtype.name} with C")
193            ggq = ggml_quants.quantize(rc, qtype)
194
195            if qtype == GGMLQuantizationType.F16:
196                pyq = pyq.view(np.uint8)
197            quant_equal = compare_tensors(pyq, ggq, qtype)
198
199            if not quant_equal:
200                logger.error(f"Quantization to {qtype.name} does not match ❌")
201            else:
202                logger.info(f"Quantization to {qtype.name} matches exactly ✅")
203
204        if has_dequantize:
205            if ggq is None and not quick:
206                logger.debug(f"Quantizing to {qtype.name} with C")
207                ggq = ggml_quants.quantize(rc, qtype)
208
209            if ggq is not None:
210                logger.debug(f"Dequantizing from {qtype.name} with Python")
211                pydq = gguf.quants.dequantize(ggq, qtype)
212                logger.debug(f"Dequantizing from {qtype.name} with C")
213                ggdq = ggml_quants.dequantize(ggq, qtype)
214
215                dequant_equal = compare_tensors(pydq, ggdq, qtype)
216
217                if not dequant_equal:
218                    logger.error(f"Dequantization from {qtype.name} does not match ❌")
219                else:
220                    logger.info(f"Dequantization from {qtype.name} matches exactly ✅")
221
222            rq_shape = gguf.quants.quant_shape_to_byte_shape((8, 1024, 1024 // 2), qtype)
223            rq = np.random.random(rq_shape).astype(np.float16).view(np.uint8)
224
225            logger.debug(f"Dequantizing random f16 data as {qtype.name} with Python")
226            pydq = gguf.quants.dequantize(rq, qtype)
227            logger.debug(f"Dequantizing random f16 data as {qtype.name} with C")
228            ggdq = ggml_quants.dequantize(rq, qtype)
229
230            dequant_equal = compare_tensors(pydq, ggdq, qtype)
231
232            if not dequant_equal:
233                logger.error(f"Dequantization from random f16 data as {qtype.name} does not match ❌")
234            else:
235                logger.info(f"Dequantization from random f16 data as {qtype.name} matches exactly ✅")
236
237
238if __name__ == "__main__":
239    parser = argparse.ArgumentParser(description="Test Python (de)quantization against the reference C implementation")
240    parser.add_argument("--libggml", type=Path, default=Path(__file__).parent.parent.parent / "build" / "bin" / "libggml.so", help="The path to libggml.so")
241    parser.add_argument("--quick", action="store_true", help="Don't quantize with C when it's not strictly necessary")
242    parser.add_argument("--type", type=str, help="The quant type to test (all by default)")
243
244    args = parser.parse_args()
245
246    logging.basicConfig(level=logging.DEBUG)
247
248    do_test(args.libggml, args.quick, GGMLQuantizationType[args.type.upper()] if args.type is not None else None)