File size: 7,011 Bytes
491cce7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | #!/usr/bin/env python3
"""Merge MTP head (safetensors) into BF16 GGUF — raw binary header rewrite."""
import argparse
import os, struct, shutil, gguf
import ml_dtypes
import numpy as np
from safetensors import safe_open
LLAMA_GGUF = os.environ.get(
"LLAMA_GGUF", "/home/wepiqx/llama.cpp/build/bin/llama-gguf"
)
NAME_MAP = {
"mtp.pre_fc_norm_embedding.weight": "blk.32.pre_fc_norm_embedding.weight",
"mtp.pre_fc_norm_hidden.weight": "blk.32.pre_fc_norm_hidden.weight",
"mtp.fc.weight": "blk.32.fc.weight",
"mtp.layers.0.input_layernorm.weight": "blk.32.input_layernorm.weight",
"mtp.layers.0.self_attn.q_proj.weight": "blk.32.attn_q.weight",
"mtp.layers.0.self_attn.k_proj.weight": "blk.32.attn_k.weight",
"mtp.layers.0.self_attn.v_proj.weight": "blk.32.attn_v.weight",
"mtp.layers.0.self_attn.o_proj.weight": "blk.32.attn_output.weight",
"mtp.layers.0.self_attn.q_norm.weight": "blk.32.self_attn_q_norm.weight",
"mtp.layers.0.self_attn.k_norm.weight": "blk.32.self_attn_k_norm.weight",
"mtp.layers.0.post_attention_layernorm.weight": "blk.32.post_attention_layernorm.weight",
"mtp.layers.0.mlp.gate_proj.weight": "blk.32.ffn_gate.weight",
"mtp.layers.0.mlp.up_proj.weight": "blk.32.ffn_up.weight",
"mtp.layers.0.mlp.down_proj.weight": "blk.32.ffn_down.weight",
"mtp.norm.weight": "blk.32.norm.weight",
}
def pack_str(s):
b = s.encode()
return struct.pack('<Q', len(b)) + b
def main():
parser = argparse.ArgumentParser(description="Merge MTP safetensors into BF16 GGUF")
parser.add_argument("--bf16", required=True, help="Source BF16 GGUF path")
parser.add_argument("--mtp", required=True, help="MTP head safetensors path")
parser.add_argument("--out", required=True, help="Output GGUF path")
args = parser.parse_args()
BF16_PATH, MTP_PATH, OUT_PATH = args.bf16, args.mtp, args.out
# --- 1. Read BF16 header bytes ---
print("Reading BF16 header...")
r = gguf.GGUFReader(BF16_PATH)
data_offset = r.data_offset
n_tensors = len(r.tensors)
print(f"Tensors: {n_tensors}, data offset: {data_offset}")
with open(BF16_PATH, 'rb') as f:
hdr = f.read(data_offset)
# --- 2. Parse KVs (from byte 24 to where TIs start) ---
pos = 24
orig_kv_count = struct.unpack('<Q', hdr[16:24])[0]
for _ in range(orig_kv_count):
klen = struct.unpack('<Q', hdr[pos:pos+8])[0]
pos += 8 + klen
vtype = struct.unpack('<I', hdr[pos:pos+4])[0]; pos += 4
if vtype == 8:
sl = struct.unpack('<Q', hdr[pos:pos+8])[0]; pos += 8 + sl
elif vtype == 9:
et = struct.unpack('<I', hdr[pos:pos+4])[0]; pos += 4
ct = struct.unpack('<Q', hdr[pos:pos+8])[0]; pos += 8
for _ in range(ct):
if et == 8:
sl = struct.unpack('<Q', hdr[pos:pos+8])[0]; pos += 8 + sl
elif et in (4,5): pos += 4
elif et == 6: pos += 4
else: pos += 4 # unknown elem, skip 4
elif vtype in (0,1,2,3,4,5,6,12): pos += 4
elif vtype in (7,10,11): pos += 8
else:
print(f"UNKNOWN vtype {vtype} at {pos-4}")
break
kv_raw_len = pos - 24
ti_start = pos # TIs start right after KVs, no padding
print(f"KV raw: {kv_raw_len} bytes, TI at: {ti_start}")
# --- 3. Parse existing TIs (between ti_start and data_offset) ---
ti_data = hdr[ti_start:data_offset]
ti_data_len = len(ti_data)
print(f"TI section: {ti_data_len} bytes")
# Find where raw TIs end (before padding)
pos2 = 0
for i in range(n_tensors):
nl = struct.unpack('<Q', ti_data[pos2:pos2+8])[0]
pos2 += 8 + nl
nd = struct.unpack('<I', ti_data[pos2:pos2+4])[0]; pos2 += 4
pos2 += 8 * nd
pos2 += 4 + 8 # type + offset
ti_raw_end = pos2
ti_pad = ti_data_len - ti_raw_end
print(f"TIs: {n_tensors} parsed, {ti_raw_end} raw + {ti_pad} pad")
# --- 4. Load MTP tensors ---
print("Loading MTP tensors...")
mtp_entries = [] # (dst_name, dims, f16_data_bytes)
with safe_open(MTP_PATH, framework='np') as f:
for src_name, dst_name in NAME_MAP.items():
arr = f.get_tensor(src_name)
data = arr.astype(np.float32).astype(np.float16).tobytes()
mtp_entries.append((dst_name, list(arr.shape), data))
# --- 5. Build MTP TIs ---
bf16_data_sz = os.path.getsize(BF16_PATH) - data_offset
mtp_offset = bf16_data_sz
mtp_ti_bytes = b''
for name, dims, data in mtp_entries:
ti = pack_str(name)
ti += struct.pack('<I', len(dims))
for d in dims: ti += struct.pack('<Q', d)
ti += struct.pack('<I', 1) # F16
ti += struct.pack('<Q', mtp_offset)
mtp_ti_bytes += ti
dpad = (32 - len(data) % 32) % 32
mtp_offset += len(data) + dpad
# --- 6. Build new header ---
raw_kv = hdr[24:24+kv_raw_len] # original KVs
nextn_kv = pack_str("nextn_predict_layers")
nextn_kv += struct.pack('<I', 4) # UINT32
nextn_kv += struct.pack('<I', 1)
new_kv = raw_kv + nextn_kv
new_ti = ti_data[:ti_raw_end] + mtp_ti_bytes # old raw TIs + MTP TIs
# Pad TIs so data section aligns to 32 bytes from file start
total_before_ti = 24 + len(new_kv)
new_ti_pad = (32 - (total_before_ti + len(new_ti)) % 32) % 32
new_ti += b'\x00' * new_ti_pad
new_tc = n_tensors + len(mtp_entries)
new_kvc = orig_kv_count + 1
new_hdr = b'GGUF'
new_hdr += struct.pack('<I', 3)
new_hdr += struct.pack('<Q', new_tc)
new_hdr += struct.pack('<Q', new_kvc)
new_hdr += new_kv + new_ti
print(f"New header: {len(new_hdr)} bytes")
print(f"TC: {new_tc}, KVC: {new_kvc}")
print(f"New data section at: {len(new_hdr)}")
# --- 7. Write output ---
with open(OUT_PATH, 'wb') as fout:
fout.write(new_hdr)
print("Copying BF16 data...")
with open(BF16_PATH, 'rb') as fb:
fb.seek(data_offset)
shutil.copyfileobj(fb, fout)
print("Writing MTP tensors...")
for name, dims, data in mtp_entries:
fout.write(data)
pad = (32 - len(data) % 32) % 32
fout.write(b'\x00' * pad)
print(f" {name}: {len(data)} bytes")
sz = os.path.getsize(OUT_PATH)
print(f"\nDone! {sz} bytes ({sz/1024/1024/1024:.2f} GB)")
# Verify
import subprocess
res = subprocess.run(
[LLAMA_GGUF, OUT_PATH, "r"],
capture_output=True, text=True, timeout=30)
for line in res.stdout.split('\n')[:25]:
if any(x in line.lower() for x in ['blk.32', 'total', 'tensor', 'arch', 'name', 'layer']):
print(line)
if res.returncode != 0:
print(f"VERIFY FAIL: {res.stderr[:300]}")
if __name__ == "__main__":
main()
|