From d005ac5aa1b52ac4c8b358965cab8df73168bead Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 11 Jul 2026 08:48:52 +0100 Subject: [PATCH] feat(model): support comfy_quant header metadata containers Newer quantized checkpoints record per-layer formats in the safetensors header _quantization_metadata instead of marker tensors. The native loader now reads the header map, re-keys it through the same prefix strip as the tensors, and transcodes it into marker tensors so both container forms share one detection path; header entries win over markers. - drop optional input_scale sidecars for marked layers - map full_precision_matrix_mult onto the sdnq per-layer matmul exclusion list - log the detection source (markers, header, both) --- pipelines/native_transformer.py | 87 ++++++++++++- test/test-native-transformer.py | 222 +++++++++++++++++++++++++++++++- 2 files changed, 303 insertions(+), 6 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index fdffd7fcc..aea2fc80c 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -274,8 +274,12 @@ def load( quant_type = model_quant.get_quant_type(quant_args) state_dict = sd_models.read_state_dict(local_file, what="transformer") + metadata_layers = read_quantization_metadata(local_file) state_dict = drop_companion_keys(state_dict, spec.ignored_prefixes, spec.cls.__name__) state_dict, detected_prefix = strip_prefix(state_dict, spec.prefixes, spec.cls.__name__) + if metadata_layers and detected_prefix: + # header metadata names mirror the file's tensor naming, so they carry the same prefix + metadata_layers = {name[len(detected_prefix):] if name.startswith(detected_prefix) else name: meta for name, meta in metadata_layers.items()} check_forbidden_markers(state_dict, spec.forbidden_markers, spec.cls.__name__, local_file) transformer_sd, sibling_sds = partition_siblings(state_dict, spec.siblings) del state_dict @@ -304,6 +308,7 @@ def load( modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict, converter_handles_quant=spec.converter_handles_quant, + metadata_layers=metadata_layers, **kwargs, ) del transformer_sd @@ -427,6 +432,64 @@ def check_forbidden_markers( ) +def read_quantization_metadata(local_file: str) -> dict[str, dict] | None: + """Read per-layer quantization info from the safetensors header metadata. + + Newer comfy_quant checkpoints record layer formats in the header + ``__metadata__`` under ``_quantization_metadata`` (a JSON string with a + ``layers`` map keyed by tensor naming) instead of per-layer marker + tensors. Returns the layers map, or ``None`` when the header carries no + such entry. Raises :class:`OverrideArchMismatch` when the entry is + present but malformed, so the caller's base-repo fallback engages. + """ + from modules.model_probe import read_safetensors_header + try: + header = read_safetensors_header(local_file) + except Exception as e: + log.debug(f'Load model: file="{local_file}" header metadata unreadable ({e})') + return None + meta = (header.get("__metadata__") or {}).get("_quantization_metadata") + if meta is None: + return None + try: + parsed = json.loads(meta) if isinstance(meta, str) else meta + layers = parsed["layers"] + except Exception as e: + raise OverrideArchMismatch( + f'Load model: file="{local_file}" native_transformer _quantization_metadata ' + f"is malformed ({type(e).__name__}: {e})" + ) from e + if not isinstance(layers, dict) or not all(isinstance(entry, dict) for entry in layers.values()): + raise OverrideArchMismatch( + f'Load model: file="{local_file}" native_transformer _quantization_metadata ' + f"layers map is malformed" + ) + return layers + + +def transcode_quant_metadata(state_dict: dict, metadata_layers: dict[str, dict]) -> tuple[dict, str]: + """Synthesize per-layer marker tensors from header quantization metadata. + + Header metadata and marker tensors describe the same per-layer format + dicts; converging on markers lets detection, converters, and remapping + handle both forms identically. Entries without a matching ``.weight`` + (companion components filtered earlier) are skipped; existing markers + are overwritten by their header entry. Returns a new dict plus the + detection source (``header`` or ``both``) for logging. + """ + sd = dict(state_dict) + had_markers = any(key.endswith(COMFY_QUANT_MARKER) for key in sd) + count = 0 + for name, meta in metadata_layers.items(): + if f"{name}.weight" not in sd: + continue + sd[f"{name}{COMFY_QUANT_MARKER}"] = torch.tensor(list(json.dumps(meta).encode("utf-8")), dtype=torch.uint8) + count += 1 + if count == 0: + return state_dict, "markers" + return sd, "both" if had_markers else "header" + + def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[dict[str, dict], str] | None: """Detect ``comfy_quant`` pre-quantized layers in a state dict. @@ -490,14 +553,19 @@ def remap_comfy_quant(state_dict: dict, marked_names: set[str]) -> dict: Renames ``.weight_scale`` to ``.scale`` (reshaped to 2-D, ``[]`` becomes ``[1, 1]``, since SDNQ transposes scales in place) and - drops the ``.comfy_quant`` markers. Returns a new dict; the input - (which may be the cached state dict) is not mutated. + drops the ``.comfy_quant`` markers plus optional + ``.input_scale`` activation-calibration sidecars (SDNQ re-derives + activation scales dynamically). Returns a new dict; the input (which may + be the cached state dict) is not mutated. """ scale_suffix = ".weight_scale" + input_scale_suffix = ".input_scale" remapped: dict = {} for key, value in state_dict.items(): if key.endswith(COMFY_QUANT_MARKER) and key[: -len(COMFY_QUANT_MARKER)] in marked_names: continue + if key.endswith(input_scale_suffix) and key[: -len(input_scale_suffix)] in marked_names: + continue if key.endswith(scale_suffix) and key[: -len(scale_suffix)] in marked_names: remapped[f"{key[: -len(scale_suffix)]}.scale"] = value.reshape(-1, 1) continue @@ -699,6 +767,9 @@ def build_component_prequantized( f"{comfy_format} expects {storage_dtype} weights but {name!r} has {found}" ) + # layers flagged full_precision_matrix_mult stay on the dequant path even + # when the user enables quantized matmul (exact-name match on the config list) + full_precision_mm = [f"{name}.weight" for name in sorted(marked_names) if marker_meta[name].get("full_precision_matrix_mult")] quantization_config = SDNQConfig( weights_dtype=weights_dtype, quantized_matmul_dtype=matmul_dtype, @@ -707,6 +778,7 @@ def build_component_prequantized( dequantize_fp32=shared.opts.sdnq_dequantize_fp32, add_skip_keys=False, modules_to_not_convert=[], + modules_to_not_use_matmul=full_precision_mm, ) quantizer = SDNQQuantizer(quantization_config, pre_quantized=True) quantizer.torch_dtype = target_dtype @@ -844,6 +916,7 @@ def build_component( modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None, converter_handles_quant: bool = False, + metadata_layers: dict[str, dict] | None = None, **kwargs, ) -> object: """Convert (if needed), instantiate, load weights, dtype-cast, quantize, @@ -851,7 +924,9 @@ def build_component( Transformer state dicts carrying ``comfy_quant`` markers dispatch to :func:`build_component_prequantized` (``quant_args`` are bypassed: the - file is already quantized); a ``converter_handles_quant`` converter runs + file is already quantized). ``metadata_layers`` (header-metadata quant + info) is transcoded into markers first, so both container forms share + one path; a ``converter_handles_quant`` converter runs before that detection, float-oriented converters after it. Under SDNQ the transformer uses the per-tensor pre-mode path in :func:`build_component_quantized` so quantization is applied in flight. @@ -866,6 +941,10 @@ def build_component( reach ``cls.from_config`` for both construction paths. """ try: + quant_source = "markers" + if component_name == "transformer" and metadata_layers: + state_dict, quant_source = transcode_quant_metadata(state_dict, metadata_layers) + if converter is not None and converter_handles_quant and component_name == "transformer": state_dict = apply_converter(converter, state_dict, cls, component_name) converter = None # consumed; must not run again on the non-comfy path below @@ -876,7 +955,7 @@ def build_component( convrot_count = sum(1 for meta in marker_meta.values() if meta.get("convrot")) log.info( f'Load model: transformer=native {component_name} quant=comfy ' - f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} keys={len(state_dict)} cls={cls.__name__}' + f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} source={quant_source} keys={len(state_dict)} cls={cls.__name__}' ) component = build_component_prequantized( component_name=component_name, diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 9996d3a1e..85f14bdd3 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -435,6 +435,121 @@ def test_remap_comfy_does_not_mutate_input(): assert set(sd.keys()) == keys_before, 'input dict must not be mutated (read_state_dict caches it)' +def test_remap_comfy_drops_input_scale(): + """Optional activation-calibration sidecars are dropped for marked layers + (SDNQ derives activation scales dynamically); unmarked ones pass through.""" + unrelated = torch.tensor(1.0) + sd = { + 'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8), + 'blocks.0.attn.wq.weight_scale': torch.tensor(0.5), + 'blocks.0.attn.wq.input_scale': torch.tensor(0.125), + 'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'), + 'other.input_scale': unrelated, + } + out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'}) + assert 'blocks.0.attn.wq.input_scale' not in out + assert out['other.input_scale'] is unrelated + + +# ============================================================ +# header _quantization_metadata +# ============================================================ + +def quant_metadata(layers: dict) -> dict: + import json + return {'_quantization_metadata': json.dumps({'format_version': '1.0', 'layers': layers})} + + +def test_read_quant_metadata_absent_returns_none(): + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + write_fixture({'a.weight': torch.zeros(2)}, fd, path) + assert nt.read_quantization_metadata(path) is None + finally: + os.unlink(path) + + +def test_read_quant_metadata_parses_layers(): + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + layers = {'blocks.0.attn.wq': {'format': 'int8_tensorwise'}} + write_fixture({'blocks.0.attn.wq.weight': torch.zeros(2)}, fd, path, metadata=quant_metadata(layers)) + assert nt.read_quantization_metadata(path) == layers + finally: + os.unlink(path) + + +def test_read_quant_metadata_malformed_raises(): + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + write_fixture({'a.weight': torch.zeros(2)}, fd, path, metadata={'_quantization_metadata': 'not json'}) + try: + nt.read_quantization_metadata(path) + raise AssertionError('expected OverrideArchMismatch') + except nt.OverrideArchMismatch as e: + assert 'malformed' in str(e) + finally: + os.unlink(path) + + +def test_read_quant_metadata_non_dict_layers_raises(): + import json + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + payload = json.dumps({'layers': {'blocks.0.attn.wq': 'int8_tensorwise'}}) + write_fixture({'a.weight': torch.zeros(2)}, fd, path, metadata={'_quantization_metadata': payload}) + try: + nt.read_quantization_metadata(path) + raise AssertionError('expected OverrideArchMismatch') + except nt.OverrideArchMismatch as e: + assert 'malformed' in str(e) + finally: + os.unlink(path) + + +def test_transcode_quant_metadata_synthesizes_markers(): + """Header entries become marker tensors for layers whose weight is present; + entries without a matching weight (companion components) are skipped.""" + sd = { + 'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8), + 'blocks.0.attn.wq.weight_scale': torch.tensor(0.5), + } + layers = { + 'blocks.0.attn.wq': {'format': 'int8_tensorwise'}, + 'text_encoder.mlp.up': {'format': 'int8_tensorwise'}, + } + out, source = nt.transcode_quant_metadata(sd, layers) + assert source == 'header' + assert 'blocks.0.attn.wq.comfy_quant' in out + assert 'text_encoder.mlp.up.comfy_quant' not in out + assert 'blocks.0.attn.wq.comfy_quant' not in sd, 'input dict must not be mutated' + marked, fmt = nt.detect_comfy_quant(out, 'Test') + assert fmt == 'int8_tensorwise' + assert marked['blocks.0.attn.wq'] == {'format': 'int8_tensorwise'} + + +def test_transcode_quant_metadata_header_wins(): + """When both forms are present the header entry overwrites the marker, + matching the reference loader's precedence.""" + sd = { + 'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8), + 'blocks.0.attn.wq.comfy_quant': comfy_marker('float8_e4m3fn'), + } + layers = {'blocks.0.attn.wq': {'format': 'int8_tensorwise', 'full_precision_matrix_mult': True}} + out, source = nt.transcode_quant_metadata(sd, layers) + assert source == 'both' + marked, fmt = nt.detect_comfy_quant(out, 'Test') + assert fmt == 'int8_tensorwise' + assert marked['blocks.0.attn.wq'].get('full_precision_matrix_mult') is True + + +def test_transcode_quant_metadata_no_matches_is_noop(): + sd = {'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8)} + out, source = nt.transcode_quant_metadata(sd, {'unrelated.layer': {'format': 'int8_tensorwise'}}) + assert source == 'markers' + assert out is sd + + # ============================================================ # is_noop_converter # ============================================================ @@ -639,9 +754,9 @@ class MockKwargsTransformer(MockMiniTransformer): return cls(dim=config['dim']) -def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str: +def write_fixture(state_dict_keys: dict, fd: int, path: str, metadata: dict | None = None) -> str: os.close(fd) - safetensors.torch.save_file(state_dict_keys, path) + safetensors.torch.save_file(state_dict_keys, path, metadata=metadata) return path @@ -1229,6 +1344,98 @@ def test_load_comfy_fp8_end_to_end(): os.unlink(path) +def test_load_comfy_metadata_end_to_end(): + """Marker-less container: quant info only in the header + _quantization_metadata, layer names carrying the file's tensor prefix. + The load path must re-key through the prefix strip and land on the same + prequantized build as the marker form.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = comfy_fixture(dim) + del raw['model.diffusion_model.in_proj.comfy_quant'] + layers = {'model.diffusion_model.in_proj': {'format': 'int8_tensorwise'}} + write_fixture(raw, fd, path, metadata=quant_metadata(layers)) + + with ComfyTestEnv(dim): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + transformer, _ = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + dtype=torch.float32, + ) + + in_proj = transformer.in_proj + assert in_proj.__class__.__name__ == 'SDNQLinear' + assert in_proj.weight.dtype == torch.int8 + assert torch.equal(in_proj.weight.detach().cpu(), raw['model.diffusion_model.in_proj.weight']) + assert transformer.out_proj.__class__ is torch.nn.Linear + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_comfy_metadata_bare_names_end_to_end(): + """Marker-less container without a tensor prefix (official Comfy-Org + exports): metadata names match the bare keys and need no re-keying.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = {key[len('model.diffusion_model.'):]: value for key, value in comfy_fixture(dim).items()} + del raw['in_proj.comfy_quant'] + raw['in_proj.input_scale'] = torch.tensor(0.125, dtype=torch.float32) + layers = {'in_proj': {'format': 'int8_tensorwise'}} + write_fixture(raw, fd, path, metadata=quant_metadata(layers)) + + with ComfyTestEnv(dim): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + transformer, _ = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + dtype=torch.float32, + ) + + assert transformer.in_proj.__class__.__name__ == 'SDNQLinear' + assert torch.equal(transformer.in_proj.weight.detach().cpu(), raw['in_proj.weight']) + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_comfy_full_precision_mm_excluded_from_matmul(): + """Layers flagged full_precision_matrix_mult land on the config's + per-layer matmul exclusion list, which apply_sdnq_options_to_model + honors when the user enables quantized matmul.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + import json + raw = comfy_fixture(dim) + payload = {'format': 'int8_tensorwise', 'full_precision_matrix_mult': True} + raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(json.dumps(payload).encode()), dtype=torch.uint8) + write_fixture(raw, fd, path) + + with ComfyTestEnv(dim): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + transformer, _ = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + dtype=torch.float32, + ) + + assert 'in_proj.weight' in transformer.quantization_config.modules_to_not_use_matmul + assert transformer.in_proj.sdnq_dequantizer.use_quantized_matmul is False + finally: + if os.path.exists(path): + os.unlink(path) + + def test_detect_comfy_mixed_formats_raises(): """One file mixing int8 and fp8 layers has no single SDNQ mapping; reject.""" sd = { @@ -1735,6 +1942,14 @@ def run_all(): test_remap_comfy_renames_and_reshapes_scale, test_remap_comfy_passes_unmarked_keys_verbatim, test_remap_comfy_does_not_mutate_input, + test_remap_comfy_drops_input_scale, + test_read_quant_metadata_absent_returns_none, + test_read_quant_metadata_parses_layers, + test_read_quant_metadata_malformed_raises, + test_read_quant_metadata_non_dict_layers_raises, + test_transcode_quant_metadata_synthesizes_markers, + test_transcode_quant_metadata_header_wins, + test_transcode_quant_metadata_no_matches_is_noop, ]: run_test(cat, fn) @@ -1803,6 +2018,9 @@ def run_all(): for fn in [ test_load_comfy_int8_end_to_end, test_load_comfy_fp8_end_to_end, + test_load_comfy_metadata_end_to_end, + test_load_comfy_metadata_bare_names_end_to_end, + test_load_comfy_full_precision_mm_excluded_from_matmul, test_detect_comfy_mixed_formats_raises, test_load_comfy_marker_dtype_mismatch_raises, test_load_comfy_unsupported_format_raises_mismatch,