Upload merge_mtp.py with huggingface_hub

merge_mtp.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""Merge the AngelSpec MTP head into a base HYV3 model.
+
+The base model (e.g. ``Hy3_opensource``) already has the full transformer body
+(layers 0-79 + embeddings + lm_head + final norm) *and* its own MTP head at
+``model.layers.80.*``. This script produces a NEW model directory that is a copy
+of the base, but whose MTP head (layer 80) is replaced by the AngelSpec MTP
+weights exported in ``mtp.safetensors``.
+
+Result = base body  +  AngelSpec MTP head.
+
+Only the 5 shards that contain layer-80 tensors are rewritten. The other 94
+shards, together with the config / tokenizer / index files, are linked (or
+copied) verbatim from the base model, so the merge is fast and cheap.
+
+Usage:
+    python merge_mtp.py \
+        --base ORIGINAL_MODEL_PATH \
+        --mtp  MTP_WEIGHT_PATH \
+        --out  OUTPUT_WEIGHT_PATH \
+        --link copy        # symlink | hardlink | copy  (how to bring over the 94 unchanged shards)
+"""
+import argparse
+import json
+import os
+import shutil
+
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+MTP_PREFIX = "model.layers.80."
+
+
+def bring_over(src, dst, mode):
+    """Materialise an unchanged shard in the output dir according to *mode*."""
+    if os.path.lexists(dst):
+        os.remove(dst)
+    if mode == "symlink":
+        os.symlink(os.path.abspath(src), dst)
+    elif mode == "hardlink":
+        os.link(src, dst)
+    elif mode == "copy":
+        shutil.copy2(src, dst)
+    else:
+        raise ValueError(f"unknown link mode: {mode}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--base", required=True, help="Base full model directory (body comes from here)")
+    ap.add_argument("--mtp", required=True, help="MTP export dir (contains mtp.safetensors)")
+    ap.add_argument("--out", required=True, help="Output directory for the merged model")
+    ap.add_argument("--link", default="symlink", choices=["symlink", "hardlink", "copy"],
+                    help="How to bring over the 94 unchanged shards (default: symlink)")
+    ap.add_argument("--prefix", default=MTP_PREFIX, help="Tensor prefix of the MTP layer")
+    args = ap.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+
+    # --- 1. Load the base index and locate the MTP (layer-80) tensors ----------
+    base_index_path = os.path.join(args.base, "model.safetensors.index.json")
+    base_index = json.load(open(base_index_path))
+    weight_map = base_index["weight_map"]
+
+    mtp_keys = sorted(k for k in weight_map if k.startswith(args.prefix))
+    if not mtp_keys:
+        raise SystemExit(f"No {args.prefix!r} tensors found in {base_index_path}")
+
+    mtp_shards = sorted({weight_map[k] for k in mtp_keys})
+    all_shards = sorted(set(weight_map.values()))
+    print(f"MTP layer has {len(mtp_keys)} tensors across {len(mtp_shards)} shards.")
+    print(f"Base model has {len(all_shards)} shards total; "
+          f"{len(all_shards) - len(mtp_shards)} will be linked verbatim.")
+
+    # --- 2. Load replacement MTP tensors from the export -----------------------
+    mtp_file = os.path.join(args.mtp, "mtp.safetensors")
+    new_mtp = {}
+    with safe_open(mtp_file, framework="pt", device="cpu") as f:
+        exported = set(f.keys())
+        missing = set(mtp_keys) - exported
+        if missing:
+            raise SystemExit(f"{len(missing)} MTP tensors missing from {mtp_file}, "
+                             f"e.g. {sorted(missing)[:3]}")
+        for k in mtp_keys:
+            new_mtp[k] = f.get_tensor(k)
+
+    # Sanity check: shapes/dtypes must match what the base expects.
+    for shard in mtp_shards:
+        with safe_open(os.path.join(args.base, shard), framework="pt", device="cpu") as f:
+            for k in [k for k in mtp_keys if weight_map[k] == shard]:
+                bt = f.get_slice(k)
+                bshape = tuple(bt.get_shape())
+                if bshape != tuple(new_mtp[k].shape):
+                    raise SystemExit(f"Shape mismatch for {k}: base {bshape} vs mtp {tuple(new_mtp[k].shape)}")
+
+    # --- 3. Bring over the unchanged shards ------------------------------------
+    for shard in all_shards:
+        if shard in mtp_shards:
+            continue
+        bring_over(os.path.join(args.base, shard), os.path.join(args.out, shard), args.link)
+    print(f"Linked {len(all_shards) - len(mtp_shards)} unchanged shards ({args.link}).")
+
+    # --- 4. Rewrite the MTP-containing shards with the new head ----------------
+    for shard in mtp_shards:
+        src = os.path.join(args.base, shard)
+        tensors = {}
+        with safe_open(src, framework="pt", device="cpu") as f:
+            for k in f.keys():
+                tensors[k] = f.get_tensor(k)
+        replaced = 0
+        for k in tensors:
+            if k in new_mtp:
+                tensors[k] = new_mtp[k]
+                replaced += 1
+        save_file(tensors, os.path.join(args.out, shard), metadata={"format": "pt"})
+        print(f"  rewrote {shard}: replaced {replaced} MTP tensors "
+              f"(kept {len(tensors) - replaced} base tensors)")
+
+    # --- 5. Copy every non-shard file (config, tokenizer, index, ...) ----------
+    for name in sorted(os.listdir(args.base)):
+        src = os.path.join(args.base, name)
+        if not os.path.isfile(src):
+            continue
+        if name.endswith(".safetensors"):   # shards handled above
+            continue
+        shutil.copy2(src, os.path.join(args.out, name))
+    print("Copied config / tokenizer / index and other metadata files.")
+
+    print(f"\nDone. Merged model written to: {args.out}")
+    if args.link == "symlink":
+        print("NOTE: unchanged shards are symlinks into the base model. Keep the base "
+              "model in place, or re-run with --link copy to make a self-contained copy.")
+
+
+if __name__ == "__main__":
+    main()