Merge 516f3ee922 into 1d1099bea0

comfy/model_patcher.py

@@ -165,6 +165,19 @@ def low_vram_patch_estimate_vram(model, key):
 
     return weight.numel() * model_dtype.itemsize * LOWVRAM_PATCH_ESTIMATE_MATH_FACTOR
 
+def _collect_quantized_weight_state_dict_keys(model):
+    """Return state-dict keys owned by each module's QuantizedTensor weight."""
+    keys = set()
+    for module_name, module in model.named_modules(remove_duplicate=False):
+        weight = getattr(module, "weight", None)
+        if not isinstance(weight, QuantizedTensor):
+            continue
+        prefix = f"{module_name}." if module_name else ""
+        weight_key = f"{prefix}weight"
+        keys.update(k for k in weight.state_dict(weight_key) if k != weight_key)
+        keys.add(f"{prefix}comfy_quant")
+    return keys
+
 def get_key_weight(model, key):
     set_func = None
     convert_func = None
@@ -812,11 +825,14 @@ class ModelPatcher:
 
     def get_key_patches(self, filter_prefix=None):
         model_sd = self.model_state_dict()
+        quantized_weight_state_dict_keys = _collect_quantized_weight_state_dict_keys(self.model)
         p = {}
         for k in model_sd:
             if filter_prefix is not None:
                 if not k.startswith(filter_prefix):
                     continue
+            if k in quantized_weight_state_dict_keys:
+                continue
             bk = self.backup.get(k, None)
             hbk = self.hook_backup.get(k, None)
             weight, set_func, convert_func = get_key_weight(self.model, k)

tests-unit/comfy_quant/test_mixed_precision.py

@@ -15,7 +15,8 @@ if not has_gpu():
     args.cpu = True
 
 from comfy import ops
-from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor
+from comfy.model_patcher import ModelPatcher
+from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, TensorCoreFP8E4M3Layout
 import comfy.utils
 
 
@@ -337,5 +338,35 @@ class TestMixedPrecisionOps(unittest.TestCase):
         self.assertEqual(saved_conf["linear_dtype"], "int8")
         self.assertNotIn("quant_group_size", saved_conf)
 
+    def test_get_key_patches_skips_only_quantized_weight_pieces(self):
+        operations = ops.mixed_precision_ops(compute_dtype=torch.float32)
+        model = torch.nn.Module()
+        model.linear = operations.Linear(4, 4, bias=False, device="cpu")
+        qdata, params = TensorCoreFP8E4M3Layout.quantize(
+            torch.ones(4, 4), scale="recalculate"
+        )
+        model.linear.quant_format = "float8_e4m3fn"
+        model.linear.layout_type = "TensorCoreFP8E4M3Layout"
+        model.linear.weight = torch.nn.Parameter(
+            QuantizedTensor(qdata, model.linear.layout_type, params),
+            requires_grad=False,
+        )
+        model.linear.input_scale = torch.nn.Parameter(
+            torch.tensor(0.125), requires_grad=False
+        )
+        model.linear_alias = model.linear
+
+        patcher = ModelPatcher(model, torch.device("cpu"), torch.device("cpu"))
+
+        self.assertEqual(
+            set(patcher.get_key_patches()),
+            {
+                "linear.weight",
+                "linear.input_scale",
+                "linear_alias.weight",
+                "linear_alias.input_scale",
+            },
+        )
+
 if __name__ == "__main__":
     unittest.main()