#!/usr/bin/env python """Repackage GLM-5.3's native MTP block (blk.45) as a standalone draft GGUF. GLM-5.3 ships a Multi-Token Prediction module inside every GGUF: blk.45, with nextn.eh_proj / enorm / hnorm / shared_head_norm plus a full MLA+DSA attention and a 144-expert MoE. Nothing executes it - transformers drops it on load, llama.cpp loads it and leaves it out of the graph. It is roughly 1.5 GiB of every quant that currently does no work. This lifts it into a model llama.cpp can actually run as a speculative draft. The repackaging is byte-preserving: quantised blocks are copied verbatim, never dequantised and requantised, so the draft is exactly the tensor that shipped in the parent file. Three renames carry the design: blk.45.* -> blk.0.* the MTP block becomes the draft's only layer blk.45.nextn.shared_head_norm -> output_norm it IS the final norm before the shared lm_head (token_embd, output) -> copied MTP has no embedding or head of its own; the reference shares the parent's, so the draft must carry a copy to stand alone The result is ~2.3 GiB at IQ3_M and pairs with exactly one parent quant. It is NOT interchangeable across quants: hidden states drift hard between quantisations (IQ3_M vs IQ4_XS cosine 0.891 against a 0.999 same-quant floor), so a draft built from one parent must be served with that parent. """ import argparse, logging, os, sys from pathlib import Path # Prefer an installed gguf, so this script is usable by anyone who has `pip install gguf`, and fall # back to a checkout beside this tree for the case where it is not installed. GGUF_PY_DIR overrides # both, for a llama.cpp checked out somewhere else entirely. try: import gguf except ImportError: # noqa: E722 _candidates = [] if os.environ.get("GGUF_PY_DIR"): _candidates.append(Path(os.environ["GGUF_PY_DIR"])) _here = Path(__file__).resolve() _candidates += [ _here.parent.parent.parent / "glm5-llama.cpp" / "gguf-py", _here.parent.parent.parent / "llama.cpp" / "gguf-py", _here.parent.parent / "llama.cpp" / "gguf-py", ] for _c in _candidates: if (_c / "gguf").is_dir(): sys.path.insert(0, str(_c)) break else: sys.exit("cannot find the gguf python package - `pip install gguf`, " "or set GGUF_PY_DIR to a llama.cpp gguf-py directory") import gguf # noqa: E402 from tqdm import tqdm # noqa: E402 logger = logging.getLogger("make_mtp_draft") SRC_ARCH = "glm5-next" DST_ARCH = "glm5-next-mtp" # Keys that describe machinery the MTP block does not have. blk.45 carries no hc_* tensors (it is # a plain pre-norm residual block, unlike the 45 hyper-connected layers ahead of it) and no ssm_* # tensors (it is MLA, not KDA), so advertising either would make the loader look for weights that # are not in the file. DROP_SUFFIXES = ( ".attention.hc.mult", ".attention.hc.sinkhorn_iters", ".attention.hc.eps", ".ssm.conv_kernel", ".ssm.gate_lower_bound", ".kda.head_dim", ".nextn_predict_layers", ) def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("src", type=Path, help="parent GGUF containing blk.45") ap.add_argument("dst", type=Path, help="draft GGUF to write") ap.add_argument("--mtp-layer", type=int, default=None, help="source block index of the MTP module (default: block_count-1)") ap.add_argument("--force", action="store_true", help="overwrite dst if it exists") ap.add_argument("--verbose", action="store_true") args = ap.parse_args() logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(levelname)s: %(message)s") if args.dst.exists() and not args.force: logger.error("%s exists; pass --force to overwrite", args.dst) sys.exit(1) reader = gguf.GGUFReader(args.src, "r") arch_f = reader.get_field("general.architecture") arch = arch_f.contents() if arch_f else None if arch != SRC_ARCH: logger.error("expected general.architecture=%s, found %r", SRC_ARCH, arch) sys.exit(1) bc = reader.get_field(f"{SRC_ARCH}.block_count") nextn = reader.get_field(f"{SRC_ARCH}.nextn_predict_layers") n_block = int(bc.contents()) n_nextn = int(nextn.contents()) if nextn else 0 if n_nextn != 1: logger.error("this script handles exactly one MTP layer, file declares %d", n_nextn) sys.exit(1) mtp_il = args.mtp_layer if args.mtp_layer is not None else n_block - 1 src_prefix = f"blk.{mtp_il}." logger.info("source block_count=%d, MTP module at blk.%d", n_block, mtp_il) writer = gguf.GGUFWriter(path=None, arch=DST_ARCH, endianess=reader.endianess) # ---- key/value metadata ------------------------------------------------------------- # Hyperparameters are re-prefixed rather than restated, so anything the parent knows about # its own MLA, MoE and indexer geometry reaches the draft without being retyped here (and # without silently drifting from the parent if the converter ever changes). n_copied = 0 for field in reader.fields.values(): name = field.name if name in ("GGUF.version", "GGUF.tensor_count", "GGUF.kv_count"): continue if name == "general.architecture": continue if any(name == SRC_ARCH + s for s in DROP_SUFFIXES): logger.debug("dropping %s (not present in the MTP block)", name) continue val_type = field.types[0] sub_type = field.types[-1] if val_type == gguf.GGUFValueType.ARRAY else None value = field.contents() if name.startswith(SRC_ARCH + "."): suffix = name[len(SRC_ARCH):] if suffix == ".block_count": value = 1 elif suffix == ".leading_dense_block_count": # The MTP block is a MoE block; there are no dense layers ahead of it here. value = 0 elif suffix == ".attention.head_count_kv": # Per-layer in the parent (0 on KDA layers, 1 on MLA). The draft has one MLA layer. value = [1] sub_type = gguf.GGUFValueType.INT32 name = DST_ARCH + suffix elif name == "general.name": value = str(value) + " MTP Draft" elif name.startswith("quantize.imatrix."): # The parent's imatrix provenance describes the parent, not this file. continue writer.add_key_value(name, value, val_type, sub_type=sub_type) n_copied += 1 writer.add_key_value(f"{DST_ARCH}.mtp.parent_block", mtp_il, gguf.GGUFValueType.UINT32) logger.info("copied %d kv pairs", n_copied) # ---- tensors ------------------------------------------------------------------------ renames: dict[str, str] = {} for t in reader.tensors: if t.name in ("token_embd.weight", "output.weight"): renames[t.name] = t.name elif t.name.startswith(src_prefix): tail = t.name[len(src_prefix):] if tail == "nextn.shared_head_norm.weight": renames[t.name] = "output_norm.weight" else: renames[t.name] = "blk.0." + tail required = { "token_embd.weight", "output.weight", "output_norm.weight", "blk.0.nextn.eh_proj.weight", "blk.0.nextn.enorm.weight", "blk.0.nextn.hnorm.weight", "blk.0.attn_norm.weight", "blk.0.attn_output.weight", "blk.0.ffn_norm.weight", "blk.0.ffn_gate_inp.weight", "blk.0.ffn_down_exps.weight", } produced = set(renames.values()) missing = required - produced if missing: logger.error("source is missing required MTP tensors: %s", sorted(missing)) sys.exit(1) keep = [t for t in reader.tensors if t.name in renames] total = 0 for t in keep: dst_name = renames[t.name] writer.add_tensor_info(dst_name, t.data.shape, t.data.dtype, t.data.nbytes, t.tensor_type) total += t.n_bytes logger.debug("%-46s -> %s", t.name, dst_name) logger.info("writing %d tensors, %.2f GiB to %s", len(keep), total / 2**30, args.dst) tmp = args.dst.with_suffix(args.dst.suffix + ".part") writer.open_output_file(tmp) writer.write_header_to_file() writer.write_kv_data_to_file() writer.write_ti_data_to_file() bar = tqdm(desc="writing", total=total, unit="B", unit_scale=True) for t in keep: writer.write_tensor_data(t.data, tensor_endianess=reader.endianess) bar.update(t.n_bytes) bar.close() writer.close() # Land the file only once it is complete on disk. A draft GGUF truncated by a power cut # would otherwise sit there looking like a valid artifact. os.replace(tmp, args.dst) logger.info("done: %s (%.2f GiB)", args.dst, args.dst.stat().st_size / 2**30) if __name__ == "__main__": main()