Merge b20113525f into 71b73e3b2b

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)
+        scale_input = torch.ones((), device=input.device, dtype=torch.float32)
-    input = torch.clamp(input, min=-448, max=448, out=input)
+        input = torch.clamp(input, min=-448, max=448, out=input)
-    input_fp8 = input.to(dtype).contiguous()
+        input_fp8 = input.to(dtype).contiguous()
-    layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
+        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)
+        quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
 
-    # Wrap weight in QuantizedTensor - this enables unified dispatch
+        # Wrap weight in QuantizedTensor - this enables unified dispatch
-    # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
+        # 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))
+        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)
+        quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
-    o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
+        o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
-
+    finally:
-    uncast_bias_weight(self, w, bias, offload_stream)
+        uncast_bias_weight(self, w, bias, offload_stream)
     if tensor_3d:
         o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))
 

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)