char-fork/llama.cpp

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

Oğuzhan Akkayagguf-py : add size guards to GGUFReader (#27188)0329fcdac

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
15.9 KiB390 linesraw
1#
2# GGUF file reading/modification support. For API usage information,
3# please see the files scripts/ for some fairly simple examples.
4#
5from __future__ import annotations
6
7import logging
8import os
9import sys
10from collections import OrderedDict
11from typing import Any, Literal, NamedTuple, TypeVar, Union
12
13import numpy as np
14import numpy.typing as npt
15
16from .quants import quant_shape_to_byte_shape
17
18if __name__ == "__main__":
19    from pathlib import Path
20
21    # Allow running file in package as a script.
22    sys.path.insert(0, str(Path(__file__).parent.parent))
23
24from gguf.constants import (
25    GGML_MAX_DIMS,
26    GGML_QUANT_SIZES,
27    GGUF_DEFAULT_ALIGNMENT,
28    GGUF_MAGIC,
29    GGUF_VERSION,
30    GGMLQuantizationType,
31    GGUFValueType,
32    GGUFEndian,
33)
34
35# limits mirroring ggml/src/gguf.cpp (not part of gguf.h)
36GGUF_MAX_STRING_LENGTH  = 1024 * 1024 * 1024
37GGUF_MAX_ARRAY_ELEMENTS = 1024 * 1024 * 1024
38
39logger = logging.getLogger(__name__)
40
41READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION]
42
43
44class ReaderField(NamedTuple):
45    # Offset to start of this field.
46    offset: int
47
48    # Name of the field (not necessarily from file data).
49    name: str
50
51    # Data parts. Some types have multiple components, such as strings
52    # that consist of a length followed by the string data.
53    parts: list[npt.NDArray[Any]] = []
54
55    # Indexes into parts that we can call the actual data. For example
56    # an array of strings will be populated with indexes to the actual
57    # string data.
58    data: list[int] = [-1]
59
60    types: list[GGUFValueType] = []
61
62    def contents(self, index_or_slice: int | slice = slice(None)) -> Any:
63        if self.types:
64            to_string = lambda x: str(x.tobytes(), encoding='utf-8') # noqa: E731
65            main_type = self.types[0]
66
67            if main_type == GGUFValueType.ARRAY:
68                sub_type = self.types[-1]
69
70                if sub_type == GGUFValueType.STRING:
71                    indices = self.data[index_or_slice]
72
73                    if isinstance(index_or_slice, int):
74                        return to_string(self.parts[indices]) # type: ignore
75                    else:
76                        return [to_string(self.parts[idx]) for idx in indices] # type: ignore
77                else:
78                    # FIXME: When/if _get_field_parts() support multi-dimensional arrays, this must do so too
79
80                    # Check if it's unsafe to perform slice optimization on data
81                    # if any(True for idx in self.data if len(self.parts[idx]) != 1):
82                    #     optim_slice = slice(None)
83                    # else:
84                    #     optim_slice = index_or_slice
85                    #     index_or_slice = slice(None)
86
87                    # if isinstance(optim_slice, int):
88                    #     return self.parts[self.data[optim_slice]].tolist()[0]
89                    # else:
90                    #     return [pv for idx in self.data[optim_slice] for pv in self.parts[idx].tolist()][index_or_slice]
91
92                    if isinstance(index_or_slice, int):
93                        return self.parts[self.data[index_or_slice]].tolist()[0]
94                    else:
95                        return [pv for idx in self.data[index_or_slice] for pv in self.parts[idx].tolist()]
96
97            if main_type == GGUFValueType.STRING:
98                return to_string(self.parts[-1])
99            else:
100                return self.parts[-1].tolist()[0]
101
102        return None
103
104
105class ReaderTensor(NamedTuple):
106    name: str
107    tensor_type: GGMLQuantizationType
108    shape: npt.NDArray[np.uint32]
109    n_elements: int
110    n_bytes: int
111    data_offset: int
112    data: npt.NDArray[Any]
113    field: ReaderField
114
115
116class GGUFReader:
117    # I - same as host, S - swapped
118    byte_order: Literal['I', 'S'] = 'I'
119    alignment: int = GGUF_DEFAULT_ALIGNMENT
120    data_offset: int
121
122    # Note: Internal helper, API may change.
123    gguf_scalar_to_np: dict[GGUFValueType, type[np.generic]] = {
124        GGUFValueType.UINT8:   np.uint8,
125        GGUFValueType.INT8:    np.int8,
126        GGUFValueType.UINT16:  np.uint16,
127        GGUFValueType.INT16:   np.int16,
128        GGUFValueType.UINT32:  np.uint32,
129        GGUFValueType.INT32:   np.int32,
130        GGUFValueType.FLOAT32: np.float32,
131        GGUFValueType.UINT64:  np.uint64,
132        GGUFValueType.INT64:   np.int64,
133        GGUFValueType.FLOAT64: np.float64,
134        GGUFValueType.BOOL:    np.bool_,
135    }
136
137    def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'):
138        self.data = np.memmap(path, mode = mode)
139        offs = 0
140
141        # Check for GGUF magic
142        if self._get(offs, np.uint32, override_order = '<')[0] != GGUF_MAGIC:
143            raise ValueError('GGUF magic invalid')
144        offs += 4
145
146        # Check GGUF version
147        temp_version = self._get(offs, np.uint32)
148        if temp_version[0] & 65535 == 0:
149            # If we get 0 here that means it's (probably) a GGUF file created for
150            # the opposite byte order of the machine this script is running on.
151            self.byte_order = 'S'
152            temp_version = temp_version.view(temp_version.dtype.newbyteorder(self.byte_order))
153        version = temp_version[0]
154        if version not in READER_SUPPORTED_VERSIONS:
155            raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle')
156        if sys.byteorder == "little":
157            # Host is little endian
158            host_endian = GGUFEndian.LITTLE
159            swapped_endian = GGUFEndian.BIG
160        else:
161            # Sorry PDP or other weird systems that don't use BE or LE.
162            host_endian = GGUFEndian.BIG
163            swapped_endian = GGUFEndian.LITTLE
164        self.endianess = swapped_endian if self.byte_order == "S" else host_endian
165        self.fields: OrderedDict[str, ReaderField] = OrderedDict()
166        self.tensors: list[ReaderTensor] = []
167        offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp_version], [0], [GGUFValueType.UINT32]))
168
169        # Check tensor count and kv count
170        temp_counts = self._get(offs, np.uint64, 2)
171        offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64]))
172        offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64]))
173        tensor_count, kv_count = temp_counts
174        if tensor_count > GGUF_MAX_ARRAY_ELEMENTS:
175            raise ValueError(f'Tensor count {tensor_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
176        if kv_count > GGUF_MAX_ARRAY_ELEMENTS:
177            raise ValueError(f'KV count {kv_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
178        offs = self._build_fields(offs, kv_count)
179
180        # Build Tensor Info Fields
181        offs, tensors_fields = self._build_tensor_info(offs, tensor_count)
182        new_align = self.fields.get('general.alignment')
183        if new_align is not None:
184            if new_align.types != [GGUFValueType.UINT32]:
185                raise ValueError('Bad type for general.alignment field')
186            self.alignment = new_align.parts[-1][0]
187            # Ensure alignment is a non-zero power of two
188            if self.alignment == 0 or (self.alignment & (self.alignment - 1)) != 0:
189                raise ValueError('Invalid alignment: must be a non-zero power of two')
190        padding = offs % self.alignment
191        if padding != 0:
192            offs += self.alignment - padding
193        self.data_offset = offs
194        self._build_tensors(offs, tensors_fields)
195
196    _DT = TypeVar('_DT', bound = npt.DTypeLike)
197
198    # Fetch a key/value metadata field by key.
199    def get_field(self, key: str) -> Union[ReaderField, None]:
200        return self.fields.get(key, None)
201
202    # Fetch a tensor from the list by index.
203    def get_tensor(self, idx: int) -> ReaderTensor:
204        return self.tensors[idx]
205
206    def _get(
207        self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None,
208    ) -> npt.NDArray[Any]:
209        count = int(count)
210        itemsize = int(np.empty([], dtype = dtype).itemsize)
211        end_offs = offset + itemsize * count
212        arr = self.data[offset:end_offs].view(dtype=dtype)[:count]
213        return arr.view(arr.dtype.newbyteorder(self.byte_order if override_order is None else override_order))
214
215    def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int:
216        if field.name in self.fields:
217            # TODO: add option to make this a warning and accept duplicate keys like below
218            raise KeyError(f'Duplicate {field.name} already in list at offset {field.offset}')
219
220            # logger.warning(f'Duplicate key {field.name} at offset {field.offset}')
221            # self.fields[field.name + '_{}'.format(field.offset)] = field
222        else:
223            self.fields[field.name] = field
224        return 0 if skip_sum else sum(int(part.nbytes) for part in field.parts)
225
226    def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]:
227        slen = self._get(offset, np.uint64)
228        if int(slen[0]) > GGUF_MAX_STRING_LENGTH:
229            raise ValueError(f'String length {int(slen[0])} exceeds maximum {GGUF_MAX_STRING_LENGTH}')
230        if offset + 8 + int(slen[0]) > self.data.nbytes:
231            raise ValueError(f'String length {int(slen[0])} exceeds remaining file size {self.data.nbytes - offset - 8}')
232        return slen, self._get(offset + 8, np.uint8, slen[0])
233
234    def _get_field_parts(
235        self, orig_offs: int, raw_type: int,
236    ) -> tuple[int, list[npt.NDArray[Any]], list[int], list[GGUFValueType]]:
237        offs = orig_offs
238        types: list[GGUFValueType] = []
239        gtype = GGUFValueType(raw_type)
240        types.append(gtype)
241        # Handle strings.
242        if gtype == GGUFValueType.STRING:
243            sparts: list[npt.NDArray[Any]] = list(self._get_str(offs))
244            size = sum(int(part.nbytes) for part in sparts)
245            return size, sparts, [1], types
246        # Check if it's a simple scalar type.
247        nptype = self.gguf_scalar_to_np.get(gtype)
248        if nptype is not None:
249            val = self._get(offs, nptype)
250            return int(val.nbytes), [val], [0], types
251        # Handle arrays.
252        if gtype == GGUFValueType.ARRAY:
253            raw_itype = self._get(offs, np.uint32)
254            offs += int(raw_itype.nbytes)
255            alen = self._get(offs, np.uint64)
256            if int(alen[0]) > GGUF_MAX_ARRAY_ELEMENTS:
257                raise ValueError(f'Array length {int(alen[0])} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}')
258            offs += int(alen.nbytes)
259            aparts: list[npt.NDArray[Any]] = [raw_itype, alen]
260            data_idxs: list[int] = []
261            # FIXME: Handle multi-dimensional arrays properly instead of flattening
262            for idx in range(alen[0]):
263                curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0])
264                if idx == 0:
265                    types += curr_types
266                idxs_offs = len(aparts)
267                aparts += curr_parts
268                data_idxs += (idx + idxs_offs for idx in curr_idxs)
269                offs += curr_size
270            return offs - orig_offs, aparts, data_idxs, types
271        # We can't deal with this one.
272        raise ValueError(f'Unknown/unhandled field type {gtype}')
273
274    def _get_tensor_info_field(self, orig_offs: int) -> ReaderField:
275        offs = orig_offs
276
277        # Get Tensor Name
278        name_len, name_data = self._get_str(offs)
279        offs += int(name_len.nbytes + name_data.nbytes)
280
281        # Get Tensor Dimensions Count
282        n_dims = self._get(offs, np.uint32)
283        offs += int(n_dims.nbytes)
284        if n_dims[0] > GGML_MAX_DIMS:
285            raise ValueError(f'Tensor dimensions count {n_dims[0]} exceeds GGML_MAX_DIMS ({GGML_MAX_DIMS})')
286
287        # Get Tensor Dimension Array
288        dims = self._get(offs, np.uint64, n_dims[0])
289        offs += int(dims.nbytes)
290
291        # Get Tensor Encoding Scheme Type
292        raw_dtype = self._get(offs, np.uint32)
293        offs += int(raw_dtype.nbytes)
294
295        # Get Tensor Offset
296        offset_tensor = self._get(offs, np.uint64)
297        offs += int(offset_tensor.nbytes)
298
299        return ReaderField(
300            orig_offs,
301            str(bytes(name_data), encoding = 'utf-8'),
302            [name_len, name_data, n_dims, dims, raw_dtype, offset_tensor],
303            [1, 3, 4, 5],
304        )
305
306    def _build_fields(self, offs: int, count: int) -> int:
307        for _ in range(count):
308            orig_offs = offs
309            kv_klen, kv_kdata = self._get_str(offs)
310            offs += int(kv_klen.nbytes + kv_kdata.nbytes)
311            raw_kv_type = self._get(offs, np.uint32)
312            offs += int(raw_kv_type.nbytes)
313            parts: list[npt.NDArray[Any]] = [kv_klen, kv_kdata, raw_kv_type]
314            idxs_offs = len(parts)
315            field_size, field_parts, field_idxs, field_types = self._get_field_parts(offs, raw_kv_type[0])
316            parts += field_parts
317            self._push_field(ReaderField(
318                orig_offs,
319                str(bytes(kv_kdata), encoding = 'utf-8'),
320                parts,
321                [idx + idxs_offs for idx in field_idxs],
322                field_types,
323            ), skip_sum = True)
324            offs += field_size
325        return offs
326
327    def _build_tensor_info(self, offs: int, count: int) -> tuple[int, list[ReaderField]]:
328        tensor_fields = []
329        for _ in range(count):
330            field = self._get_tensor_info_field(offs)
331            offs += sum(int(part.nbytes) for part in field.parts)
332            tensor_fields.append(field)
333        return offs, tensor_fields
334
335    def _build_tensors(self, start_offs: int, fields: list[ReaderField]) -> None:
336        tensors = []
337        tensor_names = set() # keep track of name to prevent duplicated tensors
338        for field in fields:
339            _name_len, name_data, _n_dims, dims, raw_dtype, offset_tensor = field.parts
340            # check if there's any tensor having same name already in the list
341            tensor_name = str(bytes(name_data), encoding = 'utf-8')
342            if tensor_name in tensor_names:
343                raise ValueError(f'Found duplicated tensor with name {tensor_name}')
344            tensor_names.add(tensor_name)
345            ggml_type = GGMLQuantizationType(raw_dtype[0])
346            # use Python ints: np.prod on uint64 wraps silently on overflow
347            n_elems = 1
348            for dim in dims.tolist():
349                n_elems *= int(dim)
350            np_dims = tuple(reversed(dims.tolist()))
351            block_size, type_size = GGML_QUANT_SIZES[ggml_type]
352            n_bytes = n_elems * type_size // block_size
353            data_offs = int(start_offs + offset_tensor[0])
354            item_type: npt.DTypeLike
355            if ggml_type == GGMLQuantizationType.F16:
356                item_count = n_elems
357                item_type = np.float16
358            elif ggml_type == GGMLQuantizationType.F32:
359                item_count = n_elems
360                item_type = np.float32
361            elif ggml_type == GGMLQuantizationType.F64:
362                item_count = n_elems
363                item_type = np.float64
364            elif ggml_type == GGMLQuantizationType.I8:
365                item_count = n_elems
366                item_type = np.int8
367            elif ggml_type == GGMLQuantizationType.I16:
368                item_count = n_elems
369                item_type = np.int16
370            elif ggml_type == GGMLQuantizationType.I32:
371                item_count = n_elems
372                item_type = np.int32
373            elif ggml_type == GGMLQuantizationType.I64:
374                item_count = n_elems
375                item_type = np.int64
376            else:
377                item_count = n_bytes
378                item_type = np.uint8
379                np_dims = quant_shape_to_byte_shape(np_dims, ggml_type)
380            tensors.append(ReaderTensor(
381                name = tensor_name,
382                tensor_type = ggml_type,
383                shape = dims,
384                n_elements = n_elems,
385                n_bytes = n_bytes,
386                data_offset = data_offs,
387                data = self._get(data_offs, item_type, item_count).reshape(np_dims),
388                field = field,
389            ))
390        self.tensors = tensors