diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index d6c95d681..aab36c7a5 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -157,6 +157,30 @@ def fetch(network_layer_name): return entry +def lookup_scores(network_layer_name): + """Cached select scores for a layer as ((s0, s1), (a0, a1)), or None. + + Score records ride the same signature-keyed entry as factors, and the + signature already pins everything the scores depend on (pair, multipliers, + stack mode and params). No hit/miss accounting: a record saves scoring and + delta assembly, not a sketch. + """ + if state['sig'] is None: + return None + t = state['store'].get(f'{network_layer_name}.sel') + if t is None: + return None + return (float(t[0]), float(t[1])), (float(t[2]), float(t[3])) + + +def store_scores(network_layer_name, scores, abs_sums): + """Persist a select-mode score record; additive to the entry, older files upgrade on their next pass.""" + if state['sig'] is None: + return + state['store'][f'{network_layer_name}.sel'] = torch.tensor([scores[0], scores[1], abs_sums[0], abs_sums[1]], dtype=torch.float64) + state['dirty'] = True + + def store(network_layer_name, up, down, energy, calibrated, rms): """Quantize-before-use: returns the dequantized round-trip the caller must apply. diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 375e8d05a..35635b227 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -460,6 +460,54 @@ def truncate_delta(self, D, dtype): return up_h, down_h, energy, rms is not None +def apply_select_cached(self, network_layer_name, wanted_names): + """Serve a select pair from cache and live factors before the walk assembles deltas. + + A cached score record plus a factor pair per network (exact factors for + factorable members, cached truncations otherwise) rebuild the segments and + the selection registration without any ``calc_updown``. Returns None when + any piece is missing; the caller assembles and ``apply_select`` recomputes + and stores. + """ + from sdnq.quant_utils import rotate_hadamard + deq = self.sdnq_dequantizer + changed = remove_factors(self) + if wanted_names == (): + return changed + if len(l.loaded_networks) != 2: + return None + dtype = deq.result_dtype + lora_factor_cache.begin_pass(wanted_names) + rec = lora_factor_cache.lookup_scores(network_layer_name) + if rec is None: + return None + pairs, notes = [], [] + for i, net in enumerate(l.loaded_networks): + module = net.modules.get(network_layer_name, None) + if module is None: + return None + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) + if factors is not None: + up_i, down_i = factors + if deq.use_hadamard: + down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype) + else: + cached = lora_factor_cache.lookup(f'{network_layer_name}#{i}') + if cached is None: + return None + up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype) + notes.append((f'{network_layer_name}#{i}', cached[2], cached[3])) + pairs.append((up_i, down_i)) + scores, abs_sums = rec + segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) + lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) + for note in notes: + lora_factor_cache.note_hit() + hosted_layers.append(note) + select_layers.append(network_layer_name) + return True + + def apply_select(self, network_layer_name, per_net, wanted_names): """Attach two networks' contributions as separate side-channel segments for per-layer selection. @@ -502,10 +550,8 @@ def apply_select(self, network_layer_name, per_net, wanted_names): up_i, down_i = lora_factor_cache.store(key, up_i, down_i, energy, calibrated, float(D.detach().float().square().mean().sqrt())) hosted_layers.append((key, energy, calibrated)) pairs.append((up_i, down_i)) - d0 = per_net[0][1].detach().to(devices.device, torch.float32) - d1 = per_net[1][1].detach().to(devices.device, torch.float32) - scores, abs_sums = lora_stack.score_pair(d0, d1, ranks[0], ranks[1]) - del d0, d1 + scores, abs_sums = lora_stack.score_pair(per_net[0][1].detach(), per_net[1][1].detach(), ranks[0], ranks[1]) + lora_factor_cache.store_scores(network_layer_name, scores, abs_sums) segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]]) lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums) select_layers.append(network_layer_name) # counted apart from the plain concat: both ride the svd channel but only one is a summed set diff --git a/modules/lora/lora_stack.py b/modules/lora/lora_stack.py index ef60bc735..acd1117fa 100644 --- a/modules/lora/lora_stack.py +++ b/modules/lora/lora_stack.py @@ -145,20 +145,43 @@ def combine(named_deltas, layer_name): def score_pair(d0, d1, rank0, rank1): - """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance.""" - abs_sums = (float(d0.abs().sum()), float(d1.abs().sum())) - if mode() == 'klora': - k = max(1, int(rank0) * int(rank1)) - s0 = float(torch.topk(d0.abs().flatten(), min(k, d0.numel()), sorted=False).values.sum()) - s1 = float(torch.topk(d1.abs().flatten(), min(k, d1.numel()), sorted=False).values.sum()) - else: - s0 = float(d0.float().square().sum()) - s1 = float(d1.float().square().sum()) - return (s0, s1), abs_sums + """Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance. + + Row-chunked fp32 interiors with fp64 accumulators and one device sync for + all four reductions. Full-tensor staging (fp32 copy, abs copy, top-k + workspace) peaks hundreds of MB per large layer, which collides with block + swapping on offloaded denoisers; chunking bounds the transient to the + chunk. The global top-K over per-chunk top-K candidates selects the same + element set as a whole-tensor top-K. + """ + k = max(1, int(rank0) * int(rank1)) if mode() == 'klora' else 0 + accs = [] + for d in (d0, d1): + score = torch.zeros((), device=d.device, dtype=torch.float64) + abs_sum = torch.zeros((), device=d.device, dtype=torch.float64) + cands = [] + for start in range(0, d.shape[0], ROW_CHUNK): + c = d[start:start + ROW_CHUNK].to(torch.float32).abs() # out-of-place abs: to() may alias a caller-owned fp32 tensor + abs_sum += c.sum(dtype=torch.float64) + if k: + flat = c.flatten() + cands.append(torch.topk(flat, min(k, flat.numel()), sorted=False).values) + else: + score += c.square().sum(dtype=torch.float64) + if k and cands: + allc = torch.cat(cands) if len(cands) > 1 else cands[0] + score = torch.topk(allc, min(k, allc.numel()), sorted=False).values.sum(dtype=torch.float64) + accs.append((score, abs_sum)) + packed = torch.stack([accs[0][0], accs[0][1], accs[1][0], accs[1][1]]).cpu() + return (float(packed[0]), float(packed[2])), (float(packed[1]), float(packed[3])) -def register_weight_pair(layer_name, module, per_net): - """Score and register a weight-kind selection pair; True when the layer is scheduled.""" +def register_weight_pair(layer_name, module, per_net, wanted_names=None): + """Score and register a weight-kind selection pair; True when the layer is scheduled. + + The scores persist in the factor cache when a pass identity is given, so a + later apply of the same configuration registers from the record alone. + """ from modules.lora import lora_common as l if per_net is None or len(per_net) != 2: return False @@ -172,11 +195,35 @@ def register_weight_pair(layer_name, module, per_net): return False names.append(net_name) ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64) - scores, abs_sums = score_pair(per_net[0][1].float(), per_net[1][1].float(), ranks[0], ranks[1]) + scores, abs_sums = score_pair(per_net[0][1], per_net[1][1], ranks[0], ranks[1]) + if wanted_names is not None: + from modules.lora import lora_factor_cache + lora_factor_cache.begin_pass(wanted_names) + lora_factor_cache.store_scores(layer_name, scores, abs_sums) register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums) return True +def register_weight_pair_cached(layer_name, module, wanted_names): + """Register a weight-kind pair from its cached score record; True when served. + + The record was stored under the same configuration signature, which pins + the loaded pair, multipliers and stack settings, so both networks are known + to target the layer and the prompt-order roles are unchanged. + """ + from modules.lora import lora_common as l + from modules.lora import lora_factor_cache + if len(l.loaded_networks) != 2: + return False + lora_factor_cache.begin_pass(wanted_names) + rec = lora_factor_cache.lookup_scores(layer_name) + if rec is None: + return False + scores, abs_sums = rec + register(layer_name, module, 'weight', scores, nets=tuple(n.name for n in l.loaded_networks), abs_sums=abs_sums) + return True + + def drop(layer_name): """Forget a layer's selection entry (its factors were removed or restored).""" if layer_name is not None and state['entries'].pop(layer_name, None) is not None: diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 366997118..f62fe431a 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -117,18 +117,20 @@ def network_activate(include=None, exclude=None): 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) - per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) - if sel_bias is None: - applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted) - if applied is not None: - if applied and component_wanted: - applied_layers.append(network_layer_name) - applied_weight += 1 - module.network_current_names = component_wanted - module.network_current_stack = stack_sig - if task is not None: - pbar.update(task, advance=1) - continue + applied = lora_sdnq.apply_select_cached(module, network_layer_name, component_wanted) # a stored score record and factor pair serve before the deltas are assembled + if applied is None: + per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) + if sel_bias is None: + applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted) + if applied is not None: + if applied and component_wanted: + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue lora_stack.warn_once('select-unridable', f'Network stack: mode={lora_stack.mode()} layer="{network_layer_name}" fallback=sum') # a pair the channel cannot carry (bias delta or malformed member) sums like any unsupported set elif getattr(module, 'sdnq_dequantizer', None) is not None: # hosting disabled: quantized layers have no side-channel to carry segments and packed backups cannot flip, so the sum paths below take the layer if any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks): @@ -137,8 +139,18 @@ def network_activate(include=None, exclude=None): sel_backup = network_backup_weights(module, network_layer_name, component_wanted, fuse) weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is not None and not isinstance(weights_backup, bool): + if lora_stack.register_weight_pair_cached(network_layer_name, module, component_wanted): # a stored score record registers without assembling the pair + backup_size += sel_backup + network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner + applied_layers.append(network_layer_name) + applied_weight += 1 + module.network_current_names = component_wanted + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True) - if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net): + if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net, component_wanted): backup_size += sel_backup # counted only when this branch keeps the layer; the fallthrough re-enters the shared backup call below, which counts it then network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner applied_layers.append(network_layer_name) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 7b4455077..61627f5f8 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -2003,6 +2003,122 @@ def test_flip_lands_before_crossover_step(): return True +def test_score_pair_chunked_precision(): + torch.manual_seed(71) + shapes = [(700, 460), (OUT_F, IN_F), (64,)] # off-chunk rows, square, and a 1-D norm delta + for shape in shapes: + d0 = (torch.randn(*shape, device=DEVICE) * 1e-2).to(torch.bfloat16) + d1 = (torch.randn(*shape, device=DEVICE) * 1e-2).to(torch.bfloat16) + for mode_name in ('klora', 'estlora'): + with select_mode(mode_name): + (s0, s1), (a0, a1) = lora_stack.score_pair(d0, d1, 8, 8) + f0, f1 = d0.to(torch.float64), d1.to(torch.float64) + ra0, ra1 = float(f0.abs().sum()), float(f1.abs().sum()) + if mode_name == 'klora': + k = 64 + rs0 = float(torch.topk(f0.abs().flatten(), min(k, f0.numel()), sorted=False).values.sum()) + rs1 = float(torch.topk(f1.abs().flatten(), min(k, f1.numel()), sorted=False).values.sum()) + else: + rs0, rs1 = float(f0.square().sum()), float(f1.square().sum()) + for got, ref, label in ((s0, rs0, 'score0'), (s1, rs1, 'score1'), (a0, ra0, 'abs0'), (a1, ra1, 'abs1')): + assert abs(got - ref) <= 1e-9 * max(abs(ref), 1e-12), f'{mode_name} {label} shape={shape}: got {got!r} ref {ref!r}' + frozen = torch.randn(300, 200, device=DEVICE, dtype=torch.float32) * 1e-2 # fp32 input aliases through to(); abs must stay out-of-place + pristine = frozen.clone() + with select_mode('klora'): + lora_stack.score_pair(frozen, frozen, 4, 4) + assert torch.equal(frozen, pristine), 'score_pair must not mutate a caller-owned fp32 delta' + return True + + +def test_select_replay_from_cache_skips_calc(): + import tempfile + layer = build_layer('uint4') + torch.manual_seed(73) + Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + with tempfile.TemporaryDirectory() as tmp: + with host_rank(32), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer), select_mode('klora'): + n1 = make_dense_net('selk1', layer, Dd1) + n2 = make_dense_net('selk2', layer, Dd2) + for net in (n1, n2): + lora_file = os.path.join(tmp, f'{net.name}.safetensors') + with open(lora_file, 'wb') as f: + f.write(b'0' * 64) + net.network_on_disk.filename = lora_file + from modules.modeldata import model_data + model_data.sd_model.sd_checkpoint_info = MockCheckpointInfo('test/cache-model') + Wdq0 = dq(layer) + with counting_calc() as calls: + activate(n1, n2) + assert calls['n'] > 0, 'a fresh select apply must assemble both deltas' + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor' + fresh_scores, fresh_abs = entry['scores'], entry['abs_sums'] + first = dq(layer) + activate() + calls['n'] = 0 + real_svd = torch.svd_lowrank + torch.svd_lowrank = raise_no_svd + try: + activate(n1, n2) + finally: + torch.svd_lowrank = real_svd + assert calls['n'] == 0, f'a select replay must not assemble deltas: calc_updown ran {calls["n"]} times' + entry = lora_stack.state['entries'].get('lora_transformer_test') + assert entry is not None and entry['kind'] == 'factor', 'the replay must re-register the selection' + assert entry['scores'] == fresh_scores and entry['abs_sums'] == fresh_abs, 'cached scores must replay exactly' + assert torch.equal(dq(layer), first), 'select replay must be bit-identical to the fresh apply' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_select_weight_replay_from_cache_skips_calc(): + import tempfile + 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.02) + lin.network_layer_name = 'lora_transformer_plainsel' + lin.network_current_names = () + torch.manual_seed(75) + Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 + W0 = lin.weight.detach().float().clone() + with tempfile.TemporaryDirectory() as tmp: + with host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=lin), select_mode('klora'): + n1 = make_dense_net('selw1', lin, Dd1) + n2 = make_dense_net('selw2', lin, Dd2) + for net in (n1, n2): + lora_file = os.path.join(tmp, f'{net.name}.safetensors') + with open(lora_file, 'wb') as f: + f.write(b'0' * 64) + net.network_on_disk.filename = lora_file + from modules.modeldata import model_data + model_data.sd_model.sd_checkpoint_info = MockCheckpointInfo('test/cache-model') + with counting_calc() as calls: + activate(n1, n2) + assert calls['n'] > 0, 'a fresh weight-kind select apply must assemble the pair' + entry = lora_stack.state['entries'].get('lora_transformer_plainsel') + assert entry is not None and entry['kind'] == 'weight' + fresh_scores = entry['scores'] + assert torch.equal(lin.weight.detach().float(), W0), 'weights stay pristine until the schedule applies a winner' + lora_stack.reset(20) + fresh_selected = lin.weight.detach().float().clone() + activate() + assert torch.equal(lin.weight.detach().float(), W0) + calls['n'] = 0 + activate(n1, n2) + assert calls['n'] == 0, f'a weight-kind select replay must not assemble the pair: calc_updown ran {calls["n"]} times' + entry = lora_stack.state['entries'].get('lora_transformer_plainsel') + assert entry is not None and entry['kind'] == 'weight', 'the replay must register from the score record' + assert entry['scores'] == fresh_scores, 'cached scores must replay exactly' + lora_stack.reset(20) # the winner recompute at schedule time still assembles its own delta, by design + assert torch.equal(lin.weight.detach().float(), fresh_selected), 'the replayed schedule must select the same winner' + activate() + assert torch.equal(lin.weight.detach().float(), W0) + return True + + CAT_COMPILE = category('compile') @@ -2194,7 +2310,8 @@ def run_tests(): test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled, 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, - test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step]: + test_select_gamma_tracks_live_entries, test_select_host_disabled_falls_back_to_sum, test_flip_lands_before_crossover_step, + test_score_pair_chunked_precision, test_select_replay_from_cache_skips_calc, test_select_weight_replay_from_cache_skips_calc]: run_test(CAT_SELECT, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: