wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-18 12:28:17 +08:00
Code Issues Actions 26 Packages Projects Releases Wiki Activity

Merge b20113525f into 71b73e3b2b

Browse Source
This commit is contained in:
liminfei-amd 2026-07-17 20:35:45 +08:00 committed by GitHub
commit 856e2dfa5f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 91 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
comfy/ops.py
View File

@ -238,6 +238,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
def post_cast(s, param_key, x, dtype, resident, update_weight):
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
fns = getattr(s, param_key + "_function", [])
requant = want_requant and param_key == "weight"
if x is None:
return None
@ -255,13 +256,13 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
if not resident and lowvram_fn is not None:
x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
x = lowvram_fn(x)
if (want_requant and len(fns) == 0 or update_weight):
if (requant and len(fns) == 0 or update_weight):
seed = comfy.utils.string_to_seed(s.seed_key)
if isinstance(orig, QuantizedTensor):
y = orig.requantize_from_float(x, scale="recalculate", stochastic_rounding=seed)
else:
y = comfy.float.stochastic_rounding(x, orig.dtype, seed=seed)
if want_requant and len(fns) == 0:
if requant and len(fns) == 0:
x = y
if update_weight:
orig.copy_(y)
@ -831,21 +832,22 @@ def fp8_linear(self, input):
return None
lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True)
scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
try:
scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
scale_input = torch.ones((), device=input.device, dtype=torch.float32)
input = torch.clamp(input, min=-448, max=448, out=input)
input_fp8 = input.to(dtype).contiguous()
layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
scale_input = torch.ones((), device=input.device, dtype=torch.float32)
input = torch.clamp(input, min=-448, max=448, out=input)
input_fp8 = input.to(dtype).contiguous()
layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
# Wrap weight in QuantizedTensor - this enables unified dispatch
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
uncast_bias_weight(self, w, bias, offload_stream)
# Wrap weight in QuantizedTensor - this enables unified dispatch
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
finally:
uncast_bias_weight(self, w, bias, offload_stream)
if tensor_3d:
o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))

74
tests-unit/comfy_quant/test_fp8_ops.py Normal file
View File

@ -0,0 +1,74 @@
import os
import sys
import unittest
from types import SimpleNamespace
from unittest import mock
import torch
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
from comfy.cli_args import args
if not torch.cuda.is_available():
args.cpu = True
from comfy import ops
class TestFP8Ops(unittest.TestCase):
def test_vbar_requantizes_weight_but_not_bias(self):
weight = torch.zeros((16, 16), dtype=torch.float8_e4m3fn)
bias = torch.zeros(16, dtype=torch.float8_e4m3fn)
patch = lambda tensor: tensor
layer = SimpleNamespace(
weight=weight,
bias=bias,
weight_function=[],
bias_function=[],
weight_lowvram_function=patch,
bias_lowvram_function=patch,
seed_key="layer",
_prefetch={
"resident": False,
"xfer_dest": torch.empty(1, dtype=torch.uint8),
"needs_cast": False,
"cast_geometry": None,
"signature": (1,),
},
)
with (
mock.patch.object(ops.comfy.memory_management, "interpret_gathered_like", return_value=[weight, bias]),
mock.patch.object(ops.comfy.float, "stochastic_rounding", side_effect=lambda tensor, dtype, seed: tensor.to(dtype)) as rounding,
):
cast_weight, cast_bias = ops.resolve_cast_module_with_vbar(
layer,
torch.float8_e4m3fn,
torch.device("cpu"),
torch.float16,
torch.float16,
True,
)
self.assertEqual(cast_weight.dtype, torch.float8_e4m3fn)
self.assertEqual(cast_bias.dtype, torch.float16)
self.assertEqual(rounding.call_count, 2)
def test_fp8_linear_unpins_vbar_when_linear_fails(self):
weight = torch.zeros((16, 16), dtype=torch.float8_e4m3fn)
bias = torch.zeros(16, dtype=torch.float8_e4m3fn)
vbar = object()
layer = SimpleNamespace(weight=weight, _v=vbar)
input_tensor = torch.zeros((2, 16), dtype=torch.float16)
offload_info = (None, torch.device("cpu"), None)
with (
mock.patch.object(ops, "cast_bias_weight", return_value=(weight, bias, offload_info)),
mock.patch.object(ops.comfy_aimdo.model_vbar, "vbar_unpin") as unpin,
mock.patch.object(torch.nn.functional, "linear", side_effect=RuntimeError("FP8 linear failed")),
):
with self.assertRaisesRegex(RuntimeError, "FP8 linear failed"):
ops.fp8_linear(layer, input_tensor)
unpin.assert_called_once_with(vbar)