upload fp8_scaled versions

convert_fp8.py

@@ -0,0 +1,148 @@
+import os
+import torch
+import argparse
+from tqdm.auto import tqdm
+from safetensors.torch import load_file, save_file
+from torch import dtype
+
+
+def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1):
+    _bits = torch.tensor(bits)
+    _mantissa_bit = torch.tensor(mantissa_bit)
+    _sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits)
+    E = _bits - _sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    return maxval
+
+
+def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1):
+    """
+    Default is E4M3.
+    """
+    bits = torch.tensor(bits)
+    mantissa_bit = torch.tensor(mantissa_bit)
+    sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits)
+    E = bits - sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    minval = - maxval
+    minval = - maxval if sign_bits == 1 else torch.zeros_like(maxval)
+    input_clamp = torch.min(torch.max(x, minval), maxval)
+    log_scales = torch.clamp(
+        (torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0)
+    log_scales = 2.0 ** (log_scales - M - bias.type(x.dtype))
+    # dequant
+    qdq_out = torch.round(input_clamp / log_scales) * log_scales
+    return qdq_out, log_scales
+
+
+def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1):
+    for i in range(len(x.shape) - 1):
+        scale = scale.unsqueeze(-1)
+    new_x = x / scale
+    quant_dequant_x, log_scales = quantize_to_fp8(
+        new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits)
+    return quant_dequant_x, scale, log_scales
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Convert safetensors to fp8 scaled",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--file",
+        type=str,
+        required=True,
+        help="Input .safetensors file to convert",
+    )
+    parser.add_argument(
+        "--base_dtype",
+        type=str,
+        default="bf16",
+        choices=["fp16", "bf16", "fp32"],
+        help="dtype to use for anything that can't be converted to fp8",
+    )
+    # parser.add_argument(
+    #     "--ban_list",
+    #     nargs="*",
+    #     default=[],
+    #     help="List of banned keys to keep in base dtype instead of converting to fp8 (zero or more strings)"
+    # )
+
+    args = parser.parse_args()
+    return args
+
+
+def main(args):
+    input_path = os.path.normpath(args.file)
+    output_path = os.path.splitext(input_path)[0] + "_fp8_scaled.safetensors"
+
+    orig_state_dict = load_file(input_path)
+    new_state_dict = {}
+
+    model_dtype: dtype = None
+
+    if args.base_dtype == "fp16":
+        model_dtype = torch.float16
+    elif args.base_dtype == "bf16":
+        model_dtype = torch.bfloat16
+    elif args.base_dtype == "fp32":
+        model_dtype = torch.float32
+    else:
+        raise Exception(f"unknown dtype: {args.base_dtype}")
+
+    # ban_list = ["text", "time", "head"]
+    # ban_list = args.ban_list
+    # ban_list = ["norm", "embedder", "pad_token", "modulation", "final_layer"]
+    # ban_list = ["norm", "embedder", "pad_token", "modulation", "final_layer", "to_q", "to_k", "to_v"] # for transformer, output will be ~6gb(bf16)    <- ex. z_image_turbo_bf16_fp8_scaled_1.safetensors
+    # ban_list = ["norm", "embedder", "pad_token", "modulation", "final_layer", "attention"] # for transformer, output will be ~8gb(bf16)   <- ex. z_image_turbo_bf16_fp8_scaled_2.safetensors
+    ban_list = ["norm", "embed_tokens"]  # for text encoder
+
+    maxval = get_fp_maxval()
+
+    for key in tqdm(orig_state_dict.keys()):
+        # decide whether to convert based on shape and banned keys
+        convert = False
+        if orig_state_dict[key].dim() == 2:
+            convert = True
+            for ban in ban_list:
+                if ban in key:
+                    convert = False
+
+        scale_key = key.rsplit(".", 1)[0] + ".scale_weight"
+
+        if convert:
+            weight = orig_state_dict[key]
+            scale = torch.max(torch.abs(weight.flatten())) / maxval
+
+            linear_weight, scale, log_scales = fp8_tensor_quant(weight, scale)
+            linear_weight = linear_weight.to(dtype=torch.float8_e4m3fn)
+
+            new_state_dict[scale_key] = scale
+            new_state_dict[key] = linear_weight
+
+        else:
+            if orig_state_dict[key].dim() == 2:
+                new_state_dict[scale_key] = torch.ones(1)
+
+            new_state_dict[key] = orig_state_dict[key].to(dtype=model_dtype)
+
+    new_state_dict["scaled_fp8"] = torch.zeros(2).to(dtype=torch.float8_e4m3fn)
+
+    save_file(new_state_dict, output_path)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    print(args.base_dtype)
+    main(args)

qwen_3_4b_bf16_fp8_scaled.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bc27405fbac59ac9998e52dfc34e66d09d0e0fd1025ead198f05b98760e3f68
+size 4411692798

z_image_turbo_bf16_fp8_scaled_1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ca63396607ef00904f3bd33e8ec3d4e362f225424f8c2e36b1aeb2d7c6390ee
+size 6293681826

z_image_turbo_bf16_fp8_scaled_2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:900e96b082a2010fc885fcd674fd10304ea0420de496fe0bedea169512679743
+size 8299083842

z_image_turbo_fp16_fp8_scaled_1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3822b59769b179d4e96e93ffd1cb968b38fadb1c882e1faec9ff4096818a1728
+size 6293681538

z_image_turbo_fp16_fp8_scaled_2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f251a69d1d6d0260c403dfaedeadffe360df785746960d27c58c045644cc3992
+size 8299083490

z_image_turbo_fp32_fp8_scaled_1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90024a47bc8f070a8916d367a25012f4a6a67d6974bee3ffe69e3db39450d823
+size 6571092026

z_image_turbo_fp32_fp8_scaled_2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d56f1cfe5409c4902d3d668717de7472d8cd6c592a06ffb1cc1dfad0210cd792
+size 12587297706