From 2306867b63e38b2b66184cdcebf1606c36f02cb0 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 27 Jul 2026 21:49:49 +0100 Subject: [PATCH] perf(lora): store hosted factors at the delta's effective rank Hosted truncation kept the full rank cap even when the spectrum ends in numerical zeros, padding exactly low-rank deltas (low-rank LyCORIS, full-family diffs) up to the cap. Slice the kept factors where cumulative capture reaches 1 - 1e-6 of the sketch total, and trim trailing all-zero columns when attaching cache entries written before the slice, so they collapse the same way without a format bump. The hosted log line reports the realized rank spread when it sits below the cap. - flat spectra keep the cap; a rank-8 delta under cap 256 stores 8 ranks - select segments follow the effective rank - suite pins the collapse, the flat-spectrum guard and the padded-entry trim --- modules/lora/lora_factor_cache.py | 34 +++++++------- modules/lora/lora_sdnq.py | 37 +++++++++++++++- test/test-sdnq-lora-factors.py | 74 ++++++++++++++++++++++++++++++- 3 files changed, 124 insertions(+), 21 deletions(-) diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index cda0dca39..2b00bbcc7 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -156,27 +156,27 @@ def fetch(network_layer_name): def store(network_layer_name, up, down, energy, calibrated, rms): - """Quantize-before-use: returns the pair the caller must apply. + """Quantize-before-use: returns the dequantized round-trip the caller must apply. - With caching inactive the inputs pass through untouched. Otherwise the - factors are stored as rowwise int8 and the dequantized round-trip comes - back, so the factors applied now and the factors a later hit replays are - the same tensors. ``rms`` is the assembled delta's rms, kept so replays - can evaluate the requantize routing rule without assembling the delta. + The factors quantize to rowwise int8 whether or not a cache entry can be + written, so the factors applied now, the factors a later hit replays, and a + cache-off apply are the same tensors (the round-trip also zeroes null-tail + columns the attach-side trim relies on). ``rms`` is the assembled delta's + rms, kept so replays can evaluate the requantize routing rule without + assembling the delta. """ - if state['sig'] is None: - return up, down up_q, up_s = quantize_rowwise(up) down_q, down_s = quantize_rowwise(down) - st = state['store'] - st[f'{network_layer_name}.up_q'] = up_q.to('cpu').contiguous() - st[f'{network_layer_name}.up_s'] = up_s.to('cpu').contiguous() - st[f'{network_layer_name}.down_q'] = down_q.to('cpu').contiguous() - st[f'{network_layer_name}.down_s'] = down_s.to('cpu').contiguous() - st[f'{network_layer_name}.energy'] = torch.tensor(float(energy)) - st[f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8) - st[f'{network_layer_name}.rms'] = torch.tensor(float(rms)) - state['dirty'] = True + if state['sig'] is not None: + st = state['store'] + st[f'{network_layer_name}.up_q'] = up_q.to('cpu').contiguous() + st[f'{network_layer_name}.up_s'] = up_s.to('cpu').contiguous() + st[f'{network_layer_name}.down_q'] = down_q.to('cpu').contiguous() + st[f'{network_layer_name}.down_s'] = down_s.to('cpu').contiguous() + st[f'{network_layer_name}.energy'] = torch.tensor(float(energy)) + st[f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8) + st[f'{network_layer_name}.rms'] = torch.tensor(float(rms)) + state['dirty'] = True return dequantize_rowwise(up_q, up_s).to(up.dtype), dequantize_rowwise(down_q, down_s).to(down.dtype) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 6aead3b12..d95849527 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -25,7 +25,9 @@ Only additive low-rank modules ride the channel exactly (plain LoRA: no DoRA, no CP ``mid``, no LyCORIS dense-bias, no ``diff_b``). On sub-8-bit formats, sets with non-factorable contributions are hosted instead: the families' own ``calc_updown`` delta is truncated to its top singular -directions and appended the same way. Truncation keeps the dominant part +directions and appended the same way, stored at the delta's effective +rank when the spectrum ends in a numerically null tail (dense-combined +plain pairs, low-rank LyCORIS). Truncation keeps the dominant part of the effect and drops an orthogonal residual, where requantize keeps only the grid extrema and adds grid-shift noise of the delta's own magnitude. When activation statistics for the checkpoint exist (see @@ -52,10 +54,12 @@ from modules.logger import log fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] +hosted_ranks: list[int] = [] routed_layers: list[str] = [] REQUANT_RATIO = 0.30 # delta rms over mean grid step above which requantize can retain the delta REQUANT_ENERGY = 0.90 # sketch capture below which truncation genuinely loses part of it +NULL_TAIL_EPS = 1e-6 # spectrum tail below this fraction of the capture is numerically null; dropping it keeps stored rank at the delta's effective rank def rank_bucket(r): """Fixed rank ladder for compiled-graph reuse: powers of two up to 256, multiples of 64 above (hosted rank plus exact members).""" @@ -84,6 +88,16 @@ def signature(): return '' if enabled() else '|quant=requantize' +def trim_null_tail(up_h, down_h): + """Cache entries stored before tail slicing carry null ranks as exact zero columns; trim to the effective rank on attach.""" + nz = (up_h != 0).any(dim=0) + if not bool(nz.all()): + k = max(1, int(nz.nonzero().max().item()) + 1) if bool(nz.any()) else 1 + if k < up_h.shape[1]: + return up_h[:, :k].contiguous(), down_h[:k].contiguous() + return up_h, down_h + + def get_module_factors(module, device, dtype, original_shape=None): """Return ``(up_eff, down)`` reproducing ``calc_updown`` exactly, or None. @@ -259,6 +273,7 @@ def apply_cached(self, network_layer_name, wanted_names): if entry is None: return None up_h, down_h, energy, calibrated, rms = entry + up_h, down_h = trim_null_tail(up_h, down_h) deq = self.sdnq_dequantizer dtype = deq.result_dtype remove_factors(self) # before the rule: the svd-channel check must see the checkpoint's own state, and a declined layer must fall through pristine @@ -283,6 +298,7 @@ def apply_cached(self, network_layer_name, wanted_names): lora_factor_cache.note_hit() append_factors(self, ups + [up_h.to(device=devices.device, dtype=dtype)], downs + [down_h.to(device=devices.device, dtype=dtype)]) hosted_layers.append((network_layer_name, energy, calibrated)) + hosted_ranks.append(int(up_h.shape[1])) return True @@ -349,8 +365,10 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if maybe_requant and energy < REQUANT_ENERGY: routed_layers.append(network_layer_name) return None + up_h, down_h = trim_null_tail(up_h, down_h) append_factors(self, ups + [up_h.to(device=devices.device, dtype=dtype)], downs + [down_h.to(device=devices.device, dtype=dtype)]) hosted_layers.append((network_layer_name, energy, calibrated)) + hosted_ranks.append(int(up_h.shape[1])) return True up_h, down_h, energy, calibrated = truncate_delta(self, D, dtype) @@ -358,8 +376,10 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if maybe_requant and energy < REQUANT_ENERGY: routed_layers.append(network_layer_name) # the stored entry memoizes the routing; replays skip the sketch return None + up_h, down_h = trim_null_tail(up_h, down_h) # the int8 roundtrip zeroes the numeric tail the eps slice keeps; fresh and replayed attaches must trim alike append_factors(self, ups + [up_h], downs + [down_h]) hosted_layers.append((network_layer_name, energy, calibrated)) + hosted_ranks.append(int(up_h.shape[1])) return True @@ -388,6 +408,14 @@ def truncate_delta(self, D, dtype): # oversampled sketch with extra power iterations lands within noise of exact svd; only the top q columns are kept U, S, V = torch.svd_lowrank(D, q=min(q + 64, *D.shape), niter=8) U, S, V = U[:, :q], S[:q], V[:, :q] + e = S.square() + total_e = e.sum() + if float(total_e) > 0: + # an exactly low-rank delta (dense-combined plain pairs, low-rank LyCORIS) fills the tail with + # numerical zeros; storing them would pad the channel to the cap for nothing + k = int((torch.cumsum(e, 0) < (1.0 - NULL_TAIL_EPS) * total_e).sum().item()) + 1 + if k < q: + U, S, V = U[:, :k], S[:k], V[:, :k] energy = float(S.square().sum() / D.square().sum().clamp(min=1e-30)) # captured fraction, in the weighted domain when calibrated up_h = (U * S).to(dtype=dtype) down_h = V.t() @@ -413,10 +441,15 @@ def report_fallbacks(): energies = sorted(e for _name, e, _c in hosted_layers) median = energies[len(energies) // 2] calibrated = sum(1 for _name, _e, c in hosted_layers if c) - log.info(f'Network load: type=LoRA quant=sdnq hosted={len(hosted_layers)} rank={int(shared.opts.lora_sdnq_host_rank)}{f" calib={calibrated}" if calibrated else ""} energy={median:.2f} min={energies[0]:.2f} non-factorable networks hosted on the svd side-channel') + ranks = '' + if len(hosted_ranks) > 0 and min(hosted_ranks) < int(shared.opts.lora_sdnq_host_rank): + rs = sorted(hosted_ranks) + ranks = f' k={rs[0]}-{rs[len(rs) // 2]}-{rs[-1]}' # realized rank spread; shown only when a spectrum collapsed below the cap + log.info(f'Network load: type=LoRA quant=sdnq hosted={len(hosted_layers)} rank={int(shared.opts.lora_sdnq_host_rank)}{ranks}{f" calib={calibrated}" if calibrated else ""} energy={median:.2f} min={energies[0]:.2f} non-factorable networks hosted on the svd side-channel') if l.debug: log.debug(f'Network load: type=LoRA quant=sdnq hosted={[(n, round(e, 3)) for n, e, _c in hosted_layers[:8]]}{"..." if len(hosted_layers) > 8 else ""}') hosted_layers.clear() + hosted_ranks.clear() if len(routed_layers) > 0: log.info(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(routed_layers)} routed=fat-delta') if l.debug: diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index d350b633c..3b745ad11 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -706,6 +706,34 @@ def test_hosted_dense_delta_beats_requant(): return True +def test_hosted_null_tail_collapses_to_effective_rank(): + layer = build_layer('uint4') + _A, _B, D = make_delta(sigma=3e-3) # exact rank-8 content in a non-factorable container + net = make_dense_net('nulltail', layer, D) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert layer.svd_up.shape[1] == 8, f'rank-8 delta under cap 256 must store 8 ranks, got {layer.svd_up.shape[1]}' + assert layer.svd_down.shape[0] == 8, f'down factor must slice with the up factor, got {layer.svd_down.shape[0]}' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.95, f'collapsing the null tail must not cost fidelity: rho={rho:.4f}' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + return True + + +def test_hosted_flat_spectrum_keeps_cap(): + layer = build_layer('uint4') + torch.manual_seed(13) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 3e-4 # full-rank gaussian: no null tail inside the cap + net = make_dense_net('flattail', layer, D) + with host_rank(64), mock_model(lin=layer): + activate(net) + assert layer.svd_up.shape[1] == 64, f'a flat spectrum must keep the full cap, got {layer.svd_up.shape[1]}' + activate() + return True + + def test_hosted_skips_int8(): layer = build_layer('int8') _A, _B, D = make_delta(sigma=3e-3) @@ -1215,6 +1243,47 @@ def test_factor_cache_invalidates_on_multiplier(): return True +def test_attach_trims_stored_null_tail(): + """Entries written before tail slicing pad the channel with null ranks: zero up + columns (and junk down rows behind them). Attach must trim to the effective rank + and replay the same resident tensors and weights as the unpadded entry.""" + import tempfile + layer = build_layer('uint4') + with tempfile.TemporaryDirectory() as tmp: + with host_rank(64), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer): + net, _D = cache_fixture(tmp, layer, name='padnet') + activate(net) + up0 = layer.svd_up.detach().clone() + down0 = layer.svd_down.detach().clone() + Wl0 = dq(layer) + activate() + cache_dir = os.path.join(tmp, 'cache') + entry = os.path.join(cache_dir, os.listdir(cache_dir)[0]) + from safetensors import safe_open + from safetensors.torch import save_file + with safe_open(entry, framework='pt', device='cpu') as f: + meta = dict(f.metadata()) + tensors = {k: f.get_tensor(k) for k in f.keys()} + for k in [k for k in tensors if k.endswith('.up_q')]: + base = k[: -len('.up_q')] + torch.manual_seed(5) + tensors[f'{base}.up_q'] = torch.cat([tensors[k], torch.zeros(tensors[k].shape[0], 64, dtype=torch.int8)], dim=1) + tensors[f'{base}.down_q'] = torch.cat([tensors[f'{base}.down_q'], torch.randint(-127, 128, (64, IN_F), dtype=torch.int8)], dim=0) + tensors[f'{base}.down_s'] = torch.cat([tensors[f'{base}.down_s'], torch.ones(64, 1)], dim=0) + save_file(tensors, entry, metadata=meta) + real_svd = torch.svd_lowrank + torch.svd_lowrank = raise_no_svd + try: + activate(net) + finally: + torch.svd_lowrank = real_svd + assert layer.svd_up.shape[1] == 64, f'attach must trim the padded tail back to the effective rank, got {layer.svd_up.shape[1]}' + assert torch.equal(layer.svd_up, up0) and torch.equal(layer.svd_down, down0), 'trimmed factors must match the unpadded entry' + assert torch.equal(dq(layer), Wl0), 'trimmed attach must materialize the same weight' + activate() + return True + + def test_factor_cache_int8_quantization(): from modules.lora import lora_factor_cache as fc torch.manual_seed(71) @@ -1500,7 +1569,7 @@ def run_tests(): test_hosted_disabled_by_option, test_hosted_transitions_and_rng_isolation, test_route_fat_dense_delta_requantizes, test_route_rule_terms_gate_both_ways, test_route_low_rank_fat_delta_stays_hosted, test_route_mixed_set_keeps_hosting, test_route_svd_checkpoint_keeps_hosting, - test_route_replay_from_cache]: + test_route_replay_from_cache, test_hosted_null_tail_collapses_to_effective_rank, test_hosted_flat_spectrum_keeps_cap]: run_test(CAT_HOST, fn) log.warning('=== Calibration ===') for fn in [test_calibrated_hosting_beats_plain, test_calibrated_low_rank_delta_survives, test_calib_option_off_matches_plain, @@ -1510,7 +1579,8 @@ def run_tests(): log.warning('=== Factor cache ===') for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, test_factor_cache_int8_quantization, test_factor_cache_disabled_at_zero, test_factor_cache_invalidates_on_calib_toggle, - test_cache_fastpath_skips_calc, test_cache_fastpath_serves_mixed_set]: + test_cache_fastpath_skips_calc, test_cache_fastpath_serves_mixed_set, + test_attach_trims_stored_null_tail]: run_test(CAT_FCACHE, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: