#!/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()