Enable more torch.compile optimization by reducing graph breaks (#243)
As title, there were 2 fundamental sources of graph breaks: 1. using a pre-existing set in a `torch.compile` region, this is a [known issue](https://github.com/pytorch/pytorch/issues/145761) in PyTorch and the fix is WIP, but it's a lot easier to change the set into a tuple here to fix the graph break. 2. using tensor subclass `GGMLTensor`. The full support of tensor subclass has been gated for `torch.compile`, but with [a recent WIP stack](https://github.com/pytorch/pytorch/pull/149792), we'll be able to remove the gating and have `torch.compile` fully support the `GGMLTensor` in this repo. The only required change here is avoid changing the `__class__` attribute, and instead using `torch.Tensor` to convert the subclass object back to Tensor.
This commit is contained in:
+1
-1
@@ -4,7 +4,7 @@ import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
TORCH_COMPATIBLE_QTYPES = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
|
||||
TORCH_COMPATIBLE_QTYPES = (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16)
|
||||
|
||||
def is_torch_compatible(tensor):
|
||||
return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES
|
||||
|
||||
@@ -6,15 +6,35 @@ import comfy.ops
|
||||
import comfy.model_management
|
||||
from .dequant import dequantize_tensor, is_quantized
|
||||
|
||||
# to avoid breaking really old pytorch versions
|
||||
if hasattr(torch, "compiler") and hasattr(torch.compiler, "disable"):
|
||||
torch_compiler_disable = torch.compiler.disable
|
||||
else:
|
||||
def torch_compiler_disable(*args, **kwargs):
|
||||
def chained_hasattr(obj, chained_attr):
|
||||
probe = obj
|
||||
for attr in chained_attr.split('.'):
|
||||
if hasattr(probe, attr):
|
||||
probe = getattr(probe, attr)
|
||||
else:
|
||||
return False
|
||||
return True
|
||||
|
||||
# A bakcward and forward compatible way to get `torch.compiler.disable`.
|
||||
def get_torch_compiler_disable_decorator():
|
||||
def dummy_decorator(*args, **kwargs):
|
||||
def noop(x):
|
||||
return x
|
||||
return noop
|
||||
|
||||
from packaging import version
|
||||
|
||||
if not chained_hasattr(torch, "compiler.disable"):
|
||||
return dummy_decorator # torch too old
|
||||
elif version.parse(torch.__version__) >= version.parse("2.8"):
|
||||
return dummy_decorator # torch compile works
|
||||
if chained_hasattr(torch, "_dynamo.config.nontraceable_tensor_subclasses"):
|
||||
return dummy_decorator # torch compile works, nightly before 2.8 release
|
||||
else:
|
||||
return torch.compiler.disable
|
||||
|
||||
torch_compiler_disable = get_torch_compiler_disable_decorator()
|
||||
|
||||
class GGMLTensor(torch.Tensor):
|
||||
"""
|
||||
Main tensor-like class for storing quantized weights
|
||||
@@ -149,7 +169,7 @@ class GGMLLayer(torch.nn.Module):
|
||||
|
||||
# prevent propagating custom tensor class
|
||||
if isinstance(weight, GGMLTensor):
|
||||
weight.__class__ = torch.Tensor
|
||||
weight = torch.Tensor(weight)
|
||||
|
||||
# apply patches
|
||||
if patch_list:
|
||||
@@ -189,7 +209,7 @@ class GGMLLayer(torch.nn.Module):
|
||||
|
||||
# non-ggml forward might still propagate custom tensor class
|
||||
if isinstance(out, GGMLTensor):
|
||||
out.__class__ = torch.Tensor
|
||||
out = torch.Tensor(out)
|
||||
return out
|
||||
|
||||
def forward_ggml_cast_weights(self, input):
|
||||
|
||||
Reference in New Issue
Block a user