fallback to F32 for 1-dimensional or very small tensors (#20)

* change fallback data type from F16 to F32

* only fallback to F32 for 1D or very small tensors

* change exception fallback type

* blacklist before small tensor check
This commit is contained in:
ddh0
2024-08-16 06:43:59 +02:00
committed by GitHub
parent de5f6e9dcc
commit 51c95d7d70
+26 -14
View File
@@ -83,22 +83,15 @@ def handle_tensors(args, writer, state_dict):
n_dims = len(data.shape)
data_qtype = args.qtype
data_shape = data.shape
# get number of parameters (AKA elements) in this tensor
n_params = 1
for dim_size in data_shape:
n_params *= dim_size
fallback = gguf.GGMLQuantizationType.F16
if n_dims == 1:
# these barely take up space, just use F32 / F16.
# also speeds up inference due to not dequantizing
data_qtype = fallback
elif n_dims == 4:
if min(data.shape[:2]) == 4: # output tensor
data_qtype = fallback
elif data.shape[-1] == 3: # 3x3 kernel
data_qtype = fallback
elif data.shape[-1] == 1: # 1x1 kernel
#data = np.squeeze(data) # don't do this
data_qtype = fallback
# other keys to keep as max precision
# keys to keep as max precision
blacklist = [
"time_embedding.",
"add_embedding.",
@@ -109,9 +102,28 @@ def handle_tensors(args, writer, state_dict):
"guidance_in.",
"final_layer.",
]
if any([x in key for x in blacklist]) and ".weight" in key:
data_qtype = fallback
if n_dims == 1:
# one-dimensional tensors should be kept in F32
# also speeds up inference due to not dequantizing
data_qtype = gguf.GGMLQuantizationType.F32
elif n_params <= 1024:
# very small tensors
data_qtype = gguf.GGMLQuantizationType.F32
elif n_dims == 4:
if min(data.shape[:2]) == 4: # output tensor
data_qtype = fallback
elif data_shape[-1] == 3: # 3x3 kernel
data_qtype = fallback
elif data_shape[-1] == 1: # 1x1 kernel
#data = np.squeeze(data) # don't do this
data_qtype = fallback
# TODO: find keys to keep in higher precision(s) / qtypes
# if "time_emb_proj.weight" in key:
# data_qtype = gguf.GGMLQuantizationType.F16