diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 1cdbb8d22..d07c0db8a 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -16,7 +16,10 @@ 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. +weight backup is needed. The side-channel storage is lossless; realized +fidelity floors at the compute dtype, because the dequantizer materializes +``base + factors`` in the result dtype and a delta below its ULP of the +base rounds exactly as it would on an unquantized model of that dtype. 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 @@ -33,7 +36,7 @@ from modules.logger import log fallback_layers: list[str] = [] -def get_module_factors(module, device, dtype): +def get_module_factors(module, device, dtype, original_shape=None): """Return ``(up_eff, down)`` reproducing ``calc_updown`` exactly, or None. ``updown = up @ down * calc_scale() * multiplier()`` for a plain linear @@ -50,6 +53,8 @@ def get_module_factors(module, device, dtype): down = module.down_model.weight if up.ndim != 2 or down.ndim != 2: return None + if original_shape is not None and (up.shape[0] != original_shape[0] or down.shape[1] != original_shape[-1]): + return None # factor_candidate skips shape checks for layers already in factor mode; recheck here so a malformed stack falls back instead of raising in cat dyn_dim = module.network.dyn_dim if dyn_dim is not None and up.shape[1] != dyn_dim: up = up[:, :dyn_dim] @@ -96,6 +101,10 @@ def remove_factors(self): if stash is None: return False svd_up, svd_down = stash + device = self.scale.device # the stash tuple does not follow module device moves; restore onto wherever the layer lives now + if svd_up is not None and svd_up.device != device: + svd_up = torch.nn.Parameter(svd_up.to(device=device), requires_grad=False) + svd_down = torch.nn.Parameter(svd_down.to(device=device), requires_grad=False) self.svd_up = svd_up self.svd_down = svd_down del self.sdnq_lora_svd_stash @@ -125,7 +134,7 @@ def apply_factors(self, network_layer_name, wanted_names, use_previous=False): module = net.modules.get(network_layer_name, None) if module is None: continue - factors = get_module_factors(module, devices.device, dtype) + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) if factors is None: return None up_eff, down = factors diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 50d5d27aa..c472da4ea 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -90,6 +90,7 @@ def network_activate(include=None, exclude=None): with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () applied_layers.clear() + lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind backup_size = 0 for component in modules.keys(): component_wanted = wanted_names if component in components else () diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 6a91380a2..4257f35fd 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -20,6 +20,9 @@ for the per-model analyzer): 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. +- Robustness: factor removal restores onto the layer's current device after + an offload-style move, and a shape-mismatched network stacked onto a + factor-mode layer downgrades to the legacy path instead of raising. All tensors are synthetic; no model files or running server required. @@ -535,6 +538,52 @@ def test_partial_coverage_layers_stay_independent(): return True +CAT_ROBUST = category('robustness') + + +def test_remove_factors_after_device_move(): + layer = build_layer('uint4', use_svd=True) # checkpoint svd correction so the stash holds real tensors + A, B, _D = make_delta() + net = make_net('mover', layer, A, B) + with mock_model(lin=layer): + Wdq0 = dq(layer) + orig_up = layer.svd_up.detach().clone() + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash') + layer.to('cpu') # offload moves registered params, never the stash tuple + activate() + assert layer.svd_up.device == layer.scale.device, f'restored svd must live on the layer device, got {layer.svd_up.device} vs {layer.scale.device}' + assert torch.equal(layer.svd_up, orig_up.to('cpu')), 'restored svd values must match the original factors' + layer.to(DEVICE) + assert torch.equal(dq(layer), Wdq0), 'round trip must restore bit-exact' + return True + + +def test_stacked_shape_mismatch_falls_back(): + from types import SimpleNamespace + layer = build_layer('uint4') + A, B, _D = make_delta() + net_good = make_net('good', layer, A, B) + torch.manual_seed(9) + A_bad = torch.randn(RANK, IN_F, device=DEVICE) * 0.01 + B_bad = torch.randn(OUT_F // 2, RANK, device=DEVICE) * 0.01 # wrong out_features for this layer + net_bad = make_net('badshape', layer, A_bad, B_bad) + prev_enl = l_common.extra_network_lora + l_common.extra_network_lora = SimpleNamespace(errors={}) # the error path reports through the extra-networks registry + try: + with mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net_good) + assert hasattr(layer, 'sdnq_lora_svd_stash') + activate(net_good, net_bad) # must not raise: a malformed stack downgrades the layer to the legacy path + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'shape-mismatched stack must leave factor mode' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact pristine' + finally: + l_common.extra_network_lora = prev_enl + return True + + def run_tests(): t0 = time.time() log.warning('=== Erasure law ===') @@ -552,6 +601,9 @@ def run_tests(): log.warning('=== Set transitions ===') for fn in [test_mixed_family_transition_restores_base, test_partial_coverage_layers_stay_independent]: run_test(CAT_TRANS, fn) + log.warning('=== Robustness ===') + for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: + run_test(CAT_ROBUST, fn) elapsed = time.time() - t0 log.warning('=== Results ===')