642 lines
23 KiB
Python
642 lines
23 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Memory-friendly checkpoint converter: inner -> outer format (v2).
|
|
|
|
Converts the HYV3 checkpoint from inner format (per-expert keys, old naming)
|
|
to outer format (fused 3D experts, new naming) shard by shard.
|
|
|
|
Handles the case where a single layer's experts may be split across
|
|
multiple shards (cross-shard experts) by deferring their fusion to a
|
|
post-processing step.
|
|
|
|
v2 improvements over v1:
|
|
- Post-processing is shard-centric (each shard read/written only once)
|
|
instead of prefix-centric (same shard read/written multiple times).
|
|
This fixes Bus error (core dump) when there are many cross-shard groups.
|
|
- Explicit memory management with gc.collect() to prevent memory bloat.
|
|
- Better progress reporting during post-processing.
|
|
|
|
Supports multi-process parallelism for faster conversion.
|
|
|
|
Usage:
|
|
# Default 8 workers
|
|
python convert_ckpt_to_outer.py \\
|
|
--input_dir pretrain_base/hf \\
|
|
--output_dir pretrain_base/hf_outer
|
|
|
|
# Custom worker count
|
|
python convert_ckpt_to_outer.py \\
|
|
--input_dir pretrain_base/hf \\
|
|
--output_dir pretrain_base/hf_outer \\
|
|
--workers 16
|
|
|
|
The script will:
|
|
1. Pre-scan index.json to detect cross-shard expert groups
|
|
2. Convert weights shard-by-shard in parallel (key rename + expert fuse)
|
|
3. Post-process cross-shard expert groups (merge from multiple shards)
|
|
- v2: shard-centric approach, each shard read/written only once
|
|
4. Copy config.json as-is (already in outer format)
|
|
5. Copy all other files (tokenizer, etc.)
|
|
6. Rebuild model.safetensors.index.json
|
|
"""
|
|
|
|
import argparse
|
|
import gc
|
|
import json
|
|
import os
|
|
import re
|
|
import signal
|
|
import shutil
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from collections import OrderedDict, defaultdict
|
|
from multiprocessing import Pool
|
|
|
|
import torch
|
|
|
|
try:
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
except ImportError:
|
|
raise ImportError("Please install safetensors: pip install safetensors")
|
|
|
|
# ============================================================================
|
|
# Signal handling for Bus error (SIGBUS) and other fatal signals
|
|
# ============================================================================
|
|
|
|
def _fatal_signal_handler(signum, frame):
|
|
"""Handle fatal signals (SIGBUS, SIGSEGV) by logging before exit.
|
|
|
|
These signals cannot be caught by try/except. This handler ensures
|
|
the error message is written to stderr (captured by nohup redirection)
|
|
before the process terminates.
|
|
"""
|
|
sig_name = signal.Signals(signum).name if hasattr(signal, 'Signals') else str(signum)
|
|
pid = os.getpid()
|
|
msg = (
|
|
f"\n[FATAL] Process {pid} received {sig_name} (signal {signum}).\n"
|
|
f"This typically indicates an out-of-memory condition during mmap I/O.\n"
|
|
f"Stack trace at time of signal:\n"
|
|
)
|
|
sys.stderr.write(msg)
|
|
traceback.print_stack(frame, file=sys.stderr)
|
|
sys.stderr.flush()
|
|
# Re-raise with default handler to get proper exit code
|
|
signal.signal(signum, signal.SIG_DFL)
|
|
os.kill(pid, signum)
|
|
|
|
|
|
def _install_signal_handlers():
|
|
"""Install handlers for SIGBUS and SIGSEGV in the current process."""
|
|
for sig in (signal.SIGBUS, signal.SIGSEGV):
|
|
try:
|
|
signal.signal(sig, _fatal_signal_handler)
|
|
except (OSError, ValueError):
|
|
# Some signals may not be available on all platforms
|
|
pass
|
|
|
|
|
|
def _pool_worker_init():
|
|
"""Initializer for multiprocessing pool workers.
|
|
|
|
Installs signal handlers so that Bus errors in worker processes
|
|
are also logged before the process dies.
|
|
"""
|
|
_install_signal_handlers()
|
|
|
|
|
|
# ============================================================================
|
|
# Key rename mapping (inner -> outer)
|
|
# ============================================================================
|
|
|
|
_KEY_RENAMES = [
|
|
("mlp.router.gate.", "mlp.gate."),
|
|
("mlp.expert_bias", "mlp.e_score_correction_bias"),
|
|
("mlp.shared_mlp.", "mlp.shared_experts."),
|
|
]
|
|
|
|
# Regex to match per-expert keys
|
|
_EXPERT_KEY_RE = re.compile(
|
|
r"^(.*\.mlp\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)\.weight$"
|
|
)
|
|
|
|
def rename_key(key: str) -> str:
|
|
"""Rename a single key from inner to outer format."""
|
|
for old_sub, new_sub in _KEY_RENAMES:
|
|
if old_sub in key:
|
|
key = key.replace(old_sub, new_sub)
|
|
break
|
|
return key
|
|
|
|
def scan_cross_shard_experts(index_path: str):
|
|
"""Pre-scan index.json to find expert groups that span multiple shards.
|
|
|
|
Returns:
|
|
cross_shard_prefixes: set of expert prefixes that span multiple shards
|
|
e.g. {"model.layers.80.mlp.experts."}
|
|
"""
|
|
with open(index_path) as f:
|
|
index = json.load(f)
|
|
wm = index["weight_map"]
|
|
|
|
# prefix -> set of shards
|
|
prefix_shards = defaultdict(set)
|
|
for key in wm:
|
|
m = _EXPERT_KEY_RE.match(key)
|
|
if m:
|
|
prefix = m.group(1)
|
|
prefix_shards[prefix].add(wm[key])
|
|
|
|
cross_shard_prefixes = set()
|
|
for prefix, shards in prefix_shards.items():
|
|
if len(shards) > 1:
|
|
cross_shard_prefixes.add(prefix)
|
|
|
|
return cross_shard_prefixes
|
|
|
|
def convert_shard(shard_path: str, cross_shard_prefixes: set = None):
|
|
"""Load a single shard, rename keys, and fuse experts.
|
|
|
|
For expert groups in cross_shard_prefixes, the per-expert keys are
|
|
kept as-is (just renamed) and returned separately as deferred items,
|
|
to be merged later in a post-processing step.
|
|
|
|
Returns:
|
|
result: OrderedDict of converted tensors (ready to save)
|
|
deferred_expert_keys: list of original expert keys that were deferred
|
|
(these are kept in result with their original per-expert naming
|
|
but with the outer rename applied, to be post-processed later)
|
|
"""
|
|
if cross_shard_prefixes is None:
|
|
cross_shard_prefixes = set()
|
|
|
|
tensors = OrderedDict()
|
|
with safe_open(shard_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
tensors[key] = f.get_tensor(key)
|
|
|
|
# Separate expert keys from non-expert keys
|
|
expert_groups = {} # prefix -> {expert_idx -> {proj_name -> tensor}}
|
|
deferred_expert_keys = [] # keys that belong to cross-shard experts
|
|
result = OrderedDict()
|
|
|
|
for key, tensor in tensors.items():
|
|
m = _EXPERT_KEY_RE.match(key)
|
|
if m:
|
|
prefix = m.group(1)
|
|
expert_idx = int(m.group(2))
|
|
proj_name = m.group(3)
|
|
|
|
if prefix in cross_shard_prefixes:
|
|
# Defer: keep the key as-is (with rename) for post-processing
|
|
new_key = rename_key(key)
|
|
result[new_key] = tensor
|
|
deferred_expert_keys.append(new_key)
|
|
else:
|
|
# Normal: collect for fusion within this shard
|
|
if prefix not in expert_groups:
|
|
expert_groups[prefix] = {}
|
|
if expert_idx not in expert_groups[prefix]:
|
|
expert_groups[prefix][expert_idx] = {}
|
|
expert_groups[prefix][expert_idx][proj_name] = tensor
|
|
else:
|
|
# Non-expert key: just rename
|
|
new_key = rename_key(key)
|
|
result[new_key] = tensor
|
|
|
|
# Fuse expert weights for each non-cross-shard layer prefix
|
|
for prefix in sorted(expert_groups.keys()):
|
|
experts = expert_groups[prefix]
|
|
num_experts = max(experts.keys()) + 1
|
|
|
|
gate_up_list = []
|
|
down_list = []
|
|
for i in range(num_experts):
|
|
if i not in experts:
|
|
raise ValueError(
|
|
f"Missing expert {i} in {prefix}. "
|
|
f"Found: {sorted(experts.keys())}"
|
|
)
|
|
exp = experts[i]
|
|
gate_up = torch.cat([exp["gate_proj"], exp["up_proj"]], dim=0)
|
|
gate_up_list.append(gate_up)
|
|
down_list.append(exp["down_proj"])
|
|
|
|
fused_gate_up = torch.stack(gate_up_list, dim=0)
|
|
fused_down = torch.stack(down_list, dim=0)
|
|
|
|
for exp in experts.values():
|
|
exp.clear()
|
|
gate_up_list.clear()
|
|
down_list.clear()
|
|
|
|
result[f"{prefix}gate_up_proj"] = fused_gate_up
|
|
result[f"{prefix}down_proj"] = fused_down
|
|
|
|
return result, deferred_expert_keys
|
|
|
|
def _process_one_shard(args_tuple):
|
|
"""Worker function: convert a single shard and save to output dir.
|
|
|
|
Args:
|
|
args_tuple: (idx, num_shards, shard_file, input_dir, output_dir, cross_shard_prefixes)
|
|
|
|
Returns:
|
|
(shard_file, key_list, shard_size, elapsed, deferred_keys)
|
|
"""
|
|
idx, num_shards, shard_file, input_dir, output_dir, cross_shard_prefixes = args_tuple
|
|
shard_path = os.path.join(input_dir, shard_file)
|
|
t0 = time.time()
|
|
|
|
converted, deferred_keys = convert_shard(shard_path, cross_shard_prefixes)
|
|
|
|
shard_size = sum(t.numel() * t.element_size() for t in converted.values())
|
|
|
|
out_shard_path = os.path.join(output_dir, shard_file)
|
|
save_file(converted, out_shard_path)
|
|
|
|
elapsed = time.time() - t0
|
|
num_keys = len(converted)
|
|
key_list = list(converted.keys())
|
|
|
|
del converted
|
|
|
|
deferred_info = ""
|
|
if deferred_keys:
|
|
deferred_info = f", Deferred={len(deferred_keys)}"
|
|
|
|
print(
|
|
f" [{idx + 1}/{num_shards}] {shard_file}: "
|
|
f"Keys={num_keys}, Size={shard_size / 1e9:.2f} GB, "
|
|
f"Time={elapsed:.1f}s{deferred_info}",
|
|
flush=True,
|
|
)
|
|
|
|
return shard_file, key_list, shard_size, elapsed, deferred_keys
|
|
|
|
|
|
def post_process_cross_shard_experts(output_dir, cross_shard_prefixes, all_deferred):
|
|
"""Merge cross-shard expert groups (v2: shard-centric approach).
|
|
|
|
Instead of iterating per-prefix (which causes the same shard to be
|
|
loaded/saved multiple times), this v2 approach:
|
|
1. Builds a mapping of which prefixes each shard is involved in
|
|
2. Collects all expert tensors from all involved shards in ONE pass
|
|
3. Fuses all prefixes
|
|
4. Writes each shard only ONCE with all its updates applied
|
|
|
|
This avoids the Bus error (core dump) caused by repeated mmap of
|
|
large files and memory bloat.
|
|
|
|
Args:
|
|
output_dir: path to output directory
|
|
cross_shard_prefixes: set of expert prefixes that span multiple shards
|
|
all_deferred: dict of {shard_file: [deferred_key, ...]}
|
|
|
|
Returns:
|
|
updated_shards: dict of {shard_file: (key_list, shard_size)} for updated shards
|
|
"""
|
|
if not cross_shard_prefixes:
|
|
return {}
|
|
|
|
print(f"\n Post-processing {len(cross_shard_prefixes)} cross-shard expert group(s)...",
|
|
flush=True)
|
|
|
|
# ----------------------------------------------------------------
|
|
# Step 1: Build mappings
|
|
# ----------------------------------------------------------------
|
|
# prefix -> ordered list of shards that contain its experts
|
|
prefix_to_shards = defaultdict(set)
|
|
# shard -> set of prefixes it is involved in
|
|
shard_to_prefixes = defaultdict(set)
|
|
|
|
for shard_file, deferred_keys in all_deferred.items():
|
|
for key in deferred_keys:
|
|
m = _EXPERT_KEY_RE.match(key)
|
|
if m:
|
|
prefix = m.group(1)
|
|
if prefix in cross_shard_prefixes:
|
|
prefix_to_shards[prefix].add(shard_file)
|
|
shard_to_prefixes[shard_file].add(prefix)
|
|
|
|
# For each prefix, decide which shard will hold the fused result
|
|
# (use the first shard alphabetically)
|
|
prefix_to_target_shard = {}
|
|
for prefix in sorted(prefix_to_shards.keys()):
|
|
target = sorted(prefix_to_shards[prefix])[0]
|
|
prefix_to_target_shard[prefix] = target
|
|
|
|
# All shards that need to be updated
|
|
all_involved_shards = set()
|
|
for shards in prefix_to_shards.values():
|
|
all_involved_shards.update(shards)
|
|
|
|
print(f" Involved shards: {len(all_involved_shards)}", flush=True)
|
|
print(f" Expert groups: {len(prefix_to_shards)}", flush=True)
|
|
|
|
# ----------------------------------------------------------------
|
|
# Step 2: Collect all expert tensors from all involved shards
|
|
# (one pass per shard)
|
|
# ----------------------------------------------------------------
|
|
# prefix -> {expert_idx -> {proj_name -> tensor}}
|
|
all_expert_data = defaultdict(dict)
|
|
# shard -> OrderedDict of non-expert keys (to be re-saved)
|
|
shard_non_expert = {}
|
|
|
|
sorted_involved = sorted(all_involved_shards)
|
|
for si, shard_file in enumerate(sorted_involved):
|
|
shard_path = os.path.join(output_dir, shard_file)
|
|
prefixes_in_shard = shard_to_prefixes[shard_file]
|
|
|
|
print(f" [{si+1}/{len(sorted_involved)}] Reading {shard_file} "
|
|
f"({len(prefixes_in_shard)} prefix(es))...", flush=True)
|
|
|
|
non_expert = OrderedDict()
|
|
with safe_open(shard_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
m = _EXPERT_KEY_RE.match(key)
|
|
if m and m.group(1) in prefixes_in_shard:
|
|
# This is a deferred expert key
|
|
prefix = m.group(1)
|
|
expert_idx = int(m.group(2))
|
|
proj_name = m.group(3)
|
|
if expert_idx not in all_expert_data[prefix]:
|
|
all_expert_data[prefix][expert_idx] = {}
|
|
all_expert_data[prefix][expert_idx][proj_name] = f.get_tensor(key)
|
|
else:
|
|
# Non-expert key: keep as-is
|
|
non_expert[key] = f.get_tensor(key)
|
|
|
|
shard_non_expert[shard_file] = non_expert
|
|
gc.collect()
|
|
|
|
# ----------------------------------------------------------------
|
|
# Step 3: Fuse all expert groups
|
|
# ----------------------------------------------------------------
|
|
# prefix -> {"gate_up_proj": tensor, "down_proj": tensor}
|
|
fused_results = {}
|
|
|
|
for pi, prefix in enumerate(sorted(all_expert_data.keys())):
|
|
expert_data = all_expert_data[prefix]
|
|
num_experts = max(expert_data.keys()) + 1
|
|
|
|
print(f" Fusing {prefix} ({num_experts} experts)...", flush=True)
|
|
|
|
gate_up_list = []
|
|
down_list = []
|
|
for i in range(num_experts):
|
|
if i not in expert_data:
|
|
raise ValueError(
|
|
f"Missing expert {i} in {prefix} after cross-shard merge. "
|
|
f"Found: {sorted(expert_data.keys())}"
|
|
)
|
|
exp = expert_data[i]
|
|
if "gate_proj" not in exp or "up_proj" not in exp:
|
|
raise ValueError(
|
|
f"Expert {i} in {prefix} missing gate_proj/up_proj. "
|
|
f"Has: {sorted(exp.keys())}"
|
|
)
|
|
if "down_proj" not in exp:
|
|
raise ValueError(
|
|
f"Expert {i} in {prefix} missing down_proj. "
|
|
f"Has: {sorted(exp.keys())}"
|
|
)
|
|
gate_up = torch.cat([exp["gate_proj"], exp["up_proj"]], dim=0)
|
|
gate_up_list.append(gate_up)
|
|
down_list.append(exp["down_proj"])
|
|
|
|
fused_gate_up = torch.stack(gate_up_list, dim=0)
|
|
fused_down = torch.stack(down_list, dim=0)
|
|
|
|
fused_results[prefix] = {
|
|
"gate_up_proj": fused_gate_up,
|
|
"down_proj": fused_down,
|
|
}
|
|
|
|
# Free per-expert data for this prefix
|
|
del gate_up_list, down_list
|
|
for exp in expert_data.values():
|
|
exp.clear()
|
|
del all_expert_data[prefix]
|
|
gc.collect()
|
|
|
|
del all_expert_data
|
|
gc.collect()
|
|
|
|
# ----------------------------------------------------------------
|
|
# Step 4: Write each involved shard ONCE with all updates applied
|
|
# ----------------------------------------------------------------
|
|
updated_shards = {}
|
|
|
|
for si, shard_file in enumerate(sorted_involved):
|
|
shard_path = os.path.join(output_dir, shard_file)
|
|
non_expert = shard_non_expert[shard_file]
|
|
|
|
# Add fused tensors for prefixes that target this shard
|
|
fused_added = []
|
|
for prefix, target_shard in prefix_to_target_shard.items():
|
|
if target_shard == shard_file and prefix in fused_results:
|
|
non_expert[f"{prefix}gate_up_proj"] = fused_results[prefix]["gate_up_proj"]
|
|
non_expert[f"{prefix}down_proj"] = fused_results[prefix]["down_proj"]
|
|
fused_added.append(prefix)
|
|
|
|
save_file(non_expert, shard_path)
|
|
shard_size = sum(t.numel() * t.element_size() for t in non_expert.values())
|
|
updated_shards[shard_file] = (list(non_expert.keys()), shard_size)
|
|
|
|
fused_info = ""
|
|
if fused_added:
|
|
fused_info = f", Fused {len(fused_added)} group(s)"
|
|
|
|
print(f" [{si+1}/{len(sorted_involved)}] Wrote {shard_file}: "
|
|
f"{len(non_expert)} keys, {shard_size / 1e9:.2f} GB{fused_info}",
|
|
flush=True)
|
|
|
|
# Free memory for this shard
|
|
del shard_non_expert[shard_file]
|
|
for prefix in fused_added:
|
|
del fused_results[prefix]
|
|
del non_expert
|
|
gc.collect()
|
|
|
|
return updated_shards
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Convert HYV3 checkpoint from inner to outer format (v2, shard-centric post-processing)."
|
|
)
|
|
parser.add_argument(
|
|
"--input_dir", type=str, required=True,
|
|
help="Path to the inner-format checkpoint directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir", type=str, required=True,
|
|
help="Path to the output outer-format checkpoint directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--workers", type=int, default=8,
|
|
help="Number of parallel worker processes (default: 8).",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
input_dir = os.path.abspath(args.input_dir)
|
|
output_dir = os.path.abspath(args.output_dir)
|
|
num_workers = args.workers
|
|
|
|
if not os.path.isdir(input_dir):
|
|
raise FileNotFoundError(f"Input directory not found: {input_dir}")
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Pre-scan for cross-shard expert groups
|
|
index_path = os.path.join(input_dir, "model.safetensors.index.json")
|
|
cross_shard_prefixes = set()
|
|
if os.path.exists(index_path):
|
|
cross_shard_prefixes = scan_cross_shard_experts(index_path)
|
|
if cross_shard_prefixes:
|
|
print(f"Detected {len(cross_shard_prefixes)} cross-shard expert group(s):")
|
|
for p in sorted(cross_shard_prefixes):
|
|
print(f" - {p}")
|
|
print()
|
|
|
|
# Get all safetensors files
|
|
shard_files = sorted(
|
|
f for f in os.listdir(input_dir) if f.endswith(".safetensors")
|
|
)
|
|
if not shard_files:
|
|
raise FileNotFoundError(f"No .safetensors files found in {input_dir}")
|
|
|
|
# Skip already-converted shards (for resumability)
|
|
# NOTE: if there are cross-shard experts, we cannot skip shards that
|
|
# contain deferred keys (they need post-processing). For simplicity,
|
|
# when cross-shard experts exist, we re-process all shards.
|
|
remaining = []
|
|
skipped = []
|
|
if cross_shard_prefixes:
|
|
# Re-process all shards when cross-shard experts exist
|
|
remaining = list(shard_files)
|
|
else:
|
|
for sf in shard_files:
|
|
out_path = os.path.join(output_dir, sf)
|
|
if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
|
|
skipped.append(sf)
|
|
else:
|
|
remaining.append(sf)
|
|
|
|
num_shards = len(shard_files)
|
|
num_workers = min(num_workers, len(remaining)) if remaining else 1
|
|
|
|
print(f"=" * 60)
|
|
print(f"HYV3 Checkpoint Converter (inner -> outer, v2)")
|
|
print(f" Input : {input_dir}")
|
|
print(f" Output : {output_dir}")
|
|
print(f" Shards : {num_shards} total, {len(skipped)} already done, {len(remaining)} to process")
|
|
print(f" Workers: {num_workers}")
|
|
if cross_shard_prefixes:
|
|
print(f" Cross-shard experts: {len(cross_shard_prefixes)} group(s) (will post-process)")
|
|
print(f"=" * 60)
|
|
|
|
t_start = time.time()
|
|
|
|
# Build task list for remaining shards
|
|
tasks = [
|
|
(i, len(remaining), sf, input_dir, output_dir, cross_shard_prefixes)
|
|
for i, sf in enumerate(remaining)
|
|
]
|
|
|
|
# Process in parallel
|
|
results = []
|
|
if tasks:
|
|
with Pool(processes=num_workers, initializer=_pool_worker_init) as pool:
|
|
results = pool.map(_process_one_shard, tasks)
|
|
|
|
# Collect deferred keys info
|
|
all_deferred = {} # shard_file -> [deferred_keys]
|
|
for shard_file, key_list, shard_size, elapsed, deferred_keys in results:
|
|
if deferred_keys:
|
|
all_deferred[shard_file] = deferred_keys
|
|
|
|
# Post-process cross-shard expert groups (v2: shard-centric)
|
|
updated_shards = {}
|
|
if cross_shard_prefixes and all_deferred:
|
|
updated_shards = post_process_cross_shard_experts(
|
|
output_dir, cross_shard_prefixes, all_deferred
|
|
)
|
|
|
|
# Build weight_map and total_size
|
|
weight_map = OrderedDict()
|
|
total_size = 0
|
|
|
|
# For skipped shards, read their keys from the output files
|
|
for sf in skipped:
|
|
out_path = os.path.join(output_dir, sf)
|
|
with safe_open(out_path, framework="pt", device="cpu") as f:
|
|
keys = list(f.keys())
|
|
for key in keys:
|
|
weight_map[key] = sf
|
|
t = f.get_tensor(key)
|
|
total_size += t.numel() * t.element_size()
|
|
|
|
# Collect results from newly converted shards
|
|
for shard_file, key_list, shard_size, elapsed, deferred_keys in results:
|
|
if shard_file in updated_shards:
|
|
# This shard was updated by post-processing
|
|
updated_key_list, updated_size = updated_shards[shard_file]
|
|
for key in updated_key_list:
|
|
weight_map[key] = shard_file
|
|
total_size += updated_size
|
|
else:
|
|
for key in key_list:
|
|
weight_map[key] = shard_file
|
|
total_size += shard_size
|
|
|
|
# Build and save index
|
|
sorted_weight_map = OrderedDict(sorted(weight_map.items()))
|
|
index = {
|
|
"metadata": {"total_size": total_size},
|
|
"weight_map": sorted_weight_map,
|
|
}
|
|
index_path_out = os.path.join(output_dir, "model.safetensors.index.json")
|
|
with open(index_path_out, "w") as f:
|
|
json.dump(index, f, indent=2)
|
|
f.write("\n")
|
|
print(f"\nSaved {index_path_out}")
|
|
|
|
# Copy non-safetensors files (config, tokenizer, etc.)
|
|
skip_suffixes = {".safetensors"}
|
|
skip_names = {"model.safetensors.index.json"}
|
|
copied = []
|
|
for fname in os.listdir(input_dir):
|
|
if fname in skip_names:
|
|
continue
|
|
if any(fname.endswith(s) for s in skip_suffixes):
|
|
continue
|
|
src = os.path.join(input_dir, fname)
|
|
dst = os.path.join(output_dir, fname)
|
|
if os.path.isfile(src):
|
|
shutil.copy2(src, dst)
|
|
copied.append(fname)
|
|
elif os.path.isdir(src):
|
|
if os.path.exists(dst):
|
|
shutil.rmtree(dst)
|
|
shutil.copytree(src, dst)
|
|
copied.append(fname + "/")
|
|
|
|
if copied:
|
|
print(f"\nCopied files: {', '.join(copied)}")
|
|
|
|
t_total = time.time() - t_start
|
|
print(f"\n{'=' * 60}")
|
|
print(f"Conversion complete!")
|
|
print(f" Total keys : {len(weight_map)}")
|
|
print(f" Total size : {total_size / 1e9:.2f} GB")
|
|
print(f" Total time : {t_total:.1f}s ({t_total / 60:.1f} min)")
|
|
print(f" Output dir : {output_dir}")
|
|
print(f"{'=' * 60}")
|
|
|
|
if __name__ == "__main__":
|
|
_install_signal_handlers()
|
|
main()
|