From b20113525fed2b138fda01c489cbf265c7628a6e Mon Sep 17 00:00:00 2001 From: liminfei Date: Fri, 17 Jul 2026 20:34:39 +0800 Subject: [PATCH] ops: keep FP8 LoRA biases in compute dtype The dynamic VBAR cast path applied want_requant to both weights and biases. For an FP8 checkpoint with a low-VRAM LoRA patch, this returned an FP8 bias to the optimized linear operation even though the activation dtype was FP16. Restrict the optional requantized return value to weights while preserving resident-storage updates for both parameters. Make the legacy FP8 linear path release its VBAR/offload resources in a finally block so backend failures do not leave pages pinned before the normal fallback runs. Add regressions for the weight/bias dtype contract and exception cleanup. Fixes #14952 Signed-off-by: liminfei --- comfy/ops.py | 32 +++++------ tests-unit/comfy_quant/test_fp8_ops.py | 74 ++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 15 deletions(-) create mode 100644 tests-unit/comfy_quant/test_fp8_ops.py 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)