vh/esh-pfi-infrastructure
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

fix(post-quant): handle sources that keep mtp.* inside a numbered shard

Browse Source 
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.
This commit is contained in:
vh
2026-08-17 17:18:01 -07:00
parent b0c2d3d1c4
commit 993421bf59
1 changed files with 32 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
services/gen-seat-mixed-quant/post_quant.py
+32 -6
View File
@@ -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")]