From cf36f879a1eac60a80965a5c5b303f72d1504762 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 06:09:13 +0100 Subject: [PATCH] fix(lora): harden select stack modes on quantized and offloaded models Select pairs now ride the svd side channel on any SDNQ linear, not only sub-8-bit ones: quantized backups are packed tensors, so the weight rewrite path cannot recompute a winner from them and left layers stripped mid-requantize. The sub-8-bit gate stays for dense hosting, where requantize retains the delta at 8 bits and above. Weight selection now only serves unquantized modules: finalize iterates a snapshot so dead entries drop cleanly, materializes balanced-offload modules before rewriting weights and skips modules with stripped or quantized weights instead of corrupting the layer. --- modules/lora/lora_sdnq.py | 14 ++++++++++---- modules/lora/lora_stack.py | 20 ++++++++++++++++++-- modules/lora/networks.py | 2 +- test/test-sdnq-lora-factors.py | 34 +++++++++++++++++++++++++++++++++- 4 files changed, 62 insertions(+), 8 deletions(-) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index df2e7d2b0..b58b26799 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -258,8 +258,8 @@ def append_factors(self, ups, downs): return segments, deq.use_quantized_matmul -def host_candidate(self, network_layer_name, wanted_names): - """True when a non-factorable set on this layer should be hosted as a truncated svd.""" +def select_candidate(self, network_layer_name, wanted_names): + """True when this layer can carry a set on the svd channel; select pairs ride it at any bit width.""" if not enabled(): return False if int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0) <= 0: @@ -268,11 +268,17 @@ def host_candidate(self, network_layer_name, wanted_names): return False if wanted_names == (): return False + return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) + + +def host_candidate(self, network_layer_name, wanted_names): + """True when a non-factorable set on this layer should be hosted as a truncated svd.""" + if not select_candidate(self, network_layer_name, wanted_names): + return False from sdnq.common import dtype_dict if dtype_dict[self.sdnq_dequantizer.weights_dtype]['num_bits'] >= 8: return False # requantize retains most of the delta at 8 bits and above; truncation would lose more than it saves - - return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) + return True def apply_cached(self, network_layer_name, wanted_names): diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index a9899b1d4..b74b25fde 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -248,12 +248,21 @@ def layer_flip_step(scores, total_steps): return total_steps +def materialize_model(): + """Weight-kind selection rewrites module weights outside the activation walk; rebuild balanced-offload modules real first (mirrors network_activate).""" + from modules import sd_models + if getattr(shared.opts, 'diffusers_offload_mode', None) == 'balanced' and getattr(shared, 'sd_model', None) is not None: + sd_models.apply_balanced_offload(shared.sd_model, force=True) + + def finalize(total_steps): """Build the inverted flip map for the pass; select-mode layers start at their step-0 winner.""" state['total_steps'] = int(total_steps) state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0 state['flips'] = {} - for layer_name, entry in state['entries'].items(): + if any(e['kind'] == 'weight' for e in state['entries'].values()): + materialize_model() + for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died flip_at = layer_flip_step(entry['scores'], state['total_steps']) initial = 1 if flip_at == 0 else 0 apply_selection(layer_name, entry, initial) @@ -298,6 +307,9 @@ def apply_selection(layer_name, entry, winner): def weight_selection(module, entry, winner): from modules.lora import lora_common as l from modules.lora.lora_apply import network_apply_weights + if getattr(module, 'sdnq_dequantizer', None) is not None: + warn_once('select-sdnq-weight', 'Network stack: flip=skipped layer=quantized') # quantized backups are packed tensors; only the segment path can flip them + return backup = getattr(module, 'network_weights_backup', None) if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy warn_once('select-nobackup', 'Network stack: flip=skipped backup=none') @@ -306,6 +318,10 @@ def weight_selection(module, entry, winner): net_module = net.modules.get(entry['layer'], None) if net is not None else None if net_module is None: return - device = module.weight.device + weight = getattr(module, 'weight', None) + if weight is None or weight.is_meta: + warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded') + return + device = weight.device updown = net_module.calc_updown(backup.to(device))[0] network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 6e442f5f3..818f3c100 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -112,7 +112,7 @@ def network_activate(include=None, exclude=None): device = group_offload_strip(sd_model, component, group_stripped) calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing if select_active and component_wanted and not network_layer_name.startswith('lora_te'): - if lora_sdnq.host_candidate(module, network_layer_name, component_wanted): # sub-8-bit SDNQ pairs ride the channel as separate segments + if lora_sdnq.select_candidate(module, network_layer_name, component_wanted): # SDNQ pairs ride the channel as separate segments at any bit width; weight rewrites cannot flip a quantized layer 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) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index d48e660c1..a94e38b68 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -1749,6 +1749,38 @@ def test_select_gated_off_when_compiled(): return True +def test_select_finalize_drops_dead_module(): + import weakref + layer = build_layer('uint4') + n1, n2, _D1, _D2 = select_pair(layer, seed0=61, seed1=62) + with mock_model(lin=layer), select_mode('klora'): + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None, 'pair must register before the module dies' + entry['module'] = weakref.ref(torch.nn.Linear(2, 2)) # referent dies immediately: simulates offload re-wraps replacing a registered module + assert entry['module']() is None + lora_stack.reset(12) + assert 'lora_transformer_test' not in lora_stack.state['entries'], 'a dead module must drop its entry without breaking finalize' + activate() + return True + + +def test_select_int8_pair_rides_segments(): + layer = build_layer('int8') + n1, n2, D1, D2 = select_pair(layer, seed0=63, seed1=64) + with mock_model(lin=layer), select_mode('klora'): + Wdq0 = dq(layer) + activate(n1, n2) + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor', 'an int8 pair must ride svd segments, not weight rewrites' + lora_stack.reset(16) + eff = dq(layer) - Wdq0 + assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.99, 'initial selection must deliver one exact per-net delta' + activate() + assert torch.equal(dq(layer), Wdq0), 'removal must restore bit-exact' + return True + + def test_select_gate_dormant_without_pair(): with select_mode('klora'): assert lora_stack.select_possible(1) is False, 'a single network must leave the fuse gate alone' @@ -2009,7 +2041,7 @@ def run_tests(): for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable, test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state, test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, - test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, + test_select_finalize_drops_dead_module, test_select_int8_pair_rides_segments, test_select_gate_dormant_without_pair, test_stale_schedule_dropped_on_reapply, test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===')