File size: 2,516 Bytes
59e9071 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | #!/usr/bin/env python3
"""
Simple FP32 to FP16 converter for safetensors files.
Converts ALL FP32 tensors to FP16 without discrimination.
USAGE:
python convert_fp16.py input.safetensors [output.safetensors]
"""
import argparse
from pathlib import Path
from typing import Dict
import torch
from safetensors.torch import load_file, save_file
def convert_fp32_to_fp16(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
"""Convert all FP32 tensors to FP16."""
converted = {}
fp32_count = 0
total_tensors = len(state_dict)
for name, tensor in state_dict.items():
if tensor.dtype == torch.float32:
converted[name] = tensor.half() # Convert to FP16
fp32_count += 1
else:
converted[name] = tensor # Keep as-is
print(f"Converted {fp32_count}/{total_tensors} tensors from FP32 to FP16")
return converted
def calculate_size_reduction(original: Dict[str, torch.Tensor], converted: Dict[str, torch.Tensor]) -> None:
"""Calculate and print size reduction."""
original_bytes = sum(t.numel() * t.element_size() for t in original.values())
converted_bytes = sum(t.numel() * t.element_size() for t in converted.values())
saved_bytes = original_bytes - converted_bytes
original_gb = original_bytes / (1024**3)
converted_gb = converted_bytes / (1024**3)
saved_gb = saved_bytes / (1024**3)
print(f"Size: {original_gb:.3f} GB → {converted_gb:.3f} GB (saved {saved_gb:.3f} GB)")
def main():
parser = argparse.ArgumentParser(description="Convert all FP32 tensors to FP16 in a safetensors file")
parser.add_argument("input", help="Input safetensors file")
parser.add_argument("output", nargs="?", help="Output safetensors file (optional)")
args = parser.parse_args()
# Generate output filename if not provided
if args.output is None:
input_path = Path(args.input)
stem = input_path.stem
if not stem.endswith("_fp16"):
stem += "_fp16"
args.output = str(input_path.with_name(stem + input_path.suffix))
print(f"Loading: {args.input}")
state_dict = load_file(args.input)
print("Converting FP32 → FP16...")
converted_state_dict = convert_fp32_to_fp16(state_dict)
calculate_size_reduction(state_dict, converted_state_dict)
print(f"Saving: {args.output}")
save_file(converted_state_dict, args.output)
print("Done!")
if __name__ == "__main__":
main() |