From cc478deb5ab236a77f62d640b0b1885635447f82 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 12 Apr 2025 05:41:26 +0200 Subject: [PATCH] Refactor 5D tensor fix logic --- tools/convert.py | 101 +++++++++++++++++++++++----------------- tools/fix_5d_tensors.py | 23 +++++---- 2 files changed, 71 insertions(+), 53 deletions(-) diff --git a/tools/convert.py b/tools/convert.py index 6c5669d..c7dd146 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -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() diff --git a/tools/fix_5d_tensors.py b/tools/fix_5d_tensors.py index 0e61d1c..0826638 100644 --- a/tools/fix_5d_tensors.py +++ b/tools/fix_5d_tensors.py @@ -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)