diff --git a/dequant.py b/dequant.py index 1750cd0..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): """ @@ -45,6 +44,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 +225,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, diff --git a/nodes.py b/nodes.py index 7957bd5..d095de2 100644 --- a/nodes.py +++ b/nodes.py @@ -103,10 +103,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) @@ -115,9 +111,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 8c12de7..a6f47c1 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,67 +60,29 @@ 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(input=x, weight=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): prefix_len = len(prefix) for k,v in state_dict.items(): - if k[prefix_len:] == "weight": - self.weight = v - elif k[prefix_len:] == "bias": - self.bias = v + if k[prefix_len] == "weight": + self.weight = torch.nn.Parameter(v, requires_grad=False) + elif k[prefix_len] == "bias" and v is not None: + self.bias = torch.nn.Parameter(v, requires_grad=False) else: missing_keys.append(k) 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 - - 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 + destination[prefix + "bias"] = self.get_weight(self.bias) def get_weight(self, tensor, dtype): if tensor is None: @@ -129,7 +91,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) @@ -146,47 +107,48 @@ 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 = getattr(input, "dtype", torch.float32) + 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) + bias = comfy.ops.cast_to(bias, bias_dtype, device, non_blocking=non_blocking, copy=False) + + 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 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 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): - 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 GroupNorm(GGMLLayer): - def __init__(self, *args, device=None, dtype=None, **kwargs): - super().__init__() - num_groups = kwargs.get("num_groups") - if num_groups is None: - if len(args) == 0 or not isinstance(args[0], int): - raise ValueError("Couldn't get num_groups in GroupNorm: missing or bad type") - num_groups = args[0] - eps = kwargs.get("eps", 1e-05) - self._forward_operation = partial(staticmethod(torch.nn.functional.group_norm), num_groups=num_groups, eps=eps) - - class LayerNorm(GGMLLayer): - def __init__(self, *args, device=None, dtype=None, **kwargs): - super().__init__() - normalized_shape = kwargs.get("normalized_shape") - if normalized_shape is None: - if len(args) == 0 or not all(isinstance(v, int) for v in args): - raise ValueError("Couldn't get normalized_shape in LayerNorm: missing or bad type") - normalized_shape = tuple(args) - eps = kwargs.get("eps", 1e-05) - self._forward_operation = partial(staticmethod(torch.nn.functional.layer_norm), normalized_shape=normalized_shape, eps=eps) + 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) + 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):