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:
Xiangxi Guo (Ryan)
2025-04-10 02:46:33 +02:00
committed by GitHub
parent bc5223b0e3
commit 0c3bdf336e
2 changed files with 28 additions and 8 deletions
+1 -1
View File
@@ -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
+27 -7
View File
@@ -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):