File size: 2,498 Bytes
4335e83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """Phase D: splice trained mtp.* tensors back into the merged BF16 checkpoint.
Creates a new checkpoint dir where safetensors shards containing replaced
tensors are rewritten and every other file is HARDLINKED (no 230GB copy).
Then quantize with quant_oq5e_mtp.py (preserve_mtp) as usual.
Usage:
python3 splice_mtp.py [--trained mtp_trained.safetensors]
[--src .../NVIDIA-Nemotron-3-Super-120B-merged3]
[--dst .../NVIDIA-Nemotron-3-Super-120B-merged3-mtpft]
"""
import argparse
import json
import os
from pathlib import Path
import mlx.core as mx
ROOT = Path(__file__).parent
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--trained", default=str(ROOT / "mtp_trained.safetensors"))
ap.add_argument(
"--src", default="/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged3")
ap.add_argument(
"--dst",
default="/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged3-mtpft")
args = ap.parse_args()
src, dst = Path(args.src), Path(args.dst)
trained = mx.load(args.trained)
print(f"{len(trained)} trained tensors")
index = json.load(open(src / "model.safetensors.index.json"))["weight_map"]
missing = [k for k in trained if k not in index]
if missing:
raise SystemExit(f"trained tensors not in checkpoint index: {missing[:5]}")
dirty_files = {index[k] for k in trained}
print(f"rewriting {len(dirty_files)} shard file(s): {sorted(dirty_files)}")
dst.mkdir(exist_ok=True)
for f in sorted(src.iterdir()):
if f.is_dir() or f.name.startswith("."):
continue
target = dst / f.name
if target.exists():
target.unlink()
if f.name in dirty_files:
shard = dict(mx.load(str(f)))
replaced = 0
for k in list(shard.keys()):
if k in trained:
if shard[k].shape != trained[k].shape:
raise SystemExit(
f"shape mismatch {k}: {shard[k].shape} vs "
f"{trained[k].shape}")
shard[k] = trained[k].astype(shard[k].dtype)
replaced += 1
mx.save_safetensors(str(target), shard,
metadata={"format": "pt"})
print(f" {f.name}: replaced {replaced} tensors")
else:
os.link(f, target)
print(f"DONE -> {dst}")
if __name__ == "__main__":
main()
|