From 9b37a1535c482e50d6de5a4dc4c87a926e6a8dc1 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 17 Jul 2026 02:55:03 +0100 Subject: [PATCH] fix(lora): exact lora application on sdnq-quantized layers Baking a lora into a quantized weight requantizes it, and on low-bit formats round-to-nearest erases sub-step deltas (uint4 retains roughly 2/group_size of the signal). Plain lora deltas now ride the sdnq svd side-channel: factors append to svd_up/svd_down with the down factor hadamard-rotated, applied by the dequantizer at full precision in every forward mode. Apply and remove are exact and take no weight backup. - non-factorable families (dora, lokr, loha, oft, cp mid, dense bias) fall back to requantize with a per-pass summary warning - native fuse now honors the quantized-model guard; fuse requantized in place on every network swap and accumulated drift - layers that fell back on a mixed set restore from backup before re-entering the factor path; untargeted quantized layers are no longer flagged - test/test-sdnq-lora-factors.py pins the erasure law, factor-path exactness, memory accounting and set transitions --- modules/lora/extra_networks_lora.py | 4 +- modules/lora/lora_apply.py | 8 +- modules/lora/lora_diffusers.py | 3 +- modules/lora/lora_load.py | 4 +- modules/lora/lora_overrides.py | 53 ++- modules/lora/lora_sdnq.py | 169 +++++++++ modules/lora/networks.py | 38 +- test/test-sdnq-lora-factors.py | 534 ++++++++++++++++++++++++++++ 8 files changed, 793 insertions(+), 20 deletions(-) create mode 100644 modules/lora/lora_sdnq.py create mode 100644 test/test-sdnq-lora-factors.py diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index 35605e95e..42d39ec72 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -254,7 +254,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): if has_changed: jobid = shared.state.begin('LoRA') if len(l.previously_loaded_networks) > 0: - log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_native else "backup"}') + log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if lora_overrides.fuse_native() else "backup"}') networks.network_deactivate(include, exclude) networks.network_activate(include, exclude) debug_log(f'Network change: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]}') @@ -267,7 +267,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): prompt(p) if has_changed and len(include) == 0: # print only once actual_method = 'native' if any(len(n.modules) > 0 for n in l.loaded_networks) else load_method - log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if shared.opts.lora_fuse_native else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') + log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if lora_overrides.fuse_native() else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') def deactivate(self, p, force=False): if len(lora_diffusers.diffuser_loaded) > 0 and (shared.opts.lora_force_reload or force): diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index c1b46c4ab..3cb8688c0 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)") -def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, wanted_names: tuple): +def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, wanted_names: tuple, fuse: bool): backup_size = 0 if len(l.loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in l.loaded_networks]): # noqa: C419 # pylint: disable=R1729 t0 = time.time() @@ -24,7 +24,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr weights_backup = getattr(self, "network_weights_backup", None) bias_backup = getattr(self, "network_bias_backup", None) if weights_backup is not None or bias_backup is not None: - if (shared.opts.lora_fuse_native and not isinstance(weights_backup, bool)) or (not shared.opts.lora_fuse_native and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly + if (fuse and not isinstance(weights_backup, bool)) or (not fuse and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly weights_backup = None bias_backup = None self.network_weights_backup = weights_backup @@ -33,7 +33,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr if weights_backup is None and wanted_names != (): # pylint: disable=C1803 weight = getattr(self, 'weight', None) self.network_weights_backup = None - if shared.opts.lora_fuse_native: + if fuse: self.network_weights_backup = True else: self.network_weights_backup = weight.clone().to(devices.cpu) @@ -53,7 +53,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr if bias_backup is None: if getattr(self, 'bias', None) is not None: - if shared.opts.lora_fuse_native: + if fuse: self.network_bias_backup = True else: bias_backup = self.bias.clone() diff --git a/modules/lora/lora_diffusers.py b/modules/lora/lora_diffusers.py index 94ad716c7..e8321f8a0 100644 --- a/modules/lora/lora_diffusers.py +++ b/modules/lora/lora_diffusers.py @@ -4,6 +4,7 @@ import diffusers from modules import shared, errors from modules.logger import log from modules.lora import network +from modules.lora import lora_overrides from modules.lora import lora_common as l @@ -54,7 +55,7 @@ def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale t0 = time.time() name = name.replace(".", "_") sd_model: diffusers.DiffusionPipeline = getattr(shared.sd_model, "pipe", shared.sd_model) - log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason={reason or "unknown"} scale={lora_scale} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}') + log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason={reason or "unknown"} scale={lora_scale} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}') if not hasattr(sd_model, 'load_lora_weights'): log.error(f'Network load: type=LoRA class={sd_model.__class__} does not implement load lora') return None diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index 7993ffe9a..b297108bd 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -149,7 +149,7 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne if l.debug: log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}') else: - log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}') + log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}') if len(matched_networks) == 0: return None lora_cache[name] = net @@ -350,7 +350,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non networks.network_activate() if len(l.loaded_networks) > 0 and l.debug: - log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}') + log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}') if recompile_model: log.info("Network load: type=LoRA recompiling model") diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index 2d27f7b88..eb629896e 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -77,13 +77,56 @@ def get_method(shorthash=''): return 'native', 'default' +# Roles a LoRA is fused into; a quantized component in any of them makes fusing unsafe. +fuse_roots = ('transformer', 'unet', 'text_encoder', 'llm_adapter') + + +def fuse_components(sd_model): + """Component names a network fuses into, matched by role prefix so numbered and reference siblings are covered.""" + names = getattr(sd_model, 'components', None) + if not isinstance(names, dict): + names = vars(sd_model) + return [name for name in names if name.startswith(fuse_roots)] + + +def is_quantized(module): + """Return True when ``module`` carries a quantization config. + + ``config.quantization_config`` is read first: SDNQ sets both it and the plain + attribute when it quantizes in place, but a checkpoint that ships pre-quantized + only reaches the plain attribute through the diffusers ConfigMixin name proxy, + which is deprecated for removal. + """ + if module is None: + return False + config = getattr(module, 'config', None) + if config is not None and getattr(config, 'quantization_config', None) is not None: + return True + return getattr(module, 'quantization_config', None) is not None + + def disable_fuse(): - if hasattr(shared.sd_model, 'quantization_config'): + """Return True when fusing a network into model weights is unsafe. + + Fusing keeps no pristine copy of the weight, so each apply and restore + round-trips it through its storage format. On quantized weights that is a + dequantize-add-requantize cycle per network swap whose error compounds. + """ + sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model) + if is_quantized(sd_model): return True - if hasattr(shared.sd_model, 'transformer') and hasattr(shared.sd_model.transformer, 'quantization_config'): + if any(is_quantized(getattr(sd_model, name, None)) for name in fuse_components(sd_model)): return True - if hasattr(shared.sd_model, 'transformer_2') and hasattr(shared.sd_model.transformer_2, 'quantization_config'): - return True - if hasattr(shared.sd_model, '_lora_partial'): + if hasattr(sd_model, '_lora_partial'): return True return shared.sd_model_type in fuse_ignore + + +def fuse_native(): + """Return True when the native apply path may fuse into model weights. + + The single source of truth for the native fuse decision: it must agree across + the backup, activate and deactivate passes, since backup mode restores from a + stored tensor while fuse mode restores by subtracting the delta. + """ + return shared.opts.lora_fuse_native and not disable_fuse() diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py new file mode 100644 index 000000000..1cdbb8d22 --- /dev/null +++ b/modules/lora/lora_sdnq.py @@ -0,0 +1,169 @@ +"""Exact LoRA application for SDNQ-quantized layers. + +Baking a LoRA into a quantized weight requantizes it: dequantize, add the +delta, re-round onto the integer grid. When the per-element delta is smaller +than half a quantization step (a rank-decomposed delta on a uint4 layer sits +at a few percent of a step), rounding erases it; what survives is the two +grid-extrema elements per quantization group (2/group_size of the signal) +plus grid-shift noise of the same norm as the delta. The optimal in-grid +representation provably retains ~0%, so no rewrite of the stored integers +can fix this. + +The exact path instead rides the SDNQ svd side-channel: the dequantizer +computes ``W = dq(q) + svd_up @ svd_down`` in the rotated domain at full +precision, in every forward mode. A LoRA delta ``B @ A`` is appended as +extra columns of ``svd_up`` and rows of ``svd_down``; because the Hadamard +rotation is block-diagonal, symmetric and self-inverse, storing ``A·H`` for +the down factor makes the round trip exact: ``(B @ (A·H)) · H = B @ A``. +Quantized weights are never touched, so apply and remove are exact and no +weight backup is needed. + +Only additive low-rank modules qualify (plain LoRA: no DoRA, no CP ``mid``, +no LyCORIS dense-bias, no ``diff_b``). Layers with any non-factorable +contribution fall back to the dequantize-add-requantize path. +""" + +import torch + +from modules import devices +from modules.lora import lora_common as l +from modules.logger import log + + +fallback_layers: list[str] = [] + + +def get_module_factors(module, device, dtype): + """Return ``(up_eff, down)`` reproducing ``calc_updown`` exactly, or None. + + ``updown = up @ down * calc_scale() * multiplier()`` for a plain linear + LoRA; the scalars fold into the up factor. ``dyn_dim`` slices ranks the + same way ``lyco_helpers.rebuild_conventional`` does. + """ + if module.__class__.__name__ != 'NetworkModuleLora': + return None + if module.dora_scale is not None or module.bias is not None or module.ex_bias is not None: + return None + if getattr(module, 'mid_model', None) is not None: + return None + up = module.up_model.weight + down = module.down_model.weight + if up.ndim != 2 or down.ndim != 2: + return None + dyn_dim = module.network.dyn_dim + if dyn_dim is not None and up.shape[1] != dyn_dim: + up = up[:, :dyn_dim] + down = down[:dyn_dim] + scalar = module.calc_scale() * module.multiplier() + up_eff = up.to(device=device, dtype=torch.float32) * scalar + return up_eff.to(dtype=dtype), down.to(device=device, dtype=dtype) + + +def factor_candidate(self, network_layer_name, wanted_names, use_previous=False): + """True when this layer should take the exact svd-append path. + + Requires an SDNQ linear layer whose active networks all contribute plain + factorable LoRA modules for this layer. An empty ``wanted_names`` is a + removal request and qualifies whenever factors are currently attached. + """ + if getattr(self, 'sdnq_dequantizer', None) is None or self.__class__.__name__ != 'SDNQLinear': + return False + if hasattr(self, 'sdnq_lora_svd_stash'): + return True + if wanted_names == (): # nothing attached, nothing to remove + return False + loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks + seen = False + for net in loaded: + module = net.modules.get(network_layer_name, None) + if module is None: + continue + seen = True + if module.__class__.__name__ != 'NetworkModuleLora': + return False + if module.dora_scale is not None or module.bias is not None or module.ex_bias is not None or getattr(module, 'mid_model', None) is not None: + return False + if module.up_model.weight.ndim != 2 or module.down_model.weight.ndim != 2: + return False + if module.up_model.weight.shape[0] != self.sdnq_dequantizer.original_shape[0] or module.down_model.weight.shape[1] != self.sdnq_dequantizer.original_shape[-1]: + return False + return seen + + +def remove_factors(self): + """Restore the layer's original svd factors; True when factors were attached.""" + stash = getattr(self, 'sdnq_lora_svd_stash', None) + if stash is None: + return False + svd_up, svd_down = stash + self.svd_up = svd_up + self.svd_down = svd_down + del self.sdnq_lora_svd_stash + return True + + +def apply_factors(self, network_layer_name, wanted_names, use_previous=False): + """Attach the active networks' LoRA factors to this layer's svd side-channel. + + Replaces any previously attached factors (multiplier changes re-enter + here with a new ``wanted_names`` signature). Returns True when the layer + changed. Falls back to the caller's requantize path by returning None + when factor extraction fails at this stage. + """ + from sdnq.quant_utils import rotate_hadamard + + changed = remove_factors(self) + if wanted_names == (): + return changed + + deq = self.sdnq_dequantizer + device = self.scale.device + dtype = deq.result_dtype + loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks + ups, downs = [], [] + for net in loaded: + module = net.modules.get(network_layer_name, None) + if module is None: + continue + factors = get_module_factors(module, devices.device, dtype) + if factors is None: + return None + up_eff, down = factors + if deq.use_hadamard: + down = rotate_hadamard(down.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype) + ups.append(up_eff) + downs.append(down) + if not ups: + return changed + + orig_up, orig_down = self.svd_up, self.svd_down + if deq.use_quantized_matmul: + # matmul layout stores factors transposed: svd_up [r, out], svd_down [in, r] + parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + [u.t() for u in ups] + parts_down = ([orig_down.to(device=devices.device, dtype=dtype)] if orig_down is not None else []) + [d.t() for d in downs] + new_up = torch.cat(parts_up, dim=0).contiguous() + new_down = torch.cat(parts_down, dim=1).contiguous() + else: + parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + ups + parts_down = ([orig_down.to(device=devices.device, dtype=dtype)] if orig_down is not None else []) + downs + new_up = torch.cat(parts_up, dim=1).contiguous() + new_down = torch.cat(parts_down, dim=0).contiguous() + + self.sdnq_lora_svd_stash = (orig_up, orig_down) + self.svd_up = torch.nn.Parameter(new_up.to(device=device), requires_grad=False) + self.svd_down = torch.nn.Parameter(new_down.to(device=device), requires_grad=False) + return True + + +def note_fallback(self, network_layer_name): + """Record a quantized layer taking the lossy requantize path (summary-logged per pass).""" + if getattr(self, 'sdnq_dequantizer', None) is not None: + fallback_layers.append(network_layer_name) + + +def report_fallbacks(): + if len(fallback_layers) > 0: + log.warning(f'Network load: type=LoRA quant=sdnq layers={len(fallback_layers)} non-factorable networks requantized in place (reduced fidelity on quantized weights)') + if l.debug: + log.debug(f'Network load: type=LoRA quant=sdnq requantized={fallback_layers[:8]}{"..." if len(fallback_layers) > 8 else ""}') + fallback_layers.clear() diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 31945c27d..50d5d27aa 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -3,6 +3,8 @@ import time import rich.progress as rp from modules.errorlimiter import limit_errors from modules.lora import lora_common as l +from modules.lora import lora_overrides +from modules.lora import lora_sdnq from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights from modules import shared, devices, sd_models from modules.logger import log, console @@ -53,6 +55,7 @@ def network_activate(include=None, exclude=None): net.unet_multiplier = pending['unet'] net.dyn_dim = pending['dyn'] t0 = time.time() + fuse = lora_overrides.fuse_native() # resolve once: backup and apply passes must agree with limit_errors("network_activate") as elimit: sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) if shared.opts.diffusers_offload_mode == "sequential": @@ -100,7 +103,20 @@ def network_activate(include=None, exclude=None): continue if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.loaded_networks): device = group_offload_strip(sd_model, component, group_stripped) - backup_size += network_backup_weights(module, network_layer_name, component_wanted) + if lora_sdnq.factor_candidate(module, network_layer_name, component_wanted): + weights_backup = getattr(module, "network_weights_backup", None) + if weights_backup is not None and not isinstance(weights_backup, bool): + network_apply_weights(module, None, None, device=device) # an earlier non-factorable set requantized this layer, restore the pristine base before attaching factors + applied = lora_sdnq.apply_factors(module, network_layer_name, component_wanted) + if applied is not None: # exact path took the layer; None falls through to requantize + if applied and component_wanted: + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + if task is not None: + pbar.update(task, advance=1) + continue + backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse) if not component_wanted: weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is None or isinstance(weights_backup, bool): # fuse mode has no tensor backup, restore stays with network_deactivate @@ -110,7 +126,9 @@ def network_activate(include=None, exclude=None): batch_updown, batch_ex_bias = None, None # restore-only pass, apply with no weights reverts to backup else: batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit) - if shared.opts.lora_fuse_native: + if batch_updown is not None: + lora_sdnq.note_fallback(module, network_layer_name) # only layers whose quantized weight actually takes a delta + if fuse: weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device) else: weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device) @@ -129,13 +147,14 @@ def network_activate(include=None, exclude=None): if task is not None and len(applied_layers) == 0: pbar.remove_task(task) # hide progress bar for no action global native_active, refused_writes # pylint: disable=global-statement + lora_sdnq.report_fallbacks() native_active = len(l.loaded_networks) > 0 refused_writes = refused l.timer.activate += time.time() - t0 if refused > 0: log.error(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} weights={applied_weight} bias={applied_bias} refused={refused} network partially applied') if l.debug and len(l.loaded_networks) > 0: - log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} refused={refused} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') + log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} refused={refused} backup={round(backup_size/1024/1024/1024, 2)} fuse={fuse}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') modules.clear() if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") @@ -146,7 +165,8 @@ def network_deactivate(include=None, exclude=None): exclude = [] if include is None: include = [] - if not shared.opts.lora_fuse_native or shared.opts.lora_force_diffusers: + fuse = lora_overrides.fuse_native() # must match network_activate: backup mode restores in its restore-only pass instead + if not fuse or shared.opts.lora_force_diffusers: return if len(l.previously_loaded_networks) == 0: return @@ -190,8 +210,14 @@ def network_deactivate(include=None, exclude=None): continue if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.previously_loaded_networks): device = group_offload_strip(sd_model, component, group_stripped) + if lora_sdnq.remove_factors(module): # exact inverse for factor-mode layers, weights were never touched + applied_layers.append(network_layer_name) + module.network_current_names = () + if task is not None: + pbar.update(task, advance=1) + continue batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit) - if shared.opts.lora_fuse_native: + if fuse: weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) else: weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True) @@ -206,7 +232,7 @@ def network_deactivate(include=None, exclude=None): if refused > 0: log.error(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} unapply={len(applied_layers)} refused={refused} network partially removed') if l.debug and len(l.previously_loaded_networks) > 0: - log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} refused={refused} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') + log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} refused={refused} fuse={fuse}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') modules.clear() if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py new file mode 100644 index 000000000..7eb0c2943 --- /dev/null +++ b/test/test-sdnq-lora-factors.py @@ -0,0 +1,534 @@ +#!/usr/bin/env python +""" +Offline unit tests for LoRA application on SDNQ-quantized layers. + +Pins two facts established on real checkpoints (see cli/lora-quant-fidelity.py +for the per-model analyzer): + +- The requantize path (dequantize + add + requantize) erases sub-step deltas + on low-bit formats: retention collapses to the ~2/group_size grid-extrema + floor on uint4, while int8 retains most of the delta. Guards against the + erasure law silently changing. +- The factor path (modules/lora/lora_sdnq.py) applies plain LoRA deltas + through the svd side-channel exactly, in both svd layouts, with exact + stacking, multiplier scaling and bit-exact removal, wired through the real + networks.network_activate / network_deactivate control flow. +- Multi-LoRA set transitions keep the base pristine: a layer that fell back + to requantize (mixed factorable/non-factorable set) restores from backup + before re-entering the factor path, layers targeted by only some of the + loaded networks stay independent, and untargeted quantized layers are not + flagged as requantized. + +All tensors are synthetic; no model files or running server required. + +Usage: + python test/test-sdnq-lora-factors.py +""" + +import os +import sys +import time +from contextlib import contextmanager + +import torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + +from modules.errors import log # pylint: disable=wrong-import-position +from modules import shared, sd_models # pylint: disable=wrong-import-position +from modules.lora import network, network_lora, lora_sdnq, networks # pylint: disable=wrong-import-position +from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position +from sdnq.quantizer import sdnq_quantize_layer, SDNQConfig # pylint: disable=wrong-import-position + +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +OUT_F, IN_F, RANK = 512, 512, 8 + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + record(cat, ok is not False, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +def build_layer(weights_dtype='uint4', use_quantized_matmul=False, seed=0): + torch.manual_seed(seed) + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.04) + cfg = SDNQConfig(weights_dtype=weights_dtype, group_size=0, hadamard_group_size=256, use_hadamard=True, + use_svd=False, use_quantized_matmul=use_quantized_matmul, dequantize_fp32=False, + quantization_device=str(DEVICE), return_device=str(DEVICE)) + layer, _ = sdnq_quantize_layer(lin, cfg, torch_dtype=torch.bfloat16, param_name='test.weight') + layer.network_layer_name = 'lora_transformer_test' + layer.network_current_names = () + return layer + + +def dq(layer): + return layer.sdnq_dequantizer(layer.weight, layer.scale, zero_point=layer.zero_point, + svd_up=layer.svd_up, svd_down=layer.svd_down, + skip_quantized_matmul=layer.sdnq_dequantizer.use_quantized_matmul, + dtype=torch.float32, skip_compile=True) + + +def make_delta(seed=1, sigma=3e-4): + torch.manual_seed(seed) + A = torch.randn(RANK, IN_F, device=DEVICE) * (sigma ** 0.5) + B = torch.randn(OUT_F, RANK, device=DEVICE) * (sigma ** 0.5) + return A, B, B @ A + + +class MockNOD: + def __init__(self, name): + self.filename = f'/tmp/{name}.safetensors' + self.name = name + self.shorthash = '' + self.sd_version = 'unknown' + + +def make_net(name, layer, A, B, te_mult=1.0, alpha=None, dora=False): + net = network.Network(name, MockNOD(name)) + net.te_multiplier = te_mult + net.unet_multiplier = [te_mult] * 3 + w = {'lora_up.weight': B.cpu(), 'lora_down.weight': A.cpu()} + if alpha is not None: + w['alpha'] = torch.tensor(float(alpha)) + if dora: + w['dora_scale'] = torch.ones(B.shape[0], 1) + nw = network.NetworkWeights(network_key=layer.network_layer_name, sd_key=layer.network_layer_name, w=w, sd_module=layer) + mod = network_lora.NetworkModuleLora(net, nw) + net.modules[layer.network_layer_name] = mod + return net + + +def rho_of(E, D): + return float(E.flatten() @ D.flatten() / D.flatten().square().sum()) + + +def requant_effective(layer, D): + """The lossy fallback path: quantize(W_dq + D) fresh with the layer's own params.""" + from sdnq.quantizer import sdnq_quantize_layer_weight + deq = layer.sdnq_dequantizer + Wdq = dq(layer) + deq2, data2 = sdnq_quantize_layer_weight(Wdq + D, layer_class_name='Linear', weights_dtype=deq.weights_dtype, + group_size=deq.group_size, hadamard_group_size=deq.hadamard_group_size, + use_hadamard=deq.use_hadamard, use_svd=False, use_quantized_matmul=False, + dequantize_fp32=False, torch_dtype=torch.bfloat16) + W2 = deq2(data2['weight'], data2['scale'], zero_point=data2['zero_point'], svd_up=None, svd_down=None, dtype=torch.float32, skip_compile=True) + return W2 - Wdq + + +# ============================================================ +# Tests - the erasure law (why the factor path exists) +# ============================================================ + +CAT_LAW = category('erasure-law') + + +def test_uint4_erases_substep_delta(): + layer = build_layer('uint4') + _A, _B, D = make_delta(sigma=2e-4) + rho = rho_of(requant_effective(layer, D), D) + group = layer.sdnq_dequantizer.group_size + floor = 2.0 / group + assert rho < 4 * floor, f'rho={rho:.4f} expected near extrema floor {floor:.4f}' + return True + + +def test_int8_retains_delta(): + layer = build_layer('int8') + _A, _B, D = make_delta(sigma=2e-4) + rho = rho_of(requant_effective(layer, D), D) + assert rho > 0.5, f'rho={rho:.4f} expected int8 to retain most of the delta' + return True + + +# ============================================================ +# Tests - factor path exactness +# ============================================================ + +CAT_FACTOR = category('factor-path') + + +def test_apply_exact_and_remove_bitexact(): + layer = build_layer('uint4') + A, B, D = make_delta() + net = make_net('one', layer, A, B) + l_common.loaded_networks.clear() + l_common.loaded_networks.append(net) + wanted = (('one', 1.0, 1.0, None),) + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) is True + Wdq0 = dq(layer) + assert lora_sdnq.apply_factors(layer, layer.network_layer_name, wanted) is True + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + assert lora_sdnq.remove_factors(layer) is True + assert torch.equal(dq(layer), Wdq0), 'remove must be bit-exact' + assert layer.svd_up is None and not hasattr(layer, 'sdnq_lora_svd_stash') + l_common.loaded_networks.clear() + return True + + +def test_multiplier_and_alpha_scaling(): + layer = build_layer('uint4') + A, B, D = make_delta() + net = make_net('one', layer, A, B, te_mult=0.5, alpha=RANK // 2) # alpha/rank = 0.5 + l_common.loaded_networks.clear() + l_common.loaded_networks.append(net) + Wdq0 = dq(layer) + lora_sdnq.apply_factors(layer, layer.network_layer_name, (('one', 0.5, 0.5, None),)) + rho = rho_of(dq(layer) - Wdq0, D) + assert abs(rho - 0.25) < 0.01, f'expected 0.5*0.5 scaling, rho={rho:.4f}' + lora_sdnq.remove_factors(layer) + l_common.loaded_networks.clear() + return True + + +def test_stacking_two_networks(): + layer = build_layer('uint4') + A1, B1, D1 = make_delta(seed=1) + A2, B2, D2 = make_delta(seed=2) + l_common.loaded_networks.clear() + l_common.loaded_networks.extend([make_net('a', layer, A1, B1), make_net('b', layer, A2, B2)]) + Wdq0 = dq(layer) + lora_sdnq.apply_factors(layer, layer.network_layer_name, (('a', 1.0, 1.0, None), ('b', 1.0, 1.0, None))) + rho = rho_of(dq(layer) - Wdq0, D1 + D2) + assert rho > 0.99, f'rho={rho:.4f}' + lora_sdnq.remove_factors(layer) + l_common.loaded_networks.clear() + return True + + +def test_matmul_layout_transposed(): + layer = build_layer('uint4', use_quantized_matmul=True) + A, B, D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + Wdq0 = dq(layer) + res = lora_sdnq.apply_factors(layer, layer.network_layer_name, (('one', 1.0, 1.0, None),)) + rho = rho_of(dq(layer) - Wdq0, D) + assert res is True and rho > 0.99, f'rho={rho:.4f}' + lora_sdnq.remove_factors(layer) + assert torch.equal(dq(layer), Wdq0) + l_common.loaded_networks.clear() + return True + + +def test_dora_falls_back(): + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('dora', layer, A, B, dora=True)) + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, (('dora', 1.0, 1.0, None),)) is False + l_common.loaded_networks.clear() + return True + + +# ============================================================ +# Tests - memory accounting across apply modes +# ============================================================ + +CAT_MEM = category('memory') + + +def tensor_bytes(t): + return t.numel() * t.element_size() if isinstance(t, torch.Tensor) else 0 + + +def test_factor_path_memory_is_factors_only(): + """Factor path: no weight/quant-state backups; added memory = the factor tensors.""" + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + lora_sdnq.apply_factors(layer, layer.network_layer_name, (('one', 1.0, 1.0, None),)) + assert getattr(layer, 'network_weights_backup', None) is None + assert not hasattr(layer, 'sdnq_dequantizer_backup') and not hasattr(layer, 'sdnq_scale_backup') + added = tensor_bytes(layer.svd_up) + tensor_bytes(layer.svd_down) + expected = RANK * (OUT_F + IN_F) * 2 # bf16 factors + assert added == expected, f'factor bytes {added} != expected {expected}' + would_be_backup = tensor_bytes(layer.weight) + tensor_bytes(layer.scale) + tensor_bytes(layer.zero_point) + assert added < would_be_backup / 4, f'factors {added}B should undercut the {would_be_backup}B backup this layer would otherwise clone' + lora_sdnq.remove_factors(layer) + assert layer.svd_up is None and layer.svd_down is None + l_common.loaded_networks.clear() + return True + + +def test_backup_mode_clones_full_quant_state(): + """Fallback in backup mode: packed weight + scale + zero_point are cloned to cpu.""" + from modules.lora.lora_apply import network_backup_weights + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('dora', layer, A, B, dora=True)) # non-factorable + reported = network_backup_weights(layer, layer.network_layer_name, (('dora', 1.0, 1.0, None),), fuse=False) + assert isinstance(layer.network_weights_backup, torch.Tensor) and layer.network_weights_backup.device.type == 'cpu' + assert hasattr(layer, 'sdnq_dequantizer_backup') and isinstance(layer.sdnq_scale_backup, torch.Tensor) + assert reported == tensor_bytes(layer.weight), f'reported {reported} != packed weight bytes {tensor_bytes(layer.weight)}' + total = reported + tensor_bytes(layer.sdnq_scale_backup) + tensor_bytes(layer.sdnq_zero_point_backup) + expected_min = OUT_F * IN_F // 2 # uint4 packs two weights per byte + assert total >= expected_min, f'backup {total}B below packed-weight floor {expected_min}B' + l_common.loaded_networks.clear() + return True + + +def test_fuse_mode_marker_takes_no_memory(): + """Fuse mode stores a boolean marker instead of tensors; guard forces backup on quantized models.""" + from modules.lora.lora_apply import network_backup_weights + from modules.lora import lora_overrides + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('dora', layer, A, B, dora=True)) + reported = network_backup_weights(layer, layer.network_layer_name, (('dora', 1.0, 1.0, None),), fuse=True) + assert layer.network_weights_backup is True and reported == 0 + assert not hasattr(layer, 'sdnq_dequantizer_backup') + + # the guard: a quantized component forces fuse off model-wide regardless of the option + class MockCfg: + quantization_config = {'quant_method': 'sdnq'} + class MockSd: + pass + sd = MockSd() + sd.transformer = torch.nn.Linear(4, 4) + sd.transformer.config = MockCfg() + from modules.modeldata import model_data + prev_model = model_data.sd_model + old_fuse = shared.opts.lora_fuse_native + try: + model_data.sd_model = sd + shared.opts.lora_fuse_native = True + assert lora_overrides.disable_fuse() is True + assert lora_overrides.fuse_native() is False + finally: + shared.opts.lora_fuse_native = old_fuse + model_data.sd_model = prev_model + l_common.loaded_networks.clear() + return True + + +# ============================================================ +# Tests - integration through networks.network_activate +# ============================================================ + +CAT_E2E = category('activate-e2e') + + +class MockHolder(torch.nn.Module): + @property + def device(self): + return DEVICE + + +@contextmanager +def mock_model(**layers): + """Install a one-component mock pipeline holding the given layers as shared.sd_model.""" + class MockPipe: + pass + class MockSd: + pass + holder = MockHolder() + for attr, lyr in layers.items(): + setattr(holder, attr, lyr) + pipe = MockPipe() + pipe.transformer = holder + sd = MockSd() + sd.pipe = pipe + from modules.modeldata import model_data + model_data.sd_model = sd + real_offload = sd_models.set_diffuser_offload + sd_models.set_diffuser_offload = lambda *a, **k: None + old_fuse = shared.opts.lora_fuse_native + shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config, so pin it instead of inheriting the running config + try: + yield + finally: + shared.opts.lora_fuse_native = old_fuse + sd_models.set_diffuser_offload = real_offload + l_common.loaded_networks.clear() + l_common.previously_loaded_networks.clear() + + +def activate(*nets): + l_common.loaded_networks.clear() + l_common.loaded_networks.extend(nets) + networks.network_activate() + + +def test_network_activate_roundtrip(): + layer = build_layer('uint4') + A, B, D = make_delta() + net = make_net('one', layer, A, B) + + with mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + assert getattr(layer, 'network_weights_backup', None) is None, 'factor path must not take weight backups' + + activate() # restore pass + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + + # fuse-mode deactivate route + layer.network_current_names = () + activate(net) + l_common.previously_loaded_networks[:] = l_common.loaded_networks + shared.opts.lora_fuse_native = True + networks.network_deactivate() + assert torch.equal(dq(layer), Wdq0), 'fuse-mode deactivate must restore bit-exact' + return True + + +# ============================================================ +# Tests - multi-LoRA set transitions between the two paths +# ============================================================ + +CAT_TRANS = category('transitions') + + +def test_mixed_family_transition_restores_base(): + layer = build_layer('uint4') + bystander = build_layer('uint4', seed=7) + bystander.network_layer_name = 'lora_transformer_bystander' + A, B, D = make_delta() + net_plain = make_net('plain', layer, A, B) + A2, B2, _ = make_delta(seed=5, sigma=3e-3) + net_dora = make_net('doranet', layer, A2, B2, dora=True) + + noted = [] + real_report = lora_sdnq.report_fallbacks + def capture_report(): + noted.append(len(lora_sdnq.fallback_layers)) + real_report() + lora_sdnq.report_fallbacks = capture_report + try: + with mock_model(lin=layer, bystander=bystander): + Wdq0 = dq(layer) + activate(net_plain) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'plain set must take the factor path' + assert noted[-1] == 0, f'untargeted quantized layers must not be flagged as requantized: noted={noted[-1]}' + + activate(net_plain, net_dora) + assert not hasattr(layer, 'sdnq_lora_svd_stash') and isinstance(layer.network_weights_backup, torch.Tensor), 'mixed set must fall back with a tensor backup' + assert not torch.equal(dq(layer), Wdq0), 'fallback must have requantized the weights' + assert noted[-1] == 1, f'exactly the requantized layer must be flagged: noted={noted[-1]}' + + activate(net_plain) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'plain-only set must re-enter the factor path' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + stash, up, down = layer.sdnq_lora_svd_stash, layer.svd_up, layer.svd_down + lora_sdnq.remove_factors(layer) + base_clean = torch.equal(dq(layer), Wdq0) + layer.sdnq_lora_svd_stash, layer.svd_up, layer.svd_down = stash, up, down + assert base_clean, 'base under factors must be restored from backup on mixed-set exit' + + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must return bit-exact pristine' + finally: + lora_sdnq.report_fallbacks = real_report + return True + + +def test_partial_coverage_layers_stay_independent(): + layer_plain = build_layer('uint4') + layer_dora = build_layer('uint4', seed=7) + layer_dora.network_layer_name = 'lora_transformer_other' + A, B, D = make_delta() + net_plain = make_net('plain', layer_plain, A, B) + A2, B2, _ = make_delta(seed=5, sigma=3e-3) + net_dora = make_net('dorafar', layer_dora, A2, B2, dora=True) + + with mock_model(lin=layer_plain, other=layer_dora): + Wdq0, Wdq0_dora = dq(layer_plain), dq(layer_dora) + activate(net_plain, net_dora) + assert hasattr(layer_plain, 'sdnq_lora_svd_stash') and getattr(layer_plain, 'network_weights_backup', None) is None, 'plain layer must stay on the factor path' + assert isinstance(getattr(layer_dora, 'network_weights_backup', None), torch.Tensor), 'dora layer must take the backup fallback' + rho = rho_of(dq(layer_plain) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + activate() + assert torch.equal(dq(layer_plain), Wdq0), 'factor layer must restore bit-exact' + assert torch.equal(dq(layer_dora), Wdq0_dora), 'fallback layer must restore bit-exact' + return True + + +def run_tests(): + t0 = time.time() + log.warning('=== Erasure law ===') + for fn in [test_uint4_erases_substep_delta, test_int8_retains_delta]: + run_test(CAT_LAW, fn) + log.warning('=== Factor path ===') + for fn in [test_apply_exact_and_remove_bitexact, test_multiplier_and_alpha_scaling, test_stacking_two_networks, test_matmul_layout_transposed, test_dora_falls_back]: + run_test(CAT_FACTOR, fn) + log.warning('=== Memory accounting ===') + for fn in [test_factor_path_memory_is_factors_only, test_backup_mode_clones_full_quant_state, test_fuse_mode_marker_takes_no_memory]: + run_test(CAT_MEM, fn) + log.warning('=== Activate integration ===') + for fn in [test_network_activate_roundtrip]: + run_test(CAT_E2E, fn) + log.warning('=== Set transitions ===') + for fn in [test_mixed_family_transition_restores_base, test_partial_coverage_layers_stay_independent]: + run_test(CAT_TRANS, fn) + + elapsed = time.time() - t0 + log.warning('=== Results ===') + total_pass = total_fail = 0 + for cat, info in results.items(): + status = 'PASS' if info['failed'] == 0 else 'FAIL' + log.info(f' {cat}: {info["passed"]} passed, {info["failed"]} failed [{status}]') + total_pass += info['passed'] + total_fail += info['failed'] + log.warning(f'Total: {total_pass} passed, {total_fail} failed in {elapsed:.2f}s') + return total_fail == 0 + + +if __name__ == '__main__': + with torch.inference_mode(): + ok = run_tests() + sys.exit(0 if ok else 1)