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

Merge upstream ops_refactor

Browse Source
This commit is contained in:
blepping
2024-09-01 03:02:59 -06:00
parent ed0702fefc 6618801f1c
commit 436ecb4913
3 changed files with 54 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dequant.py
+10 -6
View File
@@ -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,
nodes.py
-7
View File
@@ -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
ops.py
+44 -82
View File
@@ -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):