nodes/city96-ComfyUI-GGUF
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor 5D tensor fix logic

Browse Source
This commit is contained in:
City
2025-04-12 05:41:26 +02:00
parent 6de4bdba30
commit cc478deb5a
2 changed files with 71 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/convert.py
+58 -43
View File
@@ -14,13 +14,11 @@ MAX_TENSOR_DIMS = 4
class ModelTemplate:
arch = "invalid" # string describing architecture
shape_fix = False # whether to reshape tensors
ndims_fix = False # whether to save fix file for tensors exceeding max dims
keys_detect = [] # list of lists to match in state dict
keys_banned = [] # list of keys that should mark model as invalid for conversion
keys_hiprec = [] # list of keys that need to be kept in fp32 for some reason
def handle_nd_tensor(self, key, data):
raise NotImplementedError(f"Tensor detected that exceeds dims supported by C++ code! ({key} @ {data.shape})")
class ModelFlux(ModelTemplate):
arch = "flux"
keys_detect = [
@@ -47,6 +45,7 @@ class ModelAura(ModelTemplate):
class ModelHyVid(ModelTemplate):
arch = "hyvid"
ndims_fix = True
keys_detect = [
(
"double_blocks.0.img_attn_proj.weight",
@@ -54,17 +53,9 @@ class ModelHyVid(ModelTemplate):
)
]
def handle_nd_tensor(self, key, data):
# hacky but don't have any better ideas
path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here??
if os.path.isfile(path):
raise RuntimeError(f"5D tensor fix file already exists! {path}")
fsd = {key: torch.from_numpy(data)}
tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}")
save_file(fsd, path)
class ModelWan(ModelHyVid):
class ModelWan(ModelTemplate):
arch = "wan"
ndims_fix = True
keys_detect = [
(
"blocks.0.self_attn.norm_q.weight",
@@ -165,6 +156,24 @@ def strip_prefix(state_dict):
return sd
def find_main_dtype(state_dict, allow_fp32=False):
# detect most common dtype in input
dtypes = [x.dtype for x in state_dict.values()]
dtypes = {x:dtypes.count(x) for x in set(dtypes)}
main_dtype = max(dtypes, key=dtypes.get)
if main_dtype == torch.bfloat16:
ftype_name = "BF16"
ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16
elif main_dtype == torch.float32 and allow_fp32:
ftype_name = "F32"
ftype_gguf = gguf.LlamaFileType.ALL_F32
else:
ftype_name = "F16"
ftype_gguf = gguf.LlamaFileType.MOSTLY_F16
return ftype_name, ftype_gguf
def load_state_dict(path):
if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]):
state_dict = torch.load(path, map_location="cpu", weights_only=True)
@@ -179,7 +188,7 @@ def load_state_dict(path):
return strip_prefix(state_dict)
def handle_tensors(writer, state_dict, model_arch):
def handle_tensors(writer, state_dict, model_arch, allow_fp32=False):
name_lengths = tuple(sorted(
((key, len(key)) for key in state_dict.keys()),
key=lambda item: item[1],
@@ -188,9 +197,13 @@ def handle_tensors(writer, state_dict, model_arch):
if not name_lengths:
return
max_name_len = name_lengths[0][1]
if max_name_len > MAX_TENSOR_NAME_LENGTH:
bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH)
raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}")
invalid_tensors = {}
quantized_tensors = {}
for key, data in tqdm(state_dict.items()):
old_dtype = data.dtype
@@ -206,14 +219,14 @@ def handle_tensors(writer, state_dict, model_arch):
data_shape = data.shape
if old_dtype == torch.bfloat16:
data_qtype = gguf.GGMLQuantizationType.BF16
# elif old_dtype == torch.float32:
# data_qtype = gguf.GGMLQuantizationType.F32
elif old_dtype == torch.float32 and allow_fp32:
data_qtype = gguf.GGMLQuantizationType.F32
else:
data_qtype = gguf.GGMLQuantizationType.F16
# The max no. of dimensions that can be handled by the quantization code is 4
if len(data.shape) > MAX_TENSOR_DIMS:
model_arch.handle_nd_tensor(key, data)
invalid_tensors[key] = data
continue # needs to be added back later
# get number of parameters (AKA elements) in this tensor
@@ -247,38 +260,27 @@ def handle_tensors(writer, state_dict, model_arch):
try:
data = gguf.quants.quantize(data, data_qtype)
quantized_tensors[key] = data_qtype
except (AttributeError, gguf.QuantError) as e:
tqdm.write(f"falling back to F16: {e}")
data_qtype = gguf.GGMLQuantizationType.F16
data = gguf.quants.quantize(data, data_qtype)
new_name = key # do we need to rename?
quantized_tensors[key] = data_qtype
shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"
tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{new_name}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
tqdm.write(f"{f'%-{max_name_len + 4}s' % f'{key}'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
writer.add_tensor(new_name, data, raw_dtype=data_qtype)
writer.add_tensor(key, data, raw_dtype=data_qtype)
return quantized_tensors, invalid_tensors
def convert_file(path, dst_path=None, interact=True, overwrite=False):
def convert_file(path, dst_path=None, interact=True, overwrite=False, allow_fp32=False):
# load & run model detection logic
state_dict = load_state_dict(path)
model_arch = detect_arch(state_dict)
print(f"* Architecture detected from input: {model_arch.arch}")
# detect & set dtype for output file
dtypes = [x.dtype for x in state_dict.values()]
dtypes = {x:dtypes.count(x) for x in set(dtypes)}
main_dtype = max(dtypes, key=dtypes.get)
if main_dtype == torch.bfloat16:
ftype_name = "BF16"
ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16
# elif main_dtype == torch.float32:
# ftype_name = "F32"
# ftype_gguf = None
else:
ftype_name = "F16"
ftype_gguf = gguf.LlamaFileType.MOSTLY_F16
ftype_name, ftype_gguf = find_main_dtype(state_dict, allow_fp32=allow_fp32)
if dst_path is None:
dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf"
@@ -297,18 +299,31 @@ def convert_file(path, dst_path=None, interact=True, overwrite=False):
if ftype_gguf is not None:
writer.add_file_type(ftype_gguf)
handle_tensors(writer, state_dict, model_arch)
quantized_tensors, invalid_tensors = handle_tensors(writer, state_dict, model_arch, allow_fp32=allow_fp32)
if len(invalid_tensors) > 0:
if not model_arch.ndims_fix: # only applies to 5D fix for now, possibly expand to cover more cases?
raise ValueError(f"Tensor(s) detected that exceeds dims supported by C++ code! ({invalid_tensors.keys()})")
fix_path = os.path.join(
os.path.dirname(dst_path),
f"fix_5d_tensors_{model_arch.arch}.safetensors"
)
if os.path.isfile(fix_path):
raise RuntimeError(f"Tensor fix file already exists! {path}")
invalid_tensors = {k:torch.from_numpy(v.copy()) for k,v in invalid_tensors.items()}
save_file(invalid_tensors, fix_path)
print(f"\n### Warning! Fix file saved to '{fix_path}'")
print(f" you most likely need to run 'fix_5d_tensors.py' after quantization.")
else:
fix_path = None
writer.write_header_to_file(path=dst_path)
writer.write_kv_data_to_file()
writer.write_tensors_to_file(progress=True)
writer.close()
fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors"
if os.path.isfile(fix):
print(f"\n### Warning! Fix file found at '{fix}'")
print(f" you most likely need to run 'fix_5d_tensors.py' after quantization.")
return dst_path, model_arch
return dst_path, model_arch, fix_path
if __name__ == "__main__":
args = parse_args()
tools/fix_5d_tensors.py
+13 -10
View File
@@ -30,23 +30,21 @@ def get_file_type(reader):
ft = int(field.parts[field.data[-1]])
return gguf.LlamaFileType(ft)
if __name__ == "__main__":
args = get_args()
def apply_5d_fix(src, dst, fix=None, overwrite=False):
# read existing
reader = gguf.GGUFReader(args.src)
reader = gguf.GGUFReader(src)
arch = get_arch_str(reader)
file_type = get_file_type(reader)
print(f"Detected arch: '{arch}' (ftype: {str(file_type)})")
# prep fix
if args.fix is None:
args.fix = f"./fix_5d_tensors_{arch}.safetensors"
if fix is None:
fix = f"./fix_5d_tensors_{arch}.safetensors"
if not os.path.isfile(args.fix):
raise OSError(f"No 5D tensor fix file: {args.fix}")
if not os.path.isfile(fix):
raise OSError(f"No 5D tensor fix file: {fix}")
sd5d = load_file(args.fix)
sd5d = load_file(fix)
sd5d = {k:v.numpy() for k,v in sd5d.items()}
print("5D tensors:", sd5d.keys())
@@ -55,6 +53,7 @@ if __name__ == "__main__":
writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
writer.add_file_type(file_type)
global added
added = []
def add_extra_key(writer, key, data):
global added
@@ -76,7 +75,11 @@ if __name__ == "__main__":
if key not in added:
add_extra_key(writer, key, data)
writer.write_header_to_file(path=args.dst)
writer.write_header_to_file(path=dst)
writer.write_kv_data_to_file()
writer.write_tensors_to_file(progress=True)
writer.close()
if __name__ == "__main__":
args = get_args()
apply_5d_fix(args.src, args.dst, fix=args.fix, overwrite=args.overwrite)