| |
| """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 |
|
|
| 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) |
| 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}") |
|
|
