Upload quantize.py

quantize.py

@@ -0,0 +1,26 @@
+import sys
+
+from safetensors.torch import load_file, save_file
+from torch import float8_e4m3fn, float8_e5m2
+
+if len(sys.argv) != 4:
+    print("Provide input/output file names and fp8 type as either e4m3 or e5m2")
+    exit(1)
+
+if sys.argv[3] == "e4m3":
+    dt = float8_e4m3fn
+elif sys.argv[3] == "e5m2":
+    dt = float8_e5m2
+else:
+    print("Invalid quantization type, should be either e4m3 or e5m2")
+    exit(1)
+
+state_dict = load_file(sys.argv[1])
+
+for k in state_dict:
+    if "norm" in k or "bias" in "k":
+        state_dict[k] = state_dict[k].bfloat16()
+    else:
+        state_dict[k] = state_dict[k].to(dt)
+
+save_file(state_dict, sys.argv[2])