diff --git a/convert_to_safetensors.py b/convert_to_safetensors.py new file mode 100644 index 0000000..19d9633 --- /dev/null +++ b/convert_to_safetensors.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Convert torch.save quantized shards to safetensors format.""" + +import gc +import torch +from pathlib import Path +from safetensors.torch import save_file +import glob + + +def convert_shards(model_path): + output_path = Path(model_path) + shards = sorted(glob.glob(f"{model_path}/*.safetensors")) + + print(f"Converting {len(shards)} shards to safetensors format...\n") + + for i, shard_path in enumerate(shards): + print(f"Converting shard {i+1}/{len(shards)}: {shard_path.name}") + + # Load torch.save format + ckpt = torch.load(shard_path, map_location="cpu", weights_only=False) + + # Separate tensors from QuantState objects + tensors = {} + for key, value in ckpt.items(): + if isinstance(value, torch.Tensor): + tensors[key] = value + else: + # Save QuantState separately if needed + print(f" Skipping non-tensor: {key} ({type(value)})") + + # Save as safetensors + new_path = output_path / shard_path.name + save_file(tensors, new_path) + print(f" āœ“ Saved {len(tensors)} tensors") + + # Cleanup + del ckpt, tensors + gc.collect() + + print(f"\nāœ… Converted {len(shards)} shards to safetensors format") + + +if __name__ == "__main__": + convert_shards("/data/models/Ornith-1.0-35B-nf4")