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