185 lines
8.1 KiB
Python
185 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Convert a STARFlow-V checkpoint into a sharded safetensors checkpoint.
|
|
|
|
Design notes / invariants
|
|
--------------------------
|
|
* **Exact tensor names are preserved by default.** No wrapper prefix is
|
|
stripped automatically -- the official Apple ml-starflow checkpoint must
|
|
round-trip byte- and name-identically. Prefix stripping is strictly
|
|
opt-in via ``--strip-prefix`` (repeatable) and is only applied when
|
|
*every* key shares that prefix.
|
|
* **Dtype-agnostic content hashing.** ``sha256_tensor`` hashes the raw
|
|
little-endian byte content by reinterpreting storage as ``uint8``. This
|
|
is required for ``bfloat16`` (and any other dtype numpy cannot express)
|
|
and keeps fp16/fp32/int/bool checkpoints working unchanged.
|
|
* Every shard is reopened after writing and verified key/shape/dtype/
|
|
content-hash identical to the source tensors.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from typing import Dict
|
|
|
|
import torch
|
|
|
|
|
|
def sha256_file(path: str) -> str:
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def sha256_tensor(t: torch.Tensor) -> str:
|
|
# Hash the raw byte content independent of dtype. numpy() cannot
|
|
# represent bfloat16, so reinterpret the (contiguous) storage as uint8;
|
|
# this is uniform across bf16/fp16/fp32/int/bool tensors.
|
|
raw = t.detach().cpu().contiguous().reshape(-1).view(torch.uint8)
|
|
return hashlib.sha256(raw.numpy().tobytes()).hexdigest()
|
|
|
|
|
|
def load_source_state_dict(path: str) -> Dict[str, torch.Tensor]:
|
|
if path.endswith(".safetensors"):
|
|
from safetensors.torch import load_file
|
|
state_dict = load_file(path, device="cpu")
|
|
elif path.endswith(".pth") or path.endswith(".pt"):
|
|
state_dict = torch.load(path, map_location="cpu", weights_only=True)
|
|
if not isinstance(state_dict, dict):
|
|
raise ValueError(f"{path}: expected a state_dict (dict), got {type(state_dict)}")
|
|
else:
|
|
raise ValueError(f"unrecognized checkpoint format for {path}")
|
|
|
|
for k, v in state_dict.items():
|
|
if not isinstance(v, torch.Tensor):
|
|
raise ValueError(f"{path}: non-tensor value for key {k!r} ({type(v)}); refusing to convert")
|
|
return state_dict
|
|
|
|
|
|
def strip_known_prefix(state_dict: Dict[str, torch.Tensor], prefixes) -> Dict[str, torch.Tensor]:
|
|
"""Strip a wrapper prefix ONLY if the caller explicitly requested one and
|
|
every key shares it. With an empty ``prefixes`` (the default) the state
|
|
dict is returned unchanged, preserving exact tensor names."""
|
|
keys = list(state_dict.keys())
|
|
if not keys:
|
|
return state_dict
|
|
for prefix in prefixes:
|
|
if prefix and all(k.startswith(prefix) for k in keys):
|
|
print(f" stripping shared prefix {prefix!r} from all {len(keys)} keys (explicitly requested).")
|
|
return {k[len(prefix):]: v for k, v in state_dict.items()}
|
|
return state_dict
|
|
|
|
|
|
def plan_shards(state_dict: Dict[str, torch.Tensor], max_shard_bytes: int):
|
|
shards = []
|
|
current: Dict[str, torch.Tensor] = {}
|
|
current_size = 0
|
|
for key, tensor in state_dict.items():
|
|
tensor_size = tensor.numel() * tensor.element_size()
|
|
if current and current_size + tensor_size > max_shard_bytes:
|
|
shards.append(current)
|
|
current, current_size = {}, 0
|
|
current[key] = tensor
|
|
current_size += tensor_size
|
|
if current:
|
|
shards.append(current)
|
|
return shards
|
|
|
|
|
|
def convert(input_path: str, output_dir: str, max_shard_size_gb: float,
|
|
expected_sha256: str = None, strip_prefixes=()) -> None:
|
|
if expected_sha256:
|
|
print(f"Verifying source SHA-256 of {input_path} ...")
|
|
actual = sha256_file(input_path)
|
|
if actual.lower() != expected_sha256.lower():
|
|
raise RuntimeError(
|
|
f"source checksum mismatch: expected {expected_sha256}, got {actual}")
|
|
print(" OK, source checksum matches.")
|
|
|
|
print(f"Loading source checkpoint {input_path} (weights_only / tensor-only) ...")
|
|
state_dict = load_source_state_dict(input_path)
|
|
# Default: NO automatic prefix stripping -- exact names are preserved.
|
|
state_dict = strip_known_prefix(state_dict, tuple(strip_prefixes))
|
|
print(f" {len(state_dict)} tensors.")
|
|
|
|
print("Hashing source tensors (pre-save reference hashes) ...")
|
|
source_hashes = {k: sha256_tensor(v) for k, v in state_dict.items()}
|
|
source_shapes = {k: tuple(v.shape) for k, v in state_dict.items()}
|
|
source_dtypes = {k: str(v.dtype) for k, v in state_dict.items()}
|
|
|
|
max_shard_bytes = int(max_shard_size_gb * (1024 ** 3))
|
|
shard_plans = plan_shards(state_dict, max_shard_bytes)
|
|
print(f"Planned {len(shard_plans)} shard(s) (<= {max_shard_size_gb} GB each).")
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
weight_map = {}
|
|
shard_shas = {}
|
|
total_size = 0
|
|
num_shards = len(shard_plans)
|
|
for i, shard in enumerate(shard_plans, start=1):
|
|
shard_filename = f"model-{i:05d}-of-{num_shards:05d}.safetensors"
|
|
shard_path = os.path.join(output_dir, shard_filename)
|
|
to_save = {k: v.clone().contiguous() for k, v in shard.items()}
|
|
save_file(to_save, shard_path)
|
|
for k in shard:
|
|
weight_map[k] = shard_filename
|
|
total_size += shard[k].numel() * shard[k].element_size()
|
|
shard_shas[shard_filename] = sha256_file(shard_path)
|
|
print(f" wrote {shard_filename}: {len(shard)} tensors, sha256={shard_shas[shard_filename][:16]}...")
|
|
|
|
index_path = os.path.join(output_dir, "model.safetensors.index.json")
|
|
index = {
|
|
"metadata": {
|
|
"total_size": total_size,
|
|
"source_checkpoint": os.path.basename(input_path),
|
|
"source_sha256": (expected_sha256 or sha256_file(input_path)),
|
|
"shard_sha256": shard_shas,
|
|
},
|
|
"weight_map": weight_map,
|
|
}
|
|
with open(index_path, "w") as f:
|
|
json.dump(index, f, indent=2, sort_keys=True)
|
|
print(f"Wrote index: {index_path}")
|
|
|
|
print("Reopening shards and verifying key/shape/dtype/content-hash ...")
|
|
verified = 0
|
|
for shard_filename in sorted(set(weight_map.values())):
|
|
shard_path = os.path.join(output_dir, shard_filename)
|
|
reopened = load_file(shard_path, device="cpu")
|
|
for k, v in reopened.items():
|
|
if source_shapes[k] != tuple(v.shape):
|
|
raise RuntimeError(f"{k}: shape mismatch after reopen ({source_shapes[k]} vs {tuple(v.shape)})")
|
|
if source_dtypes[k] != str(v.dtype):
|
|
raise RuntimeError(f"{k}: dtype mismatch after reopen ({source_dtypes[k]} vs {v.dtype})")
|
|
got_hash = sha256_tensor(v)
|
|
if source_hashes[k] != got_hash:
|
|
raise RuntimeError(f"{k}: content hash mismatch after reopen")
|
|
verified += 1
|
|
if verified != len(state_dict):
|
|
raise RuntimeError(f"verified {verified} tensors but source has {len(state_dict)}")
|
|
print(f" OK: {verified}/{len(state_dict)} tensors verified byte-identical after reopen.")
|
|
print("Conversion complete.")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--input", required=True, help="source .pth/.pt or single .safetensors checkpoint")
|
|
parser.add_argument("--output-dir", required=True, help="directory to write sharded safetensors and index into")
|
|
parser.add_argument("--max-shard-size-gb", type=float, default=5.0)
|
|
parser.add_argument("--expected-sha256", default=None, help="optional SHA-256 to verify the source file against")
|
|
parser.add_argument("--strip-prefix", action="append", default=[],
|
|
help="wrapper prefix to strip IF shared by every key (repeatable). "
|
|
"Default: none -- exact tensor names are preserved.")
|
|
args = parser.parse_args()
|
|
|
|
convert(args.input, args.output_dir, args.max_shard_size_gb,
|
|
expected_sha256=args.expected_sha256, strip_prefixes=args.strip_prefix)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|