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:
chelsealong 2026-08-08 00:46:48 +00:00
parent 43cb4fffc8
commit 67ce9cbf1c
2 changed files with 35 additions and 2 deletions

View File

@ -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.

View File

@ -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)