From 993421bf598b4c5668b7c2d2231c123e48ab7cf0 Mon Sep 17 00:00:00 2001 From: Vuong Hoang Date: Mon, 17 Aug 2026 17:18:01 -0700 Subject: [PATCH] fix(post-quant): handle sources that keep mtp.* inside a numbered shard post_quant assumed the source ships a standalone model-mtp.safetensors, which is how JonathanColetti's grafted head is packaged. MuXodious/absolute-heresy is an unmodified full checkpoint, so its mtp.* lives in model-00012-of-00012 -- the copy silently did nothing while the index was still rewritten to point at model-mtp.safetensors, leaving 15 unresolvable tensors. Tensor counts looked correct; the checkpoint would have failed at load. The existing FAILED-CHECKS assertion caught it, which is the design working. Now extracts from the numbered shard when the standalone file is absent. Verified on the heresy build: 1968 tensors, all resolvable, 15 mtp, 333 visual, no missing shards, no orphans. --- services/gen-seat-mixed-quant/post_quant.py | 38 +++++++++++++++++---- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/services/gen-seat-mixed-quant/post_quant.py b/services/gen-seat-mixed-quant/post_quant.py index 4b39277..b6ed0d5 100644 --- a/services/gen-seat-mixed-quant/post_quant.py +++ b/services/gen-seat-mixed-quant/post_quant.py @@ -20,16 +20,42 @@ def main(): fail = [] # --- 1. MTP graft --------------------------------------------------------- + # Two source layouts exist in the wild and both must work: + # (a) a standalone `model-mtp.safetensors` -- how JonathanColetti ships its + # grafted head, so a plain file copy suffices; + # (b) mtp.* living inside a NUMBERED shard -- how MuXodious/absolute-heresy + # ships (model-00012-of-00012.safetensors), because it is an unmodified + # full checkpoint rather than a graft. + # Handling only (a) leaves the output index pointing at a `model-mtp.safetensors` + # that was never created: the checkpoint looks fine to a tensor count but every + # mtp tensor is unresolvable at load. Extract instead of copy for (b). mtp_src = os.path.join(src, "model-mtp.safetensors") mtp_dst = os.path.join(out, "model-mtp.safetensors") - if not os.path.exists(mtp_src): - fail.append(f"missing MTP shard at {mtp_src}") + if os.path.exists(mtp_dst): + print("MTP shard already present in output") + elif os.path.exists(mtp_src): + print(f"copying MTP shard ({os.path.getsize(mtp_src)/1e9:.2f} GB) ...", flush=True) + shutil.copy2(mtp_src, mtp_dst) else: - if not os.path.exists(mtp_dst): - print(f"copying MTP shard ({os.path.getsize(mtp_src)/1e9:.2f} GB) ...", flush=True) - shutil.copy2(mtp_src, mtp_dst) + # layout (b): materialise the standalone shard the index will reference + idx_path = os.path.join(src, "model.safetensors.index.json") + wm = json.load(open(idx_path))["weight_map"] + shards = sorted({wm[k] for k in wm if k.startswith("mtp")}) + if not shards: + fail.append(f"no mtp.* in {src} (neither model-mtp.safetensors nor any shard)") else: - print("MTP shard already present") + from safetensors import safe_open + from safetensors.torch import save_file + print(f"extracting mtp.* from {shards} -> model-mtp.safetensors ...", flush=True) + tensors = {} + for shard in shards: + with safe_open(os.path.join(src, shard), framework="pt") as f: + for k in f.keys(): + if k.startswith("mtp"): + tensors[k] = f.get_tensor(k) + save_file(tensors, mtp_dst, metadata={"format": "pt"}) + print(f" wrote {len(tensors)} tensors, " + f"{os.path.getsize(mtp_dst)/1e6:.1f} MB") src_idx = json.load(open(os.path.join(src, "model.safetensors.index.json"))) mtp_keys = [k for k in src_idx["weight_map"] if k.startswith("mtp")]