feat: add safetensors conversion script
This commit is contained in:
45
convert_to_safetensors.py
Normal file
45
convert_to_safetensors.py
Normal file
@@ -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")
|
||||||
Reference in New Issue
Block a user