check
Browse files- README.md +1 -1
- __pycache__/convert_diffusion_to_gguf.cpython-310.pyc +0 -0
- app.py +101 -0
- convert_diffusion_to_gguf.py +362 -0
- custom_quants.py +1802 -0
- requirements.txt +5 -0
README.md
CHANGED
@@ -1,5 +1,5 @@
|
|
1 |
---
|
2 |
-
title: Diffusers To
|
3 |
emoji: 💻
|
4 |
colorFrom: blue
|
5 |
colorTo: purple
|
|
|
1 |
---
|
2 |
+
title: Diffusers To GGUF
|
3 |
emoji: 💻
|
4 |
colorFrom: blue
|
5 |
colorTo: purple
|
__pycache__/convert_diffusion_to_gguf.cpython-310.pyc
ADDED
Binary file (10.7 kB). View file
|
|
app.py
ADDED
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
from convert_diffusion_to_gguf import SUPPORTED_ARCHS, qconfig_map, convert
|
3 |
+
from huggingface_hub import create_repo, upload_file
|
4 |
+
from argparse import Namespace
|
5 |
+
from pathlib import Path
|
6 |
+
|
7 |
+
|
8 |
+
def upload(args, outfile):
|
9 |
+
url = ""
|
10 |
+
if args.host_repo_id and args.hf_token:
|
11 |
+
repo_id = create_repo(args.host_repo_id, repo_type="model", exist_ok=True, token=args.hf_token).repo_id
|
12 |
+
info = upload_file(repo_id=repo_id, path_in_repo=str(outfile), path_or_fileobj=str(outfile), token=args.token)
|
13 |
+
url = info.commit_url
|
14 |
+
print(f"Uploaded to {url}")
|
15 |
+
|
16 |
+
return url
|
17 |
+
|
18 |
+
|
19 |
+
def go_gguf(model_repo_id, subfolder, arch, outtype, outfile_name, bigendian, verbose, host_repo_id, hf_token):
|
20 |
+
args = Namespace(
|
21 |
+
model=model_repo_id,
|
22 |
+
subfolder=subfolder,
|
23 |
+
arch=arch,
|
24 |
+
outtype=outtype,
|
25 |
+
outfile=Path(outfile_name),
|
26 |
+
bigendian=bigendian,
|
27 |
+
verbose=verbose,
|
28 |
+
host_repo_id=host_repo_id,
|
29 |
+
hf_token=hf_token,
|
30 |
+
)
|
31 |
+
convert(args)
|
32 |
+
upload(args)
|
33 |
+
|
34 |
+
|
35 |
+
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
36 |
+
gr.Markdown("<h1><center>GGUF Converter for Diffusers format model checkpoints</center></h1>")
|
37 |
+
gr.Markdown(
|
38 |
+
"Convert `diffusers` format model checkpoints from the Hub to GGUF format and optionally upload them back. Based on [this repo](https://github.com/ngxson/diffusion-to-gguf)."
|
39 |
+
)
|
40 |
+
|
41 |
+
with gr.Row():
|
42 |
+
with gr.Column(scale=1):
|
43 |
+
gr.Markdown("### 📥 Input Model")
|
44 |
+
model_repo_id = gr.Textbox(label="Model Repo ID", placeholder="e.g., Qwen/Qwen-Image")
|
45 |
+
subfolder = gr.Textbox(label="Subfolder (Optional)", placeholder="e.g., transformer")
|
46 |
+
|
47 |
+
gr.Markdown("### ⚙️ Conversion Settings")
|
48 |
+
arch = gr.Dropdown(choices=SUPPORTED_ARCHS, label="Architecture")
|
49 |
+
outtype = gr.Dropdown(choices=list(qconfig_map.keys()), label="Quantization Type", value="F16")
|
50 |
+
outfile_name = gr.Textbox(label="Output Filename", value="{ftype}.gguf")
|
51 |
+
|
52 |
+
with gr.Accordion("Advanced Settings", open=False):
|
53 |
+
bigendian = gr.Checkbox(label="Use Big Endian")
|
54 |
+
verbose = gr.Checkbox(label="Verbose Logging", value=True)
|
55 |
+
|
56 |
+
gr.Markdown("### 📤 Upload to Hub (Optional)")
|
57 |
+
host_repo_id = gr.Textbox(label="Your Hub Repo ID", placeholder="e.g., YourUsername/My-GGUFs")
|
58 |
+
hf_token = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...")
|
59 |
+
|
60 |
+
convert_btn = gr.Button("Convert & Upload", variant="primary")
|
61 |
+
|
62 |
+
with gr.Column(scale=2):
|
63 |
+
gr.Markdown("### 🚀 Result")
|
64 |
+
url_output = gr.Markdown()
|
65 |
+
|
66 |
+
gr.Examples(
|
67 |
+
examples=[
|
68 |
+
[
|
69 |
+
"black-forest-labs/FLUX.1-schnell",
|
70 |
+
"transformer",
|
71 |
+
"flux",
|
72 |
+
"Q4_0",
|
73 |
+
"flux-schnell-q4.gguf",
|
74 |
+
False,
|
75 |
+
False,
|
76 |
+
"YourUsername/MyGGUFs",
|
77 |
+
"hf_...",
|
78 |
+
],
|
79 |
+
[
|
80 |
+
"Qwen/Qwen-Image",
|
81 |
+
"transformer",
|
82 |
+
"flux",
|
83 |
+
"Q8_0",
|
84 |
+
"qwen-q4.gguf",
|
85 |
+
False,
|
86 |
+
False,
|
87 |
+
"YourUsername/MyGGUFs",
|
88 |
+
"hf_...",
|
89 |
+
],
|
90 |
+
],
|
91 |
+
inputs=[model_repo_id, subfolder, arch, outtype, outfile_name, bigendian, verbose, host_repo_id, hf_token],
|
92 |
+
)
|
93 |
+
|
94 |
+
convert_btn.click(
|
95 |
+
fn=lambda x: x,
|
96 |
+
inputs=[model_repo_id, subfolder, arch, outtype, outfile_name, bigendian, verbose, host_repo_id, hf_token],
|
97 |
+
outputs=[url_output],
|
98 |
+
)
|
99 |
+
|
100 |
+
if __name__ == "__main__":
|
101 |
+
demo.launch()
|
convert_diffusion_to_gguf.py
ADDED
@@ -0,0 +1,362 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python3
|
2 |
+
# -*- coding: utf-8 -*-
|
3 |
+
|
4 |
+
from __future__ import annotations
|
5 |
+
|
6 |
+
import logging
|
7 |
+
import argparse
|
8 |
+
import json
|
9 |
+
import safetensors.torch
|
10 |
+
import os
|
11 |
+
import sys
|
12 |
+
from pathlib import Path
|
13 |
+
from typing import Any, ContextManager, cast
|
14 |
+
from torch import Tensor
|
15 |
+
|
16 |
+
import numpy as np
|
17 |
+
import torch
|
18 |
+
import gguf
|
19 |
+
|
20 |
+
# TODO: add more:
|
21 |
+
SUPPORTED_ARCHS = ["flux", "sd3", "ltxv", "hyvid", "wan", "hidream", "qwen"]
|
22 |
+
|
23 |
+
logger = logging.getLogger(__name__)
|
24 |
+
|
25 |
+
|
26 |
+
class QuantConfig:
|
27 |
+
ftype: gguf.LlamaFileType
|
28 |
+
qtype: gguf.GGMLQuantizationType
|
29 |
+
|
30 |
+
def __init__(self, ftype: gguf.LlamaFileType, qtype: gguf.GGMLQuantizationType):
|
31 |
+
self.ftype = ftype
|
32 |
+
self.qtype = qtype
|
33 |
+
|
34 |
+
|
35 |
+
qconfig_map: dict[str, QuantConfig] = {
|
36 |
+
"F16": QuantConfig(gguf.LlamaFileType.MOSTLY_F16, gguf.GGMLQuantizationType.F16),
|
37 |
+
"BF16": QuantConfig(gguf.LlamaFileType.MOSTLY_BF16, gguf.GGMLQuantizationType.BF16),
|
38 |
+
"Q8_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q8_0, gguf.GGMLQuantizationType.Q8_0),
|
39 |
+
"Q6_K": QuantConfig(gguf.LlamaFileType.MOSTLY_Q6_K, gguf.GGMLQuantizationType.Q6_K),
|
40 |
+
"Q5_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_K_S, gguf.GGMLQuantizationType.Q5_K),
|
41 |
+
"Q5_1": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_1, gguf.GGMLQuantizationType.Q5_1),
|
42 |
+
"Q5_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q5_0, gguf.GGMLQuantizationType.Q5_0),
|
43 |
+
"Q4_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_K_S, gguf.GGMLQuantizationType.Q4_K),
|
44 |
+
"Q4_1": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_1, gguf.GGMLQuantizationType.Q4_1),
|
45 |
+
"Q4_0": QuantConfig(gguf.LlamaFileType.MOSTLY_Q4_0, gguf.GGMLQuantizationType.Q4_0),
|
46 |
+
"Q3_K_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q3_K_S, gguf.GGMLQuantizationType.Q3_K),
|
47 |
+
# "Q2_S": QuantConfig(gguf.LlamaFileType.MOSTLY_Q2_K, gguf.GGMLQuantizationType.Q2_K), # not yet supported in python
|
48 |
+
}
|
49 |
+
|
50 |
+
|
51 |
+
# tree of lazy tensors
|
52 |
+
class LazyTorchTensor(gguf.LazyBase):
|
53 |
+
_tensor_type = torch.Tensor
|
54 |
+
# to keep the type-checker happy
|
55 |
+
dtype: torch.dtype
|
56 |
+
shape: torch.Size
|
57 |
+
|
58 |
+
# only used when converting a torch.Tensor to a np.ndarray
|
59 |
+
_dtype_map: dict[torch.dtype, type] = {
|
60 |
+
torch.float16: np.float16,
|
61 |
+
torch.float32: np.float32,
|
62 |
+
}
|
63 |
+
|
64 |
+
# used for safetensors slices
|
65 |
+
# ref: https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/src/lib.rs#L1046
|
66 |
+
# TODO: uncomment U64, U32, and U16, ref: https://github.com/pytorch/pytorch/issues/58734
|
67 |
+
_dtype_str_map: dict[str, torch.dtype] = {
|
68 |
+
"F64": torch.float64,
|
69 |
+
"F32": torch.float32,
|
70 |
+
"BF16": torch.bfloat16,
|
71 |
+
"F16": torch.float16,
|
72 |
+
# "U64": torch.uint64,
|
73 |
+
"I64": torch.int64,
|
74 |
+
# "U32": torch.uint32,
|
75 |
+
"I32": torch.int32,
|
76 |
+
# "U16": torch.uint16,
|
77 |
+
"I16": torch.int16,
|
78 |
+
"U8": torch.uint8,
|
79 |
+
"I8": torch.int8,
|
80 |
+
"BOOL": torch.bool,
|
81 |
+
"F8_E4M3": torch.float8_e4m3fn,
|
82 |
+
"F8_E5M2": torch.float8_e5m2,
|
83 |
+
}
|
84 |
+
|
85 |
+
def numpy(self) -> gguf.LazyNumpyTensor:
|
86 |
+
dtype = self._dtype_map[self.dtype]
|
87 |
+
return gguf.LazyNumpyTensor(
|
88 |
+
meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(dtype, self.shape),
|
89 |
+
args=(self,),
|
90 |
+
func=(lambda s: s.numpy()),
|
91 |
+
)
|
92 |
+
|
93 |
+
@classmethod
|
94 |
+
def meta_with_dtype_and_shape(cls, dtype: torch.dtype, shape: tuple[int, ...]) -> Tensor:
|
95 |
+
return torch.empty(size=shape, dtype=dtype, device="meta")
|
96 |
+
|
97 |
+
@classmethod
|
98 |
+
def from_safetensors_slice(cls, st_slice: Any) -> Tensor:
|
99 |
+
dtype = cls._dtype_str_map[st_slice.get_dtype()]
|
100 |
+
shape: tuple[int, ...] = tuple(st_slice.get_shape())
|
101 |
+
lazy = cls(meta=cls.meta_with_dtype_and_shape(dtype, shape), args=(st_slice,), func=lambda s: s[:])
|
102 |
+
return cast(torch.Tensor, lazy)
|
103 |
+
|
104 |
+
@classmethod
|
105 |
+
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
106 |
+
del types # unused
|
107 |
+
|
108 |
+
if kwargs is None:
|
109 |
+
kwargs = {}
|
110 |
+
|
111 |
+
if func is torch.Tensor.numpy:
|
112 |
+
return args[0].numpy()
|
113 |
+
|
114 |
+
return cls._wrap_fn(func)(*args, **kwargs)
|
115 |
+
|
116 |
+
|
117 |
+
class Converter:
|
118 |
+
path_safetensors: Path
|
119 |
+
endianess: gguf.GGUFEndian
|
120 |
+
outtype: QuantConfig
|
121 |
+
outfile: Path
|
122 |
+
gguf_writer: gguf.GGUFWriter
|
123 |
+
|
124 |
+
def __init__(
|
125 |
+
self,
|
126 |
+
arch: str,
|
127 |
+
path_safetensors: Path,
|
128 |
+
endianess: gguf.GGUFEndian,
|
129 |
+
outtype: QuantConfig,
|
130 |
+
outfile: Path,
|
131 |
+
subfolder: str = None,
|
132 |
+
repo_id: str = None,
|
133 |
+
is_diffusers: bool = False,
|
134 |
+
):
|
135 |
+
self.path_safetensors = path_safetensors
|
136 |
+
self.endianess = endianess
|
137 |
+
self.outtype = outtype
|
138 |
+
self.outfile = outfile
|
139 |
+
|
140 |
+
self.gguf_writer = gguf.GGUFWriter(path=None, arch=arch, endianess=self.endianess)
|
141 |
+
self.gguf_writer.add_file_type(self.outtype.ftype)
|
142 |
+
self.gguf_writer.add_type("diffusion") # for HF hub to detect the type correctly
|
143 |
+
if repo_id:
|
144 |
+
self.gguf_writer.add_string("repo_id", repo_id)
|
145 |
+
if subfolder:
|
146 |
+
self.gguf_writer.add_string("subfolder", subfolder)
|
147 |
+
if is_diffusers:
|
148 |
+
self.gguf_writer.add_bool("is_diffusers", True)
|
149 |
+
|
150 |
+
# load tensors and process
|
151 |
+
from safetensors import safe_open
|
152 |
+
|
153 |
+
ctx = cast(ContextManager[Any], safe_open(path_safetensors, framework="pt", device="cpu"))
|
154 |
+
with ctx as model_part:
|
155 |
+
for name in model_part.keys():
|
156 |
+
data = model_part.get_slice(name)
|
157 |
+
data = LazyTorchTensor.from_safetensors_slice(data)
|
158 |
+
self.process_tensor(name, data)
|
159 |
+
|
160 |
+
def process_tensor(self, name: str, data_torch: LazyTorchTensor) -> None:
|
161 |
+
is_1d = len(data_torch.shape) == 1
|
162 |
+
current_dtype = data_torch.dtype
|
163 |
+
target_dtype = gguf.GGMLQuantizationType.F32 if is_1d else self.outtype.qtype
|
164 |
+
|
165 |
+
if data_torch.dtype not in (torch.float16, torch.float32):
|
166 |
+
data_torch = data_torch.to(torch.float32)
|
167 |
+
|
168 |
+
data = data_torch.numpy()
|
169 |
+
|
170 |
+
if current_dtype != target_dtype:
|
171 |
+
from custom_quants import quantize as custom_quantize, QuantError
|
172 |
+
|
173 |
+
try:
|
174 |
+
data = custom_quantize(data, target_dtype)
|
175 |
+
except QuantError as e:
|
176 |
+
logger.warning("%s, %s", e, "falling back to F16")
|
177 |
+
target_dtype = gguf.GGMLQuantizationType.F16
|
178 |
+
data = custom_quantize(data, target_dtype)
|
179 |
+
|
180 |
+
# reverse shape to make it similar to the internal ggml dimension order
|
181 |
+
shape = gguf.quant_shape_from_byte_shape(data.shape, target_dtype) if data.dtype == np.uint8 else data.shape
|
182 |
+
shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}"
|
183 |
+
logger.info(f"{f'%-32s' % f'{name},'} {current_dtype} --> {target_dtype.name}, shape = {shape_str}")
|
184 |
+
|
185 |
+
# add tensor to gguf
|
186 |
+
self.gguf_writer.add_tensor(name, data, raw_dtype=target_dtype)
|
187 |
+
|
188 |
+
def write(self) -> None:
|
189 |
+
self.gguf_writer.write_header_to_file(path=self.outfile)
|
190 |
+
self.gguf_writer.write_kv_data_to_file()
|
191 |
+
self.gguf_writer.write_tensors_to_file(progress=True)
|
192 |
+
self.gguf_writer.close()
|
193 |
+
|
194 |
+
|
195 |
+
# https://github.com/bghira/SimpleTuner/blob/cea2457ab063f6dedb9e697830ae68a96be90641/helpers/training/save_hooks.py#L64
|
196 |
+
def _merge_sharded_checkpoints(folder: Path):
|
197 |
+
with open(folder / "diffusion_pytorch_model.safetensors.index.json", "r") as f:
|
198 |
+
ckpt_metadata = json.load(f)
|
199 |
+
weight_map = ckpt_metadata.get("weight_map", None)
|
200 |
+
if weight_map is None:
|
201 |
+
raise KeyError("'weight_map' key not found in the shard index file.")
|
202 |
+
|
203 |
+
# Collect all unique safetensors files from weight_map
|
204 |
+
files_to_load = set(weight_map.values())
|
205 |
+
merged_state_dict = {}
|
206 |
+
|
207 |
+
# Load tensors from each unique file
|
208 |
+
for file_name in files_to_load:
|
209 |
+
part_file_path = folder / file_name
|
210 |
+
if not os.path.exists(part_file_path):
|
211 |
+
raise FileNotFoundError(f"Part file {file_name} not found.")
|
212 |
+
|
213 |
+
with safetensors.safe_open(part_file_path, framework="pt", device="cpu") as f:
|
214 |
+
for tensor_key in f.keys():
|
215 |
+
if tensor_key in weight_map:
|
216 |
+
merged_state_dict[tensor_key] = f.get_tensor(tensor_key)
|
217 |
+
|
218 |
+
return merged_state_dict
|
219 |
+
|
220 |
+
|
221 |
+
def parse_args() -> argparse.Namespace:
|
222 |
+
parser = argparse.ArgumentParser(description="Convert a flux model to GGUF")
|
223 |
+
parser.add_argument(
|
224 |
+
"--outfile",
|
225 |
+
type=Path,
|
226 |
+
default=Path("model-{ftype}.gguf"),
|
227 |
+
help="path to write to; default: 'model-{ftype}.gguf' ; note: {ftype} will be replaced by the outtype",
|
228 |
+
)
|
229 |
+
parser.add_argument(
|
230 |
+
"--outtype",
|
231 |
+
type=str,
|
232 |
+
choices=qconfig_map.keys(),
|
233 |
+
default="F16",
|
234 |
+
help="output quantization scheme",
|
235 |
+
)
|
236 |
+
parser.add_argument(
|
237 |
+
"--arch",
|
238 |
+
type=str,
|
239 |
+
choices=SUPPORTED_ARCHS,
|
240 |
+
help="output model architecture",
|
241 |
+
)
|
242 |
+
parser.add_argument(
|
243 |
+
"--bigendian",
|
244 |
+
action="store_true",
|
245 |
+
help="model is executed on big endian machine",
|
246 |
+
)
|
247 |
+
parser.add_argument(
|
248 |
+
"model",
|
249 |
+
type=Path,
|
250 |
+
help="directory containing safetensors model file",
|
251 |
+
nargs="?",
|
252 |
+
)
|
253 |
+
parser.add_argument("--cache_dir", type=Path, help="Directory to store the intermediate files when needed.")
|
254 |
+
parser.add_argument(
|
255 |
+
"--subfolder", type=Path, default=None, help="Subfolder on the HF Hub to load checkpoints from."
|
256 |
+
)
|
257 |
+
parser.add_argument(
|
258 |
+
"--verbose",
|
259 |
+
action="store_true",
|
260 |
+
help="increase output verbosity",
|
261 |
+
)
|
262 |
+
|
263 |
+
args = parser.parse_args()
|
264 |
+
if args.model is None:
|
265 |
+
parser.error("the following arguments are required: model")
|
266 |
+
if args.arch is None:
|
267 |
+
parser.error("the following arguments are required: --arch")
|
268 |
+
if args.arch not in SUPPORTED_ARCHS:
|
269 |
+
parser.error(f"Unsupported architecture: {args.arch}. Supported architectures: {', '.join(SUPPORTED_ARCHS)}")
|
270 |
+
return args
|
271 |
+
|
272 |
+
|
273 |
+
def convert(args):
|
274 |
+
if args.verbose:
|
275 |
+
logging.basicConfig(level=logging.DEBUG)
|
276 |
+
else:
|
277 |
+
logging.basicConfig(level=logging.INFO)
|
278 |
+
|
279 |
+
if not args.model.is_dir() and not args.model.is_file():
|
280 |
+
if not len(str(args.model).split("/")) == 2:
|
281 |
+
logging.error(f"Model path {args.model} does not exist.")
|
282 |
+
sys.exit(1)
|
283 |
+
|
284 |
+
is_diffusers = False
|
285 |
+
repo_id = None
|
286 |
+
merged_state_dict = None
|
287 |
+
if args.model.is_dir():
|
288 |
+
logging.info("Supplied a directory.")
|
289 |
+
files = list(args.model.glob("*.safetensors"))
|
290 |
+
n = len(files)
|
291 |
+
if n == 0:
|
292 |
+
logging.error("No safetensors files found.")
|
293 |
+
sys.exit(1)
|
294 |
+
if n == 1:
|
295 |
+
logging.info(f"Assinging {files[0]} to `args.model`")
|
296 |
+
args.model = files[0]
|
297 |
+
if n > 1:
|
298 |
+
assert args.model / "diffusion_pytorch_model.safetensors.index.json" in list(args.model.glob("*.*"))
|
299 |
+
assert args.cache_dir
|
300 |
+
merged_state_dict = _merge_sharded_checkpoints(args.model)
|
301 |
+
filepath = args.cache_dir / "merged_state_dict.safetensors"
|
302 |
+
safetensors.torch.save_file(merged_state_dict, filepath)
|
303 |
+
logging.info(f"Serialized merged state dict to {filepath}")
|
304 |
+
args.model = Path(filepath)
|
305 |
+
|
306 |
+
elif len(str(args.model).split("/")) == 2:
|
307 |
+
from huggingface_hub import snapshot_download
|
308 |
+
|
309 |
+
logging.info("Hub repo ID detected.")
|
310 |
+
allow_patterns = f"{args.subfolder}/*.*" if args.subfolder else None
|
311 |
+
local_dir = snapshot_download(repo_id=str(args.model), local_dir=args.cache_dir, allow_patterns=allow_patterns)
|
312 |
+
repo_id = str(args.model)
|
313 |
+
local_dir = Path(local_dir)
|
314 |
+
local_dir = local_dir / args.subfolder if args.subfolder else local_dir
|
315 |
+
merged_state_dict = _merge_sharded_checkpoints(local_dir)
|
316 |
+
filepath = (
|
317 |
+
args.cache_dir / "merged_state_dict.safetensors" if args.cache_dir else "merged_state_dict.safetensors"
|
318 |
+
)
|
319 |
+
safetensors.torch.save_file(merged_state_dict, filepath)
|
320 |
+
logging.info(f"Serialized merged state dict to {filepath}")
|
321 |
+
args.model = Path(filepath)
|
322 |
+
is_diffusers = True
|
323 |
+
|
324 |
+
if merged_state_dict is not None:
|
325 |
+
os.remove(filepath)
|
326 |
+
logging.info(f"Removed the intermediate {filepath}.")
|
327 |
+
|
328 |
+
if args.model.suffix != ".safetensors":
|
329 |
+
logging.error(f"Model path {args.model} is not a safetensors file.")
|
330 |
+
sys.exit(1)
|
331 |
+
|
332 |
+
if args.outfile.suffix != ".gguf":
|
333 |
+
logging.error("Output file must have .gguf extension.")
|
334 |
+
sys.exit(1)
|
335 |
+
|
336 |
+
qconfig = qconfig_map[args.outtype]
|
337 |
+
outfile = Path(str(args.outfile).format(ftype=args.outtype.upper()))
|
338 |
+
|
339 |
+
logger.info(f"Converting model in {args.model} to {outfile} with quantization {args.outtype}")
|
340 |
+
converter = Converter(
|
341 |
+
arch=args.arch,
|
342 |
+
path_safetensors=args.model,
|
343 |
+
endianess=gguf.GGUFEndian.BIG if args.bigendian else gguf.GGUFEndian.LITTLE,
|
344 |
+
outtype=qconfig,
|
345 |
+
outfile=outfile,
|
346 |
+
repo_id=repo_id,
|
347 |
+
subfolder=str(args.subfolder) if args.subfolder else None,
|
348 |
+
is_diffusers=is_diffusers,
|
349 |
+
)
|
350 |
+
converter.write()
|
351 |
+
logger.info(
|
352 |
+
f"Conversion complete. Output written to {outfile}, architecture: {args.arch}, quantization: {qconfig.qtype.name}"
|
353 |
+
)
|
354 |
+
|
355 |
+
|
356 |
+
def main() -> None:
|
357 |
+
args = parse_args()
|
358 |
+
convert(args)
|
359 |
+
|
360 |
+
|
361 |
+
if __name__ == "__main__":
|
362 |
+
main()
|
custom_quants.py
ADDED
@@ -0,0 +1,1802 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from __future__ import annotations
|
2 |
+
from abc import ABC, abstractmethod
|
3 |
+
from typing import Any, Callable, Sequence
|
4 |
+
from math import log2, ceil
|
5 |
+
|
6 |
+
from numpy.typing import DTypeLike
|
7 |
+
|
8 |
+
from gguf.constants import GGML_QUANT_SIZES, GGMLQuantizationType, QK_K
|
9 |
+
from gguf.lazy import LazyNumpyTensor
|
10 |
+
|
11 |
+
import numpy as np
|
12 |
+
|
13 |
+
|
14 |
+
def quant_shape_to_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]:
|
15 |
+
block_size, type_size = GGML_QUANT_SIZES[quant_type]
|
16 |
+
if shape[-1] % block_size != 0:
|
17 |
+
raise ValueError(
|
18 |
+
f"Quantized tensor row size ({shape[-1]}) is not a multiple of {quant_type.name} block size ({block_size})"
|
19 |
+
)
|
20 |
+
return (*shape[:-1], shape[-1] // block_size * type_size)
|
21 |
+
|
22 |
+
|
23 |
+
def quant_shape_from_byte_shape(shape: Sequence[int], quant_type: GGMLQuantizationType) -> tuple[int, ...]:
|
24 |
+
block_size, type_size = GGML_QUANT_SIZES[quant_type]
|
25 |
+
if shape[-1] % type_size != 0:
|
26 |
+
raise ValueError(
|
27 |
+
f"Quantized tensor bytes per row ({shape[-1]}) is not a multiple of {quant_type.name} type size ({type_size})"
|
28 |
+
)
|
29 |
+
return (*shape[:-1], shape[-1] // type_size * block_size)
|
30 |
+
|
31 |
+
|
32 |
+
# This is faster than np.vectorize and np.apply_along_axis because it works on more than one row at a time
|
33 |
+
def _apply_over_grouped_rows(
|
34 |
+
func: Callable[[np.ndarray], np.ndarray], arr: np.ndarray, otype: DTypeLike, oshape: tuple[int, ...]
|
35 |
+
) -> np.ndarray:
|
36 |
+
rows = arr.reshape((-1, arr.shape[-1]))
|
37 |
+
osize = 1
|
38 |
+
for dim in oshape:
|
39 |
+
osize *= dim
|
40 |
+
out = np.empty(shape=osize, dtype=otype)
|
41 |
+
# compute over groups of 16 rows (arbitrary, but seems good for performance)
|
42 |
+
n_groups = (rows.shape[0] // 16) or 1
|
43 |
+
np.concatenate([func(group).ravel() for group in np.array_split(rows, n_groups)], axis=0, out=out)
|
44 |
+
return out.reshape(oshape)
|
45 |
+
|
46 |
+
|
47 |
+
# round away from zero
|
48 |
+
# ref: https://stackoverflow.com/a/59143326/22827863
|
49 |
+
def np_roundf(n: np.ndarray) -> np.ndarray:
|
50 |
+
a = abs(n)
|
51 |
+
floored = np.floor(a)
|
52 |
+
b = floored + np.floor(2 * (a - floored))
|
53 |
+
return np.sign(n) * b
|
54 |
+
|
55 |
+
|
56 |
+
class QuantError(Exception): ...
|
57 |
+
|
58 |
+
|
59 |
+
_type_traits: dict[GGMLQuantizationType, type[__Quant]] = {}
|
60 |
+
|
61 |
+
|
62 |
+
def quantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
|
63 |
+
if qtype == GGMLQuantizationType.F32:
|
64 |
+
return data.astype(np.float32, copy=False)
|
65 |
+
elif qtype == GGMLQuantizationType.F16:
|
66 |
+
return data.astype(np.float16, copy=False)
|
67 |
+
elif (q := _type_traits.get(qtype)) is not None:
|
68 |
+
return q.quantize(data)
|
69 |
+
else:
|
70 |
+
raise NotImplementedError(f"Quantization for {qtype.name} is not yet implemented")
|
71 |
+
|
72 |
+
|
73 |
+
def dequantize(data: np.ndarray, qtype: GGMLQuantizationType) -> np.ndarray:
|
74 |
+
if qtype == GGMLQuantizationType.F32:
|
75 |
+
return data.view(np.float32)
|
76 |
+
elif qtype == GGMLQuantizationType.F16:
|
77 |
+
return data.view(np.float16).astype(np.float32)
|
78 |
+
elif (q := _type_traits.get(qtype)) is not None:
|
79 |
+
return q.dequantize(data)
|
80 |
+
else:
|
81 |
+
raise NotImplementedError(f"Dequantization for {qtype.name} is not yet implemented")
|
82 |
+
|
83 |
+
|
84 |
+
class __Quant(ABC):
|
85 |
+
qtype: GGMLQuantizationType
|
86 |
+
block_size: int
|
87 |
+
type_size: int
|
88 |
+
|
89 |
+
grid: np.ndarray[Any, np.dtype[np.float32]] | None = None
|
90 |
+
grid_shape: tuple[int, int] = (0, 0)
|
91 |
+
grid_map: tuple[int | float, ...] = ()
|
92 |
+
grid_hex: bytes | None = None
|
93 |
+
|
94 |
+
def __init__(self):
|
95 |
+
return TypeError("Quant conversion classes can't have instances")
|
96 |
+
|
97 |
+
def __init_subclass__(cls, qtype: GGMLQuantizationType) -> None:
|
98 |
+
cls.qtype = qtype
|
99 |
+
cls.block_size, cls.type_size = GGML_QUANT_SIZES[qtype]
|
100 |
+
cls.__quantize_lazy = LazyNumpyTensor._wrap_fn(
|
101 |
+
cls.__quantize_array, meta_noop=(np.uint8, cls.__shape_to_bytes)
|
102 |
+
)
|
103 |
+
cls.__dequantize_lazy = LazyNumpyTensor._wrap_fn(
|
104 |
+
cls.__dequantize_array, meta_noop=(np.float32, cls.__shape_from_bytes)
|
105 |
+
)
|
106 |
+
assert qtype not in _type_traits
|
107 |
+
_type_traits[qtype] = cls
|
108 |
+
|
109 |
+
@classmethod
|
110 |
+
def init_grid(cls):
|
111 |
+
if cls.grid is not None or cls.grid_hex is None:
|
112 |
+
return
|
113 |
+
|
114 |
+
bits_per_elem = ceil(log2(len(cls.grid_map)))
|
115 |
+
assert bits_per_elem != 0, cls.qtype.name
|
116 |
+
elems_per_byte = 8 // bits_per_elem
|
117 |
+
|
118 |
+
grid = np.frombuffer(cls.grid_hex, dtype=np.uint8)
|
119 |
+
# decode hexadecimal chars from grid
|
120 |
+
grid = grid.reshape((-1, 2))
|
121 |
+
grid = (np.where(grid > 0x40, grid + 9, grid) & 0x0F) << np.array([4, 0], dtype=np.uint8).reshape((1, 2))
|
122 |
+
grid = grid[..., 0] | grid[..., 1]
|
123 |
+
# unpack the grid values
|
124 |
+
grid = grid.reshape((-1, 1)) >> np.array(
|
125 |
+
[i for i in range(0, 8, 8 // elems_per_byte)], dtype=np.uint8
|
126 |
+
).reshape((1, elems_per_byte))
|
127 |
+
grid = (grid & ((1 << bits_per_elem) - 1)).reshape((-1, 1))
|
128 |
+
grid_map = np.array(cls.grid_map, dtype=np.float32).reshape((1, -1))
|
129 |
+
grid = np.take_along_axis(grid_map, grid, axis=-1)
|
130 |
+
cls.grid = grid.reshape((1, 1, *cls.grid_shape))
|
131 |
+
|
132 |
+
@classmethod
|
133 |
+
@abstractmethod
|
134 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
135 |
+
raise NotImplementedError
|
136 |
+
|
137 |
+
@classmethod
|
138 |
+
@abstractmethod
|
139 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
140 |
+
raise NotImplementedError
|
141 |
+
|
142 |
+
@classmethod
|
143 |
+
def quantize_rows(cls, rows: np.ndarray) -> np.ndarray:
|
144 |
+
rows = rows.astype(np.float32, copy=False)
|
145 |
+
shape = rows.shape
|
146 |
+
n_blocks = rows.size // cls.block_size
|
147 |
+
blocks = rows.reshape((n_blocks, cls.block_size))
|
148 |
+
blocks = cls.quantize_blocks(blocks)
|
149 |
+
assert blocks.dtype == np.uint8
|
150 |
+
assert blocks.shape[-1] == cls.type_size
|
151 |
+
return blocks.reshape(cls.__shape_to_bytes(shape))
|
152 |
+
|
153 |
+
@classmethod
|
154 |
+
def dequantize_rows(cls, rows: np.ndarray) -> np.ndarray:
|
155 |
+
rows = rows.view(np.uint8)
|
156 |
+
shape = rows.shape
|
157 |
+
n_blocks = rows.size // cls.type_size
|
158 |
+
blocks = rows.reshape((n_blocks, cls.type_size))
|
159 |
+
blocks = cls.dequantize_blocks(blocks)
|
160 |
+
assert blocks.dtype == np.float32
|
161 |
+
assert blocks.shape[-1] == cls.block_size
|
162 |
+
return blocks.reshape(cls.__shape_from_bytes(shape))
|
163 |
+
|
164 |
+
@classmethod
|
165 |
+
def __shape_to_bytes(cls, shape: Sequence[int]):
|
166 |
+
return quant_shape_to_byte_shape(shape, cls.qtype)
|
167 |
+
|
168 |
+
@classmethod
|
169 |
+
def __shape_from_bytes(cls, shape: Sequence[int]):
|
170 |
+
return quant_shape_from_byte_shape(shape, cls.qtype)
|
171 |
+
|
172 |
+
@classmethod
|
173 |
+
def __quantize_array(cls, array: np.ndarray) -> np.ndarray:
|
174 |
+
return _apply_over_grouped_rows(
|
175 |
+
cls.quantize_rows, arr=array, otype=np.uint8, oshape=cls.__shape_to_bytes(array.shape)
|
176 |
+
)
|
177 |
+
|
178 |
+
@classmethod
|
179 |
+
def __dequantize_array(cls, array: np.ndarray) -> np.ndarray:
|
180 |
+
cls.init_grid()
|
181 |
+
return _apply_over_grouped_rows(
|
182 |
+
cls.dequantize_rows, arr=array, otype=np.float32, oshape=cls.__shape_from_bytes(array.shape)
|
183 |
+
)
|
184 |
+
|
185 |
+
@classmethod
|
186 |
+
def __quantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any:
|
187 |
+
pass
|
188 |
+
|
189 |
+
@classmethod
|
190 |
+
def __dequantize_lazy(cls, lazy_tensor: LazyNumpyTensor, /) -> Any:
|
191 |
+
pass
|
192 |
+
|
193 |
+
@classmethod
|
194 |
+
def can_quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> bool:
|
195 |
+
return tensor.shape[-1] % cls.block_size == 0
|
196 |
+
|
197 |
+
@classmethod
|
198 |
+
def quantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray:
|
199 |
+
if not cls.can_quantize(tensor):
|
200 |
+
raise QuantError(f"Can't quantize tensor with shape {tensor.shape} to {cls.qtype.name}")
|
201 |
+
if isinstance(tensor, LazyNumpyTensor):
|
202 |
+
return cls.__quantize_lazy(tensor)
|
203 |
+
else:
|
204 |
+
return cls.__quantize_array(tensor)
|
205 |
+
|
206 |
+
@classmethod
|
207 |
+
def dequantize(cls, tensor: np.ndarray | LazyNumpyTensor) -> np.ndarray:
|
208 |
+
if isinstance(tensor, LazyNumpyTensor):
|
209 |
+
return cls.__dequantize_lazy(tensor)
|
210 |
+
else:
|
211 |
+
return cls.__dequantize_array(tensor)
|
212 |
+
|
213 |
+
|
214 |
+
class BF16(__Quant, qtype=GGMLQuantizationType.BF16):
|
215 |
+
@classmethod
|
216 |
+
# same as ggml_compute_fp32_to_bf16 in ggml-impl.h
|
217 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
218 |
+
n = blocks.view(np.uint32)
|
219 |
+
# force nan to quiet
|
220 |
+
n = np.where((n & 0x7FFFFFFF) > 0x7F800000, (n & np.uint32(0xFFFF0000)) | np.uint32(64 << 16), n)
|
221 |
+
# round to nearest even
|
222 |
+
n = (np.uint64(n) + (0x7FFF + ((n >> 16) & 1))) >> 16
|
223 |
+
return n.astype(np.uint16).view(np.uint8)
|
224 |
+
|
225 |
+
@classmethod
|
226 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
227 |
+
return (blocks.view(np.int16).astype(np.int32) << 16).view(np.float32)
|
228 |
+
|
229 |
+
|
230 |
+
class Q4_0(__Quant, qtype=GGMLQuantizationType.Q4_0):
|
231 |
+
@classmethod
|
232 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
233 |
+
n_blocks = blocks.shape[0]
|
234 |
+
|
235 |
+
imax = abs(blocks).argmax(axis=-1, keepdims=True)
|
236 |
+
max = np.take_along_axis(blocks, imax, axis=-1)
|
237 |
+
|
238 |
+
d = max / -8
|
239 |
+
with np.errstate(divide="ignore"):
|
240 |
+
id = np.where(d == 0, 0, 1 / d)
|
241 |
+
# FIXME: Q4_0's reference rounding is cursed and depends on FMA
|
242 |
+
qs = (
|
243 |
+
np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(8.5), dtype=np.float32)
|
244 |
+
.astype(np.uint8)
|
245 |
+
.clip(0, 15)
|
246 |
+
)
|
247 |
+
|
248 |
+
qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
|
249 |
+
qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
|
250 |
+
|
251 |
+
d = d.astype(np.float16).view(np.uint8)
|
252 |
+
|
253 |
+
return np.concatenate([d, qs], axis=-1)
|
254 |
+
|
255 |
+
@classmethod
|
256 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
257 |
+
n_blocks = blocks.shape[0]
|
258 |
+
|
259 |
+
d, qs = np.hsplit(blocks, [2])
|
260 |
+
|
261 |
+
d = d.view(np.float16).astype(np.float32)
|
262 |
+
|
263 |
+
qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape(
|
264 |
+
(1, 1, 2, 1)
|
265 |
+
)
|
266 |
+
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.int8) - np.int8(8)
|
267 |
+
|
268 |
+
return d * qs.astype(np.float32)
|
269 |
+
|
270 |
+
|
271 |
+
class Q4_1(__Quant, qtype=GGMLQuantizationType.Q4_1):
|
272 |
+
@classmethod
|
273 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
274 |
+
n_blocks = blocks.shape[0]
|
275 |
+
|
276 |
+
max = blocks.max(axis=-1, keepdims=True)
|
277 |
+
min = blocks.min(axis=-1, keepdims=True)
|
278 |
+
|
279 |
+
d = (max - min) / 15
|
280 |
+
with np.errstate(divide="ignore"):
|
281 |
+
id = np.where(d == 0, 0, 1 / d)
|
282 |
+
qs = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 15)
|
283 |
+
|
284 |
+
qs = qs.reshape((n_blocks, 2, cls.block_size // 2))
|
285 |
+
qs = qs[..., 0, :] | (qs[..., 1, :] << np.uint8(4))
|
286 |
+
|
287 |
+
d = d.astype(np.float16).view(np.uint8)
|
288 |
+
m = min.astype(np.float16).view(np.uint8)
|
289 |
+
|
290 |
+
return np.concatenate([d, m, qs], axis=-1)
|
291 |
+
|
292 |
+
@classmethod
|
293 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
294 |
+
n_blocks = blocks.shape[0]
|
295 |
+
|
296 |
+
d, rest = np.hsplit(blocks, [2])
|
297 |
+
m, qs = np.hsplit(rest, [2])
|
298 |
+
|
299 |
+
d = d.view(np.float16).astype(np.float32)
|
300 |
+
m = m.view(np.float16).astype(np.float32)
|
301 |
+
|
302 |
+
qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape(
|
303 |
+
(1, 1, 2, 1)
|
304 |
+
)
|
305 |
+
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1)).astype(np.float32)
|
306 |
+
|
307 |
+
return (d * qs) + m
|
308 |
+
|
309 |
+
|
310 |
+
class Q5_0(__Quant, qtype=GGMLQuantizationType.Q5_0):
|
311 |
+
@classmethod
|
312 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
313 |
+
n_blocks = blocks.shape[0]
|
314 |
+
|
315 |
+
imax = abs(blocks).argmax(axis=-1, keepdims=True)
|
316 |
+
max = np.take_along_axis(blocks, imax, axis=-1)
|
317 |
+
|
318 |
+
d = max / -16
|
319 |
+
with np.errstate(divide="ignore"):
|
320 |
+
id = np.where(d == 0, 0, 1 / d)
|
321 |
+
# FIXME: Q5_0's reference rounding is cursed and depends on FMA
|
322 |
+
q = (
|
323 |
+
np.trunc((np.float64(blocks) * np.float64(id)) + np.float64(16.5), dtype=np.float32)
|
324 |
+
.astype(np.uint8)
|
325 |
+
.clip(0, 31)
|
326 |
+
)
|
327 |
+
|
328 |
+
qs = q.reshape((n_blocks, 2, cls.block_size // 2))
|
329 |
+
qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
|
330 |
+
|
331 |
+
qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
|
332 |
+
|
333 |
+
d = d.astype(np.float16).view(np.uint8)
|
334 |
+
|
335 |
+
return np.concatenate([d, qh, qs], axis=-1)
|
336 |
+
|
337 |
+
@classmethod
|
338 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
339 |
+
n_blocks = blocks.shape[0]
|
340 |
+
|
341 |
+
d, rest = np.hsplit(blocks, [2])
|
342 |
+
qh, qs = np.hsplit(rest, [4])
|
343 |
+
|
344 |
+
d = d.view(np.float16).astype(np.float32)
|
345 |
+
qh = qh.view(np.uint32)
|
346 |
+
|
347 |
+
qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
|
348 |
+
ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape(
|
349 |
+
(1, 1, 2, 1)
|
350 |
+
)
|
351 |
+
qh = (qh & np.uint32(0x01)).astype(np.uint8)
|
352 |
+
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
|
353 |
+
|
354 |
+
qs = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(16)
|
355 |
+
|
356 |
+
return d * qs.astype(np.float32)
|
357 |
+
|
358 |
+
|
359 |
+
class Q5_1(__Quant, qtype=GGMLQuantizationType.Q5_1):
|
360 |
+
@classmethod
|
361 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
362 |
+
n_blocks = blocks.shape[0]
|
363 |
+
|
364 |
+
max = blocks.max(axis=-1, keepdims=True)
|
365 |
+
min = blocks.min(axis=-1, keepdims=True)
|
366 |
+
|
367 |
+
d = (max - min) / 31
|
368 |
+
with np.errstate(divide="ignore"):
|
369 |
+
id = np.where(d == 0, 0, 1 / d)
|
370 |
+
q = np.trunc((blocks - min) * id + np.float32(0.5), dtype=np.float32).astype(np.uint8).clip(0, 31)
|
371 |
+
|
372 |
+
qs = q.reshape((n_blocks, 2, cls.block_size // 2))
|
373 |
+
qs = (qs[..., 0, :] & np.uint8(0x0F)) | (qs[..., 1, :] << np.uint8(4))
|
374 |
+
|
375 |
+
qh = np.packbits(q.reshape((n_blocks, 1, 32)) >> np.uint8(4), axis=-1, bitorder="little").reshape(n_blocks, 4)
|
376 |
+
|
377 |
+
d = d.astype(np.float16).view(np.uint8)
|
378 |
+
m = min.astype(np.float16).view(np.uint8)
|
379 |
+
|
380 |
+
return np.concatenate([d, m, qh, qs], axis=-1)
|
381 |
+
|
382 |
+
@classmethod
|
383 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
384 |
+
n_blocks = blocks.shape[0]
|
385 |
+
|
386 |
+
d, rest = np.hsplit(blocks, [2])
|
387 |
+
m, rest = np.hsplit(rest, [2])
|
388 |
+
qh, qs = np.hsplit(rest, [4])
|
389 |
+
|
390 |
+
d = d.view(np.float16).astype(np.float32)
|
391 |
+
m = m.view(np.float16).astype(np.float32)
|
392 |
+
qh = qh.view(np.uint32)
|
393 |
+
|
394 |
+
qh = qh.reshape((n_blocks, 1)) >> np.array([i for i in range(32)], dtype=np.uint32).reshape((1, 32))
|
395 |
+
ql = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape(
|
396 |
+
(1, 1, 2, 1)
|
397 |
+
)
|
398 |
+
qh = (qh & np.uint32(0x01)).astype(np.uint8)
|
399 |
+
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1))
|
400 |
+
|
401 |
+
qs = (ql | (qh << np.uint8(4))).astype(np.float32)
|
402 |
+
|
403 |
+
return (d * qs) + m
|
404 |
+
|
405 |
+
|
406 |
+
class Q8_0(__Quant, qtype=GGMLQuantizationType.Q8_0):
|
407 |
+
@classmethod
|
408 |
+
# Implementation of Q8_0 with bit-exact same results as reference implementation in ggml-quants.c
|
409 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
410 |
+
d = abs(blocks).max(axis=1, keepdims=True) / 127
|
411 |
+
with np.errstate(divide="ignore"):
|
412 |
+
id = np.where(d == 0, 0, 1 / d)
|
413 |
+
qs = np_roundf(blocks * id)
|
414 |
+
|
415 |
+
# (n_blocks, 2)
|
416 |
+
d = d.astype(np.float16).view(np.uint8)
|
417 |
+
# (n_blocks, block_size)
|
418 |
+
qs = qs.astype(np.int8).view(np.uint8)
|
419 |
+
|
420 |
+
return np.concatenate([d, qs], axis=1)
|
421 |
+
|
422 |
+
@classmethod
|
423 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
424 |
+
d, x = np.split(blocks, [2], axis=1)
|
425 |
+
d = d.view(np.float16).astype(np.float32)
|
426 |
+
x = x.view(np.int8).astype(np.float32)
|
427 |
+
|
428 |
+
return x * d
|
429 |
+
|
430 |
+
|
431 |
+
class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K):
|
432 |
+
@classmethod
|
433 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
434 |
+
n_blocks = blocks.shape[0]
|
435 |
+
|
436 |
+
scales, rest = np.hsplit(blocks, [QK_K // 16])
|
437 |
+
qs, rest = np.hsplit(rest, [QK_K // 4])
|
438 |
+
d, dmin = np.hsplit(rest, [2])
|
439 |
+
|
440 |
+
d = d.view(np.float16).astype(np.float32)
|
441 |
+
dmin = dmin.view(np.float16).astype(np.float32)
|
442 |
+
|
443 |
+
# (n_blocks, 16, 1)
|
444 |
+
dl = (d * (scales & 0xF).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
|
445 |
+
ml = (dmin * (scales >> 4).astype(np.float32)).reshape((n_blocks, QK_K // 16, 1))
|
446 |
+
|
447 |
+
shift = np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
448 |
+
|
449 |
+
qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & np.uint8(3)
|
450 |
+
|
451 |
+
qs = qs.reshape((n_blocks, QK_K // 16, 16)).astype(np.float32)
|
452 |
+
|
453 |
+
qs = dl * qs - ml
|
454 |
+
|
455 |
+
return qs.reshape((n_blocks, -1))
|
456 |
+
|
457 |
+
|
458 |
+
class Q3_K(__Quant, qtype=GGMLQuantizationType.Q3_K):
|
459 |
+
@classmethod
|
460 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
461 |
+
"""
|
462 |
+
Quantizes a numpy array of floats into Q3_K format.
|
463 |
+
Vectorized implementation of the C++ reference code.
|
464 |
+
"""
|
465 |
+
n_blocks = blocks.shape[0]
|
466 |
+
sub_blocks = blocks.reshape((n_blocks, 16, 16))
|
467 |
+
|
468 |
+
# --- Vectorized make_qx_quants logic for per-sub-block scales ---
|
469 |
+
nmax_data = 4 # Quantization range for data: [-4, 3]
|
470 |
+
|
471 |
+
flat_sub_blocks = sub_blocks.reshape(-1, 16)
|
472 |
+
weights_data = flat_sub_blocks * flat_sub_blocks # rmse_type=1 uses w=x*x
|
473 |
+
|
474 |
+
# Find max absolute values for each sub-block
|
475 |
+
abs_sub_blocks = np.abs(flat_sub_blocks)
|
476 |
+
max_indices = np.argmax(abs_sub_blocks, axis=-1, keepdims=True)
|
477 |
+
max_vals = np.take_along_axis(flat_sub_blocks, max_indices, axis=-1)
|
478 |
+
|
479 |
+
# Iteratively find the best scale for each sub-block
|
480 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
481 |
+
initial_iscale = np.where(max_vals == 0, 0, -nmax_data / max_vals)
|
482 |
+
|
483 |
+
# Initial calculation (is=0)
|
484 |
+
l = np_roundf(flat_sub_blocks * initial_iscale).clip(-nmax_data, nmax_data - 1)
|
485 |
+
sumlx = np.sum(weights_data * flat_sub_blocks * l, axis=-1)
|
486 |
+
suml2 = np.sum(weights_data * l * l, axis=-1)
|
487 |
+
|
488 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
489 |
+
current_scales = np.divide(sumlx, suml2, out=np.zeros_like(sumlx), where=suml2 != 0)
|
490 |
+
|
491 |
+
best_scores = current_scales * sumlx
|
492 |
+
best_scales = current_scales.copy()
|
493 |
+
|
494 |
+
# Iterative search over potential iscale adjustments
|
495 |
+
for is_ in range(-9, 10):
|
496 |
+
if is_ == 0:
|
497 |
+
continue
|
498 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
499 |
+
iscale_try = -(nmax_data + 0.1 * is_) / max_vals
|
500 |
+
iscale_try[max_vals == 0] = 0
|
501 |
+
|
502 |
+
l_try = np_roundf(flat_sub_blocks * iscale_try).clip(-nmax_data, nmax_data - 1)
|
503 |
+
sumlx_try = np.sum(weights_data * flat_sub_blocks * l_try, axis=-1)
|
504 |
+
suml2_try = np.sum(weights_data * l_try * l_try, axis=-1)
|
505 |
+
|
506 |
+
improvement_mask = (suml2_try > 0) & (sumlx_try * sumlx_try * suml2 > best_scores * suml2_try)
|
507 |
+
if np.any(improvement_mask):
|
508 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
509 |
+
scales_try = np.divide(sumlx_try, suml2_try, out=np.zeros_like(sumlx_try), where=suml2_try != 0)
|
510 |
+
best_scores[improvement_mask] = (scales_try * sumlx_try)[improvement_mask]
|
511 |
+
best_scales[improvement_mask] = scales_try[improvement_mask]
|
512 |
+
# Update suml2 for the next comparison in the loop
|
513 |
+
suml2[improvement_mask] = suml2_try[improvement_mask]
|
514 |
+
|
515 |
+
scales = best_scales.reshape(n_blocks, 16)
|
516 |
+
|
517 |
+
# --- Vectorized logic to quantize the scales themselves ---
|
518 |
+
nmax_scales = 32 # Quantization range for scales: [-32, 31]
|
519 |
+
abs_scales = np.abs(scales)
|
520 |
+
max_scale_indices = np.argmax(abs_scales, axis=-1, keepdims=True)
|
521 |
+
max_scale_vals = np.take_along_axis(scales, max_scale_indices, axis=-1)
|
522 |
+
|
523 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
524 |
+
iscale_s = np.where(max_scale_vals == 0, 0, -nmax_scales / max_scale_vals)
|
525 |
+
|
526 |
+
l_s = np_roundf(scales * iscale_s).clip(-nmax_scales, nmax_scales - 1)
|
527 |
+
d_val = np.divide(
|
528 |
+
np.sum(scales * l_s, axis=-1, keepdims=True),
|
529 |
+
np.sum(l_s * l_s, axis=-1, keepdims=True),
|
530 |
+
out=np.zeros((n_blocks, 1)),
|
531 |
+
where=np.sum(l_s * l_s, axis=-1, keepdims=True) != 0,
|
532 |
+
)
|
533 |
+
|
534 |
+
# Pack the 6-bit quantized scales into 12 bytes
|
535 |
+
l = (l_s + 32).astype(np.uint8)
|
536 |
+
scales_packed = np.zeros((n_blocks, 12), dtype=np.uint8)
|
537 |
+
l_low = l & 0x0F
|
538 |
+
l_high = (l >> 4) & 0x03
|
539 |
+
scales_packed[:, 0:8] = l_low[:, 0:8] | (l_low[:, 8:16] << 4)
|
540 |
+
l_high_reshaped = l_high.reshape(n_blocks, 4, 4).transpose(0, 2, 1)
|
541 |
+
packed_high_bits = (
|
542 |
+
l_high_reshaped[:, :, 0]
|
543 |
+
| (l_high_reshaped[:, :, 1] << 2)
|
544 |
+
| (l_high_reshaped[:, :, 2] << 4)
|
545 |
+
| (l_high_reshaped[:, :, 3] << 6)
|
546 |
+
)
|
547 |
+
scales_packed[:, 8:12] = packed_high_bits
|
548 |
+
d = d_val.astype(np.float16).view(np.uint8)
|
549 |
+
|
550 |
+
# --- Re-quantize data with final scales and pack ---
|
551 |
+
sc_dequant = (l.astype(np.int8) - 32).astype(np.float32)
|
552 |
+
d_eff = (d_val * sc_dequant).reshape(n_blocks, 16, 1)
|
553 |
+
|
554 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
555 |
+
l_data_float = np.divide(sub_blocks, d_eff, out=np.zeros_like(sub_blocks), where=d_eff != 0)
|
556 |
+
|
557 |
+
l_data = (np.clip(np_roundf(l_data_float), -4, 3) + 4).astype(np.uint8)
|
558 |
+
l_data = l_data.reshape(n_blocks, 256)
|
559 |
+
|
560 |
+
# hmask stores the 3rd bit
|
561 |
+
hmask_values = (l_data > 3).reshape(n_blocks, 8, 32).transpose(0, 2, 1)
|
562 |
+
hmask = np.packbits(hmask_values, axis=-1, bitorder="little").reshape(n_blocks, -1)
|
563 |
+
|
564 |
+
# qs stores the lower 2 bits
|
565 |
+
l_data[l_data > 3] -= 4
|
566 |
+
l_data_low = (l_data & 0x03).reshape(n_blocks, 2, 4, 32)
|
567 |
+
qs_parts = (
|
568 |
+
l_data_low[:, :, 0, :]
|
569 |
+
| (l_data_low[:, :, 1, :] << 2)
|
570 |
+
| (l_data_low[:, :, 2, :] << 4)
|
571 |
+
| (l_data_low[:, :, 3, :] << 6)
|
572 |
+
)
|
573 |
+
qs = qs_parts.reshape(n_blocks, 64)
|
574 |
+
|
575 |
+
return np.concatenate([hmask, qs, scales_packed, d], axis=1)
|
576 |
+
|
577 |
+
@classmethod
|
578 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
579 |
+
n_blocks = blocks.shape[0]
|
580 |
+
|
581 |
+
hmask, rest = np.hsplit(blocks, [QK_K // 8])
|
582 |
+
qs, rest = np.hsplit(rest, [QK_K // 4])
|
583 |
+
scales, d = np.hsplit(rest, [12])
|
584 |
+
|
585 |
+
d = d.view(np.float16).astype(np.float32)
|
586 |
+
|
587 |
+
# The scales are packed at 6-bit each in this pattern:
|
588 |
+
# 0: IIIIAAAA
|
589 |
+
# 1: JJJJBBBB
|
590 |
+
# 2: KKKKCCCC
|
591 |
+
# 3: LLLLDDDD
|
592 |
+
# 4: MMMMEEEE
|
593 |
+
# 5: NNNNFFFF
|
594 |
+
# 6: OOOOGGGG
|
595 |
+
# 7: PPPPHHHH
|
596 |
+
# 8: MMIIEEAA
|
597 |
+
# 9: NNJJFFBB
|
598 |
+
# 10: OOKKGGCC
|
599 |
+
# 11: PPLLHHDD
|
600 |
+
lscales, hscales = np.hsplit(scales, [8])
|
601 |
+
lscales = lscales.reshape((n_blocks, 1, 8)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 2, 1))
|
602 |
+
lscales = lscales.reshape((n_blocks, 16))
|
603 |
+
hscales = hscales.reshape((n_blocks, 1, 4)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 4, 1))
|
604 |
+
hscales = hscales.reshape((n_blocks, 16))
|
605 |
+
scales = (lscales & np.uint8(0x0F)) | ((hscales & np.uint8(0x03)) << np.uint8(4))
|
606 |
+
scales = (scales.astype(np.int8) - np.int8(32)).astype(np.float32)
|
607 |
+
|
608 |
+
dl = (d * scales).reshape((n_blocks, 16, 1))
|
609 |
+
|
610 |
+
ql = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
611 |
+
qh = hmask.reshape(n_blocks, -1, 1, 32) >> np.array([i for i in range(8)], dtype=np.uint8).reshape(
|
612 |
+
(1, 1, 8, 1)
|
613 |
+
)
|
614 |
+
ql = ql.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(3)
|
615 |
+
qh = qh.reshape((n_blocks, 16, QK_K // 16)) & np.uint8(1)
|
616 |
+
qh = qh ^ np.uint8(1) # strangely, the offset is zero when the bitmask is 1
|
617 |
+
q = (ql.astype(np.int8) - (qh << np.uint8(2)).astype(np.int8)).astype(np.float32)
|
618 |
+
|
619 |
+
return (dl * q).reshape((n_blocks, QK_K))
|
620 |
+
|
621 |
+
|
622 |
+
class Q4_K(__Quant, qtype=GGMLQuantizationType.Q4_K):
|
623 |
+
K_SCALE_SIZE = 12
|
624 |
+
QK_K = QK_K # Block size
|
625 |
+
|
626 |
+
@classmethod
|
627 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
628 |
+
"""
|
629 |
+
Quantizes a numpy array of floats into Q4_K format.
|
630 |
+
Vectorized implementation inspired by the C++ reference code.
|
631 |
+
"""
|
632 |
+
if blocks.shape[-1] % cls.QK_K != 0:
|
633 |
+
raise ValueError(
|
634 |
+
f"The last dimension of the input array must be a multiple of {cls.QK_K}, but got {blocks.shape[-1]}"
|
635 |
+
)
|
636 |
+
|
637 |
+
n_blocks = blocks.size // cls.QK_K
|
638 |
+
sub_blocks = blocks.reshape((n_blocks, 8, 32))
|
639 |
+
|
640 |
+
# --- Vectorized make_qkx2_quants logic ---
|
641 |
+
nmax = 15
|
642 |
+
rmin = -1.0
|
643 |
+
rdelta = 0.1
|
644 |
+
nstep = 20
|
645 |
+
|
646 |
+
# Calculate weights for all sub-blocks
|
647 |
+
sum_x2 = np.sum(sub_blocks * sub_blocks, axis=-1, keepdims=True)
|
648 |
+
# Use np.maximum to avoid sqrt of negative number due to float precision
|
649 |
+
av_x = np.sqrt(np.maximum(0, sum_x2 / 32.0))
|
650 |
+
weights = av_x + np.abs(sub_blocks)
|
651 |
+
sum_w = np.sum(weights, axis=-1, keepdims=True)
|
652 |
+
sum_x = np.sum(weights * sub_blocks, axis=-1, keepdims=True)
|
653 |
+
|
654 |
+
# Initial guess for scales and mins
|
655 |
+
min_v = np.min(sub_blocks, axis=-1, keepdims=True)
|
656 |
+
max_v = np.max(sub_blocks, axis=-1, keepdims=True)
|
657 |
+
min_v[min_v > 0] = 0.0
|
658 |
+
|
659 |
+
max_minus_min = max_v - min_v
|
660 |
+
|
661 |
+
# Handle cases where all values in a sub-block are the same
|
662 |
+
is_flat = max_minus_min < 1e-8
|
663 |
+
max_minus_min[is_flat] = 1.0 # Avoid division by zero
|
664 |
+
|
665 |
+
with np.errstate(divide="ignore"):
|
666 |
+
iscale = nmax / max_minus_min
|
667 |
+
scale = 1.0 / iscale
|
668 |
+
scale[is_flat] = 0.0
|
669 |
+
|
670 |
+
l_current = np_roundf(iscale * (sub_blocks - min_v)).clip(0, nmax).astype(np.uint8)
|
671 |
+
diff = scale * l_current + min_v - sub_blocks
|
672 |
+
best_mse = np.sum(weights * (diff * diff), axis=-1)
|
673 |
+
|
674 |
+
scale_best = scale.squeeze(-1)
|
675 |
+
min_best = min_v.squeeze(-1)
|
676 |
+
|
677 |
+
# Iterative search loop over all sub-blocks at once
|
678 |
+
for is_ in range(nstep + 1):
|
679 |
+
with np.errstate(divide="ignore"):
|
680 |
+
current_iscale = (rmin + rdelta * is_ + nmax) / max_minus_min
|
681 |
+
current_iscale[is_flat] = 0.0
|
682 |
+
|
683 |
+
l_aux = np_roundf(current_iscale * (sub_blocks - min_v)).clip(0, nmax).astype(np.uint8)
|
684 |
+
|
685 |
+
w_l = weights * l_aux
|
686 |
+
sum_l = np.sum(w_l, axis=-1, keepdims=True)
|
687 |
+
sum_l2 = np.sum(w_l * l_aux, axis=-1, keepdims=True)
|
688 |
+
sum_xl = np.sum(w_l * sub_blocks, axis=-1, keepdims=True)
|
689 |
+
|
690 |
+
D = sum_w * sum_l2 - sum_l * sum_l
|
691 |
+
|
692 |
+
valid_D_mask = D > 0
|
693 |
+
# Use np.where for safe division, filling invalid entries with 0
|
694 |
+
this_scale = np.divide((sum_w * sum_xl - sum_x * sum_l), D, out=np.zeros_like(D), where=valid_D_mask)
|
695 |
+
this_min = np.divide((sum_l2 * sum_x - sum_l * sum_xl), D, out=np.zeros_like(D), where=valid_D_mask)
|
696 |
+
|
697 |
+
# Handle case where candidate min > 0
|
698 |
+
min_gt_zero_mask = valid_D_mask & (this_min > 0)
|
699 |
+
if np.any(min_gt_zero_mask):
|
700 |
+
recalc_scale = np.divide(sum_xl, sum_l2, out=np.zeros_like(sum_xl), where=sum_l2 > 0)
|
701 |
+
this_scale = np.where(min_gt_zero_mask, recalc_scale, this_scale)
|
702 |
+
this_min = np.where(min_gt_zero_mask, 0.0, this_min)
|
703 |
+
|
704 |
+
# Calculate current MSE
|
705 |
+
diff = this_scale * l_aux + this_min - sub_blocks
|
706 |
+
current_mse = np.sum(weights * (diff * diff), axis=-1)
|
707 |
+
|
708 |
+
# Update best values where MSE has improved
|
709 |
+
improvement_mask = valid_D_mask.squeeze(-1) & (current_mse < best_mse)
|
710 |
+
if np.any(improvement_mask):
|
711 |
+
best_mse[improvement_mask] = current_mse[improvement_mask]
|
712 |
+
scale_best[improvement_mask] = this_scale.squeeze(-1)[improvement_mask]
|
713 |
+
min_best[improvement_mask] = this_min.squeeze(-1)[improvement_mask]
|
714 |
+
|
715 |
+
scales_all = scale_best
|
716 |
+
mins_all = -min_best
|
717 |
+
# --- End of vectorized search ---
|
718 |
+
|
719 |
+
# Find block-level d and dmin
|
720 |
+
max_scale_per_block = np.max(scales_all, axis=1, keepdims=True)
|
721 |
+
max_min_per_block = np.max(mins_all, axis=1, keepdims=True)
|
722 |
+
|
723 |
+
# Quantize and pack scales and mins
|
724 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
725 |
+
inv_scale = np.where(max_scale_per_block == 0, 0, 63.0 / max_scale_per_block)
|
726 |
+
inv_min = np.where(max_min_per_block == 0, 0, 63.0 / max_min_per_block)
|
727 |
+
|
728 |
+
ls = np.clip(np_roundf(scales_all * inv_scale), 0, 63).astype(np.uint8)
|
729 |
+
lm = np.clip(np_roundf(mins_all * inv_min), 0, 63).astype(np.uint8)
|
730 |
+
|
731 |
+
scales_packed = np.zeros((n_blocks, cls.K_SCALE_SIZE), dtype=np.uint8)
|
732 |
+
scales_packed[:, 0:4] = ls[:, 0:4] & 0x3F
|
733 |
+
scales_packed[:, 4:8] = lm[:, 0:4] & 0x3F
|
734 |
+
scales_packed[:, 8:12] = (ls[:, 4:8] & 0x0F) | ((lm[:, 4:8] & 0x0F) << 4)
|
735 |
+
scales_packed[:, 0:4] |= (ls[:, 4:8] >> 4) << 6
|
736 |
+
scales_packed[:, 4:8] |= (lm[:, 4:8] >> 4) << 6
|
737 |
+
|
738 |
+
# Store block-level d and dmin
|
739 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
740 |
+
d_val = np.where(max_scale_per_block == 0, 0, max_scale_per_block / 63.0)
|
741 |
+
dmin_val = np.where(max_min_per_block == 0, 0, max_min_per_block / 63.0)
|
742 |
+
|
743 |
+
d = d_val.reshape(n_blocks, 1).astype(np.float16).view(np.uint8)
|
744 |
+
dmin = dmin_val.reshape(n_blocks, 1).astype(np.float16).view(np.uint8)
|
745 |
+
|
746 |
+
# Re-quantize the actual data
|
747 |
+
d_eff = (d_val * ls.astype(np.float32)).reshape(n_blocks, 8, 1)
|
748 |
+
m_eff = (dmin_val * lm.astype(np.float32)).reshape(n_blocks, 8, 1)
|
749 |
+
|
750 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
751 |
+
L_float = np.divide(sub_blocks + m_eff, d_eff, out=np.zeros_like(sub_blocks), where=d_eff != 0)
|
752 |
+
|
753 |
+
L = np.clip(np_roundf(L_float), 0, 15).astype(np.uint8)
|
754 |
+
|
755 |
+
# Pack the 4-bit quantized data
|
756 |
+
L_reshaped = L.reshape((n_blocks, cls.QK_K // 64, 2, 32))
|
757 |
+
L_low = L_reshaped[:, :, 0, :].reshape(n_blocks, -1)
|
758 |
+
L_high = L_reshaped[:, :, 1, :].reshape(n_blocks, -1)
|
759 |
+
qs = L_low | (L_high << 4)
|
760 |
+
|
761 |
+
# Assemble and return the final block
|
762 |
+
return np.concatenate([d, dmin, scales_packed, qs], axis=1)
|
763 |
+
|
764 |
+
@staticmethod
|
765 |
+
def get_scale_min(scales: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
766 |
+
n_blocks = scales.shape[0]
|
767 |
+
s = scales.view(np.uint8).reshape(n_blocks, Q4_K.K_SCALE_SIZE)
|
768 |
+
|
769 |
+
sc = np.zeros((n_blocks, 8), dtype=np.uint8)
|
770 |
+
m = np.zeros((n_blocks, 8), dtype=np.uint8)
|
771 |
+
|
772 |
+
sc[:, 0:4] = s[:, 0:4] & 0x3F
|
773 |
+
m[:, 0:4] = s[:, 4:8] & 0x3F
|
774 |
+
|
775 |
+
sc[:, 4:8] = (s[:, 8:12] & 0x0F) | ((s[:, 0:4] >> 6) << 4)
|
776 |
+
m[:, 4:8] = (s[:, 8:12] >> 4) | ((s[:, 4:8] >> 6) << 4)
|
777 |
+
|
778 |
+
return sc, m
|
779 |
+
|
780 |
+
@classmethod
|
781 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
782 |
+
n_blocks = blocks.shape[0]
|
783 |
+
|
784 |
+
d, rest = np.hsplit(blocks, [2])
|
785 |
+
dmin, rest = np.hsplit(rest, [2])
|
786 |
+
scales, qs = np.hsplit(rest, [cls.K_SCALE_SIZE])
|
787 |
+
|
788 |
+
d = d.view(np.float16).astype(np.float32)
|
789 |
+
dmin = dmin.view(np.float16).astype(np.float32)
|
790 |
+
|
791 |
+
sc, m = cls.get_scale_min(scales)
|
792 |
+
|
793 |
+
d_eff = (d * sc.astype(np.float32)).reshape((n_blocks, 8, 1))
|
794 |
+
dm_eff = (dmin * m.astype(np.float32)).reshape((n_blocks, 8, 1))
|
795 |
+
|
796 |
+
# Unpack 4-bit values and arrange back into sub-blocks
|
797 |
+
qs_reshaped = qs.reshape(n_blocks, QK_K // 64, 32)
|
798 |
+
qs_unpacked = np.empty((n_blocks, 8, 32), dtype=np.float32)
|
799 |
+
qs_unpacked[:, [0, 2, 4, 6], :] = qs_reshaped & 0x0F
|
800 |
+
qs_unpacked[:, [1, 3, 5, 7], :] = qs_reshaped >> 4
|
801 |
+
|
802 |
+
return (d_eff * qs_unpacked - dm_eff).reshape((n_blocks, QK_K))
|
803 |
+
|
804 |
+
|
805 |
+
class Q5_K(__Quant, qtype=GGMLQuantizationType.Q5_K):
|
806 |
+
@classmethod
|
807 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
808 |
+
"""
|
809 |
+
Quantizes a numpy array of floats into Q5_K format.
|
810 |
+
Vectorized implementation of the C++ reference code.
|
811 |
+
"""
|
812 |
+
if blocks.shape[-1] % QK_K != 0:
|
813 |
+
raise ValueError(
|
814 |
+
f"The last dimension of the input array must be a multiple of {QK_K}, but got {blocks.shape[-1]}"
|
815 |
+
)
|
816 |
+
|
817 |
+
n_blocks = blocks.size // QK_K
|
818 |
+
sub_blocks = blocks.reshape((n_blocks, 8, 32))
|
819 |
+
|
820 |
+
# --- Vectorized make_qkx3_quants logic for 5 bits ---
|
821 |
+
nmax = 31
|
822 |
+
nstep = 36
|
823 |
+
rmin = -0.9
|
824 |
+
rdelta = 0.05
|
825 |
+
|
826 |
+
# Calculate weights for all sub-blocks
|
827 |
+
sum_x2 = np.sum(sub_blocks * sub_blocks, axis=-1, keepdims=True)
|
828 |
+
av_x = np.sqrt(np.maximum(0, 2 * sum_x2 / QK_K)) # sigma calculation from C++
|
829 |
+
weights = av_x + np.abs(sub_blocks)
|
830 |
+
sum_w = np.sum(weights, axis=-1, keepdims=True)
|
831 |
+
sum_x = np.sum(weights * sub_blocks, axis=-1, keepdims=True)
|
832 |
+
|
833 |
+
min_v = np.min(sub_blocks, axis=-1, keepdims=True)
|
834 |
+
max_v = np.max(sub_blocks, axis=-1, keepdims=True)
|
835 |
+
min_v[min_v > 0] = 0.0
|
836 |
+
|
837 |
+
max_minus_min = max_v - min_v
|
838 |
+
is_flat = max_minus_min < 1e-8
|
839 |
+
max_minus_min[is_flat] = 1.0
|
840 |
+
|
841 |
+
# Initial mse for comparison
|
842 |
+
with np.errstate(divide="ignore"):
|
843 |
+
iscale_initial = nmax / max_minus_min
|
844 |
+
scale_initial = 1.0 / iscale_initial
|
845 |
+
scale_initial[is_flat] = 0.0
|
846 |
+
l_initial = np_roundf(iscale_initial * (sub_blocks - min_v)).clip(0, nmax).astype(np.uint8)
|
847 |
+
diff = scale_initial * l_initial + min_v - sub_blocks
|
848 |
+
best_mse = np.sum(weights * (diff * diff), axis=-1)
|
849 |
+
|
850 |
+
scale_best = scale_initial.squeeze(-1)
|
851 |
+
min_best = min_v.squeeze(-1)
|
852 |
+
|
853 |
+
# Iterative search
|
854 |
+
for is_ in range(nstep + 1):
|
855 |
+
with np.errstate(divide="ignore"):
|
856 |
+
current_iscale = (rmin + rdelta * is_ + nmax) / max_minus_min
|
857 |
+
current_iscale[is_flat] = 0.0
|
858 |
+
|
859 |
+
l_aux = np_roundf(current_iscale * (sub_blocks - min_v)).clip(0, nmax).astype(np.uint8)
|
860 |
+
w_l = weights * l_aux
|
861 |
+
sum_l = np.sum(w_l, axis=-1, keepdims=True)
|
862 |
+
sum_l2 = np.sum(w_l * l_aux, axis=-1, keepdims=True)
|
863 |
+
sum_xl = np.sum(w_l * sub_blocks, axis=-1, keepdims=True)
|
864 |
+
|
865 |
+
D = sum_w * sum_l2 - sum_l * sum_l
|
866 |
+
valid_D_mask = D > 0
|
867 |
+
this_scale = np.divide((sum_w * sum_xl - sum_x * sum_l), D, out=np.zeros_like(D), where=valid_D_mask)
|
868 |
+
this_min = np.divide((sum_l2 * sum_x - sum_l * sum_xl), D, out=np.zeros_like(D), where=valid_D_mask)
|
869 |
+
|
870 |
+
min_gt_zero_mask = valid_D_mask & (this_min > 0)
|
871 |
+
if np.any(min_gt_zero_mask):
|
872 |
+
recalc_scale = np.divide(sum_xl, sum_l2, out=np.zeros_like(sum_xl), where=sum_l2 > 0)
|
873 |
+
this_scale = np.where(min_gt_zero_mask, recalc_scale, this_scale)
|
874 |
+
this_min = np.where(min_gt_zero_mask, 0.0, this_min)
|
875 |
+
|
876 |
+
diff = this_scale * l_aux + this_min - sub_blocks
|
877 |
+
current_mse = np.sum(weights * (diff * diff), axis=-1)
|
878 |
+
improvement_mask = valid_D_mask.squeeze(-1) & (current_mse < best_mse)
|
879 |
+
if np.any(improvement_mask):
|
880 |
+
best_mse[improvement_mask] = current_mse[improvement_mask]
|
881 |
+
scale_best[improvement_mask] = this_scale.squeeze(-1)[improvement_mask]
|
882 |
+
min_best[improvement_mask] = this_min.squeeze(-1)[improvement_mask]
|
883 |
+
|
884 |
+
scales_all = scale_best
|
885 |
+
mins_all = -min_best
|
886 |
+
|
887 |
+
# --- Quantize and pack scales/mins (identical to Q4_K) ---
|
888 |
+
max_scale_per_block = np.max(scales_all, axis=1, keepdims=True)
|
889 |
+
max_min_per_block = np.max(mins_all, axis=1, keepdims=True)
|
890 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
891 |
+
inv_scale = np.where(max_scale_per_block == 0, 0, 63.0 / max_scale_per_block)
|
892 |
+
inv_min = np.where(max_min_per_block == 0, 0, 63.0 / max_min_per_block)
|
893 |
+
ls = np.clip(np_roundf(scales_all * inv_scale), 0, 63).astype(np.uint8)
|
894 |
+
lm = np.clip(np_roundf(mins_all * inv_min), 0, 63).astype(np.uint8)
|
895 |
+
|
896 |
+
scales_packed = np.zeros((n_blocks, Q4_K.K_SCALE_SIZE), dtype=np.uint8)
|
897 |
+
scales_packed[:, 0:4] = ls[:, 0:4] & 0x3F
|
898 |
+
scales_packed[:, 4:8] = lm[:, 0:4] & 0x3F
|
899 |
+
scales_packed[:, 8:12] = (ls[:, 4:8] & 0x0F) | ((lm[:, 4:8] & 0x0F) << 4)
|
900 |
+
scales_packed[:, 0:4] |= (ls[:, 4:8] >> 4) << 6
|
901 |
+
scales_packed[:, 4:8] |= (lm[:, 4:8] >> 4) << 6
|
902 |
+
|
903 |
+
# --- Store block-level d and dmin (identical to Q4_K) ---
|
904 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
905 |
+
d_val = np.where(max_scale_per_block == 0, 0, max_scale_per_block / 63.0)
|
906 |
+
dmin_val = np.where(max_min_per_block == 0, 0, max_min_per_block / 63.0)
|
907 |
+
d = d_val.reshape(n_blocks, 1).astype(np.float16).view(np.uint8)
|
908 |
+
dmin = dmin_val.reshape(n_blocks, 1).astype(np.float16).view(np.uint8)
|
909 |
+
|
910 |
+
# --- Re-quantize the actual data to 5 bits ---
|
911 |
+
d_eff = (d_val * ls.astype(np.float32)).reshape(n_blocks, 8, 1)
|
912 |
+
m_eff = (dmin_val * lm.astype(np.float32)).reshape(n_blocks, 8, 1)
|
913 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
914 |
+
L_float = np.divide(sub_blocks + m_eff, d_eff, out=np.zeros_like(sub_blocks), where=d_eff != 0)
|
915 |
+
L = np.clip(np_roundf(L_float), 0, 31).astype(np.uint8)
|
916 |
+
|
917 |
+
# --- Pack the 5-bit quantized data into qh and qs ---
|
918 |
+
# qh (high bits)
|
919 |
+
h = (L > 15).astype(np.uint8)
|
920 |
+
h_reshaped = h.reshape(n_blocks, 8, 32).transpose(0, 2, 1)
|
921 |
+
bit_shifts = 2 ** np.arange(8, dtype=np.uint8).reshape(1, 1, 8)
|
922 |
+
qh = np.sum(h_reshaped * bit_shifts, axis=-1).astype(np.uint8)
|
923 |
+
|
924 |
+
# qs (low bits)
|
925 |
+
L[L > 15] -= 16
|
926 |
+
l_reshaped = L.reshape(n_blocks, 8, 32)
|
927 |
+
part1 = l_reshaped[:, ::2, :].reshape(n_blocks, -1)
|
928 |
+
part2 = l_reshaped[:, 1::2, :].reshape(n_blocks, -1)
|
929 |
+
qs = part1 | (part2 << 4)
|
930 |
+
|
931 |
+
return np.concatenate([d, dmin, scales_packed, qh, qs], axis=1)
|
932 |
+
|
933 |
+
@classmethod
|
934 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
935 |
+
n_blocks = blocks.shape[0]
|
936 |
+
|
937 |
+
d, rest = np.hsplit(blocks, [2])
|
938 |
+
dmin, rest = np.hsplit(rest, [2])
|
939 |
+
scales, rest = np.hsplit(rest, [Q4_K.K_SCALE_SIZE])
|
940 |
+
qh, qs = np.hsplit(rest, [QK_K // 8])
|
941 |
+
|
942 |
+
d = d.view(np.float16).astype(np.float32)
|
943 |
+
dmin = dmin.view(np.float16).astype(np.float32)
|
944 |
+
|
945 |
+
sc, m = Q4_K.get_scale_min(scales)
|
946 |
+
|
947 |
+
d_eff = (d * sc.astype(np.float32)).reshape((n_blocks, -1, 1))
|
948 |
+
dm_eff = (dmin * m.astype(np.float32)).reshape((n_blocks, -1, 1))
|
949 |
+
|
950 |
+
# Unpack high bits (qh)
|
951 |
+
bit_shifts = 2 ** np.arange(8, dtype=np.uint8).reshape(1, 1, 8)
|
952 |
+
qh_unpacked = (qh[:, :, np.newaxis] & bit_shifts) != 0
|
953 |
+
qh_unpacked = qh_unpacked.transpose(0, 2, 1).reshape(n_blocks, -1, 32)
|
954 |
+
|
955 |
+
# Unpack low bits (qs)
|
956 |
+
ql_unpacked = np.empty((n_blocks, 8, 32), dtype=np.uint8)
|
957 |
+
qs_reshaped = qs.reshape(n_blocks, 4, 32)
|
958 |
+
ql_unpacked[:, ::2, :] = qs_reshaped & 0x0F
|
959 |
+
ql_unpacked[:, 1::2, :] = qs_reshaped >> 4
|
960 |
+
|
961 |
+
# Combine high and low bits and dequantize
|
962 |
+
q = (ql_unpacked + (qh_unpacked * 16)).astype(np.float32)
|
963 |
+
return (d_eff * q - dm_eff).reshape((n_blocks, QK_K))
|
964 |
+
|
965 |
+
|
966 |
+
class Q6_K(__Quant, qtype=GGMLQuantizationType.Q6_K):
|
967 |
+
@classmethod
|
968 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
969 |
+
"""
|
970 |
+
Quantizes a numpy array of floats into Q6_K format.
|
971 |
+
Vectorized implementation of the C++ reference code.
|
972 |
+
"""
|
973 |
+
n_blocks = blocks.shape[0]
|
974 |
+
# Reshape for sub-block processing
|
975 |
+
sub_blocks = blocks.reshape(n_blocks * 16, 16)
|
976 |
+
|
977 |
+
# --- Vectorized `make_qx_quants` for all sub-blocks to find initial scales ---
|
978 |
+
nmax_data = 32 # For Q6_K, data range is [-32, 31]
|
979 |
+
|
980 |
+
# Weights are x*x for the reference implementation
|
981 |
+
weights_data = sub_blocks * sub_blocks
|
982 |
+
|
983 |
+
# Find max absolute values for each sub-block to determine the initial scale
|
984 |
+
abs_sub_blocks = np.abs(sub_blocks)
|
985 |
+
max_indices = np.argmax(abs_sub_blocks, axis=-1, keepdims=True)
|
986 |
+
max_vals = np.take_along_axis(sub_blocks, max_indices, axis=-1)
|
987 |
+
|
988 |
+
is_zero_mask = np.abs(max_vals) < 1e-15
|
989 |
+
|
990 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
991 |
+
initial_iscale = np.where(is_zero_mask, 0, -nmax_data / max_vals)
|
992 |
+
|
993 |
+
# Use np.round for round-half-to-even, matching C's nearest_int
|
994 |
+
l = np.round(sub_blocks * initial_iscale).clip(-nmax_data, nmax_data - 1)
|
995 |
+
sumlx = np.sum(weights_data * sub_blocks * l, axis=-1)
|
996 |
+
suml2 = np.sum(weights_data * l * l, axis=-1)
|
997 |
+
|
998 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
999 |
+
scales_cand = np.divide(sumlx, suml2, out=np.zeros_like(sumlx), where=suml2 != 0)
|
1000 |
+
best_scores = scales_cand * sumlx
|
1001 |
+
best_l = l.copy()
|
1002 |
+
|
1003 |
+
# Iterative search over potential iscale adjustments
|
1004 |
+
for is_ in range(-9, 10):
|
1005 |
+
if is_ == 0:
|
1006 |
+
continue
|
1007 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
1008 |
+
iscale_try = np.where(is_zero_mask, 0, -(nmax_data + 0.1 * is_) / max_vals)
|
1009 |
+
|
1010 |
+
l_try = np.round(sub_blocks * iscale_try).clip(-nmax_data, nmax_data - 1)
|
1011 |
+
sumlx_try = np.sum(weights_data * sub_blocks * l_try, axis=-1)
|
1012 |
+
suml2_try = np.sum(weights_data * l_try * l_try, axis=-1)
|
1013 |
+
|
1014 |
+
improvement_mask = (suml2_try > 0) & (sumlx_try * sumlx_try * suml2 > best_scores * suml2_try)
|
1015 |
+
if np.any(improvement_mask):
|
1016 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
1017 |
+
new_best_scores = np.divide(sumlx_try * sumlx_try, suml2_try, where=suml2_try > 0)
|
1018 |
+
best_scores[improvement_mask] = new_best_scores[improvement_mask]
|
1019 |
+
best_l[improvement_mask] = l_try[improvement_mask]
|
1020 |
+
suml2[improvement_mask] = suml2_try[improvement_mask]
|
1021 |
+
|
1022 |
+
# Recompute final best scales from the best quants (best_l)
|
1023 |
+
sumlx_final = np.sum(weights_data * sub_blocks * best_l, axis=-1)
|
1024 |
+
suml2_final = np.sum(weights_data * best_l * best_l, axis=-1)
|
1025 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
1026 |
+
scales = np.divide(sumlx_final, suml2_final, out=np.zeros_like(sumlx_final), where=suml2_final != 0)
|
1027 |
+
|
1028 |
+
scales[np.all(sub_blocks == 0, axis=-1)] = 0.0
|
1029 |
+
scales = scales.reshape(n_blocks, 16)
|
1030 |
+
|
1031 |
+
# --- Quantize the scales themselves ---
|
1032 |
+
abs_scales = np.abs(scales)
|
1033 |
+
max_abs_scale_indices = np.argmax(abs_scales, axis=-1, keepdims=True)
|
1034 |
+
max_scale_vals = np.take_along_axis(scales, max_abs_scale_indices, axis=-1)
|
1035 |
+
|
1036 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
1037 |
+
is_zero_mask = np.abs(max_scale_vals) < 1e-15
|
1038 |
+
iscale_s = np.where(is_zero_mask, 0, -128.0 / max_scale_vals)
|
1039 |
+
d_val = np.where(is_zero_mask, 0, max_scale_vals / -128.0)
|
1040 |
+
|
1041 |
+
quantized_scales = np.round(scales * iscale_s).clip(-128, 127).astype(np.int8)
|
1042 |
+
d = d_val.astype(np.float16).view(np.uint8)
|
1043 |
+
|
1044 |
+
# --- Re-quantize original data with final scales ---
|
1045 |
+
d_sub = d_val * quantized_scales.astype(np.float32)
|
1046 |
+
d_sub_reshaped = d_sub.reshape(n_blocks, 16, 1)
|
1047 |
+
|
1048 |
+
sub_blocks_reshaped = blocks.reshape(n_blocks, 16, 16)
|
1049 |
+
with np.errstate(divide="ignore", invalid="ignore"):
|
1050 |
+
l_float = np.divide(
|
1051 |
+
sub_blocks_reshaped, d_sub_reshaped, out=np.zeros_like(sub_blocks_reshaped), where=d_sub_reshaped != 0
|
1052 |
+
)
|
1053 |
+
|
1054 |
+
l_final = np.round(l_float).clip(-32, 31).astype(np.int8)
|
1055 |
+
L = (l_final + 32).astype(np.uint8).reshape(n_blocks, 256)
|
1056 |
+
|
1057 |
+
# --- Pack the 6-bit quantized data ---
|
1058 |
+
L_reshaped = L.reshape(n_blocks, 2, 4, 32)
|
1059 |
+
L_low = L_reshaped & 0xF
|
1060 |
+
L_high = L_reshaped >> 4
|
1061 |
+
|
1062 |
+
# Pack lower 4 bits into ql
|
1063 |
+
ql = np.empty((n_blocks, 128), dtype=np.uint8)
|
1064 |
+
ql[:, 0:32] = L_low[:, 0, 0, :] | (L_low[:, 0, 2, :] << 4)
|
1065 |
+
ql[:, 32:64] = L_low[:, 0, 1, :] | (L_low[:, 0, 3, :] << 4)
|
1066 |
+
ql[:, 64:96] = L_low[:, 1, 0, :] | (L_low[:, 1, 2, :] << 4)
|
1067 |
+
ql[:, 96:128] = L_low[:, 1, 1, :] | (L_low[:, 1, 3, :] << 4)
|
1068 |
+
|
1069 |
+
# Pack higher 2 bits into qh
|
1070 |
+
qh_packed = (
|
1071 |
+
L_high[:, :, 0, :] | (L_high[:, :, 1, :] << 2) | (L_high[:, :, 2, :] << 4) | (L_high[:, :, 3, :] << 6)
|
1072 |
+
)
|
1073 |
+
qh = qh_packed.reshape(n_blocks, -1)
|
1074 |
+
|
1075 |
+
# Final assembly: view scales as uint8 before concatenating
|
1076 |
+
return np.concatenate([ql, qh, quantized_scales.view(np.uint8), d], axis=1)
|
1077 |
+
|
1078 |
+
@classmethod
|
1079 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1080 |
+
n_blocks = blocks.shape[0]
|
1081 |
+
|
1082 |
+
ql, rest = np.hsplit(blocks, [QK_K // 2])
|
1083 |
+
qh, rest = np.hsplit(rest, [QK_K // 4])
|
1084 |
+
scales, d = np.hsplit(rest, [QK_K // 16])
|
1085 |
+
|
1086 |
+
scales = scales.view(np.int8).astype(np.float32)
|
1087 |
+
d = d.view(np.float16).astype(np.float32)
|
1088 |
+
d = (d * scales).reshape((n_blocks, QK_K // 16, 1))
|
1089 |
+
|
1090 |
+
ql = ql.reshape((n_blocks, -1, 1, 64)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
1091 |
+
ql = (ql & np.uint8(0x0F)).reshape((n_blocks, -1, 32))
|
1092 |
+
qh = qh.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
1093 |
+
qh = (qh & np.uint8(0x03)).reshape((n_blocks, -1, 32))
|
1094 |
+
q = (ql | (qh << np.uint8(4))).astype(np.int8) - np.int8(32)
|
1095 |
+
q = q.reshape((n_blocks, QK_K // 16, -1)).astype(np.float32)
|
1096 |
+
|
1097 |
+
return (d * q).reshape((n_blocks, QK_K))
|
1098 |
+
|
1099 |
+
|
1100 |
+
class TQ1_0(__Quant, qtype=GGMLQuantizationType.TQ1_0):
|
1101 |
+
@classmethod
|
1102 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1103 |
+
n_blocks = blocks.shape[0]
|
1104 |
+
|
1105 |
+
d = abs(blocks).max(axis=-1, keepdims=True)
|
1106 |
+
with np.errstate(divide="ignore"):
|
1107 |
+
id = np.where(d == 0, 0, 1 / d)
|
1108 |
+
qs = np_roundf(blocks * id)
|
1109 |
+
qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8)
|
1110 |
+
|
1111 |
+
qs0, qs1, qh = qs[..., : (32 * 5)], qs[..., (32 * 5) : (48 * 5)], qs[..., (48 * 5) :]
|
1112 |
+
qs0 = qs0.reshape((n_blocks, -1, 5, 32)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1))
|
1113 |
+
qs0 = np.sum(qs0, axis=-2).reshape((n_blocks, -1))
|
1114 |
+
qs1 = qs1.reshape((n_blocks, -1, 5, 16)) * np.array([81, 27, 9, 3, 1], dtype=np.uint8).reshape((1, 1, 5, 1))
|
1115 |
+
qs1 = np.sum(qs1, axis=-2).reshape((n_blocks, -1))
|
1116 |
+
qh = qh.reshape((n_blocks, -1, 4, 4)) * np.array([81, 27, 9, 3], dtype=np.uint8).reshape((1, 1, 4, 1))
|
1117 |
+
qh = np.sum(qh, axis=-2).reshape((n_blocks, -1))
|
1118 |
+
qs = np.concatenate([qs0, qs1, qh], axis=-1)
|
1119 |
+
qs = (qs.astype(np.uint16) * 256 + (243 - 1)) // 243
|
1120 |
+
|
1121 |
+
qs = qs.astype(np.uint8)
|
1122 |
+
d = d.astype(np.float16).view(np.uint8)
|
1123 |
+
|
1124 |
+
return np.concatenate([qs, d], axis=-1)
|
1125 |
+
|
1126 |
+
@classmethod
|
1127 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1128 |
+
n_blocks = blocks.shape[0]
|
1129 |
+
|
1130 |
+
qs, rest = np.hsplit(blocks, [(QK_K - 4 * QK_K // 64) // 5])
|
1131 |
+
qh, d = np.hsplit(rest, [QK_K // 64])
|
1132 |
+
|
1133 |
+
d = d.view(np.float16).astype(np.float32)
|
1134 |
+
|
1135 |
+
qs0, qs1 = qs[..., :32], qs[..., 32:]
|
1136 |
+
qs0 = qs0.reshape((n_blocks, -1, 1, 32)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1))
|
1137 |
+
qs0 = qs0.reshape((n_blocks, -1))
|
1138 |
+
qs1 = qs1.reshape((n_blocks, -1, 1, 16)) * np.array([1, 3, 9, 27, 81], dtype=np.uint8).reshape((1, 1, 5, 1))
|
1139 |
+
qs1 = qs1.reshape((n_blocks, -1))
|
1140 |
+
qh = qh.reshape((n_blocks, -1, 1, 4)) * np.array([1, 3, 9, 27], dtype=np.uint8).reshape((1, 1, 4, 1))
|
1141 |
+
qh = qh.reshape((n_blocks, -1))
|
1142 |
+
qs = np.concatenate([qs0, qs1, qh], axis=-1)
|
1143 |
+
qs = ((qs.astype(np.uint16) * 3) >> 8).astype(np.int8) - np.int8(1)
|
1144 |
+
|
1145 |
+
return d * qs.astype(np.float32)
|
1146 |
+
|
1147 |
+
|
1148 |
+
class TQ2_0(__Quant, qtype=GGMLQuantizationType.TQ2_0):
|
1149 |
+
@classmethod
|
1150 |
+
def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1151 |
+
n_blocks = blocks.shape[0]
|
1152 |
+
|
1153 |
+
d = abs(blocks).max(axis=-1, keepdims=True)
|
1154 |
+
with np.errstate(divide="ignore"):
|
1155 |
+
id = np.where(d == 0, 0, 1 / d)
|
1156 |
+
qs = np_roundf(blocks * id)
|
1157 |
+
qs = (qs.astype(np.int8) + np.int8(1)).astype(np.uint8)
|
1158 |
+
|
1159 |
+
qs = qs.reshape((n_blocks, -1, 4, 32)) << np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
1160 |
+
qs = qs[..., 0, :] | qs[..., 1, :] | qs[..., 2, :] | qs[..., 3, :]
|
1161 |
+
qs = qs.reshape((n_blocks, -1))
|
1162 |
+
|
1163 |
+
d = d.astype(np.float16).view(np.uint8)
|
1164 |
+
|
1165 |
+
return np.concatenate([qs, d], axis=-1)
|
1166 |
+
|
1167 |
+
@classmethod
|
1168 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1169 |
+
n_blocks = blocks.shape[0]
|
1170 |
+
|
1171 |
+
qs, d = np.hsplit(blocks, [QK_K // 4])
|
1172 |
+
|
1173 |
+
d = d.view(np.float16).astype(np.float32)
|
1174 |
+
|
1175 |
+
qs = qs.reshape((n_blocks, -1, 1, 32)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4, 1))
|
1176 |
+
qs = (qs & 0x03).reshape((n_blocks, -1)).astype(np.int8) - np.int8(1)
|
1177 |
+
|
1178 |
+
return d * qs.astype(np.float32)
|
1179 |
+
|
1180 |
+
|
1181 |
+
class IQ2_XXS(__Quant, qtype=GGMLQuantizationType.IQ2_XXS):
|
1182 |
+
ksigns: bytes = (
|
1183 |
+
b"\x00\x81\x82\x03\x84\x05\x06\x87\x88\x09\x0a\x8b\x0c\x8d\x8e\x0f"
|
1184 |
+
b"\x90\x11\x12\x93\x14\x95\x96\x17\x18\x99\x9a\x1b\x9c\x1d\x1e\x9f"
|
1185 |
+
b"\xa0\x21\x22\xa3\x24\xa5\xa6\x27\x28\xa9\xaa\x2b\xac\x2d\x2e\xaf"
|
1186 |
+
b"\x30\xb1\xb2\x33\xb4\x35\x36\xb7\xb8\x39\x3a\xbb\x3c\xbd\xbe\x3f"
|
1187 |
+
b"\xc0\x41\x42\xc3\x44\xc5\xc6\x47\x48\xc9\xca\x4b\xcc\x4d\x4e\xcf"
|
1188 |
+
b"\x50\xd1\xd2\x53\xd4\x55\x56\xd7\xd8\x59\x5a\xdb\x5c\xdd\xde\x5f"
|
1189 |
+
b"\x60\xe1\xe2\x63\xe4\x65\x66\xe7\xe8\x69\x6a\xeb\x6c\xed\xee\x6f"
|
1190 |
+
b"\xf0\x71\x72\xf3\x74\xf5\xf6\x77\x78\xf9\xfa\x7b\xfc\x7d\x7e\xff"
|
1191 |
+
)
|
1192 |
+
|
1193 |
+
# iq2xxs_grid, but with each byte of the original packed in 2 bits,
|
1194 |
+
# by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
|
1195 |
+
grid_shape = (256, 8)
|
1196 |
+
grid_map = (0x08, 0x19, 0x2B)
|
1197 |
+
grid_hex = (
|
1198 |
+
b"00000200050008000a00110014002000220028002a0041004400500058006100"
|
1199 |
+
b"6400800082008a00a20001010401100115014001840198010002020222028202"
|
1200 |
+
b"010404041004210424044004420448046004810484049004a404000502050805"
|
1201 |
+
b"200546056905800591050906100640068406a406000805080808140828084108"
|
1202 |
+
b"440850085208880804094009020a140a01100410101021104010601084109010"
|
1203 |
+
b"951000110811201150115a118011241245120014081420142514491480141815"
|
1204 |
+
b"6215001616160118041810184018811800190519a019511a002002200a204420"
|
1205 |
+
b"6120802082202921482100220222012404241024402456240025412564259026"
|
1206 |
+
b"082820289428442a014004401040184021402440404048405640604081408440"
|
1207 |
+
b"9040004120416141804185410142104248425642684200440844204480449944"
|
1208 |
+
b"124524450046014804481048404845480049584961498249454a904a00500850"
|
1209 |
+
b"1150195020508050885004514251a4519152905492540a550156545600581158"
|
1210 |
+
b"195864584059085a046010604060686000615561186260620064056410651265"
|
1211 |
+
b"84654268008002800a8041808280048118814081118201840484108415844084"
|
1212 |
+
b"608400854685948509864086608602880489118a0490109024904090a1901691"
|
1213 |
+
b"8091459200942294449451958198209902a050a085a009a100a218a450a804a9"
|
1214 |
+
)
|
1215 |
+
|
1216 |
+
@classmethod
|
1217 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1218 |
+
n_blocks = blocks.shape[0]
|
1219 |
+
|
1220 |
+
d, qs = np.hsplit(blocks, [2])
|
1221 |
+
|
1222 |
+
d = d.view(np.float16).astype(np.float32)
|
1223 |
+
|
1224 |
+
qs = qs.view(np.uint32).reshape(n_blocks, -1, 2)
|
1225 |
+
|
1226 |
+
db = d * (np.float32(0.5) + (qs[..., 1] >> 28).astype(np.float32)) * np.float32(0.25)
|
1227 |
+
db = db.reshape((n_blocks, -1, 1, 1))
|
1228 |
+
|
1229 |
+
# get the sign indices and unpack the bits
|
1230 |
+
signs = qs[..., 1].reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4))
|
1231 |
+
ksigns = np.frombuffer(cls.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128))
|
1232 |
+
signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1))
|
1233 |
+
signs = np.take_along_axis(ksigns, signs, axis=-1)
|
1234 |
+
signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape(
|
1235 |
+
(1, 1, 1, 8)
|
1236 |
+
)
|
1237 |
+
signs = signs & np.uint8(0x01)
|
1238 |
+
signs = np.where(signs == 0, np.float32(1), np.float32(-1))
|
1239 |
+
signs = signs.reshape((n_blocks, -1, 4, 8))
|
1240 |
+
|
1241 |
+
assert cls.grid is not None
|
1242 |
+
grid = np.take_along_axis(cls.grid, qs[..., 0].copy().view(np.uint8).reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1243 |
+
grid = grid.reshape((n_blocks, -1, 4, 8))
|
1244 |
+
|
1245 |
+
return (db * grid * signs).reshape((n_blocks, -1))
|
1246 |
+
|
1247 |
+
|
1248 |
+
class IQ2_XS(__Quant, qtype=GGMLQuantizationType.IQ2_XS):
|
1249 |
+
# iq2xs_grid, but with each byte of the original packed in 2 bits,
|
1250 |
+
# by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
|
1251 |
+
grid_shape = (512, 8)
|
1252 |
+
grid_map = (0x08, 0x19, 0x2B)
|
1253 |
+
grid_hex = (
|
1254 |
+
b"00000200050008000a0011001400160019002000220025002800410044004600"
|
1255 |
+
b"49005000520055005800610064008000820085008800910094009900a0000101"
|
1256 |
+
b"04010601090110011201150118011a0121012401400142014501480151015401"
|
1257 |
+
b"6001680181018401900100020202050208021102140220024102440250025502"
|
1258 |
+
b"80028a0201040404060409041004120415041804210424044004420445044804"
|
1259 |
+
b"5104540456046004810484049004000502050505080511051405200541054405"
|
1260 |
+
b"500561058005010604061006260640064206840600080208050808080a081108"
|
1261 |
+
b"14082008250841084408500858088008a008aa08010904091009400981098909"
|
1262 |
+
b"000a200a280a960aa00a01100410061009101010121015101810211024104010"
|
1263 |
+
b"4210451048105110541060106a10811084109010001102110511081111111411"
|
1264 |
+
b"2011411144115011801194119611011204120612101240126012001402140514"
|
1265 |
+
b"0814111414142014411444144914501464148014011504151015401500161416"
|
1266 |
+
b"49160118041810181218401854188618001905196619511aa91a002002200520"
|
1267 |
+
b"08200a201120142020204120442050208020a020012104211021402148216521"
|
1268 |
+
b"002222228022a82201240424102429244024002541255225992501261a26a626"
|
1269 |
+
b"002808280a28202855288828a22868299029082a202a822a882a8a2a01400440"
|
1270 |
+
b"0640094010401240154018402140244040404240454048404a40514054406040"
|
1271 |
+
b"6540814084409040004102410541084111411441204141414441504180418541"
|
1272 |
+
b"a241014204421042124229424042004402440544084411441444194420444144"
|
1273 |
+
b"4444504480449444014504451045244540459a4500460a464446504601480448"
|
1274 |
+
b"1048404845485448624800491149444950496949044a00500250055008501150"
|
1275 |
+
b"145020502850415044505050805001510451105115514051425100524452aa52"
|
1276 |
+
b"0154045410542154405460548154a154005508558055885521566856a1560058"
|
1277 |
+
b"14584158505899581a5940594259855a0160046010604060546062608660a960"
|
1278 |
+
b"006124624a62926200641664106540654565a46501686a682569066a546a626a"
|
1279 |
+
b"00800280058008801180148020802a8041804480508080808280a880aa800181"
|
1280 |
+
b"0481068110814081518159810082208280828282a082a8820184048410841284"
|
1281 |
+
b"158440846084898400854485a58518866a860088088825885a8880888288a888"
|
1282 |
+
b"0689228a808a888a968aa88a0190049010904090569084900091229164915692"
|
1283 |
+
b"89920094059444945094589429959095929541965198a6984999159a609a00a0"
|
1284 |
+
b"02a008a00aa020a02aa0a0a051a159a1a6a100a202a208a22aa280a2a0a240a4"
|
1285 |
+
b"95a465a698a60aa820a822a828a8a0a8a8a804a984a986a928aa2aaa91aaaaaa"
|
1286 |
+
)
|
1287 |
+
|
1288 |
+
@classmethod
|
1289 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1290 |
+
n_blocks = blocks.shape[0]
|
1291 |
+
|
1292 |
+
d, rest = np.hsplit(blocks, [2])
|
1293 |
+
qs, scales = np.hsplit(rest, [2 * QK_K // 8])
|
1294 |
+
|
1295 |
+
d = d.view(np.float16).astype(np.float32)
|
1296 |
+
qs = qs.view(np.uint16)
|
1297 |
+
|
1298 |
+
scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
|
1299 |
+
scales = (scales & 0x0F).reshape((n_blocks, -1))
|
1300 |
+
db = d * (np.float32(0.5) + scales) * np.float32(0.25)
|
1301 |
+
db = db.reshape((n_blocks, -1, 1, 1))
|
1302 |
+
|
1303 |
+
# get the sign indices and unpack the bits
|
1304 |
+
signs = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape(1, 1, 128)
|
1305 |
+
signs = np.take_along_axis(signs, (qs >> 9).reshape((n_blocks, -1, 1)), axis=-1)
|
1306 |
+
signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
|
1307 |
+
signs = signs & np.uint8(0x01)
|
1308 |
+
signs = np.where(signs == 0, np.float32(1), np.float32(-1))
|
1309 |
+
signs = signs.reshape((n_blocks, -1, 2, 8))
|
1310 |
+
|
1311 |
+
assert cls.grid is not None
|
1312 |
+
grid = np.take_along_axis(cls.grid, (qs & np.uint16(511)).reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1313 |
+
grid = grid.reshape((n_blocks, -1, 2, 8))
|
1314 |
+
|
1315 |
+
return (db * grid * signs).reshape((n_blocks, -1))
|
1316 |
+
|
1317 |
+
|
1318 |
+
class IQ2_S(__Quant, qtype=GGMLQuantizationType.IQ2_S):
|
1319 |
+
# iq2s_grid, but with each byte of the original packed in 2 bits,
|
1320 |
+
# by mapping 0x08 to 0, 0x19 to 1, and 0x2b to 2.
|
1321 |
+
grid_shape = (1024, 8)
|
1322 |
+
grid_map = (0x08, 0x19, 0x2B)
|
1323 |
+
grid_hex = (
|
1324 |
+
b"00000200050008000a0011001400160019002000220025002800410044004600"
|
1325 |
+
b"490050005200550058006100640066006900800082008500880091009400a000"
|
1326 |
+
b"a500aa0001010401060109011001120115011801210124014001420145014801"
|
1327 |
+
b"510154015601590160016501680181018401900192019501a101a40100020202"
|
1328 |
+
b"050208021102140220022a02410244024602490250025502800285028a029402"
|
1329 |
+
b"a202010404040604090410041204150418042104240426042904400442044504"
|
1330 |
+
b"48044a0451045404560459046004620465048104840486048904900495049804"
|
1331 |
+
b"a104a40400050205050508050a05110514051605190520052505280541054405"
|
1332 |
+
b"46054905500552055505580561056405800582058505880591059405a0050106"
|
1333 |
+
b"0406060609061006150640064506480651065406600681068406900600080208"
|
1334 |
+
b"050808081108140816081908200825082a084108440846084908500852085508"
|
1335 |
+
b"580861086408800885089408aa08010904091009120915091809210940094509"
|
1336 |
+
b"480951095409600981099009000a110a140a220a280a2a0a500a990a01100410"
|
1337 |
+
b"0610091010101210151018102110241026104010421045104810511054105610"
|
1338 |
+
b"59106010621065106810811084108610901095109810a110a410001102110511"
|
1339 |
+
b"08110a1111111411161119112011221125112811411144114611491150115211"
|
1340 |
+
b"5511581161116411801182118511881191119411011204120912101215122112"
|
1341 |
+
b"2412401245125112541281128412901200140214051408141114141416141914"
|
1342 |
+
b"2014251428144114441446144914501452145514581461146414801482148514"
|
1343 |
+
b"881491149414a014011504150615091510151215151518152115241540154215"
|
1344 |
+
b"4515481551155415601581158415901500160516081611161416201641164416"
|
1345 |
+
b"50168016aa160118041806180918101815181818211840184218451848185118"
|
1346 |
+
b"541860188118841800190219051908191119141920194119441950196919a219"
|
1347 |
+
b"041a101a401a561a00200220052008201120142016201920202025202a204120"
|
1348 |
+
b"4420502052205520642080208a209420aa200121042110211221152121214021"
|
1349 |
+
b"4221452151215421602181218421902100220a22222228222a22442250228822"
|
1350 |
+
b"8a22a82201240424062409241024152418242124242440244224452448245124"
|
1351 |
+
b"5424602481248424902400250525082511251425202541254425502566258025"
|
1352 |
+
b"0126042610264026592600280528112814284128442850288a28aa2801290429"
|
1353 |
+
b"102995290a2a222a642a882a8a2a014004400640094010401240154018401a40"
|
1354 |
+
b"21402440264040404240454048404a4051405440564059406040624065408140"
|
1355 |
+
b"8440904095409840a140a4400041024105410841114114411641194120412241"
|
1356 |
+
b"2541414144414641494150415241554158416141644180418241854188419141"
|
1357 |
+
b"9441a04101420442104212421542184224424042454248425142544260428142"
|
1358 |
+
b"844200440244054408440a441144144416441944204422442544284441444444"
|
1359 |
+
b"46444944504452445544584461446444804482448544884491449444a0440145"
|
1360 |
+
b"0445064509451045124515451845214524454045424545454845514554456045"
|
1361 |
+
b"6a4581458445904500460246054608461146144620464146444650468046a546"
|
1362 |
+
b"0148044809481048124815481848214824484048424845484848514854486048"
|
1363 |
+
b"84489048004902490549084911491449204941494449504980499649014a044a"
|
1364 |
+
b"104a404a00500250055008501150145016501950205022502550285041504450"
|
1365 |
+
b"4650495050505250555058506150645080508250855088509150945001510451"
|
1366 |
+
b"0651095110511251155118512151245140514251455148515151545160518151"
|
1367 |
+
b"8451905100520552085211521452205241524452505269528052015404540654"
|
1368 |
+
b"0954105412541554185421542454405442544554485451545454605481548454"
|
1369 |
+
b"9054005502550555085511551455205541554455505580550156045610562656"
|
1370 |
+
b"405600580258055808581158145820584158445850585a588058015904591059"
|
1371 |
+
b"4059005a195a855aa85a01600460066010601260156018602160246040604560"
|
1372 |
+
b"4860516054606060846090600061026105610861116114612061416144615061"
|
1373 |
+
b"806199610462106240625662a162006405640864116414642064416444645064"
|
1374 |
+
b"806401650465106540654a656865926500669466016804681068656898680069"
|
1375 |
+
b"2a69426aa16a0080028005800880118014801980208025804180448050805280"
|
1376 |
+
b"5580588061808080858091809480018104810981108112811581188121812481"
|
1377 |
+
b"408142814581488151815481818184819081a981008205820a82118214824182"
|
1378 |
+
b"4482508201840484068409841084128415841884218440844284458448845184"
|
1379 |
+
b"5484608481848484908400850285058508851185148520854185448550858085"
|
1380 |
+
b"8a85018604861086298640860088058811881488418844885088a28801890489"
|
1381 |
+
b"40896589228a588a5a8a828aa28a019004900990109012901590189024904090"
|
1382 |
+
b"4290459048905190549060908190849090900091059111911491419144915091"
|
1383 |
+
b"5a910192049210924092a6920094029405940894119414942094419444945094"
|
1384 |
+
b"8094969401950495109540959895a19500964696649601980498109826984098"
|
1385 |
+
b"a998009949995299909a00a005a00aa014a022a02aa041a044a050a0a2a0aaa0"
|
1386 |
+
b"40a165a102a20aa222a228a22aa282a288a28aa2a8a201a404a410a440a489a4"
|
1387 |
+
b"a4a400a519a551a60aa828a8a2a854a986a908aa0aaa20aa22aa28aa88aaaaaa"
|
1388 |
+
)
|
1389 |
+
|
1390 |
+
@classmethod
|
1391 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1392 |
+
n_blocks = blocks.shape[0]
|
1393 |
+
|
1394 |
+
d, rest = np.hsplit(blocks, [2])
|
1395 |
+
qs, rest = np.hsplit(rest, [QK_K // 8])
|
1396 |
+
signs, rest = np.hsplit(rest, [QK_K // 8])
|
1397 |
+
qh, scales = np.hsplit(rest, [QK_K // 32])
|
1398 |
+
|
1399 |
+
d = d.view(np.float16).astype(np.float32)
|
1400 |
+
|
1401 |
+
scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
|
1402 |
+
scales = (scales & 0x0F).reshape((n_blocks, -1))
|
1403 |
+
db = d * (np.float32(0.5) + scales) * np.float32(0.25)
|
1404 |
+
db = db.reshape((n_blocks, -1, 1, 1))
|
1405 |
+
|
1406 |
+
# unpack the sign bits
|
1407 |
+
signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
|
1408 |
+
signs = signs & np.uint8(0x01)
|
1409 |
+
signs = np.where(signs == 0, np.float32(1), np.float32(-1))
|
1410 |
+
signs = signs.reshape((n_blocks, -1, 2, 8))
|
1411 |
+
|
1412 |
+
qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape((1, 1, 4))
|
1413 |
+
qs = qs.astype(np.uint16) | ((qh & 0x03).astype(np.uint16) << 8).reshape((n_blocks, -1))
|
1414 |
+
|
1415 |
+
assert cls.grid is not None
|
1416 |
+
grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1417 |
+
grid = grid.reshape((n_blocks, -1, 2, 8))
|
1418 |
+
|
1419 |
+
return (db * grid * signs).reshape((n_blocks, -1))
|
1420 |
+
|
1421 |
+
|
1422 |
+
class IQ3_XXS(__Quant, qtype=GGMLQuantizationType.IQ3_XXS):
|
1423 |
+
grid_shape = (256, 4)
|
1424 |
+
grid_map = (0x04, 0x0C, 0x14, 0x1C, 0x24, 0x2C, 0x34, 0x3E)
|
1425 |
+
grid_hex = (
|
1426 |
+
b"0000020004001100130017002000220031004200730075000101030110011201"
|
1427 |
+
b"2101250130013201410154017001000202020402110220022202310233023702"
|
1428 |
+
b"5102570275020103070310031203250370031304370444045704730475040105"
|
1429 |
+
b"0705320552053506640610071407160743076107011003101010121021102310"
|
1430 |
+
b"3010321034104710501000110211111120112211011203121012121221123012"
|
1431 |
+
b"7212001302132013311346136613011405145014201524154615711505162217"
|
1432 |
+
b"4017002002201120132020202220262031204220012103210521102112212121"
|
1433 |
+
b"3021632167217021002202221122172220222222372240225522012310231423"
|
1434 |
+
b"7023742335245324032527254125742501270327162745270130103012302130"
|
1435 |
+
b"2330503065307230003102312031313144314631013203321032253252327232"
|
1436 |
+
b"1133333330344734723400350635223555351436363663363337603704401740"
|
1437 |
+
b"3540374053405740744120423742404260426642074345430444514464442545"
|
1438 |
+
b"4345704505471047124730471250415070500051065126515551145232527252"
|
1439 |
+
b"0253535310542354275472540255315550562457425724604460466064602161"
|
1440 |
+
b"6161176264623063366344640565526533660367216703700570077010703270"
|
1441 |
+
b"5270267140711272457252720073157333736073217441740075027524753076"
|
1442 |
+
)
|
1443 |
+
|
1444 |
+
@classmethod
|
1445 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1446 |
+
n_blocks = blocks.shape[0]
|
1447 |
+
|
1448 |
+
d, rest = np.hsplit(blocks, [2])
|
1449 |
+
qs, scales = np.hsplit(rest, [QK_K // 4])
|
1450 |
+
|
1451 |
+
d = d.view(np.float16).astype(np.float32)
|
1452 |
+
scales = scales.view(np.uint32)
|
1453 |
+
|
1454 |
+
db = d * (np.float32(0.5) + (scales >> 28).astype(np.float32)) * np.float32(0.5)
|
1455 |
+
db = db.reshape((n_blocks, -1, 1, 1))
|
1456 |
+
|
1457 |
+
# get the sign indices and unpack the bits
|
1458 |
+
signs = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 7, 14, 21], dtype=np.uint32).reshape((1, 1, 4))
|
1459 |
+
ksigns = np.frombuffer(IQ2_XXS.ksigns, dtype=np.uint8).reshape((1, 1, 1, 128))
|
1460 |
+
signs = (signs & np.uint32(0x7F)).reshape((n_blocks, -1, 4, 1))
|
1461 |
+
signs = np.take_along_axis(ksigns, signs, axis=-1)
|
1462 |
+
signs = signs.reshape((n_blocks, -1, 4, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape(
|
1463 |
+
(1, 1, 1, 8)
|
1464 |
+
)
|
1465 |
+
signs = signs & np.uint8(0x01)
|
1466 |
+
signs = np.where(signs == 0, np.float32(1), np.float32(-1))
|
1467 |
+
signs = signs.reshape((n_blocks, -1, 4, 8))
|
1468 |
+
|
1469 |
+
assert cls.grid is not None
|
1470 |
+
grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1471 |
+
grid = grid.reshape((n_blocks, -1, 4, 8))
|
1472 |
+
|
1473 |
+
return (db * grid * signs).reshape((n_blocks, -1))
|
1474 |
+
|
1475 |
+
|
1476 |
+
class IQ3_S(__Quant, qtype=GGMLQuantizationType.IQ3_S):
|
1477 |
+
grid_shape = (512, 4)
|
1478 |
+
grid_map = (0x01, 0x03, 0x05, 0x07, 0x09, 0x0B, 0x0D, 0x0F)
|
1479 |
+
grid_hex = (
|
1480 |
+
b"0000010002000500070010001100120014001600200021002500330040004200"
|
1481 |
+
b"4500470051005300600062007100740077000001010102010401100111011501"
|
1482 |
+
b"2001230127013101350144016101650172010002010205020702100213021602"
|
1483 |
+
b"2102250230023402420245024702510253027002730203031103150320032203"
|
1484 |
+
b"3103330336034403500352036703710375030004130417042104240432044004"
|
1485 |
+
b"4304510470040205040520052205260533054105450547056605730506061106"
|
1486 |
+
b"1306310652067106000702070407200722072607330750075407001001100210"
|
1487 |
+
b"0410101011101310151017102010221031103410361054105610611072100011"
|
1488 |
+
b"0111031106111011141121113011331141115011521170117611001212121512"
|
1489 |
+
b"1712201224123212401243125512601272120113041307131013131321132713"
|
1490 |
+
b"3013341341136213701303140514121414143114331442144614501454140115"
|
1491 |
+
b"1015131521153015321551152016241627164416461601170317101712172117"
|
1492 |
+
b"3517411762177017002001200320052007201020122014201620212023202720"
|
1493 |
+
b"3020322041204320452050205220672070207320752000210221102113211721"
|
1494 |
+
b"2221252131213421422151210122042207222122232230223722412253225722"
|
1495 |
+
b"7122742200230223052311232223242331233323422350236623012407242024"
|
1496 |
+
b"2324322435244124722475240425112522253725402553257025002602260726"
|
1497 |
+
b"2126552661260527112726273027432750270230113013301530173022303130"
|
1498 |
+
b"3330353042304430473051306330713001310331053114312131233140316031"
|
1499 |
+
b"7231763100321232203232323432503201331033143321332333273330334133"
|
1500 |
+
b"4333473355337333033411341634223431345234603464340135103512352535"
|
1501 |
+
b"3235443556357335163641360137033720372237353700400440124020402440"
|
1502 |
+
b"2740324041405040704002410741114113412241304135414341514155410142"
|
1503 |
+
b"0342104215422142334240425742624270420443114313432043224331433543"
|
1504 |
+
b"0044024424443744404471440545074521456245134634466046104715473047"
|
1505 |
+
b"4347514702501050145022504050445047505250665074500151035105511251"
|
1506 |
+
b"2151325172510052115223523052365253520253075310532753445351536553"
|
1507 |
+
b"7353015404542054325446541255265551555355425602570457225711601360"
|
1508 |
+
b"1560316033606060006120612761646112623462426255626262706200631463"
|
1509 |
+
b"2163406325644364626400650365346560650566406611671367007004700770"
|
1510 |
+
b"2070227036704070547062700271117124714371457101720472107216722172"
|
1511 |
+
b"3072517202733273357353730174057413742074507422754275027631760077"
|
1512 |
+
)
|
1513 |
+
|
1514 |
+
@classmethod
|
1515 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1516 |
+
n_blocks = blocks.shape[0]
|
1517 |
+
|
1518 |
+
d, rest = np.hsplit(blocks, [2])
|
1519 |
+
qs, rest = np.hsplit(rest, [QK_K // 4])
|
1520 |
+
qh, rest = np.hsplit(rest, [QK_K // 32])
|
1521 |
+
signs, scales = np.hsplit(rest, [QK_K // 8])
|
1522 |
+
|
1523 |
+
d = d.view(np.float16).astype(np.float32)
|
1524 |
+
|
1525 |
+
scales = scales.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
|
1526 |
+
scales = (scales & 0x0F).reshape((n_blocks, -1))
|
1527 |
+
db = d * (1 + 2 * scales)
|
1528 |
+
db = db.reshape((n_blocks, -1, 1, 1))
|
1529 |
+
|
1530 |
+
# unpack the sign bits
|
1531 |
+
signs = signs.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8).reshape((1, 1, 8))
|
1532 |
+
signs = signs & np.uint8(0x01)
|
1533 |
+
signs = np.where(signs == 0, np.float32(1), np.float32(-1))
|
1534 |
+
signs = signs.reshape((n_blocks, -1, 4, 8))
|
1535 |
+
|
1536 |
+
qh = qh.reshape((n_blocks, -1, 1)) >> np.array([i for i in range(8)], dtype=np.uint8)
|
1537 |
+
qh = (qh & 0x01).astype(np.uint16).reshape((n_blocks, -1))
|
1538 |
+
qs = qs.astype(np.uint16) | (qh << 8)
|
1539 |
+
|
1540 |
+
assert cls.grid is not None
|
1541 |
+
grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1542 |
+
grid = grid.reshape((n_blocks, -1, 4, 8))
|
1543 |
+
|
1544 |
+
return (db * grid * signs).reshape((n_blocks, -1))
|
1545 |
+
|
1546 |
+
|
1547 |
+
class IQ1_S(__Quant, qtype=GGMLQuantizationType.IQ1_S):
|
1548 |
+
# iq1s_grid, with each byte packed into 2 bits
|
1549 |
+
# -1, 0, 1 <=> 0, 1, 2
|
1550 |
+
grid_shape = (2048, 8)
|
1551 |
+
grid_map = (-1, 0, 1)
|
1552 |
+
grid_hex = (
|
1553 |
+
b"00000200050008000a00110015002000220028002a0045005100540056006500"
|
1554 |
+
b"8000820088008a009500a000a200a800aa000401050111011401160119011a01"
|
1555 |
+
b"2501410146014901520155015a0161016401660168018501910194019601a501"
|
1556 |
+
b"0002020208020a0215022002220228022a024502510259026402690280028202"
|
1557 |
+
b"88028a02910295029902a002a202a802aa021104140416042504410449045504"
|
1558 |
+
b"5a046404650491049904a5040105040505050605150518051a05290540054505"
|
1559 |
+
b"4a0550055105540555055605590560056205650568056a058105910595059805"
|
1560 |
+
b"9a05a105a405a505a605a9051406190641064406500652065506580660066106"
|
1561 |
+
b"6606690685069106940699060008020808080a0815082008220828082a084508"
|
1562 |
+
b"5108560865088008820888088a089508a008a208a808aa080509110914091909"
|
1563 |
+
b"2409250941095009510955096109640969099109940996099909a509000a020a"
|
1564 |
+
b"080a0a0a150a200a220a280a2a0a450a510a590a610a650a800a820a850a880a"
|
1565 |
+
b"8a0a950aa00aa20aa80aaa0a1010111014101910241025104110441050105510"
|
1566 |
+
b"58106110641065106910911094109610a110a510011104110611091110111211"
|
1567 |
+
b"1511181121112411291145114a11501151115211541155115611591160116511"
|
1568 |
+
b"841192119511a111a41111121412161225124012461249125212551258125a12"
|
1569 |
+
b"641266128512911294129612a512011406140914141415141814191421142614"
|
1570 |
+
b"41144514461448144a1451145414551456145914621465146814841489149014"
|
1571 |
+
b"94149514981499149a14a114a414a514a914021505150a151115141515151615"
|
1572 |
+
b"191520152215251528152a154115441545154615511552155415551556155915"
|
1573 |
+
b"5a1561156415651566156915801582158415851588158a159015911594159515"
|
1574 |
+
b"961599159a15a015a215a51501160416051606161516161618161a1621162616"
|
1575 |
+
b"401642164416451648164a165116551656165816591661166416651668166916"
|
1576 |
+
b"6a1686168a1692169516a416a916111816182518411844184618491850185518"
|
1577 |
+
b"58185a1860186118641866186918851891189418a5181019121915191a192119"
|
1578 |
+
b"25194219441945194819511954195519561959195a19601965196a1989199119"
|
1579 |
+
b"921995199819a119a619a919091a161a241a261a441a461a491a501a521a551a"
|
1580 |
+
b"581a611a661a691a851a911a961a9a1a0020022008200a201520202022202520"
|
1581 |
+
b"28202a20452051205920612065208020822088208a209520a020a220a520a820"
|
1582 |
+
b"aa2005211121142119212521422144214921552158215a216121642165216621"
|
1583 |
+
b"8521902196219921a521012208220a22112215222022222228222a2245225122"
|
1584 |
+
b"562259226522812288228a2291229522a022a222a822aa220524142416241924"
|
1585 |
+
b"252444244524462449245224552458245a2466248524912494249924a124a524"
|
1586 |
+
b"0925152521252925402545254825512554255525592562256525682589259025"
|
1587 |
+
b"9425952598259a25a125a425a625a92505261026122619262526412649265526"
|
1588 |
+
b"6026612669268426862690269a260028022808280a2815282028222828282a28"
|
1589 |
+
b"45285128542865288028822888288a28a028a228a828aa280929112914291929"
|
1590 |
+
b"2529462949295229552961296429662969298529902996299929a429a529002a"
|
1591 |
+
b"022a082a0a2a202a222a282a2a2a452a512a562a592a652a802a822a882a8a2a"
|
1592 |
+
b"952aa02aa22aa82aaa2a054011401640254049405240554058405a4061406440"
|
1593 |
+
b"664094409940a140a6400041014104410641094112411541164118411a412141"
|
1594 |
+
b"26412941454148414a41514154415541564159415a41654168416a4181418441"
|
1595 |
+
b"8641904192419541a041a141a241054211421442164225424142524255425a42"
|
1596 |
+
b"6442694289429442a5420144154419442944454448444a445144544455445644"
|
1597 |
+
b"61446244654468446a44814486448944904492449544a044a144a94401450245"
|
1598 |
+
b"05450a4511451445154516451945204525452a45414544454545464549455045"
|
1599 |
+
b"5145544555455645584559456145644565456645694582458445854588459145"
|
1600 |
+
b"94459545964599459a45a545a845aa450146054609461446154618461a462146"
|
1601 |
+
b"2446294640464246454648465046514652465546564659466246654668468146"
|
1602 |
+
b"85468a4694469546a146a446a6460548114815481a4825484248494850485548"
|
1603 |
+
b"5848614864486648694885489148944896489948a5480149054906490a491049"
|
1604 |
+
b"144915491849214924492649404945494a495149524954495549564959496049"
|
1605 |
+
b"6249654966496a49864989499249954996499849a149a449a649a949164a444a"
|
1606 |
+
b"464a494a554a584a5a4a644a694a944aa54a0150045005500650095012501550"
|
1607 |
+
b"1a50215024502950405045504850515054505550565059506550685086508950"
|
1608 |
+
b"95509850a050a150a650a9500551085109510a51115114511551165118511951"
|
1609 |
+
b"20512551265128512a5141514451455146514951505151515251545155515651"
|
1610 |
+
b"585159515a51615164516551665169518251855191519451955196519951a051"
|
1611 |
+
b"a551aa5101520652125215521a5221522452425245524a525152545255525652"
|
1612 |
+
b"595262526552855290529252955299529a52a452045405541154145415541654"
|
1613 |
+
b"185419542154255428542a54415444544554465449544a545054515454545554"
|
1614 |
+
b"5654585459545a54615462546454655466546954805488548a54915494549554"
|
1615 |
+
b"96549954a154a454a554aa540155025504550555065509551055115512551455"
|
1616 |
+
b"1555165519551a55215524552555265529554055415542554455455546554855"
|
1617 |
+
b"4955505551555255545555555655585559555a55605561556455655566556855"
|
1618 |
+
b"69556a5581558455855589558a559055915594559555965598559955a155a455"
|
1619 |
+
b"a555a655a9550056015602560456065608560956115614561556185619562056"
|
1620 |
+
b"2156225624562556265628562956415645564656485649564a56505651565256"
|
1621 |
+
b"545655565656585659565a566156645665566956825685568656885689568a56"
|
1622 |
+
b"915695569a56a256a556a656a856a95604580558065809581058155818582158"
|
1623 |
+
b"2a58455848584a58515854585558565858585958605862586458655882588958"
|
1624 |
+
b"9058925895589858a158a9580159025905590a59115914591559165919592559"
|
1625 |
+
b"41594459455946594959505951595259545955595659585959595a5961596459"
|
1626 |
+
b"655966596959815985598959915994599559965998599959a559045a085a155a"
|
1627 |
+
b"1a5a205a255a265a295a455a485a495a515a555a565a585a595a625a655a685a"
|
1628 |
+
b"6a5a815a8a5a925a955a965a985a9a5aa15a0560146016601960256044605060"
|
1629 |
+
b"5560566058605a60616064606660696081609660a56001610461066109611261"
|
1630 |
+
b"15612161226126612961456149615161556156615961656166616a6184618a61"
|
1631 |
+
b"92619561a161a661a96111621662196240624162466255625662586260628562"
|
1632 |
+
b"91629662a56211641264156416641a6421642664296440644264456448644a64"
|
1633 |
+
b"516454645564566459645a646064626465648464856489649064926494649564"
|
1634 |
+
b"966498649a64a164a464a964056508650a651165156516651965446545654665"
|
1635 |
+
b"496550655165546555655665596561656465656566656965866589658a659165"
|
1636 |
+
b"9565966599659a65a265a565a665a86502660966156620662666286629664066"
|
1637 |
+
b"456648664a66516654665566566658665a666066656668668066826685668a66"
|
1638 |
+
b"9466966698669966a066a466a666aa661668196825684168526855685a686168"
|
1639 |
+
b"6968856891689868a66801690469106915692169246926692969406941694569"
|
1640 |
+
b"4669486951695469556956695969606965696a69826984698a699569a169a469"
|
1641 |
+
b"a569a969116a166a186a416a446a496a506a556a586a5a6a646a656a696a866a"
|
1642 |
+
b"946a986a9a6aa66a0080028008800a802080228028802a804580508051805480"
|
1643 |
+
b"5680598065808080828088808a809580a080a280a880aa800581118114811681"
|
1644 |
+
b"1981258141814481498150815281558156815881598164816681698185818981"
|
1645 |
+
b"948196819981a5810082028208820a8215822082228228822a82518254825982"
|
1646 |
+
b"65828082828288828a829582a082a282a882aa82148419844184448451845584"
|
1647 |
+
b"5a846184648469849484998401850985128515851a8526852985408541854585"
|
1648 |
+
b"4885518554855585568559855a856585668568856a8581858485868589859085"
|
1649 |
+
b"928595859885a68511861686198625864186448649864a865086558659865a86"
|
1650 |
+
b"618666866a86858691869a86a4860088028808880a8815882088228828882a88"
|
1651 |
+
b"41884588518854885988658869888088828888888a889588a088a288a888aa88"
|
1652 |
+
b"05890689118914891689258941894489468949895089528955895a8961896489"
|
1653 |
+
b"858996899989a589008a028a088a0a8a158a208a228a288a2a8a458a518a548a"
|
1654 |
+
b"568a808a828a888a8a8a958aa08aa28aa88aaa8a059011901690189019902590"
|
1655 |
+
b"419046904990559058905a9069906a9085909190949096909990a59001910491"
|
1656 |
+
b"069109911091159118911a912191249126912991409145915091519154915591"
|
1657 |
+
b"569159916291659184918691929195919891a191a491a691a991059211921492"
|
1658 |
+
b"19922592449246924992509252925592589266926992859294929692a9920194"
|
1659 |
+
b"04940694109415941894269440944a9451945494559456945894599460946194"
|
1660 |
+
b"62946594849486949294949495949894a194a9940095059508950a9510951195"
|
1661 |
+
b"14951595169519952195259529952a9541954495459546954995509551955295"
|
1662 |
+
b"549555955695589559955a956195649565956695699581958595889591959295"
|
1663 |
+
b"94959595969599959a95a095a295a595a895aa95019604961096159619962096"
|
1664 |
+
b"2696299645964896499651965296559656965996659668968296849689968a96"
|
1665 |
+
b"929694969596a496a696a9960598169819982598419846985098529855985698"
|
1666 |
+
b"5a98649865988598919896989998a59804990699099910991299159918991a99"
|
1667 |
+
b"209921992499269940994299459948994a995199549955995699599962996599"
|
1668 |
+
b"66996a99819984999099929995999a99a199a699059a159a259a449a469a499a"
|
1669 |
+
b"509a559a589a619a859a919a949a959a969a00a002a008a00aa015a020a022a0"
|
1670 |
+
b"28a02aa045a051a054a056a059a080a082a088a08aa095a0a0a0a2a0a8a0aaa0"
|
1671 |
+
b"05a109a111a114a116a119a11aa146a149a151a155a158a15aa161a164a185a1"
|
1672 |
+
b"90a192a196a199a102a208a20aa210a219a222a228a22aa245a251a256a259a2"
|
1673 |
+
b"65a280a282a288a28aa295a2a0a2a2a2a8a2aaa219a425a441a444a450a454a4"
|
1674 |
+
b"55a458a45aa461a465a466a468a469a485a406a509a510a512a515a518a526a5"
|
1675 |
+
b"29a542a545a551a554a555a556a559a565a56aa581a584a585a586a589a592a5"
|
1676 |
+
b"95a598a505a611a616a61aa621a625a644a646a64aa652a655a656a658a660a6"
|
1677 |
+
b"62a686a690a695a696a699a6a1a6a4a6a6a600a802a808a80aa820a822a828a8"
|
1678 |
+
b"2aa851a854a856a859a880a882a888a88aa895a8a0a8a2a8a8a8aaa805a914a9"
|
1679 |
+
b"19a921a925a941a950a955a95aa961a966a969a990a996a900aa02aa08aa0aaa"
|
1680 |
+
b"20aa22aa28aa2aaa51aa54aa56aa80aa82aa88aa8aaa95aaa0aaa2aaa8aaaaaa"
|
1681 |
+
)
|
1682 |
+
|
1683 |
+
delta = np.float32(0.125)
|
1684 |
+
|
1685 |
+
@classmethod
|
1686 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1687 |
+
n_blocks = blocks.shape[0]
|
1688 |
+
|
1689 |
+
d, rest = np.hsplit(blocks, [2])
|
1690 |
+
qs, qh = np.hsplit(rest, [QK_K // 8])
|
1691 |
+
|
1692 |
+
d = d.view(np.float16).astype(np.float32)
|
1693 |
+
qh = qh.view(np.uint16)
|
1694 |
+
|
1695 |
+
dl = d * (2 * ((qh >> 12) & 7) + 1)
|
1696 |
+
dl = dl.reshape((n_blocks, -1, 1, 1))
|
1697 |
+
delta = np.where((qh & np.uint16(0x8000)) == 0, cls.delta, -cls.delta)
|
1698 |
+
delta = delta.reshape((n_blocks, -1, 1, 1))
|
1699 |
+
|
1700 |
+
qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4))
|
1701 |
+
qs = qs.astype(np.uint16) | ((qh & 7) << 8).reshape((n_blocks, -1))
|
1702 |
+
|
1703 |
+
assert cls.grid is not None
|
1704 |
+
grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1705 |
+
grid = grid.reshape((n_blocks, -1, 4, 8))
|
1706 |
+
|
1707 |
+
return (dl * (grid + delta)).reshape((n_blocks, -1))
|
1708 |
+
|
1709 |
+
|
1710 |
+
class IQ1_M(__Quant, qtype=GGMLQuantizationType.IQ1_M):
|
1711 |
+
grid_shape = IQ1_S.grid_shape
|
1712 |
+
grid_map = IQ1_S.grid_map
|
1713 |
+
grid_hex = IQ1_S.grid_hex
|
1714 |
+
|
1715 |
+
delta = IQ1_S.delta
|
1716 |
+
|
1717 |
+
# Okay *this* type is weird. It's the only one which stores the f16 scales in multiple parts.
|
1718 |
+
@classmethod
|
1719 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1720 |
+
n_blocks = blocks.shape[0]
|
1721 |
+
|
1722 |
+
qs, rest = np.hsplit(blocks, [QK_K // 8])
|
1723 |
+
qh, scales = np.hsplit(rest, [QK_K // 16])
|
1724 |
+
|
1725 |
+
# The f16 scale is packed across multiple bytes
|
1726 |
+
scales = scales.view(np.uint16)
|
1727 |
+
d = (scales.reshape((n_blocks, 4)) & np.uint16(0xF000)) >> np.array([12, 8, 4, 0], dtype=np.uint16).reshape(
|
1728 |
+
(1, 4)
|
1729 |
+
)
|
1730 |
+
d = d[..., 0] | d[..., 1] | d[..., 2] | d[..., 3]
|
1731 |
+
d = d.view(np.float16).astype(np.float32).reshape((n_blocks, 1))
|
1732 |
+
|
1733 |
+
scales = scales.reshape(n_blocks, -1, 1) >> np.array([0, 3, 6, 9], dtype=np.uint16).reshape((1, 1, 4))
|
1734 |
+
scales = (scales & 0x07).reshape((n_blocks, -1))
|
1735 |
+
dl = d * (2 * scales + 1)
|
1736 |
+
dl = dl.reshape((n_blocks, -1, 2, 1, 1))
|
1737 |
+
|
1738 |
+
qh = qh.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
|
1739 |
+
qs = qs.astype(np.uint16) | ((qh & 0x07).astype(np.uint16) << 8).reshape((n_blocks, -1))
|
1740 |
+
|
1741 |
+
delta = np.where(qh & 0x08 == 0, cls.delta, -cls.delta)
|
1742 |
+
delta = delta.reshape((n_blocks, -1, 2, 2, 1))
|
1743 |
+
|
1744 |
+
assert cls.grid is not None
|
1745 |
+
grid = np.take_along_axis(cls.grid, qs.reshape((n_blocks, -1, 1, 1)), axis=-2)
|
1746 |
+
grid = grid.reshape((n_blocks, -1, 2, 2, 8))
|
1747 |
+
|
1748 |
+
return (dl * (grid + delta)).reshape((n_blocks, -1))
|
1749 |
+
|
1750 |
+
|
1751 |
+
class IQ4_NL(__Quant, qtype=GGMLQuantizationType.IQ4_NL):
|
1752 |
+
kvalues = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
|
1753 |
+
|
1754 |
+
@classmethod
|
1755 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1756 |
+
n_blocks = blocks.shape[0]
|
1757 |
+
|
1758 |
+
d, qs = np.hsplit(blocks, [2])
|
1759 |
+
|
1760 |
+
d = d.view(np.float16).astype(np.float32)
|
1761 |
+
|
1762 |
+
qs = qs.reshape((n_blocks, -1, 1, cls.block_size // 2)) >> np.array([0, 4], dtype=np.uint8).reshape(
|
1763 |
+
(1, 1, 2, 1)
|
1764 |
+
)
|
1765 |
+
|
1766 |
+
qs = (qs & np.uint8(0x0F)).reshape((n_blocks, -1, 1))
|
1767 |
+
|
1768 |
+
kvalues = np.array(cls.kvalues, dtype=np.int8).reshape(1, 1, 16)
|
1769 |
+
qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1))
|
1770 |
+
|
1771 |
+
return d * qs
|
1772 |
+
|
1773 |
+
|
1774 |
+
class IQ4_XS(__Quant, qtype=GGMLQuantizationType.IQ4_XS):
|
1775 |
+
@classmethod
|
1776 |
+
def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray:
|
1777 |
+
n_blocks = blocks.shape[0]
|
1778 |
+
|
1779 |
+
d, rest = np.hsplit(blocks, [2])
|
1780 |
+
scales_h, rest = np.hsplit(rest, [2])
|
1781 |
+
scales_l, qs = np.hsplit(rest, [QK_K // 64])
|
1782 |
+
|
1783 |
+
d = d.view(np.float16).astype(np.float32)
|
1784 |
+
scales_h = scales_h.view(np.uint16)
|
1785 |
+
|
1786 |
+
scales_l = scales_l.reshape((n_blocks, -1, 1)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2))
|
1787 |
+
scales_h = scales_h.reshape((n_blocks, 1, -1)) >> np.array(
|
1788 |
+
[2 * i for i in range(QK_K // 32)], dtype=np.uint16
|
1789 |
+
).reshape((1, -1, 1))
|
1790 |
+
scales_l = scales_l.reshape((n_blocks, -1)) & np.uint8(0x0F)
|
1791 |
+
scales_h = scales_h.reshape((n_blocks, -1)).astype(np.uint8) & np.uint8(0x03)
|
1792 |
+
|
1793 |
+
scales = (scales_l | (scales_h << np.uint8(4))).astype(np.int8) - np.int8(32)
|
1794 |
+
dl = (d * scales.astype(np.float32)).reshape((n_blocks, -1, 1))
|
1795 |
+
|
1796 |
+
qs = qs.reshape((n_blocks, -1, 1, 16)) >> np.array([0, 4], dtype=np.uint8).reshape((1, 1, 2, 1))
|
1797 |
+
qs = qs.reshape((n_blocks, -1, 32, 1)) & np.uint8(0x0F)
|
1798 |
+
|
1799 |
+
kvalues = np.array(IQ4_NL.kvalues, dtype=np.int8).reshape((1, 1, 1, -1))
|
1800 |
+
qs = np.take_along_axis(kvalues, qs, axis=-1).astype(np.float32).reshape((n_blocks, -1, 32))
|
1801 |
+
|
1802 |
+
return (dl * qs).reshape((n_blocks, -1))
|
requirements.txt
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
numpy
|
3 |
+
safetensors
|
4 |
+
huggingface-hub
|
5 |
+
gguf
|