diff --git a/tools/convert.py b/tools/convert.py index 7b44300..748391b 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -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