add convert_minicpm_mixed.py

convert_minicpm_mixed.py

@@ -0,0 +1,247 @@
+"""convert_minicpm_mixed.py - Mixed-precision converter: MiniCPM5-2B bf16 safetensors -> LAL GPQ8.
+
+Mixed-precision scheme (user requirement: 4-bit main body, part 8-bit, part 3-value):
+  - embed_tokens / lm_head        -> Q8 per-row (qtype=1)  [8-bit, quality-critical]
+  - layernorms / model.norm       -> F32 (qtype=0)
+  - layers 0 and 41 (edges)       -> Q8_0 (qtype=3)        [8-bit, sensitive first/last layers]
+  - layers 1..40 q/k/v/o/down     -> Q4_K (qtype=5)        [4-bit main body]
+  - layers 1..7,32..40 gate/up    -> Q4_K (qtype=5)
+  - layers 8..31 gate/up          -> TERNARY 3-value (qtype=6)  {-s,0,+s}, 2-bit codes
+
+Formats:
+  Q4_K    : 144B / 256 elems (fp16 d + fp16 dmin + 12B 6-bit packed scales/mins + 128B qs)
+  Q8_0    : 34B  / 32  elems (fp16 scale + 32 int8)
+  Q8_row  : int8 + per-row fp32 scale (stored in trailing scale array)
+  TERNARY : 18B  / 64  elems (fp16 scale + 16B codes: 8B bit0 stream + 8B bit1 stream)
+            code t in {-1,0,+1}: bit0=(t==+1), bit1=(t==-1); kernel computes (S0-S1)
+"""
+import struct, json, os, sys, time
+import numpy as np
+
+SRC = 'C:/models/model-00000-of-00001.safetensors'
+OUT = 'C:/models/minicpm_mixed.bin'
+LOG = 'C:/models/convert_status.log'
+TERN_COS_MIN = 0.8970  # per-matrix cosine gate: below this, fall back to Q4_K
+
+TERN_RANGE = (8, 31)     # inclusive layer range for 3-value gate/up candidates
+EDGE_LAYERS = (0, 41)    # layers fully Q8_0 (8-bit)
+n_tern_stats = []
+# top-8 candidates by measured cosine (highest fidelity 3-value matrices):
+TERN_KEEP = {
+    'model.layers.31.mlp.up_proj.weight',   # cos=0.8989
+    'model.layers.10.mlp.up_proj.weight',   # cos=0.8986
+    'model.layers.11.mlp.up_proj.weight',   # cos=0.8984
+    'model.layers.10.mlp.gate_proj.weight', # cos=0.8981
+    'model.layers.19.mlp.up_proj.weight',   # cos=0.8980
+    'model.layers.30.mlp.up_proj.weight',   # cos=0.8979
+    'model.layers.31.mlp.gate_proj.weight', # cos=0.8979
+    'model.layers.18.mlp.up_proj.weight',   # cos=0.8978
+}
+
+log = open(LOG, 'w', buffering=1)
+def L(msg):
+    log.write(msg + '\n')
+
+# ---------- dtype conversions ----------
+def to_f32(raw, dtype):
+    if dtype == 'BF16':
+        u16 = np.frombuffer(raw, dtype=np.uint16)
+        return (u16.astype(np.uint32) << 16).view(np.float32)
+    if dtype == 'F32':
+        return np.frombuffer(raw, dtype=np.float32).copy()
+    if dtype == 'F16':
+        return np.frombuffer(raw, dtype=np.float16).astype(np.float32)
+    raise ValueError('unsupported dtype ' + dtype)
+
+# ---------- quantizers ----------
+def q8_per_row(w2d):
+    """Q8 per-row: int8 + per-row fp32 scale -> (bytes, scales)"""
+    w = w2d.astype(np.float32)
+    s = np.abs(w).max(axis=1) / 127.0
+    s = np.where(s < 1e-12, 1e-12, s).astype(np.float32)
+    q = np.clip(np.round(w / s[:, None]), -127, 127).astype(np.int8)
+    return q.tobytes(), s.tobytes()
+
+def q8_0(w2d):
+    """Q8_0 blocks: per-32 fp16 scale + 32 int8 -> 34B/32elem"""
+    w = w2d.astype(np.float32)
+    out_dim, in_dim = w.shape
+    assert in_dim % 32 == 0
+    nb = in_dim // 32
+    wb = w.reshape(out_dim, nb, 32)
+    amax = np.abs(wb).max(axis=2) / 127.0
+    amax = np.where(amax < 1e-12, 1e-12, amax).astype(np.float32)
+    q = np.clip(np.round(wb / amax[:, :, None]), -127, 127).astype(np.int8)
+    blocks = np.zeros((out_dim, nb, 34), dtype=np.uint8)
+    h16 = amax.astype(np.float16).view(np.uint8).reshape(out_dim, nb, 2)
+    blocks[:, :, 0:2] = h16
+    blocks[:, :, 2:34] = q.view(np.uint8).reshape(out_dim, nb, 32)
+    return blocks.tobytes()
+
+def q4_k(w2d):
+    """Q4_K superblocks: 144B/256elem (llama.cpp-compatible, matches lal kernel)"""
+    w = w2d.astype(np.float32)
+    out_dim, in_dim = w.shape
+    assert in_dim % 256 == 0
+    n_super = in_dim // 256
+    wb = w.reshape(out_dim, n_super, 8, 32)
+    sub_min = wb.min(axis=3)
+    sub_max = wb.max(axis=3)
+    sub_scale = (sub_max - sub_min) / 15.0
+    sub_scale = np.where(sub_scale < 1e-8, 1e-8, sub_scale)
+    max_scale = np.where(sub_scale.max(axis=2) < 1e-8, 1e-8, sub_scale.max(axis=2)).astype(np.float32)
+    max_min = np.where(np.abs(sub_min).max(axis=2) < 1e-8, 1e-8, np.abs(sub_min).max(axis=2)).astype(np.float32)
+    d, dmin = max_scale, max_min
+    sc6 = np.clip(np.round(sub_scale / d[:, :, None] * 63), 0, 63).astype(np.uint8)
+    m6 = np.clip(np.round(np.abs(sub_min) / dmin[:, :, None] * 63), 0, 63).astype(np.uint8)
+    act_s = d[:, :, None] * sc6.astype(np.float32) / 63.0
+    act_m = dmin[:, :, None] * m6.astype(np.float32) / 63.0
+    q = np.clip(np.round((wb + act_m[:, :, :, None]) / (act_s[:, :, :, None] + 1e-8)), 0, 15).astype(np.uint8)
+    qf = q.reshape(out_dim, n_super, 256)
+    qs = np.zeros((out_dim, n_super, 128), dtype=np.uint8)
+    for sub in range(8):
+        for i in range(16):
+            qs[:, :, sub * 16 + i] = qf[:, :, sub * 32 + 2 * i] | (qf[:, :, sub * 32 + 2 * i + 1] << 4)
+    # 16 x 6-bit values (8 scales + 8 mins) -> continuous 96-bit little-endian
+    # bitstream -> 12 bytes (llama.cpp Q4_K packing).
+    flat = np.zeros((out_dim * n_super, 16), dtype=np.uint64)
+    flat[:, 0:8] = sc6.reshape(-1, 8)
+    flat[:, 8:16] = m6.reshape(-1, 8)
+    outb = np.zeros((flat.shape[0], 12), dtype=np.uint8)
+    for j in range(16):
+        v = flat[:, j]
+        bp = j * 6
+        for b in range(6):
+            bit = bp + b
+            byte_i, bit_i = bit // 8, bit % 8
+            outb[:, byte_i] |= (((v >> np.uint64(b)) & np.uint64(1)) << np.uint64(bit_i)).astype(np.uint8)
+    blocks = np.zeros((out_dim, n_super, 144), dtype=np.uint8)
+    blocks[:, :, 0:2] = d.astype(np.float16).view(np.uint8).reshape(out_dim, n_super, 2)
+    blocks[:, :, 2:4] = dmin.astype(np.float16).view(np.uint8).reshape(out_dim, n_super, 2)
+    blocks[:, :, 4:16] = outb.reshape(out_dim, n_super, 12)
+    blocks[:, :, 16:144] = qs
+    return blocks.tobytes()
+
+def ternary(w2d):
+    """3-value blocks: 18B/64elem = fp16 scale + 8B bit0 + 8B bit1.
+    Structure: TWN threshold delta = 0.7*mean|w| picks the +/-1 support.
+    Scale: per-block least-squares s* = sum(w*t)/sum(t*t), the minimizer of
+    ||w - s*t||^2 -- restores the energy that plain absmean destroys (PTQ)."""
+    w = w2d.astype(np.float32)
+    out_dim, in_dim = w.shape
+    assert in_dim % 64 == 0
+    nb = in_dim // 64
+    wb = w.reshape(out_dim, nb, 64)
+    aabs = np.abs(wb)
+    delta = 0.7 * aabs.mean(axis=2, keepdims=True)
+    t = np.where(wb > delta, 1, np.where(wb < -delta, -1, 0)).astype(np.float32)
+    denom = (t * t).sum(axis=2, keepdims=True)
+    s = np.where(denom > 0, (wb * t).sum(axis=2, keepdims=True) / np.maximum(denom, 1.0), 0.0)
+    s = s.astype(np.float32)
+    ti = t.astype(np.int8)
+    bit0 = (ti == 1).astype(np.uint8).reshape(out_dim, nb, 64)
+    bit1 = (ti == -1).astype(np.uint8).reshape(out_dim, nb, 64)
+    blocks = np.zeros((out_dim, nb, 18), dtype=np.uint8)
+    blocks[:, :, 0:2] = s.astype(np.float16).view(np.uint8).reshape(out_dim, nb, 2)
+    blocks[:, :, 2:10] = np.packbits(bit0, axis=2, bitorder='little').reshape(out_dim, nb, 8)
+    blocks[:, :, 10:18] = np.packbits(bit1, axis=2, bitorder='little').reshape(out_dim, nb, 8)
+    return blocks.tobytes()
+
+# ---------- main ----------
+def qtype_for(key, shape):
+    if key == 'model.embed_tokens.weight' or key == 'lm_head.weight':
+        return 1
+    if key == 'model.norm.weight' or 'layernorm' in key:
+        return 0
+    import re
+    m = re.search(r'layers\.(\d+)\.', key)
+    li = int(m.group(1))
+    if li in EDGE_LAYERS:
+        return 3
+    if ('gate_proj' in key or 'up_proj' in key) and TERN_RANGE[0] <= li <= TERN_RANGE[1] and key in TERN_KEEP:
+        return 66  # candidate: decide after cosine test
+    return 5
+
+def tern_cos(w2d):
+    """weight-space cosine between bf16 W and its 3-value reconstruction"""
+    out_dim, in_dim = w2d.shape
+    blocks = ternary(w2d)
+    nb = in_dim // 64
+    b = np.frombuffer(blocks, dtype=np.uint8).reshape(out_dim, nb, 18)
+    s = b[:, :, 0:2].copy().view(np.float16).astype(np.float32).reshape(out_dim, nb, 1)
+    def expand(ab):
+        bits = np.zeros((ab.shape[0], ab.shape[1], 64), dtype=np.int8)
+        for bi in range(8):
+            for bit in range(8):
+                bits[:, :, bi * 8 + bit] = ((ab[:, :, bi] >> bit) & 1).astype(np.int8)
+        return bits
+    T = (expand(b[:, :, 2:10]) - expand(b[:, :, 10:18])).astype(np.float32)
+    Wh = (T * s).reshape(out_dim, in_dim)
+    W = w2d.reshape(out_dim, in_dim)
+    return float((W * Wh).sum() / (np.sqrt((W * W).sum()) * np.sqrt((Wh * Wh).sum()) + 1e-12))
+
+# ---------- main ----------
+def main():
+    t0 = time.time()
+    with open(SRC, 'rb') as f:
+        hlen = struct.unpack('<Q', f.read(8))[0]
+        hdr = json.loads(f.read(hlen).decode('utf-8'))
+        data_start = 8 + hlen
+    meta = [(k, v) for k, v in hdr.items() if k != '__metadata__']
+    L('TENSORS %d' % len(meta))
+    stats = {}
+    out = open(OUT, 'wb')
+    out.write(b'GPQ8')
+    out.write(struct.pack('<I', len(meta)))
+    src = open(SRC, 'rb')
+    for idx, (key, m) in enumerate(meta):
+        shape, dtype, off = m['shape'], m['dtype'], m['data_offsets']
+        src.seek(data_start + off[0])
+        raw = src.read(off[1] - off[0])
+        f32 = to_f32(raw, dtype)
+        qt = qtype_for(key, shape)
+        scales_payload = b''
+        if len(shape) == 2:
+            w2d = f32.reshape(shape[0], shape[1])
+        else:
+            w2d = f32.reshape(shape)
+        if qt == 66:
+            c = tern_cos(w2d)
+            qt = 6 if c >= TERN_COS_MIN else 5
+            n_tern_stats.append((key, c))
+        if qt == 0:
+            data, n_scale = f32.astype(np.float32).tobytes(), 0
+        elif qt == 1:
+            data, sc = q8_per_row(w2d)
+            n_scale, scales_payload = shape[0], sc  # per-row fp32 scales
+        elif qt == 3:
+            data, n_scale = q8_0(w2d), 0
+        elif qt == 5:
+            data, n_scale = q4_k(w2d), 0
+        elif qt == 6:
+            data, n_scale = ternary(w2d), 0
+        kb = key.encode('utf-8')
+        out.write(struct.pack('<I', len(kb))); out.write(kb)
+        out.write(struct.pack('<I', len(shape)))
+        for d in shape: out.write(struct.pack('<I', d))
+        out.write(struct.pack('<B', qt))
+        out.write(struct.pack('<Q', len(data)))
+        out.write(data)
+        out.write(struct.pack('<I', n_scale))
+        if scales_payload:
+            out.write(scales_payload)
+        stats[qt] = stats.get(qt, 0) + len(data)
+        if idx % 20 == 0:
+            L('[%d/%d] %s qtype=%d %.1fMB %.0fs' % (idx, len(meta), key, qt, len(data) / 1048576, time.time() - t0))
+    out.close(); src.close()
+    total = os.path.getsize(OUT)
+    L('DONE total=%.2f GB (%.0f MB f32-eq saved) in %.0fs' % (total / 1073741824, total / 1048576, time.time() - t0))
+    for key, c in n_tern_stats:
+        L('TERN %s cos=%.4f' % (key, c))
+    kept = sum(1 for _, c in n_tern_stats if c >= TERN_COS_MIN)
+    L('TERN_SUMMARY candidates=%d kept=%d fell_back=%d' % (len(n_tern_stats), kept, len(n_tern_stats) - kept))
+    for qt in sorted(stats):
+        L('  qtype %d: %.1f MB' % (qt, stats[qt] / 1048576))
+
+if __name__ == '__main__':
+    main()