Fix quant metadata marker collision handling
This commit is contained in:
Alan Li
parent
20467f61ee
commit
50af1a3563
|
|
@ -1401,15 +1401,8 @@ def detect_layer_quantization(state_dict, prefix):
|
|||
return None
|
||||
|
||||
def _resolve_quant_metadata_layer_key(state_dict, layer_key, model_prefix):
|
||||
# _quantization_metadata layer keys come from two different checkpoint
|
||||
# export conventions in the wild: some store them with the full
|
||||
# diffusion-model prefix (e.g. "model.diffusion_model.proj_in"), others
|
||||
# already strip it (e.g. "proj_in"). convert_old_quants() can also be
|
||||
# called before or after model_prefix is stripped from state_dict
|
||||
# (comfy/sd.py calls it up to twice around that strip, see #13328), so a
|
||||
# single fixed convention silently mismatches one case or the other and
|
||||
# the layer never gets wrapped in a QuantizedTensor (see #11864). Match
|
||||
# against state_dict's real keys instead of assuming a convention.
|
||||
# Metadata layer keys may use the full diffusion-model prefix or the
|
||||
# already-stripped form, so resolve them against the state dict.
|
||||
if "{}.weight".format(layer_key) in state_dict:
|
||||
return layer_key
|
||||
if model_prefix:
|
||||
|
|
@ -1479,8 +1472,9 @@ def convert_old_quants(state_dict, model_prefix="", metadata={}):
|
|||
for k, v in layers.items():
|
||||
resolved_key = _resolve_quant_metadata_layer_key(state_dict, k, model_prefix)
|
||||
marker_key = "{}.comfy_quant".format(resolved_key)
|
||||
if marker_key not in state_dict: # idempotent: convert_old_quants may run twice on the same checkpoint
|
||||
state_dict[marker_key] = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8)
|
||||
marker = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8)
|
||||
if marker_key not in state_dict or not torch.equal(state_dict[marker_key], marker):
|
||||
state_dict[marker_key] = marker
|
||||
|
||||
return state_dict, metadata
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ if not has_gpu():
|
|||
from comfy import ops
|
||||
from comfy.quant_ops import QuantizedTensor
|
||||
import comfy.utils
|
||||
import comfy.model_detection as model_detection
|
||||
|
||||
|
||||
def marker_json(state_dict, key):
|
||||
|
|
@ -188,9 +187,29 @@ class TestConvertOldQuantsPrefixAware(unittest.TestCase):
|
|||
"proj_in.weight_scale": torch.tensor(1.0),
|
||||
}
|
||||
out_sd1, _ = comfy.utils.convert_old_quants(dict(state_dict), model_prefix="", metadata=dict(metadata))
|
||||
out_sd2, _ = comfy.utils.convert_old_quants(out_sd1, model_prefix="", metadata=dict(metadata))
|
||||
self.assertEqual(set(out_sd1.keys()), set(out_sd2.keys()))
|
||||
self.assertTrue(torch.equal(out_sd1["proj_in.comfy_quant"], out_sd2["proj_in.comfy_quant"]))
|
||||
keys_before = set(out_sd1.keys())
|
||||
original_marker = out_sd1["proj_in.comfy_quant"]
|
||||
marker_before = original_marker.clone()
|
||||
out_sd2, _ = comfy.utils.convert_old_quants(dict(out_sd1), model_prefix="", metadata=dict(metadata))
|
||||
self.assertEqual(keys_before, set(out_sd2.keys()))
|
||||
self.assertTrue(torch.equal(marker_before, out_sd2["proj_in.comfy_quant"]))
|
||||
self.assertIs(original_marker, out_sd2["proj_in.comfy_quant"])
|
||||
|
||||
def test_conflicting_marker_is_replaced_with_current_metadata(self):
|
||||
layer_quant_config = {"proj_in": {"format": "float8_e4m3fn"}}
|
||||
old_marker = torch.tensor(list(json.dumps({"format": "nvfp4"}).encode("utf-8")), dtype=torch.uint8)
|
||||
state_dict = {
|
||||
"proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn),
|
||||
"proj_in.weight_scale": torch.tensor(1.0),
|
||||
"proj_in.comfy_quant": old_marker,
|
||||
}
|
||||
out_sd, _ = comfy.utils.convert_old_quants(
|
||||
state_dict,
|
||||
model_prefix="",
|
||||
metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})},
|
||||
)
|
||||
self.assertEqual(marker_json(out_sd, "proj_in.comfy_quant"), {"format": "float8_e4m3fn"})
|
||||
self.assertIsNot(old_marker, out_sd["proj_in.comfy_quant"])
|
||||
|
||||
# ---- extra: functional end-to-end, proving the fixed layer actually loads as QuantizedTensor ----
|
||||
def test_functional_load_after_prefix_mismatch_fix(self):
|
||||
|
|
@ -217,62 +236,5 @@ class TestConvertOldQuantsPrefixAware(unittest.TestCase):
|
|||
self.assertEqual(model.layer1.weight._params.scale.item(), 2.0)
|
||||
|
||||
|
||||
class TestKnownResidualGap(unittest.TestCase):
|
||||
"""Documents a related but DISTINCT failure mode this PR does not close,
|
||||
found while writing the regression tests above. Kept as an
|
||||
expectedFailure so it stays visible instead of silently passing or
|
||||
breaking CI.
|
||||
|
||||
load_diffusion_model_state_dict() passes model_prefix="" (a string
|
||||
literal, not the real diffusion-model prefix) to both of its
|
||||
convert_old_quants() calls -- see comfy/sd.py, unchanged by this PR since
|
||||
editing that call site is the path #13328 was closed for. When a
|
||||
checkpoint's real weight keys carry NO wrapper prefix at all but its
|
||||
_quantization_metadata layer keys DO carry a "model.diffusion_model."
|
||||
style prefix, convert_old_quants() (both before and after this fix) has
|
||||
no way to recognize the mismatch on the first call with model_prefix="",
|
||||
so its fallback write -- required to be byte-for-byte identical to
|
||||
today's behavior for the "neither convention matches" case -- pollutes
|
||||
state_dict with spurious "model.diffusion_model.*.comfy_quant" keys.
|
||||
comfy.model_detection.unet_prefix_from_state_dict() (a different
|
||||
function, out of this PR's scope) then falsely detects
|
||||
"model.diffusion_model." as the prefix from those spurious keys alone,
|
||||
and state_dict_prefix_replace(..., filter_keys=True) strips using that
|
||||
wrong prefix, discarding every real .weight/.weight_scale tensor.
|
||||
|
||||
This is pre-existing: it reproduces identically with and without this
|
||||
PR's fix (verified manually), because the root cause here is
|
||||
unet_prefix_from_state_dict()/the sd.py call site, not the marker-key
|
||||
resolution this PR changes. Closing it would require either passing the
|
||||
real prefix into convert_old_quants() from sd.py (a call-site change) or
|
||||
hardening unet_prefix_from_state_dict() to ignore .comfy_quant keys --
|
||||
both outside "only touch convert_old_quants()".
|
||||
"""
|
||||
|
||||
@unittest.expectedFailure
|
||||
def test_empty_model_prefix_cannot_prevent_real_prefix_detection_poisoning(self):
|
||||
prefix = "model.diffusion_model."
|
||||
num_layers = 10 # unet_prefix_from_state_dict requires > 5 matches
|
||||
sd = {}
|
||||
layers_meta = {}
|
||||
for i in range(num_layers):
|
||||
local = f"block{i}"
|
||||
sd[f"{local}.weight"] = torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn)
|
||||
sd[f"{local}.weight_scale"] = torch.tensor(1.0)
|
||||
layers_meta[f"{prefix}{local}"] = {"format": "float8_e4m3fn"}
|
||||
metadata = {"_quantization_metadata": json.dumps({"layers": layers_meta})}
|
||||
|
||||
# exact two-call dance from comfy/sd.py::load_diffusion_model_state_dict
|
||||
sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata)
|
||||
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
|
||||
temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)
|
||||
if len(temp_sd) > 0:
|
||||
sd = temp_sd
|
||||
sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata)
|
||||
|
||||
remaining_weights = [k for k in sd if k.endswith(".weight")]
|
||||
self.assertEqual(len(remaining_weights), num_layers)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
|
|||
Loading…
Reference in New Issue