From 826bd1804cc67e59bb23a9cbcfdf2ec65304b984 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 30 Aug 2024 03:25:57 +0200 Subject: [PATCH 1/5] Make ops less jank Also possibly faster. Ignore the massively increased system memory usage. --- ops.py | 81 +++++++++++++++++++++------------------------------------- 1 file changed, 29 insertions(+), 52 deletions(-) diff --git a/ops.py b/ops.py index 88ca894..d78f982 100644 --- a/ops.py +++ b/ops.py @@ -1,6 +1,6 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import gguf import torch -from functools import partial import comfy.ops from .dequant import dequantize_tensor @@ -60,35 +60,12 @@ class GGMLLayer(torch.nn.Module): dequant_dtype = None patch_dtype = None - def __init__(self, *args, **kwargs): - super().__init__() - self.weight = GGMLTensor(1, tensor_type=None, tensor_shape=None) - self.bias = None - - def _forward_operation(self, x, weight, bias): - raise NotImplementedError - - def forward(self, x): - # lowvram hack - device = None - if self.weight.device != x.device: - device = self.weight.device - self.to(x.device) - - weight, bias = self.get_weights(x.dtype) - x = self._forward_operation(x, weight, bias=bias) - del weight, bias - - if device: - self.to(device) - return x - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): for k,v in state_dict.items(): if k[len(prefix):] == "weight": - self.weight = v - elif k[len(prefix):] == "bias": - self.bias = v + self.weight = torch.nn.Parameter(v, requires_grad=False) + elif k[len(prefix):] == "bias" and v is not None: + self.bias = torch.nn.Parameter(v, requires_grad=False) else: missing_keys.append(k) @@ -109,18 +86,6 @@ class GGMLLayer(torch.nn.Module): if bias is not None: destination[f"{prefix}bias"] = weight - def _apply(self, fn): - if self.weight is not None: - try: - self.weight = fn(self.weight) - except TypeError: - # TODO: Figure out why this happens - pass - if self.bias is not None: - self.bias = fn(self.bias) - super()._apply(fn) - return self - def get_weight(self, tensor, dtype): if tensor is None: return @@ -128,7 +93,6 @@ class GGMLLayer(torch.nn.Module): # consolidate and load patches to GPU in async patch_list = [] device = tensor.device - t_move = lambda x: x.to(device) if torch.is_tensor(x) else x for function, patches, key in getattr(tensor, "patches", []): patch_list += move_patch_to_cuda(patches, device) @@ -145,25 +109,38 @@ class GGMLLayer(torch.nn.Module): weight = function(patch_list, weight, key, patch_dtype) return weight - def get_weights(self, dtype=torch.float16): - weight = self.get_weight(self.weight, dtype) - bias = self.get_weight(self.bias, dtype) - return (weight, bias) + def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): + if input is not None: + if dtype is None: + dtype = input.dtype + if bias_dtype is None: + bias_dtype = dtype + if device is None: + device = input.device + bias = None + non_blocking = comfy.model_management.device_supports_non_blocking(device) + if s.bias is not None: + bias = s.get_weight(s.bias.to(device), dtype) if hasattr(s, "get_weight") else s.bias + bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) + + weight = s.get_weight(s.weight.to(device), dtype) if hasattr(s, "get_weight") else s.weight + weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False) + return weight, bias class GGMLOps(comfy.ops.manual_cast): """ Dequantize weights on the fly before doing the compute """ - class Linear(GGMLLayer): - _forward_operation = staticmethod(torch.nn.functional.linear) - - class Conv2d(GGMLLayer): - def __init__(self, *args, device=None, dtype=None, **kwargs): - super().__init__() - _ = kwargs.pop("kernel_size", None) - self._forward_operation = partial(staticmethod(torch.nn.functional.conv2d), **kwargs) + class Linear(GGMLLayer, comfy.ops.manual_cast.Linear): + def forward_comfy_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return torch.nn.functional.linear(input, weight, bias) + class Conv2d(GGMLLayer, comfy.ops.manual_cast.Conv2d): + def forward_comfy_cast_weights(self, input): + weight, bias = self.cast_bias_weight(input) + return self._conv_forward(input, weight, bias) def move_patch_to_cuda(item, device): if isinstance(item, torch.Tensor): From f40b149bc5117582828c399a382d60c9781896e6 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 30 Aug 2024 03:39:57 +0200 Subject: [PATCH 2/5] Add BF16 support --- dequant.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/dequant.py b/dequant.py index 1750cd0..454bd8c 100644 --- a/dequant.py +++ b/dequant.py @@ -45,6 +45,10 @@ def split_block_dims(blocks, *args): dims = list(args) + [n_max - sum(args)] return torch.split(blocks, dims, dim=1) +# Full weights # +def dequantize_blocks_BF16(blocks, block_size, type_size, dtype=None): + return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32) + # Legacy Quants # def dequantize_blocks_Q8_0(blocks, block_size, type_size, dtype=None): d, x = split_block_dims(blocks, 2) @@ -222,6 +226,7 @@ def dequantize_blocks_Q2_K(blocks, block_size, type_size, dtype=None): return qs.reshape((n_blocks, -1)) dequantize_functions = { + gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1, gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0, From ec24e75d89402a2f4d113d01f0e255a006632b31 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 30 Aug 2024 03:40:15 +0200 Subject: [PATCH 3/5] Remove unused check --- ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ops.py b/ops.py index d78f982..dbdaed9 100644 --- a/ops.py +++ b/ops.py @@ -121,10 +121,10 @@ class GGMLLayer(torch.nn.Module): bias = None non_blocking = comfy.model_management.device_supports_non_blocking(device) if s.bias is not None: - bias = s.get_weight(s.bias.to(device), dtype) if hasattr(s, "get_weight") else s.bias + bias = s.get_weight(s.bias.to(device), dtype) bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) - weight = s.get_weight(s.weight.to(device), dtype) if hasattr(s, "get_weight") else s.weight + weight = s.get_weight(s.weight.to(device), dtype) weight = comfy.ops.cast_to(weight, dtype, device, non_blocking=non_blocking, copy=False) return weight, bias From c231239fdfb6e3c484f091c6ea004f767ee2ea1f Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 30 Aug 2024 03:47:52 +0200 Subject: [PATCH 4/5] cleanup Perf difference from `torch.tensor(tensor.data)` doesn't seem realistically measurable and the warning is annoying. --- dequant.py | 11 +++++------ ops.py | 13 +++++-------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/dequant.py b/dequant.py index 454bd8c..3688d7f 100644 --- a/dequant.py +++ b/dequant.py @@ -4,20 +4,19 @@ import torch from tqdm import tqdm def dequantize_tensor(tensor, dtype=None, dequant_dtype=None): - data = torch.tensor(tensor.data) qtype = getattr(tensor, "tensor_type", None) - oshape = getattr(tensor, "tensor_shape", data.shape) + oshape = getattr(tensor, "tensor_shape", tensor.shape) if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16): - return data.to(dtype) + return tensor.to(dtype) elif qtype in dequantize_functions: dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype - return dequantize(data, qtype, oshape, dtype=dequant_dtype).to(dtype) + return dequantize(tensor, qtype, oshape, dtype=dequant_dtype).to(dtype) else: # this is incredibly slow tqdm.write(f"Falling back to numpy dequant for qtype: {qtype}") - new = gguf.quants.dequantize(data.cpu().numpy(), qtype) - return torch.from_numpy(new).to(data.device, dtype=dtype) + new = gguf.quants.dequantize(tensor.cpu().numpy(), qtype) + return torch.from_numpy(new).to(tensor.device, dtype=dtype) def dequantize(data, qtype, oshape, dtype=None): """ diff --git a/ops.py b/ops.py index dbdaed9..43f6299 100644 --- a/ops.py +++ b/ops.py @@ -71,20 +71,17 @@ class GGMLLayer(torch.nn.Module): def _save_to_state_dict(self, destination, prefix, keep_vars): # This is a fake state dict for vram estimation - if self.weight is not None: - weight = torch.zeros_like(self.weight, device=torch.device("meta")) - destination[f"{prefix}weight"] = weight + weight = torch.zeros_like(self.weight, device=torch.device("meta")) + destination[prefix + "weight"] = weight if self.bias is not None: bias = torch.zeros_like(self.bias, device=torch.device("meta")) - destination[f"{prefix}bias"] = bias + destination[prefix + "bias"] = bias return # This would return the actual state dict - weight, bias = self.get_weights() - if weight is not None: - destination[f"{prefix}weight"] = weight + destination[prefix + "weight"] = self.get_weight(self.weight) if bias is not None: - destination[f"{prefix}bias"] = weight + destination[prefix + "bias"] = self.get_weight(self.bias) def get_weight(self, tensor, dtype): if tensor is None: From 6618801f1c1ee6c4203e4a76b87d17f3a2a3d3eb Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Fri, 30 Aug 2024 06:25:19 +0200 Subject: [PATCH 5/5] Don't dequantize token embed on load --- nodes.py | 7 ------- ops.py | 12 +++++++++++- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/nodes.py b/nodes.py index 70823f4..9b7bede 100644 --- a/nodes.py +++ b/nodes.py @@ -78,10 +78,6 @@ clip_sd_map = { "ffn_gate": "layer.1.DenseReluDense.wi_0", "ffn_norm": "layer.1.layer_norm", } -# weights that should be dequantized on load -clip_sd_dequant = { - "shared.weight", -} def gguf_clip_loader(path): raw_sd = gguf_sd_loader(path) @@ -90,9 +86,6 @@ def gguf_clip_loader(path): for k,v in raw_sd.items(): for s,d in clip_sd_map.items(): k = k.replace(s,d) - if k in clip_sd_dequant: - v = dequantize_tensor(v, torch.float32).to(torch.float16) - v = GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape) sd[k] = v return sd diff --git a/ops.py b/ops.py index 43f6299..bc17c3a 100644 --- a/ops.py +++ b/ops.py @@ -109,7 +109,7 @@ class GGMLLayer(torch.nn.Module): def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None): if input is not None: if dtype is None: - dtype = input.dtype + dtype = getattr(input, "dtype", torch.float32) if bias_dtype is None: bias_dtype = dtype if device is None: @@ -139,6 +139,16 @@ class GGMLOps(comfy.ops.manual_cast): weight, bias = self.cast_bias_weight(input) return self._conv_forward(input, weight, bias) + class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding): + def forward_comfy_cast_weights(self, input, out_dtype=None): + output_dtype = out_dtype + if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16: + out_dtype = None + weight, bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype) + return torch.nn.functional.embedding( + input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse + ).to(dtype=output_dtype) + def move_patch_to_cuda(item, device): if isinstance(item, torch.Tensor): return item.to(device, non_blocking=True)