389 lines
14 KiB
Python
389 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
"""gguf_retype.py — retype selected tensors of a GGUF in place-of-copy, leaving
|
|
every other tensor byte-identical.
|
|
|
|
Why this exists: `llama-quantize --allow-requantize` rebuilds *every* tensor
|
|
(dequantise -> requantise), so a one-tensor change cannot be isolated from the
|
|
rounding it does to the other 700. Phase 2 of the batched-decode campaign needs
|
|
exactly two tensors moved Q5_K -> Q4_K (`blk.0/blk.1.ffn_down_exps.weight`, the
|
|
only two layers whose `ffn_down_exps` is not a type `gpu_moe_setup` can group)
|
|
with everything else preserved bit-for-bit.
|
|
|
|
How: the GGUF header + KV block + tensor-info block are copied verbatim; only
|
|
the `type` (u32) and `offset` (u64) fields of the info records are patched, both
|
|
fixed-width, so the header length is unchanged. Tensor data is streamed through
|
|
in info order, copied byte-for-byte except for the retyped tensors, which are
|
|
dequantised to f32 and requantised with ggml's own `ggml_quantize_chunk` (the
|
|
same code path `llama-quantize` uses, no imatrix).
|
|
|
|
Requires: libggml-base.so from a llama.cpp build (exports `ggml_quantize_chunk`
|
|
and `dequantize_row_*`). numpy is NOT required.
|
|
|
|
Usage:
|
|
# list tensors (name, type, dims, offset, bytes)
|
|
python3 tools/gguf_retype.py --in model.gguf --dump [--filter ffn_down_exps]
|
|
|
|
# retype
|
|
python3 tools/gguf_retype.py --in model.gguf --out model.retyped.gguf \
|
|
--ggml-lib /path/to/libggml-base.so \
|
|
--retype blk.0.ffn_down_exps.weight=Q4_K \
|
|
--retype blk.1.ffn_down_exps.weight=Q4_K
|
|
|
|
# re-parse the output and diff the tensor table against the input
|
|
python3 tools/gguf_retype.py --in model.gguf --verify model.retyped.gguf
|
|
|
|
@status Working; used for Phase 2 (saber blk.0/blk.1 ffn_down_exps Q5_K -> Q4_K)
|
|
@issues Only the quant types listed in GGML_TYPES below are known to this tool.
|
|
@todo None
|
|
"""
|
|
import argparse
|
|
import ctypes
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
# --- ggml type table: name -> (enum value, block size in elems, bytes/block) --
|
|
GGML_TYPES = {
|
|
"F32": (0, 1, 4),
|
|
"F16": (1, 1, 2),
|
|
"Q4_0": (2, 32, 18),
|
|
"Q4_1": (3, 32, 20),
|
|
"Q5_0": (6, 32, 22),
|
|
"Q5_1": (7, 32, 24),
|
|
"Q8_0": (8, 32, 34),
|
|
"Q8_1": (9, 32, 36),
|
|
"Q2_K": (10, 256, 84),
|
|
"Q3_K": (11, 256, 110),
|
|
"Q4_K": (12, 256, 144),
|
|
"Q5_K": (13, 256, 176),
|
|
"Q6_K": (14, 256, 210),
|
|
"Q8_K": (15, 256, 292),
|
|
"BF16": (30, 1, 2),
|
|
}
|
|
TYPE_BY_ID = {v[0]: k for k, v in GGML_TYPES.items()}
|
|
|
|
# GGUF metadata value types
|
|
(U8, I8, U16, I16, U32, I32, F32V, BOOL, STRING, ARRAY, U64, I64, F64) = range(13)
|
|
_FIXED = {U8: 1, I8: 1, U16: 2, I16: 2, U32: 4, I32: 4, F32V: 4, BOOL: 1,
|
|
U64: 8, I64: 8, F64: 8}
|
|
|
|
|
|
class Reader:
|
|
def __init__(self, buf):
|
|
self.b = buf
|
|
self.p = 0
|
|
|
|
def read_u32(self):
|
|
v = struct.unpack_from("<I", self.b, self.p)[0]; self.p += 4; return v
|
|
|
|
def read_u64(self):
|
|
v = struct.unpack_from("<Q", self.b, self.p)[0]; self.p += 8; return v
|
|
|
|
def read_string(self):
|
|
n = self.read_u64()
|
|
v = self.b[self.p:self.p + n]; self.p += n
|
|
return v.decode("utf-8", "replace")
|
|
|
|
def skip_value(self, t):
|
|
if t in _FIXED:
|
|
self.p += _FIXED[t]
|
|
elif t == STRING:
|
|
# NB: `self.p += self.read_u64()` would drop the length prefix —
|
|
# augmented assignment reads self.p before the RHS advances it.
|
|
n = self.read_u64()
|
|
self.p += n
|
|
elif t == ARRAY:
|
|
et = self.read_u32()
|
|
n = self.read_u64()
|
|
if et in _FIXED:
|
|
self.p += _FIXED[et] * n
|
|
elif et == STRING:
|
|
for _ in range(n):
|
|
ln = self.read_u64()
|
|
self.p += ln
|
|
else:
|
|
raise ValueError(f"nested array type {et}")
|
|
else:
|
|
raise ValueError(f"unknown value type {t}")
|
|
|
|
|
|
class Tensor:
|
|
__slots__ = ("name", "dims", "type", "offset", "nbytes",
|
|
"type_pos", "off_pos", "index")
|
|
|
|
@property
|
|
def nelem(self):
|
|
n = 1
|
|
for d in self.dims:
|
|
n *= d
|
|
return n
|
|
|
|
|
|
def quant_row_bytes(type_name, ne0):
|
|
_, blck, tsz = GGML_TYPES[type_name]
|
|
if ne0 % blck:
|
|
raise ValueError(f"row length {ne0} not a multiple of {type_name} "
|
|
f"block {blck}")
|
|
return ne0 // blck * tsz
|
|
|
|
|
|
def tensor_total_bytes(type_name, dims):
|
|
nrows = 1
|
|
for d in dims[1:]:
|
|
nrows *= d
|
|
return quant_row_bytes(type_name, dims[0]) * nrows
|
|
|
|
|
|
def pad_to_alignment(offset, align):
|
|
return (offset + align - 1) // align * align
|
|
|
|
|
|
def parse_gguf_metadata(path):
|
|
"""Parse header + KV + tensor infos. Returns (header_bytes, tensors,
|
|
data_start, alignment, kv)."""
|
|
with open(path, "rb") as f:
|
|
head = f.read(64 * 1024 * 1024) # metadata is far under this
|
|
if head[:4] != b"GGUF":
|
|
raise ValueError("not a GGUF file")
|
|
r = Reader(head)
|
|
r.p = 4
|
|
version = r.read_u32()
|
|
n_tensors = r.read_u64()
|
|
n_kv = r.read_u64()
|
|
kv = {}
|
|
for _ in range(n_kv):
|
|
key = r.read_string()
|
|
t = r.read_u32()
|
|
vpos = r.p
|
|
if t == U32:
|
|
kv[key] = struct.unpack_from("<I", head, vpos)[0]
|
|
elif t == STRING:
|
|
save = r.p
|
|
kv[key] = r.read_string()
|
|
r.p = save
|
|
r.skip_value(t)
|
|
tensors = []
|
|
for i in range(n_tensors):
|
|
t = Tensor()
|
|
t.index = i
|
|
t.name = r.read_string()
|
|
nd = r.read_u32()
|
|
t.dims = [r.read_u64() for _ in range(nd)]
|
|
t.type_pos = r.p
|
|
tid = r.read_u32()
|
|
t.type = TYPE_BY_ID.get(tid, f"UNK{tid}")
|
|
t.off_pos = r.p
|
|
t.offset = r.read_u64()
|
|
t.nbytes = tensor_total_bytes(t.type, t.dims)
|
|
tensors.append(t)
|
|
align = int(kv.get("general.alignment", 32))
|
|
data_start = pad_to_alignment(r.p, align)
|
|
return head[:data_start], tensors, data_start, align, version, kv
|
|
|
|
|
|
def cmd_dump(args):
|
|
header, tensors, data_start, align, version, kv = parse_gguf_metadata(args.inp)
|
|
print(f"gguf v{version} tensors={len(tensors)} align={align} "
|
|
f"data_start={data_start} file={os.path.getsize(args.inp)}")
|
|
counts = {}
|
|
for t in tensors:
|
|
counts[t.type] = counts.get(t.type, 0) + 1
|
|
if args.filter and args.filter not in t.name:
|
|
continue
|
|
print(f" {t.name:44s} {t.type:5s} dims={t.dims} "
|
|
f"off={t.offset} bytes={t.nbytes}")
|
|
print("type histogram: " + ", ".join(f"{k}={v}" for k, v in
|
|
sorted(counts.items())))
|
|
# self-check: recompute offsets with the standard packing rule
|
|
off = 0
|
|
bad = 0
|
|
for t in tensors:
|
|
if t.offset != off:
|
|
bad += 1
|
|
if bad < 5:
|
|
print(f" OFFSET MISMATCH {t.name}: file={t.offset} calc={off}")
|
|
off = pad_to_alignment(off + t.nbytes, align)
|
|
print(f"offset-model check: {bad} mismatches; "
|
|
f"calc end={data_start + off} actual={os.path.getsize(args.inp)}")
|
|
|
|
|
|
def load_ggml_library(libpath):
|
|
lib = ctypes.CDLL(libpath)
|
|
lib.ggml_quantize_chunk.restype = ctypes.c_size_t
|
|
lib.ggml_quantize_chunk.argtypes = [
|
|
ctypes.c_int, ctypes.POINTER(ctypes.c_float), ctypes.c_void_p,
|
|
ctypes.c_int64, ctypes.c_int64, ctypes.c_int64,
|
|
ctypes.POINTER(ctypes.c_float)]
|
|
return lib
|
|
|
|
|
|
def dequant_row_fn(lib, type_name):
|
|
sym = f"dequantize_row_{type_name.lower()}"
|
|
# ggml keeps the K-quant suffix capitalised: dequantize_row_q5_K
|
|
if type_name.endswith("_K"):
|
|
sym = f"dequantize_row_{type_name[:-2].lower()}_K"
|
|
fn = getattr(lib, sym)
|
|
fn.restype = None
|
|
fn.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_float),
|
|
ctypes.c_int64]
|
|
return fn
|
|
|
|
|
|
def cmd_retype(args):
|
|
header, tensors, data_start, align, version, kv = parse_gguf_metadata(args.inp)
|
|
want = {}
|
|
for spec in args.retype:
|
|
name, _, ty = spec.partition("=")
|
|
if ty not in GGML_TYPES:
|
|
sys.exit(f"unknown target type {ty}")
|
|
want[name] = ty
|
|
by_name = {t.name: t for t in tensors}
|
|
for name in want:
|
|
if name not in by_name:
|
|
sys.exit(f"tensor not in model: {name}")
|
|
lib = load_ggml_library(args.ggml_lib)
|
|
|
|
# 1. rebuild the info block: new types + repacked offsets
|
|
hdr = bytearray(header)
|
|
off = 0
|
|
plan = []
|
|
for t in tensors:
|
|
new_type = want.get(t.name, t.type)
|
|
nbytes = tensor_total_bytes(new_type, t.dims)
|
|
plan.append((t, new_type, off, nbytes))
|
|
struct.pack_into("<I", hdr, t.type_pos, GGML_TYPES[new_type][0])
|
|
struct.pack_into("<Q", hdr, t.off_pos, off)
|
|
off = pad_to_alignment(off + nbytes, align)
|
|
total = data_start + off
|
|
print(f"[retype] {len(want)} tensor(s); out size {total} "
|
|
f"({total / 2**30:.3f} GiB) vs in {os.path.getsize(args.inp)}")
|
|
|
|
# 2. stream the data section
|
|
CHUNK = 8 << 20
|
|
with open(args.inp, "rb") as fi, open(args.out, "wb") as fo:
|
|
fo.write(hdr)
|
|
written = 0
|
|
for t, new_type, new_off, nbytes in plan:
|
|
if written < new_off: # inter-tensor padding
|
|
fo.write(b"\0" * (new_off - written))
|
|
written = new_off
|
|
fi.seek(data_start + t.offset)
|
|
if new_type == t.type:
|
|
left = t.nbytes
|
|
while left:
|
|
b = fi.read(min(CHUNK, left))
|
|
if len(b) != min(CHUNK, left):
|
|
sys.exit(f"short read on {t.name}")
|
|
fo.write(b)
|
|
left -= len(b)
|
|
written += t.nbytes
|
|
else:
|
|
n_per_row = t.dims[0]
|
|
nrows = t.nelem // n_per_row
|
|
src_rs = quant_row_bytes(t.type, n_per_row)
|
|
dst_rs = quant_row_bytes(new_type, n_per_row)
|
|
deq = dequant_row_fn(lib, t.type)
|
|
rows_per_batch = max(1, (4 << 20) // (n_per_row * 4))
|
|
fbuf = (ctypes.c_float * (rows_per_batch * n_per_row))()
|
|
dbuf = ctypes.create_string_buffer(rows_per_batch * dst_rs)
|
|
print(f"[retype] {t.name}: {t.type} -> {new_type}, "
|
|
f"{nrows} rows x {n_per_row} "
|
|
f"({t.nbytes} -> {nbytes} bytes)")
|
|
done = 0
|
|
while done < nrows:
|
|
nb = min(rows_per_batch, nrows - done)
|
|
raw = fi.read(nb * src_rs)
|
|
if len(raw) != nb * src_rs:
|
|
sys.exit(f"short read on {t.name}")
|
|
deq(ctypes.cast(ctypes.c_char_p(raw), ctypes.c_void_p),
|
|
fbuf, nb * n_per_row)
|
|
got = lib.ggml_quantize_chunk(
|
|
GGML_TYPES[new_type][0], fbuf,
|
|
ctypes.cast(dbuf, ctypes.c_void_p),
|
|
0, nb, n_per_row, None)
|
|
if got != nb * dst_rs:
|
|
sys.exit(f"quantize_chunk returned {got}, "
|
|
f"expected {nb * dst_rs}")
|
|
fo.write(dbuf.raw[:nb * dst_rs])
|
|
done += nb
|
|
written += nbytes
|
|
if written < off:
|
|
fo.write(b"\0" * (off - written))
|
|
print(f"[retype] wrote {args.out} ({os.path.getsize(args.out)} bytes)")
|
|
|
|
|
|
def cmd_verify(args):
|
|
ha, ta, dsa, aa, va, _ = parse_gguf_metadata(args.inp)
|
|
hb, tb, dsb, ab, vb, _ = parse_gguf_metadata(args.verify)
|
|
ok = True
|
|
if len(ta) != len(tb) or dsa != dsb or aa != ab or va != vb:
|
|
print(f"HEADER DIFF: n={len(ta)}/{len(tb)} data_start={dsa}/{dsb} "
|
|
f"align={aa}/{ab} ver={va}/{vb}")
|
|
ok = False
|
|
fa = open(args.inp, "rb")
|
|
fb = open(args.verify, "rb")
|
|
changed = 0
|
|
for x, y in zip(ta, tb):
|
|
if x.name != y.name or x.dims != y.dims:
|
|
print(f"TENSOR DIFF: {x.name} vs {y.name}")
|
|
ok = False
|
|
continue
|
|
if x.type != y.type:
|
|
changed += 1
|
|
print(f" retyped {x.name}: {x.type} -> {y.type} "
|
|
f"({x.nbytes} -> {y.nbytes} bytes)")
|
|
continue
|
|
if x.nbytes != y.nbytes:
|
|
print(f"SIZE DIFF: {x.name}")
|
|
ok = False
|
|
continue
|
|
# byte-compare the payload
|
|
fa.seek(dsa + x.offset)
|
|
fb.seek(dsb + y.offset)
|
|
left = x.nbytes
|
|
while left:
|
|
n = min(8 << 20, left)
|
|
if fa.read(n) != fb.read(n):
|
|
print(f"BYTE DIFF in unchanged tensor {x.name}")
|
|
ok = False
|
|
break
|
|
left -= n
|
|
fa.close()
|
|
fb.close()
|
|
# metadata block: identical except the patched type/offset fields
|
|
diffs = sum(1 for i in range(min(len(ha), len(hb))) if ha[i] != hb[i])
|
|
print(f"metadata bytes differing: {diffs} "
|
|
f"(expected: type+offset fields only)")
|
|
print(f"retyped tensors: {changed}")
|
|
print("VERIFY: " + ("PASS" if ok else "FAIL"))
|
|
return 0 if ok else 1
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--in", dest="inp", required=True)
|
|
ap.add_argument("--out")
|
|
ap.add_argument("--dump", action="store_true")
|
|
ap.add_argument("--filter", default=None)
|
|
ap.add_argument("--verify", default=None, help="output file to check")
|
|
ap.add_argument("--retype", action="append", default=[],
|
|
help="name=TYPE, repeatable")
|
|
ap.add_argument("--ggml-lib", default=os.environ.get("GGML_LIB"),
|
|
help="path to libggml-base.so (default: $GGML_LIB)")
|
|
a = ap.parse_args()
|
|
if a.dump:
|
|
return cmd_dump(a)
|
|
if a.verify:
|
|
return cmd_verify(a)
|
|
if not a.out or not a.retype:
|
|
ap.error("need --out and at least one --retype (or --dump/--verify)")
|
|
if not a.ggml_lib:
|
|
ap.error("retype needs --ggml-lib or GGML_LIB (libggml-base.so)")
|
|
return cmd_retype(a)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main() or 0)
|