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
This commit is contained in:
parent
43cb4fffc8
commit
67ce9cbf1c
10
comfy/ops.py
10
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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue