From 5d5ede5c473ce8c7e9747c9237bd6b01b88f3abf Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 11 Jul 2026 06:40:07 +0100 Subject: [PATCH] feat(model): load convrot comfy_quant checkpoints via sdnq hadamard ConvRot is the regular Hadamard rotation SDNQ implements: identical construction, normalization, axis, and dequant order. Per-layer markers map onto the dequantizer's use_hadamard and hadamard_group_size; group sizes must be powers of 4 and divide in_features, else base-repo fallback. Detection carries per-layer metadata since files mix plain and rotated layers. --- pipelines/native_transformer.py | 58 +++++++----- test/test-native-transformer.py | 159 ++++++++++++++++++++++++++++---- 2 files changed, 176 insertions(+), 41 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 2fbf3f797..fdffd7fcc 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -427,21 +427,25 @@ def check_forbidden_markers( ) -def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] | None: +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. ComfyUI-format quantized checkpoints mark each quantized layer with a ``.comfy_quant`` uint8 tensor whose bytes are a JSON object naming the storage format, alongside a ``.weight_scale`` tensor. Returns - the set of marked module names and the format string, or ``None`` when no - markers are present. Raises :class:`OverrideArchMismatch` for malformed - markers, unsupported formats, or a file mixing formats across layers, so - the caller's base-repo fallback engages. + a mapping of marked module names to their parsed marker JSON plus the + file's format string, or ``None`` when no markers are present. ConvRot + markers (``convrot``/``convrot_groupsize``) are accepted per layer: the + rotation is the same regular Hadamard SDNQ implements, so flagged layers + map onto ``use_hadamard`` (regular Hadamards only exist for power-of-4 + group sizes). Raises :class:`OverrideArchMismatch` for + malformed markers, unsupported formats or group sizes, or a file mixing + formats across layers, so the caller's base-repo fallback engages. """ marker_keys = [key for key in state_dict if key.endswith(COMFY_QUANT_MARKER)] if not marker_keys: return None - marked: set[str] = set() + marked: dict[str, dict] = {} formats: set[str] = set() for key in marker_keys: name = key[: -len(COMFY_QUANT_MARKER)] @@ -454,14 +458,14 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] f"for {name!r} is malformed ({type(e).__name__}: {e})" ) from e if meta.get("convrot"): - # ConvRot-flagged layers may require an inverse rotation at - # runtime; adopting them as plain tensors would corrupt outputs. - log.error(f'Load model: type={type_name} quant=comfy format={fmt} convrot=true not supported') - raise OverrideArchMismatch( - f"Load model: type={type_name} native_transformer comfy_quant layer " - f"{name!r} carries a convrot rotation flag, not supported" - ) - marked.add(name) + group_size = int(meta.get("convrot_groupsize", 256)) + is_pow4 = group_size >= 4 and (group_size & (group_size - 1)) == 0 and (group_size.bit_length() & 1) == 1 + if not is_pow4: + raise OverrideArchMismatch( + f"Load model: type={type_name} native_transformer comfy_quant layer " + f"{name!r} convrot_groupsize={group_size} is not a power of 4" + ) + marked[name] = meta formats.add(fmt) unsupported = sorted(formats - set(COMFY_QUANT_FORMATS)) if unsupported: @@ -644,7 +648,7 @@ def build_component_prequantized( state_dict: dict, config: dict, cls: type, - marked_names: set[str], + marker_meta: dict[str, dict], comfy_format: str, dtype, acceptable_missing: tuple[str, ...], @@ -656,7 +660,9 @@ def build_component_prequantized( The supported formats are subsets of SDNQ's symmetric quantization: same storage layout (unpacked 8-bit ``[out, in]``), same dequant math - (``weight * scale``, no zero point), so tensors are adopted bit-exact. + (``weight * scale``, no zero point), so tensors are adopted bit-exact; + convrot layers map onto SDNQ's Hadamard support, whose identical regular + Hadamard construction lets the dequantizer undo the stored rotation. The file dictates which layers are quantized, independent of the user's quantization settings, and floating-point SDNQ params are not cast to the target dtype (the fp32 scales must survive). Layers are assembled in @@ -679,6 +685,7 @@ def build_component_prequantized( matmul_dtype = "int8" if dtype_dict[weights_dtype]["is_integer"] else "float8_e4m3fn" storage_dtype = dtype_dict[weights_dtype]["storage_dtype"] target_dtype = dtype if dtype is not None else devices.dtype + marked_names = set(marker_meta) sd = remap_comfy_quant(state_dict, marked_names) # Civitai relabels these containers freely; trust the marker only as far @@ -722,6 +729,14 @@ def build_component_prequantized( f"Load model: transformer=native {component_name} comfy_quant marked " f"module {name!r} is {linear.__class__.__name__}, expected Linear" ) + # convrot weights are stored rotated by SDNQ's regular Hadamard; the dequantizer undoes it + use_hadamard = bool(marker_meta[name].get("convrot")) + hadamard_group_size = int(marker_meta[name].get("convrot_groupsize", 256)) if use_hadamard else 256 + if use_hadamard and linear.in_features % hadamard_group_size != 0: + raise OverrideArchMismatch( + f"Load model: transformer=native {component_name} comfy_quant marked " + f"module {name!r} convrot_groupsize={hadamard_group_size} does not divide in_features={linear.in_features}" + ) linear.sdnq_dequantizer = SDNQDequantizer( result_dtype=target_dtype, result_shape=None, @@ -730,14 +745,14 @@ def build_component_prequantized( quantized_weight_shape=torch.Size((linear.out_features, linear.in_features)), weights_dtype=weights_dtype, quantized_matmul_dtype=matmul_dtype, - hadamard_group_size=256, + hadamard_group_size=hadamard_group_size, group_size=-1, svd_rank=32, svd_steps=8, use_quantized_matmul=False, re_quantize_for_matmul=False, use_stochastic_rounding=False, - use_hadamard=False, + use_hadamard=use_hadamard, layer_class_name="Linear", ) wrapped = get_sdnq_wrapper_class(linear, dequant_forward) @@ -857,17 +872,18 @@ def build_component( comfy_quant = detect_comfy_quant(state_dict, cls.__name__) if component_name == "transformer" else None if comfy_quant is not None: - marked_names, comfy_format = comfy_quant + marker_meta, comfy_format = comfy_quant + 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(marked_names)} keys={len(state_dict)} cls={cls.__name__}' + f'format={comfy_format} layers={len(marker_meta)} convrot={convrot_count} keys={len(state_dict)} cls={cls.__name__}' ) component = build_component_prequantized( component_name=component_name, state_dict=state_dict, config=config, cls=cls, - marked_names=marked_names, + marker_meta=marker_meta, comfy_format=comfy_format, dtype=dtype, acceptable_missing=acceptable_missing, diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index a8cd14687..9996d3a1e 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -323,7 +323,8 @@ def test_detect_comfy_valid_markers(): detected = nt.detect_comfy_quant(sd, 'Test') assert detected is not None marked, fmt = detected - assert marked == {'blocks.0.attn.wq', 'blocks.0.mlp.up'} + assert set(marked) == {'blocks.0.attn.wq', 'blocks.0.mlp.up'} + assert marked['blocks.0.attn.wq'] == {'format': 'int8_tensorwise'} assert fmt == 'int8_tensorwise' @@ -351,20 +352,37 @@ def test_detect_comfy_malformed_marker_raises(): assert 'malformed' in str(e) -def test_detect_comfy_convrot_flag_raises(): - """Markers carrying convrot=true may need runtime inverse rotation; adopting - them as plain tensors would corrupt outputs, so they must be rejected.""" +def test_detect_comfy_convrot_accepted(): + """ConvRot markers map onto SDNQ's Hadamard support; detection carries the + per-layer fields through instead of rejecting.""" import json - payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256}) + payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256, 'per_row': True}) sd = { 'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8), 'blocks.0.attn.wq.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8), + 'blocks.0.mlp.up.weight': torch.zeros((4, 4), dtype=torch.int8), + 'blocks.0.mlp.up.comfy_quant': comfy_marker('int8_tensorwise'), + } + marked, fmt = nt.detect_comfy_quant(sd, 'Test') + assert fmt == 'int8_tensorwise' + assert marked['blocks.0.attn.wq'].get('convrot') is True + assert marked['blocks.0.attn.wq'].get('convrot_groupsize') == 256 + assert not marked['blocks.0.mlp.up'].get('convrot') + + +def test_detect_comfy_convrot_bad_groupsize_raises(): + """Regular Hadamards only exist for power-of-4 sizes; any other convrot + group size cannot be a compatible rotation.""" + import json + payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 8}) + sd = { + 'blocks.0.attn.wq.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8), } try: nt.detect_comfy_quant(sd, 'Test') raise AssertionError('expected OverrideArchMismatch') except nt.OverrideArchMismatch as e: - assert 'convrot' in str(e) + assert 'power of 4' in str(e) def test_detect_comfy_marker_missing_format_field_raises(): @@ -1446,25 +1464,26 @@ def test_load_fused_bf16_end_to_end(): os.unlink(path) -def test_load_fused_convrot_falls_back(): - """The one real Ideogram 4 civitai file is convrot-flagged: conversion - succeeds, then detection rejects it into the base-repo fallback.""" +def test_load_fused_convrot_end_to_end(): + """Fused qkv with a convrot marker: the converter duplicates the marker to + the three split layers, each of which loads with hadamard geometry. Row + slicing commutes with the in-features rotation, so slices stay lossless.""" import json as json_mod fd, path = tempfile.mkstemp(suffix='.safetensors') try: - dim = 8 + dim, group_size = 8, 4 raw = fused_fixture(dim, quantized=True) - payload = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256, 'per_row': True}) + payload = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': group_size, 'per_row': True}) raw['model.diffusion_model.layers.0.attention.qkv.comfy_quant'] = torch.tensor(list(payload.encode()), dtype=torch.uint8) + raw['model.diffusion_model.layers.0.attention.qkv.weight_scale'] = torch.rand(3 * dim, 1, dtype=torch.float32) write_fixture(raw, fd, path) with ComfyTestEnv(dim): - raised = False - try: - nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}) - except nt.OverrideArchMismatch as e: - raised = True - assert 'convrot' in str(e) - assert raised, 'expected OverrideArchMismatch' + transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32) + attn = transformer.layers[0].attention + for name in ('to_q', 'to_k', 'to_v'): + dq = getattr(attn, name).sdnq_dequantizer + assert dq.use_hadamard is True, f'{name} lost the convrot flag' + assert dq.hadamard_group_size == group_size finally: if os.path.exists(path): os.unlink(path) @@ -1534,6 +1553,103 @@ def test_load_transformer_secondary_slot_syncs_tracker(): os.unlink(path) +def regular_hadamard(n: int) -> torch.Tensor: + """Regular Hadamard construction the ConvRot format and SDNQ share: + 4x4 base, Kronecker recursion, 1/sqrt(n) normalization. Symmetric and + involutory, so rotation and inverse are the same matrix.""" + h4 = torch.tensor([[1., 1., 1., -1.], [1., 1., -1., 1.], [1., -1., 1., 1.], [-1., 1., 1., 1.]]) + h = h4 + size = 4 + while size < n: + h = torch.kron(h, h4) + size *= 4 + return h / (n ** 0.5) + + +def comfy_convrot_quantize(weight: torch.Tensor, group_size: int): + """Reference ConvRot int8 quantizer (mirrors the ComfyUI runtime): rotate + grouped in-features by the regular Hadamard, then row-wise symmetric int8.""" + h = regular_hadamard(group_size) + out_f, in_f = weight.shape + rotated = (weight.reshape(out_f, -1, group_size) @ h.T).reshape(out_f, in_f) + scale = rotated.abs().amax(dim=-1, keepdim=True) / 127 + q = rotated.div(scale).round().clamp(-128, 127).to(torch.int8) + return q, scale + + +def test_load_comfy_convrot_end_to_end(): + """ConvRot parity: a layer quantized with the reference ConvRot math + loads as an SDNQ hadamard layer whose dequant matches the reference dequant + and recovers the original weight within int8 quantization error. The plain + second layer proves per-layer mixing within one file.""" + import json as json_mod + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim, group_size = 8, 4 + original = torch.randn(dim, dim) + q, scale = comfy_convrot_quantize(original, group_size) + marker = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': group_size, 'per_row': True}) + raw = { + 'model.diffusion_model.in_proj.weight': q, + 'model.diffusion_model.in_proj.weight_scale': scale, + 'model.diffusion_model.in_proj.comfy_quant': torch.tensor(list(marker.encode()), dtype=torch.uint8), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim, dtype=torch.float16), + 'model.diffusion_model.out_proj.weight': torch.randint(-128, 127, (dim, dim), dtype=torch.int8), + 'model.diffusion_model.out_proj.weight_scale': torch.tensor(0.03125, dtype=torch.float32), + 'model.diffusion_model.out_proj.comfy_quant': comfy_marker('int8_tensorwise'), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16), + } + 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) + + in_proj = transformer.in_proj + dq = in_proj.sdnq_dequantizer + assert dq.use_hadamard is True + assert dq.hadamard_group_size == group_size + assert tuple(in_proj.scale.shape) == (dim, 1) + + dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None).detach().cpu() + h = regular_hadamard(group_size) + reference = ((q.float() * scale).reshape(dim, -1, group_size) @ h.T).reshape(dim, dim) + assert torch.allclose(dequantized, reference, atol=1e-5), 'SDNQ hadamard dequant diverges from the ConvRot reference' + assert torch.allclose(dequantized, original, atol=0.15), 'dequant does not recover the original weight' + + # plain layer in the same file stays rotation-free + assert transformer.out_proj.sdnq_dequantizer.use_hadamard is False + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_comfy_convrot_nondivisible_falls_back(): + """A convrot group size that does not divide in_features cannot be undone; + the file must fall back to the base repo.""" + import json as json_mod + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = comfy_fixture(dim) + marker = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 16}) + raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(marker.encode()), dtype=torch.uint8) + write_fixture(raw, fd, path) + + with ComfyTestEnv(dim): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + raised = False + try: + nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={}) + except nt.OverrideArchMismatch as e: + raised = True + assert 'does not divide' in str(e) + assert raised, 'expected OverrideArchMismatch' + finally: + if os.path.exists(path): + os.unlink(path) + + def test_build_component_comfy_preempts_sdnq_fresh_quant(): """When SDNQ on-load quant settings are active (quant_type=SDNQConfig), a comfy_quant file must still take the pre-quantized path: fresh quant of @@ -1613,7 +1729,8 @@ def run_all(): test_detect_comfy_valid_markers, test_detect_comfy_unsupported_format_raises, test_detect_comfy_malformed_marker_raises, - test_detect_comfy_convrot_flag_raises, + test_detect_comfy_convrot_accepted, + test_detect_comfy_convrot_bad_groupsize_raises, test_detect_comfy_marker_missing_format_field_raises, test_remap_comfy_renames_and_reshapes_scale, test_remap_comfy_passes_unmarked_keys_verbatim, @@ -1690,6 +1807,8 @@ def run_all(): test_load_comfy_marker_dtype_mismatch_raises, test_load_comfy_unsupported_format_raises_mismatch, test_load_comfy_marker_for_unknown_module_raises_mismatch, + test_load_comfy_convrot_end_to_end, + test_load_comfy_convrot_nondivisible_falls_back, test_load_transformer_syncs_loaded_unet, test_load_transformer_secondary_slot_syncs_tracker, test_build_component_comfy_preempts_sdnq_fresh_quant, @@ -1710,7 +1829,7 @@ def run_all(): test_ideogram4_converter_does_not_mutate_input, test_load_fused_comfy_end_to_end, test_load_fused_bf16_end_to_end, - test_load_fused_convrot_falls_back, + test_load_fused_convrot_end_to_end, ]: run_test(cat, fn)