From 67ce9cbf1cc672430a3ea0cac89ea73779f78d99 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Sat, 8 Aug 2026 00:46:48 +0000 Subject: [PATCH 1/3] Fix crash loading quantized checkpoints with all-NUL comfy_quant markers Some quantizers mark unquantized layers in a mixed-precision checkpoint with an all-NUL comfy_quant placeholder tensor instead of omitting the key. json.loads() misdetects the leading NUL bytes as UTF-32 and raises UnicodeDecodeError; once that decode is patched around, the resulting empty {} config still fails with "Unknown quantization format" because the loader always requires an explicit format for any present marker. Treat an all-NUL marker as if comfy_quant were absent so the layer loads as a plain unquantized weight, matching the official MiniMax H3 nvfp4 AWQ text encoder. Fixes #15400 --- comfy/ops.py | 10 +++++-- .../comfy_quant/test_mixed_precision.py | 27 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/comfy/ops.py b/comfy/ops.py index 14599997b..314ba005b 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1151,7 +1151,10 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat layer_conf = state_dict.pop(f"{prefix}comfy_quant", None) if layer_conf is not None: - layer_conf = json.loads(layer_conf.numpy().tobytes()) + raw_conf = layer_conf.numpy().tobytes() + # Some quantizers mark unquantized layers with an all-NUL comfy_quant + # placeholder instead of omitting it; treat that the same as absent. + layer_conf = json.loads(raw_conf) if raw_conf.strip(b"\x00") else None if layer_conf is None: module.weight = torch.nn.Parameter(weight.to(device=device, dtype=compute_dtype), requires_grad=False) @@ -1587,7 +1590,10 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec weight_key = f"{prefix}weight" layer_conf = state_dict.pop(f"{prefix}comfy_quant", None) if layer_conf is not None: - layer_conf = json.loads(layer_conf.numpy().tobytes()) + raw_conf = layer_conf.numpy().tobytes() + # Some quantizers mark unquantized layers with an all-NUL comfy_quant + # placeholder instead of omitting it; treat that the same as absent. + layer_conf = json.loads(raw_conf) if raw_conf.strip(b"\x00") else None # Only fp8 and int8_tensorwise support per-row dequant via index select. # Block-scaled formats (NVFP4, MXFP8) can't do per-row lookup efficiently. diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index 7bbc96616..f697eaf8c 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -228,6 +228,33 @@ class TestMixedPrecisionOps(unittest.TestCase): with self.assertRaises(KeyError): model.load_state_dict(state_dict, strict=False) + def test_all_nul_comfy_quant_marker_loads_as_unquantized(self): + """Some quantizers mark unquantized layers with an all-NUL comfy_quant + placeholder instead of omitting the key; it must load as plain weight, + not crash decoding it as JSON or raise for a missing format.""" + state_dict = { + "layer1.weight": torch.randn(20, 10, dtype=torch.bfloat16), + "layer1.bias": torch.randn(20, dtype=torch.bfloat16), + "layer1.comfy_quant": torch.zeros(29, dtype=torch.uint8), + "layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16), + "layer2.bias": torch.randn(30, dtype=torch.bfloat16), + "layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16), + "layer3.bias": torch.randn(40, dtype=torch.bfloat16), + } + + model = SimpleModel(operations=ops.mixed_precision_ops({})) + model.load_state_dict(state_dict, strict=False) + + self.assertNotIsInstance(model.layer1.weight, QuantizedTensor) + + for layer in [model.layer1, model.layer2, model.layer3]: + layer.weight_function = [] + layer.bias_function = [] + + input_tensor = torch.randn(5, 10, dtype=torch.bfloat16) + output = model(input_tensor) + self.assertEqual(output.shape, (5, 40)) + def test_int8_convrot_metadata_loads_into_params(self): """ConvRot metadata must reach TensorWiseINT8Layout params.""" torch.manual_seed(123) From 6fc556d8416e4cec60c67ae4e2a48658cad9632d Mon Sep 17 00:00:00 2001 From: chelsealong Date: Sat, 8 Aug 2026 00:53:55 +0000 Subject: [PATCH 2/3] Infer quant format from weight dtype when comfy_quant omits it The all-NUL-marker fix only covered layers with no comfy_quant config. Layers with valid, non-empty comfy_quant JSON that carry a weight_scale but no "format" key (e.g. a q_proj-style layer in the MiniMax H3 nvfp4 AWQ checkpoint) still hit "Unknown quantization format for layer ...". Fall back to inferring the format from the on-disk weight dtype (int8 -> int8_tensorwise, float8_e4m3fn -> float8_e4m3fn, uint8 -> nvfp4) in both the shared _load_quantized_module path and the mixed-precision Embedding's own inline handling, matching the workaround documented in the issue. --- comfy/ops.py | 16 ++++++ .../comfy_quant/test_mixed_precision.py | 53 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/comfy/ops.py b/comfy/ops.py index 314ba005b..da415ff48 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1118,6 +1118,16 @@ def _quantized_apply(module, fn, recurse=True): return module +# Some quantizers write a comfy_quant payload with a weight_scale but no explicit +# "format" (e.g. the MiniMax H3 nvfp4 AWQ checkpoint). Infer the format from the +# on-disk weight dtype in that case, matching the storage dtypes each format uses. +_QUANT_FORMAT_BY_WEIGHT_DTYPE = { + torch.int8: "int8_tensorwise", + torch.float8_e4m3fn: "float8_e4m3fn", + torch.uint8: "nvfp4", +} + + def _load_quantized_module(module, super_load, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, load_extra_params=False): """Shared _load_from_state_dict body for quantized-weight modules. @@ -1160,6 +1170,8 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat module.weight = torch.nn.Parameter(weight.to(device=device, dtype=compute_dtype), requires_grad=False) else: module.quant_format = layer_conf.get("format", None) + if module.quant_format is None and f"{prefix}weight_scale" in state_dict: + module.quant_format = _QUANT_FORMAT_BY_WEIGHT_DTYPE.get(weight.dtype) module._full_precision_mm_config = layer_conf.get("full_precision_matrix_mult", False) if not module._full_precision_mm: module._full_precision_mm = module._full_precision_mm_config @@ -1598,6 +1610,10 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec # Only fp8 and int8_tensorwise support per-row dequant via index select. # Block-scaled formats (NVFP4, MXFP8) can't do per-row lookup efficiently. quant_format = layer_conf.get("format") if layer_conf is not None else None + if quant_format is None and layer_conf is not None and f"{prefix}weight_scale" in state_dict: + _stored_weight = state_dict.get(weight_key) + if _stored_weight is not None: + quant_format = _QUANT_FORMAT_BY_WEIGHT_DTYPE.get(_stored_weight.dtype) manually_loaded_keys = [] if quant_format in ("float8_e4m3fn", "float8_e5m2", "int8_tensorwise") and weight_key in state_dict: diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index f697eaf8c..f43058e13 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -255,6 +255,59 @@ class TestMixedPrecisionOps(unittest.TestCase): output = model(input_tensor) self.assertEqual(output.shape, (5, 40)) + def test_formatless_scaled_comfy_quant_infers_format_from_dtype(self): + """Some quantizers write a comfy_quant payload with a weight_scale but no + "format" key (e.g. a q_proj-style layer in a MiniMax H3 nvfp4 AWQ + checkpoint). The loader must infer the format from the on-disk weight + dtype instead of raising "Unknown quantization format".""" + state_dict = { + "layer1.weight": torch.randint(-128, 127, (20, 10), dtype=torch.int8), + "layer1.comfy_quant": torch.tensor(list(json.dumps({}).encode("utf-8")), dtype=torch.uint8), + "layer1.weight_scale": torch.ones(20), + "layer1.bias": torch.randn(20, dtype=torch.bfloat16), + "layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16), + "layer2.bias": torch.randn(30, dtype=torch.bfloat16), + "layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16), + "layer3.bias": torch.randn(40, dtype=torch.bfloat16), + } + + model = SimpleModel(operations=ops.mixed_precision_ops({})) + model.load_state_dict(state_dict, strict=False) + + self.assertIsInstance(model.layer1.weight, QuantizedTensor) + self.assertEqual(model.layer1.quant_format, "int8_tensorwise") + + for layer in [model.layer1, model.layer2, model.layer3]: + layer.weight_function = [] + layer.bias_function = [] + + input_tensor = torch.randn(5, 10, dtype=torch.bfloat16) + output = model(input_tensor) + self.assertEqual(output.shape, (5, 40)) + + def test_formatless_scaled_comfy_quant_embedding_infers_format_from_dtype(self): + """Same formatless-but-scaled scenario, but for the Embedding load path, + which has its own inline comfy_quant handling separate from + _load_quantized_module.""" + operations = ops.mixed_precision_ops({}) + + class EmbModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.emb = operations.Embedding(100, 20, device="cpu", dtype=torch.bfloat16) + + state_dict = { + "emb.weight": torch.randint(-128, 127, (100, 20), dtype=torch.int8), + "emb.comfy_quant": torch.tensor(list(json.dumps({}).encode("utf-8")), dtype=torch.uint8), + "emb.weight_scale": torch.ones(100), + } + + model = EmbModel() + model.load_state_dict(state_dict, strict=False) + + self.assertIsInstance(model.emb.weight, QuantizedTensor) + self.assertEqual(model.emb.quant_format, "int8_tensorwise") + def test_int8_convrot_metadata_loads_into_params(self): """ConvRot metadata must reach TensorWiseINT8Layout params.""" torch.manual_seed(123) From 9601f4742b413eb1d02f8494e04331d24d470ab0 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Sat, 8 Aug 2026 01:10:41 +0000 Subject: [PATCH 3/3] Reset stale quant state on reload and reject inferred NVFP4 embeddings Address CodeRabbit review on #15410: clear quant_format/layout_type when a module reloads an unquantized weight (previously stale state made Linear.forward take the quantized path against a plain Parameter), and raise instead of silently loading raw bytes when an Embedding's inferred format is NVFP4, which the embedding path can't dequantize. --- comfy/ops.py | 4 ++ .../comfy_quant/test_mixed_precision.py | 72 +++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/comfy/ops.py b/comfy/ops.py index da415ff48..840ed176b 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1167,6 +1167,8 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat layer_conf = json.loads(raw_conf) if raw_conf.strip(b"\x00") else None if layer_conf is None: + module.quant_format = None + module.layout_type = None module.weight = torch.nn.Parameter(weight.to(device=device, dtype=compute_dtype), requires_grad=False) else: module.quant_format = layer_conf.get("format", None) @@ -1614,6 +1616,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec _stored_weight = state_dict.get(weight_key) if _stored_weight is not None: quant_format = _QUANT_FORMAT_BY_WEIGHT_DTYPE.get(_stored_weight.dtype) + if quant_format == "nvfp4": + raise ValueError(f"NVFP4 embedding format is unsupported for layer {prefix.rstrip('.')}") manually_loaded_keys = [] if quant_format in ("float8_e4m3fn", "float8_e5m2", "int8_tensorwise") and weight_key in state_dict: diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index f43058e13..07486131e 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -308,6 +308,78 @@ class TestMixedPrecisionOps(unittest.TestCase): self.assertIsInstance(model.emb.weight, QuantizedTensor) self.assertEqual(model.emb.quant_format, "int8_tensorwise") + def test_reload_unquantized_resets_stale_quant_state(self): + """A module that previously loaded a quantized checkpoint must clear + quant_format/layout_type when reloaded with an unquantized checkpoint, + so forward doesn't take the stale quantized path against what is now + a plain Parameter.""" + layer_quant_config = { + "layer1": { + "format": "float8_e4m3fn", + "params": {} + } + } + fp8_weight = torch.randn(20, 10, dtype=torch.float32).to(torch.float8_e4m3fn) + state_dict1 = { + "layer1.weight": fp8_weight, + "layer1.bias": torch.randn(20, dtype=torch.bfloat16), + "layer1.weight_scale": torch.tensor(2.0, dtype=torch.float32), + "layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16), + "layer2.bias": torch.randn(30, dtype=torch.bfloat16), + "layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16), + "layer3.bias": torch.randn(40, dtype=torch.bfloat16), + } + state_dict1, _ = comfy.utils.convert_old_quants(state_dict1, metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}) + + model = SimpleModel(operations=ops.mixed_precision_ops({})) + model.load_state_dict(state_dict1, strict=False) + self.assertIsInstance(model.layer1.weight, QuantizedTensor) + + # Reload layer1 with a plain (unquantized) weight, no comfy_quant key. + state_dict2 = { + "layer1.weight": torch.randn(20, 10, dtype=torch.bfloat16), + "layer1.bias": torch.randn(20, dtype=torch.bfloat16), + "layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16), + "layer2.bias": torch.randn(30, dtype=torch.bfloat16), + "layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16), + "layer3.bias": torch.randn(40, dtype=torch.bfloat16), + } + model.load_state_dict(state_dict2, strict=False) + + self.assertNotIsInstance(model.layer1.weight, QuantizedTensor) + self.assertIsNone(model.layer1.quant_format) + self.assertIsNone(model.layer1.layout_type) + + for layer in [model.layer1, model.layer2, model.layer3]: + layer.weight_function = [] + layer.bias_function = [] + + input_tensor = torch.randn(5, 10, dtype=torch.bfloat16) + output = model(input_tensor) + self.assertEqual(output.shape, (5, 40)) + + def test_formatless_scaled_comfy_quant_embedding_rejects_nvfp4(self): + """A formatless comfy_quant payload that infers nvfp4 from a uint8 + weight dtype must raise, since the embedding load path has no + per-row dequant support for NVFP4; it must not silently load the + raw quantized bytes as an ordinary embedding weight.""" + operations = ops.mixed_precision_ops({}) + + class EmbModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.emb = operations.Embedding(100, 20, device="cpu", dtype=torch.bfloat16) + + state_dict = { + "emb.weight": torch.randint(0, 255, (100, 20), dtype=torch.uint8), + "emb.comfy_quant": torch.tensor(list(json.dumps({}).encode("utf-8")), dtype=torch.uint8), + "emb.weight_scale": torch.ones(100), + } + + model = EmbModel() + with self.assertRaises(ValueError): + model.load_state_dict(state_dict, strict=False) + def test_int8_convrot_metadata_loads_into_params(self): """ConvRot metadata must reach TensorWiseINT8Layout params.""" torch.manual_seed(123)