diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index d70b42bf8..b5a521b1a 100644 --- a/comfy/model_patcher.py +++ b/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) diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index 7bbc96616..3771f849d 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/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()