diff --git a/comfy/ops.py b/comfy/ops.py index 13c2604fb..1d3344975 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -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])) diff --git a/tests-unit/comfy_quant/test_fp8_ops.py b/tests-unit/comfy_quant/test_fp8_ops.py new file mode 100644 index 000000000..6e19b147a --- /dev/null +++ b/tests-unit/comfy_quant/test_fp8_ops.py @@ -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)