File size: 1,921 Bytes
f53c54f | 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 | #!/usr/bin/env python3
"""Operator-fusion at the weight level: merge FFN w1 (gate) and w3 (up) into one w13 tensor.
For a Linear weight stored as [in, out] (ggml ne = [in, out]), out-rows are contiguous blocks,
so concatenating raw quantized bytes along the output dimension yields a valid quantized
tensor of [in, out1 + out2]. Output columns [0, out1) equal the original w1 output and
[out1, out1 + out2) equal w3 output, bit-for-bit, because each output element is the same dot product.
"""
import sys
import gguf
import numpy as np
src, dst = sys.argv[1], sys.argv[2]
reader = gguf.GGUFReader(src)
tensors = {t.name: t for t in reader.tensors}
prefixes = sorted({n[: -len("feed_forward.w1.weight")] for n in tensors if n.endswith("feed_forward.w1.weight")})
fused_names = {}
for p in prefixes:
w1, w3 = tensors[p + "feed_forward.w1.weight"], tensors[p + "feed_forward.w3.weight"]
assert w1.tensor_type == w3.tensor_type, p
assert list(w1.shape) == list(w3.shape), p
fused_names[p + "feed_forward.w1.weight"] = p + "feed_forward.w13.weight"
fused_names[p + "feed_forward.w3.weight"] = None # consumed
writer = gguf.GGUFWriter(dst, arch="")
n_fused = 0
for t in reader.tensors:
if t.name in fused_names and fused_names[t.name] is None:
continue
if t.name in fused_names:
p = t.name[: -len("feed_forward.w1.weight")]
w1 = tensors[p + "feed_forward.w1.weight"]
w3 = tensors[p + "feed_forward.w3.weight"]
data = np.concatenate([np.asarray(w1.data), np.asarray(w3.data)], axis=0) # [out1+out2, row_bytes]
writer.add_tensor(fused_names[t.name], data, raw_dtype=t.tensor_type)
n_fused += 1
continue
writer.add_tensor(t.name, np.asarray(t.data), raw_dtype=t.tensor_type)
writer.write_header_to_file()
writer.write_kv_data_to_file()
writer.write_tensors_to_file()
writer.close()
print(f"fused {n_fused} FFN pairs -> {dst}")
|