From a2daf9027c89a6c189bacb563a6598d0dedabc50 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 00:03:03 +0100 Subject: [PATCH 1/8] feat(lora): disk cache for hosted svd factors Hosting a non-factorable set costs one truncated svd per layer on every fresh apply. The factors are deterministic in the checkpoint, the loaded set, the host rank and the calibration statistics, so they persist under data/lora-factor-cache keyed by that identity and replay bit-identically on later applies. lora_sdnq_host_cache sets the disk budget in GB (0 disables), least-recently-used entries are evicted past it. With the cost paid once per configuration, svd subspace iterations rise from 2 to 4: about two thirds of the captured-energy gap to an exact decomposition for +45% one-time compute; an exact svd measures 200-400x slower at these shapes and is not viable. - modules/lora/lora_factor_cache.py: signature, store, flush, eviction - lora_sdnq.apply_hosted: fetch before computing, store after, hits in the load summary - cli/lora-quant-fidelity.py: matching niter for the hosted mirror - test/test-sdnq-lora-factors.py: factor-cache category, 3 tests --- cli/lora-quant-fidelity.py | 2 +- modules/lora/lora_factor_cache.py | 167 ++++++++++++++++++++++++++++++ modules/lora/lora_sdnq.py | 34 ++++-- modules/ui_definitions.py | 1 + test/test-sdnq-lora-factors.py | 102 ++++++++++++++++++ ui/locale/locale_en.json | 3 +- 6 files changed, 299 insertions(+), 10 deletions(-) create mode 100644 modules/lora/lora_factor_cache.py diff --git a/cli/lora-quant-fidelity.py b/cli/lora-quant-fidelity.py index 74bd214cd..df2c64126 100644 --- a/cli/lora-quant-fidelity.py +++ b/cli/lora-quant-fidelity.py @@ -315,7 +315,7 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None): Dw = D * rms if rms is not None else D with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []): torch.manual_seed(0) - U, S, V = torch.svd_lowrank(Dw, q=q, niter=2) + U, S, V = torch.svd_lowrank(Dw, q=q, niter=4) Dk = (U * S) @ V.t() if rms is not None: Dk = Dk / rms diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py new file mode 100644 index 000000000..220877c04 --- /dev/null +++ b/modules/lora/lora_factor_cache.py @@ -0,0 +1,167 @@ +"""Disk cache for hosted svd factors. + +Hosting a non-factorable adapter set costs one truncated svd per targeted +layer (tens of ms each, seconds per file) every time the set is applied +fresh. The resulting factors are deterministic in the checkpoint, the loaded +set (files, multipliers, dyn_dim), the host rank and the calibration +statistics, so they are cached on disk keyed by exactly that identity and +replayed bit-identically on the next apply of the same configuration. + +One safetensors file per configuration under ``models/lora-factor-cache``, +holding every hosted layer's post-rotation factor pair. Files are named by +the model and network set with an identity-hash suffix, and the exact +signature is embedded in the file metadata. The +``lora_sdnq_host_cache`` option is the size budget in GB (0 disables); +least-recently-used entries are evicted past the budget. Any doubt about +identity (unknown checkpoint, unreadable lora file, signature mismatch) +disables caching for the pass rather than risking a stale hit. +""" + +import os +import json +import hashlib + +import torch + +from modules import paths, shared +from modules.lora import lora_common as l +from modules.logger import log + + +cache_root = os.path.join(paths.models_path, 'lora-factor-cache') +state = {'wn': None, 'sig': None, 'path': None, 'store': {}, 'dirty': False, 'hits': 0, 'misses': 0} + + +def budget_gb(): + try: + return float(getattr(shared.opts, 'lora_sdnq_host_cache', 0) or 0) + except Exception: + return 0.0 + + +def signature(wanted_names): + """Content identity of a hosted-apply configuration, or None when caching is unsafe.""" + from modules.lora import lora_calib + model_name = lora_calib.checkpoint_name(getattr(shared, 'sd_model', None)) + if model_name is None: + return None + calib_path = lora_calib.calib_file(model_name) + parts = { + 'model': model_name, + 'rank': int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0), + 'calib': int(os.path.getmtime(calib_path)) if os.path.isfile(calib_path) else None, + 'nets': [], + } + for name, te, unet, dyn in wanted_names: + net = next((n for n in l.loaded_networks if n.name == name), None) + filename = getattr(getattr(net, 'network_on_disk', None), 'filename', None) + try: + st = os.stat(filename) + except Exception: + return None + parts['nets'].append([name, repr(te), repr(unet), repr(dyn), filename, int(st.st_mtime), st.st_size]) + return parts + + +def label(parts): + """Filename prefix from the model and net names, so the cache folder reads without tooling.""" + names = [parts['model'].replace('\\', '/').split('/')[-1]] + [n[0] for n in parts['nets']] + text = '-'.join(names) + text = ''.join(c if c.isalnum() or c in '._-' else '-' for c in text) + return text[:96] + + +def begin_pass(wanted_names): + """Bind the pass to its cache entry; identity-memoized on the wanted_names tuple.""" + if wanted_names is state['wn']: + return + state['wn'] = wanted_names + state.update(sig=None, path=None, dirty=False) + state['store'] = {} + if budget_gb() <= 0 or wanted_names == (): + return + parts = signature(wanted_names) + if parts is None: + return + sig = json.dumps(parts, sort_keys=True) + key = hashlib.sha256(sig.encode()).hexdigest()[:24] + path = os.path.join(cache_root, f'{label(parts)}-{key}.safetensors') + store = {} + if os.path.isfile(path): + try: + from safetensors import safe_open + with safe_open(path, framework='pt', device='cpu') as f: + if (f.metadata() or {}).get('sig') == sig: + for k in f.keys(): + store[k] = f.get_tensor(k) + os.utime(path, None) # freshness for LRU eviction + except Exception as e: + log.debug(f'Network cache: read failed path="{path}" {e}') + store = {} + state.update(sig=sig, path=path) + state['store'] = store + log.debug(f'Network cache: entry="{path}" keys={len(store)}') + + +def fetch(network_layer_name): + """Cached (up, down, energy, calibrated) for a layer, or None.""" + if state['sig'] is None: + return None + up = state['store'].get(f'{network_layer_name}.up') + down = state['store'].get(f'{network_layer_name}.down') + energy = state['store'].get(f'{network_layer_name}.energy') + calib = state['store'].get(f'{network_layer_name}.calib') + if up is None or down is None or energy is None or calib is None: + state['misses'] += 1 + return None + state['hits'] += 1 + return up, down, float(energy), bool(calib) + + +def put(network_layer_name, up, down, energy, calibrated): + if state['sig'] is None: + return + state['store'][f'{network_layer_name}.up'] = up.detach().to('cpu').contiguous() + state['store'][f'{network_layer_name}.down'] = down.detach().to('cpu').contiguous() + state['store'][f'{network_layer_name}.energy'] = torch.tensor(float(energy)) + state['store'][f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8) + state['dirty'] = True + + +def evict(): + budget = budget_gb() * 2**30 + try: + files = [os.path.join(cache_root, f) for f in os.listdir(cache_root) if f.endswith('.safetensors')] + sizes = {p: os.path.getsize(p) for p in files} + except Exception: + return + total = sum(sizes.values()) + for p in sorted(files, key=os.path.getmtime): + if total <= budget: + break + if p == state['path']: + continue # never evict the entry of the live pass + try: + os.remove(p) + total -= sizes[p] + except Exception: + pass + + +def flush(): + """Persist a dirty pass store; returns (hits, misses) since the last flush.""" + hits, misses = state['hits'], state['misses'] + state['hits'] = state['misses'] = 0 + if not state['dirty'] or state['path'] is None: + return hits, misses + state['dirty'] = False + try: + from safetensors.torch import save_file + os.makedirs(cache_root, exist_ok=True) + tmp = state['path'] + '.tmp' + save_file(state['store'], tmp, metadata={'sig': state['sig']}) + os.replace(tmp, state['path']) + evict() + except Exception as e: + log.warning(f'Network cache: write failed path="{state["path"]}" {e}') + return hits, misses diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 63597111b..51f91044b 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -38,7 +38,7 @@ remains. import torch from modules import devices, shared -from modules.lora import lora_calib +from modules.lora import lora_calib, lora_factor_cache from modules.lora import lora_common as l from modules.logger import log @@ -218,9 +218,11 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa appended exactly so they never compete with the hosted remainder for rank. When per-checkpoint activation statistics exist (``lora_calib``), input channels are weighted by their RMS before truncation so the kept - directions minimize output error rather than weight error. Returns None - when the delta cannot ride the channel (wrong shape); the caller falls - back to requantize. + directions minimize output error rather than weight error. Computed + factors are disk-cached per configuration (``lora_factor_cache``) and + replayed bit-identically on later applies. Returns None when the delta + cannot ride the channel (wrong shape); the caller falls back to + requantize. """ from sdnq.quant_utils import rotate_hadamard @@ -231,7 +233,11 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa if updown is None or updown.ndim != 2 or tuple(updown.shape) != tuple(deq.original_shape): return None dtype = deq.result_dtype - D = updown.detach().to(devices.device, torch.float32) + cached = None + if not use_previous: + lora_factor_cache.begin_pass(wanted_names) + cached = lora_factor_cache.fetch(network_layer_name) + D = None if cached is not None else updown.detach().to(devices.device, torch.float32) ups, downs = [], [] loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks @@ -243,12 +249,19 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa if factors is None: continue up_eff, down = factors - D = D.sub_(up_eff.to(torch.float32) @ down.to(torch.float32)) # factorable members ride exactly; host only the remainder + if D is not None: + D = D.sub_(up_eff.to(torch.float32) @ down.to(torch.float32)) # factorable members ride exactly; host only the remainder 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 cached is not None: + up_h, down_h, energy, calibrated = cached + 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)) + return True + cap = int(shared.opts.lora_sdnq_host_rank) q = min(cap, *D.shape) rms = lora_calib.rms_for(self) @@ -261,7 +274,7 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa # svd_lowrank draws random projections; fork so user generation seeds are untouched and re-applies are deterministic with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []): torch.manual_seed(0) - U, S, V = torch.svd_lowrank(D, q=q, niter=2) + U, S, V = torch.svd_lowrank(D, q=q, niter=4) 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() @@ -269,7 +282,9 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa down_h = down_h / rms # unscale in the original input basis, before any rotation if deq.use_hadamard: down_h = rotate_hadamard(down_h, group_size=deq.hadamard_group_size) - append_factors(self, ups + [up_h], downs + [down_h.to(dtype=dtype)]) + down_h = down_h.to(dtype=dtype) + lora_factor_cache.put(network_layer_name, up_h, down_h, energy, rms is not None) + append_factors(self, ups + [up_h], downs + [down_h]) hosted_layers.append((network_layer_name, energy, rms is not None)) return True @@ -281,6 +296,9 @@ def note_fallback(self, network_layer_name): def report_fallbacks(): + hits, misses = lora_factor_cache.flush() + if hits > 0 or misses > 0: + log.info(f'Network load: type=LoRA quant=sdnq cache hits={hits} misses={misses}') if len(hosted_layers) > 0: energies = sorted(e for _name, e, _c in hosted_layers) median = energies[len(energies) // 2] diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index d59970138..a873ea68d 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -707,6 +707,7 @@ def create_settings(cmd_opts): "lora_sdnq_apply": OptionInfo("exact", "LoRA quantized apply method", gr.Radio, {"choices": ["exact", "requantize"]}), "lora_sdnq_host_rank": OptionInfo(256, "LoRA quantized host rank", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 32}), "lora_sdnq_host_calib": OptionInfo(True, "LoRA quantized host calibration"), + "lora_sdnq_host_cache": OptionInfo(10, "LoRA quantized host cache", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), "lora_meta_sep": OptionInfo("

Metadata

", "", gr.HTML), "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info to metadata"), diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 1c0ccb4f6..95c6ddb35 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -34,6 +34,9 @@ for the per-model analyzer): plain truncation bit-exact, and the capture hooks accumulate, persist and reload statistics correctly, gated by option, format width and model compile. +- Factor cache: hosted factors replay bit-identically from the disk cache + without re-running the svd, a configuration change (multiplier) misses + and writes a separate entry, and budget 0 writes nothing. All tensors are synthetic; no model files or running server required. @@ -897,6 +900,102 @@ def test_calib_capture_gates(): return True +CAT_FCACHE = category('factor-cache') + + +@contextmanager +def host_cache(gb, root): + from modules.lora import lora_factor_cache + old_gb = getattr(shared.opts, 'lora_sdnq_host_cache', 0) + old_root = lora_factor_cache.cache_root + shared.opts.lora_sdnq_host_cache = gb + lora_factor_cache.cache_root = root + lora_factor_cache.state.update(wn=None, sig=None, path=None, dirty=False, hits=0, misses=0) + lora_factor_cache.state['store'] = {} + try: + yield lora_factor_cache + finally: + shared.opts.lora_sdnq_host_cache = old_gb + lora_factor_cache.cache_root = old_root + lora_factor_cache.state.update(wn=None, sig=None, path=None, dirty=False, hits=0, misses=0) + lora_factor_cache.state['store'] = {} + + +def cache_fixture(tmp, layer, name='cachenet', sigma=3e-4, seed=61): + """Dense net whose on-disk file exists (signature needs a stat-able path) plus a mock checkpoint identity.""" + torch.manual_seed(seed) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * sigma + net = make_dense_net(name, layer, D) + lora_file = os.path.join(tmp, f'{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') + return net, D + + +def raise_no_svd(*_args, **_kwargs): + raise AssertionError('svd must not run on a cache hit') + + +def test_factor_cache_roundtrip_bitexact(): + 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) + Wdq0 = dq(layer) + activate(net) + first_up = layer.svd_up.detach().clone() + first_down = layer.svd_down.detach().clone() + activate() # pass end flushed the entry; unload restores the base + files = os.listdir(os.path.join(tmp, 'cache')) + assert len(files) == 1, f'one cache entry expected, got {files}' + real_svd = torch.svd_lowrank + torch.svd_lowrank = raise_no_svd + try: + activate(net) # same configuration: must replay from disk without touching the svd + finally: + torch.svd_lowrank = real_svd + assert torch.equal(layer.svd_up, first_up), 'cache hit must replay bit-identical up factors' + assert torch.equal(layer.svd_down, first_down), 'cache hit must replay bit-identical down factors' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + return True + + +def test_factor_cache_invalidates_on_multiplier(): + 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) + activate(net) + up_full = layer.svd_up.detach().clone() + activate() + net.te_multiplier = 0.7 + net.unet_multiplier = [0.7] * 3 + activate(net) # different multiplier: different signature, fresh svd, second entry + assert not torch.equal(layer.svd_up, up_full), 'multiplier change must produce different factors' + activate() + files = os.listdir(os.path.join(tmp, 'cache')) + assert len(files) == 2, f'two cache entries expected, got {files}' + return True + + +def test_factor_cache_disabled_at_zero(): + import tempfile + layer = build_layer('uint4') + with tempfile.TemporaryDirectory() as tmp: + with host_rank(64), host_cache(0, os.path.join(tmp, 'cache')), mock_model(lin=layer): + net, _D = cache_fixture(tmp, layer) + activate(net) + activate() + assert not os.path.isdir(os.path.join(tmp, 'cache')), 'budget 0 must write nothing' + return True + + CAT_ROBUST = category('robustness') @@ -970,6 +1069,9 @@ def run_tests(): for fn in [test_calibrated_hosting_beats_plain, test_calibrated_low_rank_delta_survives, test_calib_option_off_matches_plain, test_calib_capture_persist_roundtrip, test_calib_capture_gates]: run_test(CAT_CALIB, fn) + log.warning('=== Factor cache ===') + for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, test_factor_cache_disabled_at_zero]: + run_test(CAT_FCACHE, fn) log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: run_test(CAT_ROBUST, fn) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 479d44eed..135283560 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -864,9 +864,10 @@ {"id":"","label":"LoRA native apply to text encoder","localized":"","hint":"","ui":"settings_lora"}, {"id":"","label":"LoRA native fuse with model","localized":"","hint":"Merge LoRA into the model for lower memory usage.

Warning: After removing or switching a LoRA, you may still see its style in generated images. To get a clean model, reload it from the model selector.","ui":"settings_lora"}, {"id":"","label":"LoRA diffusers fuse with model","localized":"","hint":"Merge LoRA into the model for lower memory usage and torch.compile compatibility.

Warning: After removing or switching a LoRA, you may still see its style in generated images. To get a clean model, reload it from the model selector.","ui":"settings_lora"}, - {"id":"","label":"LoRA quantized apply method","localized":"","hint":"How networks are applied to SDNQ-quantized model weights:
- exact: adapters are carried alongside the quantized weights at full precision; apply and removal are exact and the quantized weights are never modified. The carried factors take additional VRAM, growing with adapter rank, size and count
- requantize: adapters are merged into the quantized weights, matching the behavior of earlier releases. Uses no additional VRAM (a weight backup for network removal is held in system RAM); on models quantized below 8 bits rounding typically loses much of the adapter effect, with strong adapters retaining more

With requantize selected, the host rank and calibration options below have no effect.

Default is exact.","ui":"settings_lora"}, + {"id":"","label":"LoRA quantized apply method","localized":"","hint":"How networks are applied to SDNQ-quantized model weights:
- exact: adapters are carried alongside the quantized weights at full precision; apply and removal are exact and the quantized weights are never modified. The carried factors take additional VRAM, growing with adapter rank, size and count
- requantize: adapters are merged into the quantized weights, matching the behavior of earlier releases. Uses no additional VRAM (a weight backup for network removal is held in system RAM); on models quantized below 8 bits rounding typically loses much of the adapter effect, with strong adapters retaining more

With requantize selected, the host rank, calibration and cache options below have no effect.

Default is exact.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host rank","localized":"","hint":"Maximum rank used to carry non-factorable adapter types (LoKR, LoHA, OFT, DoRA) on the side-channel of SDNQ models quantized below 8 bits, where merging would erase most of the adapter. Higher values keep more of the adapter at proportionally higher memory cost. Set to 0 to disable and merge into the quantized weights instead.","ui":"settings_lora"}, {"id":"","label":"LoRA quantized host calibration","localized":"","hint":"Collection of per-channel activation statistics during normal generations on models quantized below 8 bits, cached per checkpoint. When available, the statistics weight side-channel hosting of non-factorable adapter types toward the channels carrying the most activation energy, improving delivered fidelity at the same host rank.","ui":"settings_lora"}, + {"id":"","label":"LoRA quantized host cache","localized":"","hint":"Disk budget in GB for caching the side-channel factors computed when hosting non-factorable adapter types on quantized models. A cached configuration skips the truncation math on reapply; least recently used entries are evicted past the budget. Set to 0 to disable.","ui":"settings_lora"}, {"id":"","label":"LoRA auto-apply tags","localized":"","hint":"Automatically add trigger words/tags from LoRA metadata to your prompt.
Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.
Set to 0 to disable, -1 to add all available tags.","ui":"settings_lora"}, {"id":"","label":"LoRA memory cache","localized":"","hint":"How many LoRAs to keep in network for future use before requiring reloading from storage","ui":"settings_lora"}, {"id":"","label":"LoRA add hash info to metadata","localized":"","hint":"Include LoRA file hashes in generated image metadata.
Useful for reproducibility and tracking which exact LoRA versions were used.","ui":"settings_lora"}, From 03dfddaa96cc835228a12cb690e319c1ab4d7976 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 00:25:49 +0100 Subject: [PATCH 2/8] perf(lora): int8 storage for cached hosted factors Rowwise int8 with fp32 scales halves cache entries; measured in output space on real hosted deltas the roundtrip is fidelity-free (within 0.0003 of fp32 factors, bf16 storage likewise). Factors are quantized before first use and the dequantized roundtrip is what the apply attaches, so a fresh compute and a later cache hit stay bit-identical; the entry format is versioned and pre-int8 entries reload as misses. - lora_factor_cache: quantize_rowwise/dequantize_rowwise, store returns the applied pair, fmt guard on read - test/test-sdnq-lora-factors.py: int8 quantization test, entry-size assertion in the roundtrip test --- modules/lora/lora_factor_cache.py | 85 +++++++++++++++++++++---------- modules/lora/lora_sdnq.py | 2 +- test/test-sdnq-lora-factors.py | 20 +++++++- 3 files changed, 78 insertions(+), 29 deletions(-) diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index 220877c04..a5ef4041a 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -8,13 +8,17 @@ statistics, so they are cached on disk keyed by exactly that identity and replayed bit-identically on the next apply of the same configuration. One safetensors file per configuration under ``models/lora-factor-cache``, -holding every hosted layer's post-rotation factor pair. Files are named by -the model and network set with an identity-hash suffix, and the exact -signature is embedded in the file metadata. The -``lora_sdnq_host_cache`` option is the size budget in GB (0 disables); -least-recently-used entries are evicted past the budget. Any doubt about -identity (unknown checkpoint, unreadable lora file, signature mismatch) -disables caching for the pass rather than risking a stale hit. +holding every hosted layer's post-rotation factor pair as rowwise int8 +with fp32 scales (measured fidelity-free in output space, half the bytes +of bf16). Files are named by the model and network set with an +identity-hash suffix, and the exact signature is embedded in the file +metadata. Factors are quantized before first use: ``store`` returns the +dequantized round-trip for the caller to apply, so a fresh compute and a +later cache hit attach bit-identical tensors. The ``lora_sdnq_host_cache`` +option is the size budget in GB (0 disables); least-recently-used entries +are evicted past the budget. Any doubt about identity (unknown checkpoint, +unreadable lora file, signature mismatch) disables caching for the pass +rather than risking a stale hit. """ import os @@ -86,46 +90,73 @@ def begin_pass(wanted_names): sig = json.dumps(parts, sort_keys=True) key = hashlib.sha256(sig.encode()).hexdigest()[:24] path = os.path.join(cache_root, f'{label(parts)}-{key}.safetensors') - store = {} + entries = {} if os.path.isfile(path): try: from safetensors import safe_open with safe_open(path, framework='pt', device='cpu') as f: - if (f.metadata() or {}).get('sig') == sig: + meta = f.metadata() or {} + if meta.get('sig') == sig and meta.get('fmt') == '2': for k in f.keys(): - store[k] = f.get_tensor(k) + entries[k] = f.get_tensor(k) os.utime(path, None) # freshness for LRU eviction except Exception as e: log.debug(f'Network cache: read failed path="{path}" {e}') - store = {} + entries = {} state.update(sig=sig, path=path) - state['store'] = store - log.debug(f'Network cache: entry="{path}" keys={len(store)}') + state['store'] = entries + log.debug(f'Network cache: entry="{path}" keys={len(entries)}') + + +def quantize_rowwise(t): + t32 = t.detach().to(torch.float32) + scale = t32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) / 127.0 + q = (t32 / scale).round().clamp(-127, 127).to(torch.int8) + return q, scale + + +def dequantize_rowwise(q, scale): + # int8 * fp32 with a single fp32 rounding: identical on any device, so hit and miss replay the same values + return q.to(torch.float32) * scale def fetch(network_layer_name): - """Cached (up, down, energy, calibrated) for a layer, or None.""" + """Cached (up, down, energy, calibrated) for a layer, or None; factors return as fp32.""" if state['sig'] is None: return None - up = state['store'].get(f'{network_layer_name}.up') - down = state['store'].get(f'{network_layer_name}.down') - energy = state['store'].get(f'{network_layer_name}.energy') - calib = state['store'].get(f'{network_layer_name}.calib') - if up is None or down is None or energy is None or calib is None: + st = state['store'] + up_q, up_s = st.get(f'{network_layer_name}.up_q'), st.get(f'{network_layer_name}.up_s') + down_q, down_s = st.get(f'{network_layer_name}.down_q'), st.get(f'{network_layer_name}.down_s') + energy = st.get(f'{network_layer_name}.energy') + calib = st.get(f'{network_layer_name}.calib') + if up_q is None or up_s is None or down_q is None or down_s is None or energy is None or calib is None: state['misses'] += 1 return None state['hits'] += 1 - return up, down, float(energy), bool(calib) + return dequantize_rowwise(up_q, up_s), dequantize_rowwise(down_q, down_s), float(energy), bool(calib) -def put(network_layer_name, up, down, energy, calibrated): +def store(network_layer_name, up, down, energy, calibrated): + """Quantize-before-use: returns the pair 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. + """ if state['sig'] is None: - return - state['store'][f'{network_layer_name}.up'] = up.detach().to('cpu').contiguous() - state['store'][f'{network_layer_name}.down'] = down.detach().to('cpu').contiguous() - state['store'][f'{network_layer_name}.energy'] = torch.tensor(float(energy)) - state['store'][f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8) + 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) state['dirty'] = True + return dequantize_rowwise(up_q, up_s).to(up.dtype), dequantize_rowwise(down_q, down_s).to(down.dtype) def evict(): @@ -159,7 +190,7 @@ def flush(): from safetensors.torch import save_file os.makedirs(cache_root, exist_ok=True) tmp = state['path'] + '.tmp' - save_file(state['store'], tmp, metadata={'sig': state['sig']}) + save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': '2'}) os.replace(tmp, state['path']) evict() except Exception as e: diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 51f91044b..8d5330e09 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -283,7 +283,7 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa if deq.use_hadamard: down_h = rotate_hadamard(down_h, group_size=deq.hadamard_group_size) down_h = down_h.to(dtype=dtype) - lora_factor_cache.put(network_layer_name, up_h, down_h, energy, rms is not None) + up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, rms is not None) append_factors(self, ups + [up_h], downs + [down_h]) hosted_layers.append((network_layer_name, energy, rms is not None)) return True diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 95c6ddb35..58028801c 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -952,6 +952,9 @@ def test_factor_cache_roundtrip_bitexact(): activate() # pass end flushed the entry; unload restores the base files = os.listdir(os.path.join(tmp, 'cache')) assert len(files) == 1, f'one cache entry expected, got {files}' + bf16_bytes = (first_up.numel() + first_down.numel()) * 2 + entry_bytes = os.path.getsize(os.path.join(tmp, 'cache', files[0])) + assert entry_bytes < bf16_bytes * 0.62 + 8192, f'int8 entry must be about half the bf16 factor bytes: {entry_bytes} vs {bf16_bytes}' real_svd = torch.svd_lowrank torch.svd_lowrank = raise_no_svd try: @@ -984,6 +987,20 @@ def test_factor_cache_invalidates_on_multiplier(): return True +def test_factor_cache_int8_quantization(): + from modules.lora import lora_factor_cache as fc + torch.manual_seed(71) + t = torch.randn(64, 128, device=DEVICE) * torch.logspace(-3, 0, 64, device=DEVICE)[:, None] # rows spanning magnitudes + q, s = fc.quantize_rowwise(t) + assert q.dtype == torch.int8 + dq = fc.dequantize_rowwise(q, s) + err = (dq - t).abs().max(dim=1).values + assert bool((err <= s.squeeze(1) * 0.51).all()), 'rowwise int8 error must stay within half a step' + cos = torch.nn.functional.cosine_similarity(dq.flatten(), t.flatten(), dim=0) + assert float(cos) > 0.99995, f'int8 roundtrip cosine {float(cos):.6f}' + return True + + def test_factor_cache_disabled_at_zero(): import tempfile layer = build_layer('uint4') @@ -1070,7 +1087,8 @@ def run_tests(): test_calib_capture_persist_roundtrip, test_calib_capture_gates]: run_test(CAT_CALIB, fn) log.warning('=== Factor cache ===') - for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, test_factor_cache_disabled_at_zero]: + 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]: run_test(CAT_FCACHE, fn) log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: From 74e68b58ce440ef3b29eb930de3e4f106d75b942 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 01:25:03 +0100 Subject: [PATCH 3/8] perf(lora): oversample the hosted-factor sketch to near-exact svd Sketch hosted deltas at rank+64 with eight power iterations and keep the top rank columns; this lands within noise of the exact decomposition at roughly twice a sketch cost the factor cache pays once per configuration. Bump the cache format so narrower-sketch entries reload as misses. --- cli/lora-quant-fidelity.py | 4 ++-- modules/lora/lora_factor_cache.py | 4 ++-- modules/lora/lora_sdnq.py | 4 +++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/cli/lora-quant-fidelity.py b/cli/lora-quant-fidelity.py index df2c64126..1dc41bfe5 100644 --- a/cli/lora-quant-fidelity.py +++ b/cli/lora-quant-fidelity.py @@ -315,8 +315,8 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None): Dw = D * rms if rms is not None else D with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []): torch.manual_seed(0) - U, S, V = torch.svd_lowrank(Dw, q=q, niter=4) - Dk = (U * S) @ V.t() + U, S, V = torch.svd_lowrank(Dw, q=min(q + 64, *D.shape), niter=8) + Dk = (U[:, :q] * S[:q]) @ V[:, :q].t() if rms is not None: Dk = Dk / rms base16 = W_dq.to(torch.bfloat16).float() diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index a5ef4041a..6e7e46a0b 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -96,7 +96,7 @@ def begin_pass(wanted_names): from safetensors import safe_open with safe_open(path, framework='pt', device='cpu') as f: meta = f.metadata() or {} - if meta.get('sig') == sig and meta.get('fmt') == '2': + if meta.get('sig') == sig and meta.get('fmt') == '3': for k in f.keys(): entries[k] = f.get_tensor(k) os.utime(path, None) # freshness for LRU eviction @@ -190,7 +190,7 @@ def flush(): from safetensors.torch import save_file os.makedirs(cache_root, exist_ok=True) tmp = state['path'] + '.tmp' - save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': '2'}) + save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': '3'}) os.replace(tmp, state['path']) evict() except Exception as e: diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 8d5330e09..b2a91de6c 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -274,7 +274,9 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa # svd_lowrank draws random projections; fork so user generation seeds are untouched and re-applies are deterministic with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []): torch.manual_seed(0) - U, S, V = torch.svd_lowrank(D, q=q, niter=4) + # 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] 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() From ff94565da879aa185274a0906ad03ddc103dc3e7 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 18 Jul 2026 01:43:26 +0100 Subject: [PATCH 4/8] perf(lora): pad side-channel factor ranks to fixed buckets The compiled dequant specializes per factor rank, so each distinct lora set shape paid a compile stall on switch. Pad appended factors to a power-of-two rank ladder (multiples of 64 past the hosted cap) with zero fill: switches inside a bucket reuse the compiled graph, and common trained ranks land on their bucket exactly so padding is usually a no-op. Regression tests pin the factor add inside the single compiled graph and the bucket reuse. --- modules/lora/lora_sdnq.py | 23 ++++++++ test/test-sdnq-lora-factors.py | 104 +++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index b2a91de6c..5635f8495 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -46,6 +46,22 @@ from modules.logger import log fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] +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).""" + if r <= 8: + return 8 + if r <= 256: + return 1 << (r - 1).bit_length() + return -(-r // 64) * 64 + + +def pad_rank(t, dim, bucket): + if t.shape[dim] >= bucket: + return t + shape = list(t.shape) + shape[dim] = bucket - t.shape[dim] + return torch.cat([t, t.new_zeros(shape)], dim=dim) + def enabled(): """True while the exact svd-channel machinery may take quantized layers; the requantize choice routes every layer to the legacy weight-rewrite path.""" @@ -187,6 +203,13 @@ def append_factors(self, ups, downs): 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() + from sdnq.common import use_torch_compile + if use_torch_compile: + # the compiled dequant specializes per factor rank; pad to a fixed bucket so set switches inside a bucket reuse the graph (zero columns contribute exactly nothing) + dim_up, dim_down = (0, 1) if deq.use_quantized_matmul else (1, 0) + bucket = rank_bucket(new_up.shape[dim_up]) + new_up = pad_rank(new_up, dim_up, bucket) + new_down = pad_rank(new_down, dim_down, bucket) 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) diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 58028801c..4c2c17ce6 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -37,6 +37,11 @@ for the per-model analyzer): - Factor cache: hosted factors replay bit-identically from the disk cache without re-running the svd, a configuration change (multiplier) misses and writes a separate entry, and budget 0 writes nothing. +- Compile: the factor add runs inside the single compiled dequant graph + (fullgraph, no breaks) and matches the eager result; factor ranks pad to + a fixed bucket ladder so set switches inside a bucket reuse the compiled + graph while a novel bucket compiles exactly once, and padding changes + the dequantized weight by nothing beyond reduction-order ulp. All tensors are synthetic; no model files or running server required. @@ -1013,6 +1018,102 @@ def test_factor_cache_disabled_at_zero(): return True +CAT_COMPILE = category('compile') + + +def dq_compiled(layer): + # the production entry: skip_compile left at its default so the shared compiled dequant runs + 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) + + +def graph_stats(): + from torch._dynamo.utils import counters + return int(counters['stats']['unique_graphs']), sum(counters['graph_break'].values()) + + +def test_factor_add_inside_compiled_graph(): + from sdnq.common import use_torch_compile + if not use_torch_compile: + return True # compile disabled at sdnq import (no triton); nothing to pin + import torch._dynamo + from torch._dynamo.utils import counters + layer = build_layer('uint4') + A, B, _D = make_delta() + dtype = layer.sdnq_dequantizer.result_dtype + torch._dynamo.reset() + counters.clear() + lora_sdnq.append_factors(layer, [B.to(dtype)], [A.to(dtype)]) + W_c = dq_compiled(layer) + graphs, breaks = graph_stats() + assert breaks == 0, f'graph breaks in the compiled dequant: {breaks}' + assert graphs == 1, f'factor-bearing dequant must be one compiled region, got {graphs} graphs' + W_e = dq(layer) + assert torch.allclose(W_c, W_e, rtol=1e-3, atol=1e-4), f'compiled vs eager dequant diverged, max {float((W_c - W_e).abs().max()):.3e}' + lora_sdnq.remove_factors(layer) + return True + + +def test_rank_bucket_graph_reuse(): + from sdnq.common import use_torch_compile + if not use_torch_compile: + return True + import torch._dynamo + from torch._dynamo.utils import counters + import sdnq.common as sdnq_common + layer = build_layer('uint4', use_hadamard=False) + dtype = layer.sdnq_dequantizer.result_dtype + torch.manual_seed(13) + mk = lambda r: (torch.randn(OUT_F, r, device=DEVICE, dtype=dtype) * 0.01, torch.randn(r, IN_F, device=DEVICE, dtype=dtype) * 0.01) + B8, A8 = mk(8) + B6, A6 = mk(6) + B24, A24 = mk(24) + + torch._dynamo.reset() + counters.clear() + lora_sdnq.append_factors(layer, [B8], [A8]) + assert layer.svd_up.shape[1] == 8, f'rank 8 must bucket to 8, got {layer.svd_up.shape[1]}' + dq_compiled(layer) + g_first, _ = graph_stats() + + lora_sdnq.remove_factors(layer) + lora_sdnq.append_factors(layer, [B6], [A6]) + assert layer.svd_up.shape[1] == 8, f'rank 6 must pad to bucket 8, got {layer.svd_up.shape[1]}' + assert float(layer.svd_up[:, 6:].abs().sum()) == 0.0, 'pad columns must be exact zeros' + dq_compiled(layer) + g_same, _ = graph_stats() + assert g_same == g_first, f'same bucket must reuse the graph: {g_first} -> {g_same}' + + lora_sdnq.remove_factors(layer) + lora_sdnq.append_factors(layer, [B24], [A24]) + assert layer.svd_up.shape[1] == 32, f'rank 24 must pad to bucket 32, got {layer.svd_up.shape[1]}' + dq_compiled(layer) + g_novel, _ = graph_stats() + assert g_novel == g_first + 1, f'novel bucket must compile exactly one new graph: {g_first} -> {g_novel}' + + lora_sdnq.remove_factors(layer) + lora_sdnq.append_factors(layer, [B8], [A8]) + dq_compiled(layer) + g_back, _ = graph_stats() + assert g_back == g_novel, f'returning to a seen bucket must be free: {g_novel} -> {g_back}' + + W_padded = dq(layer) + lora_sdnq.remove_factors(layer) + old_flag = sdnq_common.use_torch_compile + sdnq_common.use_torch_compile = False + try: + lora_sdnq.append_factors(layer, [B8], [A8]) + assert layer.svd_up.shape[1] == 8 + W_unpadded = dq(layer) + finally: + sdnq_common.use_torch_compile = old_flag + lora_sdnq.remove_factors(layer) + assert torch.allclose(W_padded, W_unpadded, rtol=0.0, atol=1e-6), f'padding must be inert beyond reduction-order ulp, max {float((W_padded - W_unpadded).abs().max()):.3e}' + return True + + CAT_ROBUST = category('robustness') @@ -1090,6 +1191,9 @@ def run_tests(): 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]: run_test(CAT_FCACHE, fn) + log.warning('=== Compile ===') + for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: + run_test(CAT_COMPILE, fn) log.warning('=== Robustness ===') for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: run_test(CAT_ROBUST, fn) From e33965fc3b348cd7f0913eceacd43d0fac59095a Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 01:15:23 +0100 Subject: [PATCH 5/8] fix(lora): bound calibration capture and cover unet denoisers Persist now fires either when every hooked layer reaches the token quota or at a fixed denoiser-forward deadline; layers under a small token floor are omitted and stay on plain truncation. The all-done barrier alone never fired on models whose modulation or pooled projections see a few tokens per forward, so hooks stayed registered forever and statistics recollected every session. eligible_modules walks the transformer or the unet, so unet checkpoints collect statistics at all. - lora_sdnq: drop the unused use_previous parameter; the factor cache store now always runs inside a begin_pass - tests: deadline persist, token-floor omission, unet root walk --- modules/lora/lora_calib.py | 66 +++++++++++++----- modules/lora/lora_sdnq.py | 28 ++++---- test/test-sdnq-lora-factors.py | 118 ++++++++++++++++++++++++++++++++- 3 files changed, 179 insertions(+), 33 deletions(-) diff --git a/modules/lora/lora_calib.py b/modules/lora/lora_calib.py index 032148944..622fa50e9 100644 --- a/modules/lora/lora_calib.py +++ b/modules/lora/lora_calib.py @@ -13,11 +13,15 @@ MLP down projections whose inputs carry the largest outlier channels. Statistics come from the model's own forwards: when a sub-8-bit SDNQ model loads and no calibration is cached for it, streaming sum-of-squares hooks attach to its quantized linears, accumulate during normal generations, -persist once enough tokens are seen, and go inert. Cached statistics load -at model load and sit on each layer as ``sdnq_calib_rms``; the hosting path -reads them through ``rms_for``. Capture is skipped when the model is -compiled (hooks would break the graph) and everything is gated by the -``lora_sdnq_host_calib`` option. +persist, and go inert. Persist fires when every layer reaches the token +quota, or at a bounded number of denoiser forwards for models where some +projections take pooled or modulation vectors (a few tokens per forward) +and could never reach an absolute quota; layers still under a small token +floor at the deadline are omitted and stay on plain truncation. Cached +statistics load at model load and sit on each layer as ``sdnq_calib_rms``; +the hosting path reads them through ``rms_for``. Capture is skipped when +the model is compiled (hooks would break the graph) and everything is +gated by the ``lora_sdnq_host_calib`` option. """ import os @@ -40,12 +44,15 @@ class CaptureState(TypedDict): model: Optional[str] recs: dict[str, CaptureRecord] handles: list[torch.utils.hooks.RemovableHandle] + forwards: int complete: bool TOKENS_DONE = 65536 +FORWARDS_DEADLINE = 48 # ~2 generations; token-rich layers normally finish their quota well inside it +TOKENS_FLOOR = 32 # below this mass the rms estimate is noise; the layer is omitted and stays on plain truncation calib_root = os.path.join(paths.models_path, 'calibration') -capture: CaptureState = {'model': None, 'recs': {}, 'handles': [], 'complete': False} +capture: CaptureState = {'model': None, 'recs': {}, 'handles': [], 'forwards': 0, 'complete': False} def enabled(): @@ -62,14 +69,20 @@ def checkpoint_name(sd_model): return getattr(info, 'name', None) +def denoiser_root(sd_model): + """The model's denoiser component, transformer first, unet otherwise.""" + root = getattr(sd_model, 'transformer', None) + return root if root is not None else getattr(sd_model, 'unet', None) + + def eligible_modules(sd_model): - """Sub-8-bit 2-D SDNQ linears of the model's transformer: the layers hosting applies to.""" - transformer = getattr(sd_model, 'transformer', None) - if transformer is None: + """Sub-8-bit 2-D SDNQ linears of the model's denoiser: the layers hosting applies to.""" + root = denoiser_root(sd_model) + if root is None: return [] from sdnq.common import dtype_dict out = [] - for name, m in transformer.named_modules(): + for name, m in root.named_modules(): deq = getattr(m, 'sdnq_dequantizer', None) if deq is None or len(deq.original_shape) != 2: continue @@ -85,9 +98,24 @@ def detach_capture(): capture['handles'].clear() capture['recs'].clear() capture['model'] = None + capture['forwards'] = 0 capture['complete'] = False +def deadline_hook(module, hook_args): # pylint: disable=unused-argument + """Count denoiser forwards and close capture at the deadline. + + Layers taking pooled or modulation vectors see a few tokens per forward + and can never reach the token quota; a global forward count bounds + capture for them and for modules the generation path never runs. + """ + if capture['complete']: + return + capture['forwards'] += 1 + if capture['forwards'] >= FORWARDS_DEADLINE: + persist() + + def hook_for(rec, in_features): def hook(module, hook_args): # pylint: disable=unused-argument if rec['done'] or capture['complete']: @@ -111,10 +139,12 @@ def hook_for(rec, in_features): def persist(): - """Write completed statistics and stamp them onto the layers. + """Write accumulated statistics and stamp them onto the layers. - Runs from the last completing hook, inside a forward; the write is a few - MB once per checkpoint ever. Handles stay registered but inert until the + Runs from the last hook to complete its quota or from the forward + deadline, inside a forward; the write is a few MB once per checkpoint + ever. Layers under the token floor are omitted rather than saved with + meaningless statistics. Handles stay registered but inert until the next safe point removes them (hook removal here would mutate the hook dict the forward is iterating). """ @@ -124,15 +154,20 @@ def persist(): from safetensors.torch import save_file tensors, min_n = {}, None for name, rec in capture['recs'].items(): - rms = (rec['ss'] / max(rec['n'], 1)).sqrt().float().cpu().contiguous().clone() + if rec['ss'] is None or rec['n'] < TOKENS_FLOOR: + continue + rms = (rec['ss'] / rec['n']).sqrt().float().cpu().contiguous().clone() tensors[name] = rms rec['m'].sdnq_calib_rms = rms min_n = rec['n'] if min_n is None else min(min_n, rec['n']) + if not tensors: + log.warning(f'Network calibration: model="{capture["model"]}" no layer reached {TOKENS_FLOOR} tokens; nothing saved') + return path = calib_file(capture['model']) try: os.makedirs(calib_root, exist_ok=True) save_file(tensors, path, metadata={'version': '1', 'model': capture['model'], 'tokens': str(min_n)}) - log.info(f'Network calibration: model="{capture["model"]}" layers={len(tensors)} tokens={min_n} saved="{path}"') + log.info(f'Network calibration: model="{capture["model"]}" layers={len(tensors)}/{len(capture["recs"])} tokens={min_n} saved="{path}"') except Exception as e: log.warning(f'Network calibration: save failed path="{path}" {e}') @@ -172,6 +207,7 @@ def on_model_loaded(sd_model): if 'Model' in (getattr(shared.opts, 'cuda_compile', None) or []): return # hooks inside a compiled module graph-break or misbehave; skip capture entirely capture['model'] = name + capture['handles'].append(denoiser_root(sd_model).register_forward_pre_hook(deadline_hook)) for mod_name, m in modules_list: rec = {'m': m, 'ss': None, 'n': 0, 'done': False} capture['recs'][mod_name] = rec diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 5635f8495..3f9b9a0a0 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -101,7 +101,7 @@ def get_module_factors(module, device, dtype, original_shape=None): return up_eff.to(dtype=dtype), down.to(device=device, dtype=dtype) -def factor_candidate(self, network_layer_name, wanted_names, use_previous=False): +def factor_candidate(self, network_layer_name, wanted_names): """True when this layer should take the exact svd-append path. Requires an SDNQ linear layer whose active networks all contribute plain @@ -116,9 +116,8 @@ def factor_candidate(self, network_layer_name, wanted_names, use_previous=False) 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: + for net in l.loaded_networks: module = net.modules.get(network_layer_name, None) if module is None: continue @@ -150,7 +149,7 @@ def remove_factors(self): return True -def apply_factors(self, network_layer_name, wanted_names, use_previous=False): +def apply_factors(self, network_layer_name, wanted_names): """Attach the active networks' LoRA factors to this layer's svd side-channel. Replaces any previously attached factors (multiplier changes re-enter @@ -166,9 +165,9 @@ def apply_factors(self, network_layer_name, wanted_names, use_previous=False): deq = self.sdnq_dequantizer dtype = deq.result_dtype - loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks + ups, downs = [], [] - for net in loaded: + for net in l.loaded_networks: module = net.modules.get(network_layer_name, None) if module is None: continue @@ -215,7 +214,7 @@ def append_factors(self, ups, downs): self.svd_down = torch.nn.Parameter(new_down.to(device=device), requires_grad=False) -def host_candidate(self, network_layer_name, wanted_names, use_previous=False): +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 enabled(): return False @@ -228,11 +227,11 @@ def host_candidate(self, network_layer_name, wanted_names, use_previous=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 - loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks - return any(net.modules.get(network_layer_name, None) is not None for net in loaded) + + return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) -def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=False): +def apply_hosted(self, network_layer_name, updown, wanted_names): """Host a set's delta on the svd channel: exact factors for factorable members, the top-k singular directions of the remainder for the rest. @@ -256,15 +255,12 @@ def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=Fa if updown is None or updown.ndim != 2 or tuple(updown.shape) != tuple(deq.original_shape): return None dtype = deq.result_dtype - cached = None - if not use_previous: - lora_factor_cache.begin_pass(wanted_names) - cached = lora_factor_cache.fetch(network_layer_name) + lora_factor_cache.begin_pass(wanted_names) + cached = lora_factor_cache.fetch(network_layer_name) D = None if cached is not None else updown.detach().to(devices.device, torch.float32) ups, downs = [], [] - loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks - for net in loaded: + for net in l.loaded_networks: module = net.modules.get(network_layer_name, None) if module is None: continue diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 4c2c17ce6..53872f021 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -858,7 +858,7 @@ def test_calib_capture_persist_roundtrip(): lora_calib.calib_root = tmp lora_calib.TOKENS_DONE = 2048 lora_calib.on_model_loaded(sd) - assert len(lora_calib.capture['handles']) == 2, 'both sub-8-bit linears must hook' + assert len(lora_calib.capture['handles']) == 3, 'both sub-8-bit linears plus the root forward counter must hook' torch.manual_seed(51) scale = torch.linspace(0.1, 4.0, IN_F, device=DEVICE) xs = [] @@ -905,6 +905,119 @@ def test_calib_capture_gates(): return True +class MockCalibDenoiser(torch.nn.Module): + """Denoiser whose forward feeds one token-rich linear and one token-starved one, like a DiT block beside its modulation projection.""" + def __init__(self, rich, starved, starved_tokens): + super().__init__() + self.rich = rich + self.starved = starved + self.starved_tokens = starved_tokens + + def forward(self, x): + self.rich(x) + self.starved(x[:self.starved_tokens]) + return x + + +def test_calib_deadline_persists_starved_layers(): + import tempfile + from safetensors import safe_open + from modules.lora import lora_calib + rich = build_layer('uint4', seed=45) + starved = build_layer('uint4', seed=46) + root = MockCalibDenoiser(rich, starved, starved_tokens=8) + sd = MockCalibSd('test/calib-deadline') + sd.transformer = root + old = (lora_calib.calib_root, lora_calib.TOKENS_DONE, lora_calib.FORWARDS_DEADLINE) + with tempfile.TemporaryDirectory() as tmp, host_calib(True): + try: + lora_calib.calib_root = tmp + lora_calib.TOKENS_DONE = 2048 + lora_calib.FORWARDS_DEADLINE = 6 + lora_calib.on_model_loaded(sd) + assert len(lora_calib.capture['handles']) == 3, 'two layer hooks plus the root forward counter must attach' + torch.manual_seed(52) + xs = [] + for i in range(6): + x = torch.randn(1024, IN_F, device=DEVICE).to(torch.bfloat16) + if i < 5: # the deadline fires at the start of the sixth forward, before its layer hooks run + xs.append(x[:8].float()) + root(x) + assert lora_calib.capture['complete'], 'the forward deadline must close capture' + assert lora_calib.capture['forwards'] == 6, f'root counter must track denoiser forwards, got {lora_calib.capture["forwards"]}' + path = lora_calib.calib_file('test/calib-deadline') + assert os.path.isfile(path), 'deadline persist must write the statistics file' + assert getattr(starved, 'sdnq_calib_rms', None) is not None, 'the starved layer must carry statistics' + expected = torch.cat(xs).square().mean(dim=0).sqrt().cpu() + assert torch.allclose(starved.sdnq_calib_rms, expected, rtol=1e-3, atol=1e-5), 'starved rms must match exactly the tokens it saw' + with safe_open(path, framework='pt', device='cpu') as f: + assert set(f.keys()) == {'rich', 'starved'}, f'both layers must persist, got {sorted(f.keys())}' + assert f.metadata()['tokens'] == '40', f'metadata must report the weakest saved layer, got {f.metadata()["tokens"]}' + finally: + lora_calib.calib_root, lora_calib.TOKENS_DONE, lora_calib.FORWARDS_DEADLINE = old + lora_calib.detach_capture() + return True + + +def test_calib_deadline_omits_subfloor_layers(): + import tempfile + from safetensors import safe_open + from modules.lora import lora_calib + rich = build_layer('uint4', seed=48) + starved = build_layer('uint4', seed=49) + root = MockCalibDenoiser(rich, starved, starved_tokens=2) # 2 tokens x 5 counted forwards = 10, under the floor of 32 + sd = MockCalibSd('test/calib-subfloor') + sd.transformer = root + old = (lora_calib.calib_root, lora_calib.TOKENS_DONE, lora_calib.FORWARDS_DEADLINE) + with tempfile.TemporaryDirectory() as tmp, host_calib(True): + try: + lora_calib.calib_root = tmp + lora_calib.TOKENS_DONE = 2048 + lora_calib.FORWARDS_DEADLINE = 6 + lora_calib.on_model_loaded(sd) + torch.manual_seed(53) + for _ in range(6): + root(torch.randn(1024, IN_F, device=DEVICE).to(torch.bfloat16)) + assert lora_calib.capture['complete'], 'the forward deadline must close capture' + path = lora_calib.calib_file('test/calib-subfloor') + with safe_open(path, framework='pt', device='cpu') as f: + assert set(f.keys()) == {'rich'}, f'a layer under the token floor must be omitted, got {sorted(f.keys())}' + assert getattr(starved, 'sdnq_calib_rms', None) is None, 'an omitted layer must not carry statistics' + del rich.sdnq_calib_rms + lora_calib.on_model_loaded(sd) # second load takes the cached path with the partial file + assert len(lora_calib.capture['handles']) == 0, 'a partial file still counts as cached; capture must not re-attach' + assert getattr(rich, 'sdnq_calib_rms', None) is not None, 'the saved layer must reload from the partial file' + assert getattr(starved, 'sdnq_calib_rms', None) is None, 'the omitted layer must stay on plain truncation after reload' + finally: + lora_calib.calib_root, lora_calib.TOKENS_DONE, lora_calib.FORWARDS_DEADLINE = old + lora_calib.detach_capture() + return True + + +def test_calib_unet_root_walk(): + import tempfile + from modules.lora import lora_calib + layer = build_layer('uint4', seed=47) + sd = MockCalibSd('test/calib-unet', lin=layer) + sd.unet = sd.transformer + sd.transformer = None + mods = lora_calib.eligible_modules(sd) + assert [n for n, _ in mods] == ['lin'], f'the unet root must be walked when no transformer exists, got {[n for n, _ in mods]}' + both = MockCalibSd('test/calib-both') + both.unet = sd.unet + assert lora_calib.eligible_modules(both) == [], 'a transformer root wins even when it holds no eligible linears' + old_root = lora_calib.calib_root + with tempfile.TemporaryDirectory() as tmp, host_calib(True): + try: + lora_calib.calib_root = tmp + lora_calib.on_model_loaded(sd) + assert len(lora_calib.capture['handles']) == 2, 'the layer hook plus the root counter must attach on a unet model' + finally: + lora_calib.calib_root = old_root + lora_calib.detach_capture() + return True + + CAT_FCACHE = category('factor-cache') @@ -1185,7 +1298,8 @@ def run_tests(): 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, - test_calib_capture_persist_roundtrip, test_calib_capture_gates]: + test_calib_capture_persist_roundtrip, test_calib_capture_gates, test_calib_deadline_persists_starved_layers, + test_calib_deadline_omits_subfloor_layers, test_calib_unet_root_walk]: run_test(CAT_CALIB, fn) log.warning('=== Factor cache ===') for fn in [test_factor_cache_roundtrip_bitexact, test_factor_cache_invalidates_on_multiplier, From 179978a55533f5ded6e89bff8ad148f6ea468ff4 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 03:25:43 +0100 Subject: [PATCH 6/8] feat(lora): route fat non-factorable deltas back to requantize Hosting truncates every non-factorable set on sub-8-bit layers, but a delta large against the grid step whose truncation capture is low is retained better by the grid than by the rank cap. apply_hosted now returns such layers to the requantize path when rms(delta)/mean(step) exceeds 0.30 and the sketch capture falls below 0.90, thresholds sized on 610 calibrated modules across krea2 and anima. Both terms must agree: a thin delta rounds away on the grid however low its capture, and a low-rank delta hosts exactly however fat it is. - scoped to pure non-factorable sets; factorable members, dense-combined deltas and svd-channel checkpoints keep hosting - factor cache entries memoize the decision through their stored capture, so replays route without re-running the sketch - routed layers log as info apart from the forced-fallback warning - the fidelity CLI applies the same rule so its reports track the loader - seven routing tests, constants module-level and test-overridable --- cli/lora-quant-fidelity.py | 38 ++++++---- modules/lora/lora_sdnq.py | 80 +++++++++++++++++---- test/test-sdnq-lora-factors.py | 128 ++++++++++++++++++++++++++++++++- 3 files changed, 217 insertions(+), 29 deletions(-) diff --git a/cli/lora-quant-fidelity.py b/cli/lora-quant-fidelity.py index 1dc41bfe5..8aba5675c 100644 --- a/cli/lora-quant-fidelity.py +++ b/cli/lora-quant-fidelity.py @@ -252,7 +252,7 @@ class Bf16Repo: return f.get_tensor(key) -def analyze_module(W_dq, deq_params, mods, calib_rms=None): +def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None): """Return fidelity metrics for one quantized module and the adapters targeting it. Deltas come from each module's production calc_updown and sum the way the @@ -260,6 +260,9 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None): is measured as applied. A module is factor-path eligible only when every contribution is a plain additive lora. With ``calib_rms``, hosting mirrors the calibrated production path and its rho is scored in the weighted norm. + With ``step_live`` (the layer's own pre-add scale), the production routing + rule applies: a delta fat against the grid whose truncation capture is low + reports the requantize path, the way the loader would route it. """ D = None for mod in mods: @@ -316,17 +319,23 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None): with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []): torch.manual_seed(0) U, S, V = torch.svd_lowrank(Dw, q=min(q + 64, *D.shape), niter=8) - Dk = (U[:, :q] * S[:q]) @ V[:, :q].t() - if rms is not None: - Dk = Dk / rms - base16 = W_dq.to(torch.bfloat16).float() - realized = (W_dq.to(torch.bfloat16) + Dk.to(torch.bfloat16)).float() - base16 - if rms is not None: # weighted norm: the diagonal-covariance output-error proxy the calibrated truncation optimizes - Dr = D * rms - applied_rho = float((realized * rms).flatten() @ Dr.flatten() / Dr.square().sum()) - else: - applied_rho = float(realized.flatten() @ D.flatten() / nD.square()) - hosted = True + energy = float(S[:q].square().sum() / Dw.square().sum().clamp(min=1e-30)) + routed = False + if step_live is not None and not deq_params.get('use_svd', False): + sr = float(D.square().mean().sqrt() / step_live.float().mean()) + routed = sr > lora_sdnq.REQUANT_RATIO and energy < lora_sdnq.REQUANT_ENERGY + if not routed: # the loader routes fat, genuinely-truncated deltas back to requantize + Dk = (U[:, :q] * S[:q]) @ V[:, :q].t() + if rms is not None: + Dk = Dk / rms + base16 = W_dq.to(torch.bfloat16).float() + realized = (W_dq.to(torch.bfloat16) + Dk.to(torch.bfloat16)).float() - base16 + if rms is not None: # weighted norm: the diagonal-covariance output-error proxy the calibrated truncation optimizes + Dr = D * rms + applied_rho = float((realized * rms).flatten() @ Dr.flatten() / Dr.square().sum()) + else: + applied_rho = float(realized.flatten() @ D.flatten() / nD.square()) + hosted = True return dict(rank=getattr(mods[0], 'dim', None), rms_delta=float(D.pow(2).mean().sqrt()), rms_weight=float(W_dq.pow(2).mean().sqrt()), step_ratio=step_ratio, crossers=crossers, requant_rho=rho, requant_resid=resid, factor_eligible=factor_eligible, hosted=hosted, applied_rho=applied_rho, @@ -398,6 +407,7 @@ def main(): skip_quantized_matmul=deq.use_quantized_matmul, dtype=torch.float32, skip_compile=True).to(device) params = dict(weights_dtype=deq.weights_dtype, group_size=deq.group_size, hadamard_group_size=deq.hadamard_group_size, use_hadamard=deq.use_hadamard, use_svd=layer.svd_up is not None, svd_rank=deq.svd_rank, svd_steps=deq.svd_steps) + step_live = layer.scale.detach().to(device) sd_module = layer else: W = bf16_repo.get(f'{path}.weight') @@ -412,16 +422,18 @@ def main(): if args.dtype == 'bf16': W_dq = W.to(device, torch.bfloat16).float() params = dict(weights_dtype='bf16', group_size=0, hadamard_group_size=0, use_hadamard=False) + step_live = None else: deq0, data0 = sdnq_quantize_layer_weight(W.to(device, torch.float32), layer_class_name='Linear', weights_dtype=args.dtype, group_size=args.group, hadamard_group_size=args.hadamard_group, use_hadamard=args.hadamard_group > 0, use_svd=False, use_quantized_matmul=False, dequantize_fp32=False, torch_dtype=torch.bfloat16) W_dq = deq0(data0['weight'], data0['scale'], zero_point=data0['zero_point'], svd_up=None, svd_down=None, dtype=torch.float32, skip_compile=True) params = dict(weights_dtype=args.dtype, group_size=deq0.group_size, hadamard_group_size=deq0.hadamard_group_size, use_hadamard=deq0.use_hadamard) + step_live = data0['scale'].detach() sd_module = make_stub(W.shape) try: mods = [build_module(fam, path, w, net, sd_module) for fam, w in entries] - row = analyze_module(W_dq, params, mods, calib_rms=calib_stats.get(lname)) + row = analyze_module(W_dq, params, mods, calib_rms=calib_stats.get(lname), step_live=step_live) except Exception as e: # a family the tool cannot rebuild must not read as a clean module failed.append(f'{path}: {type(e).__name__}: {e}') del W_dq diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 3f9b9a0a0..5a341c137 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -33,6 +33,13 @@ magnitude. When activation statistics for the checkpoint exist (see error instead of weight error. At 8 bits and above requantize retains most of the delta, so hosting is skipped there and the requantize path remains. + +A small tail of deltas inverts the tradeoff: when the delta is large +against the grid step AND the truncation genuinely cuts it, requantize +retains more than hosting drops, and the layer routes back to the +requantize path (``REQUANT_RATIO``/``REQUANT_ENERGY``). Both terms must +agree: a thin delta rounds away on the grid however low its capture, and +a low-rank delta hosts exactly however fat it is. """ import torch @@ -45,6 +52,10 @@ from modules.logger import log fallback_layers: list[str] = [] hosted_layers: list[tuple[str, float, bool]] = [] +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 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).""" @@ -243,8 +254,8 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): directions minimize output error rather than weight error. Computed factors are disk-cached per configuration (``lora_factor_cache``) and replayed bit-identically on later applies. Returns None when the delta - cannot ride the channel (wrong shape); the caller falls back to - requantize. + cannot ride the channel (wrong shape) or when the routing rule prefers + the grid for it; the caller falls back to requantize. """ from sdnq.quant_utils import rotate_hadamard @@ -255,19 +266,33 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if updown is None or updown.ndim != 2 or tuple(updown.shape) != tuple(deq.original_shape): return None dtype = deq.result_dtype - lora_factor_cache.begin_pass(wanted_names) - cached = lora_factor_cache.fetch(network_layer_name) - D = None if cached is not None else updown.detach().to(devices.device, torch.float32) - ups, downs = [], [] + members = [] for net in l.loaded_networks: module = net.modules.get(network_layer_name, None) if module is None: continue factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) - if factors is None: - continue - up_eff, down = factors + if factors is not None: + members.append(factors) + + # requantize keeps a delta the grid can resolve and that truncation would genuinely + # cut: both terms must agree, since a thin delta rounds away on the grid however + # low its capture, and a low-rank delta hosts exactly however fat it is. Scoped to + # sets the side-channel would otherwise carry whole: factorable members ride + # exactly. + maybe_requant = len(members) == 0 and self.svd_up is None + if maybe_requant: + step = float(self.scale.detach().float().mean()) + rms = float(updown.detach().float().square().mean().sqrt()) + maybe_requant = step > 0 and rms / step > REQUANT_RATIO + + lora_factor_cache.begin_pass(wanted_names) + cached = lora_factor_cache.fetch(network_layer_name) + D = None if cached is not None else updown.detach().to(devices.device, torch.float32) + + ups, downs = [], [] + for up_eff, down in members: if D is not None: D = D.sub_(up_eff.to(torch.float32) @ down.to(torch.float32)) # factorable members ride exactly; host only the remainder if deq.use_hadamard: @@ -277,10 +302,33 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if cached is not None: up_h, down_h, energy, calibrated = cached + if maybe_requant and energy < REQUANT_ENERGY: + routed_layers.append(network_layer_name) + return None 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)) return True + up_h, down_h, energy, calibrated = truncate_delta(self, D, dtype) + up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, calibrated) + 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 + append_factors(self, ups + [up_h], downs + [down_h]) + hosted_layers.append((network_layer_name, energy, calibrated)) + return True + + +def truncate_delta(self, D, dtype): + """Truncate one dense fp32 delta to hosted factors in the layer's channel layout; consumes ``D``. + + Calibration-weighted when statistics exist; the sketch is oversampled past + the kept rank so the truncation sits within noise of exact svd. Returns + ``(up_h, down_h, energy, calibrated)`` with the down factor rotated into the + layer's hadamard domain. + """ + from sdnq.quant_utils import rotate_hadamard + deq = self.sdnq_dequantizer cap = int(shared.opts.lora_sdnq_host_rank) q = min(cap, *D.shape) rms = lora_calib.rms_for(self) @@ -304,15 +352,12 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): if deq.use_hadamard: down_h = rotate_hadamard(down_h, group_size=deq.hadamard_group_size) down_h = down_h.to(dtype=dtype) - up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, rms is not None) - append_factors(self, ups + [up_h], downs + [down_h]) - hosted_layers.append((network_layer_name, energy, rms is not None)) - return True + return up_h, down_h, energy, rms is not None 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: + """Record a quantized layer taking the requantize path (summary-logged per pass); layers the routing rule sent there are counted apart.""" + if getattr(self, 'sdnq_dequantizer', None) is not None and network_layer_name not in routed_layers: fallback_layers.append(network_layer_name) @@ -328,6 +373,11 @@ def report_fallbacks(): 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() + 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: + log.debug(f'Network load: type=LoRA quant=sdnq routed={routed_layers[:8]}{"..." if len(routed_layers) > 8 else ""}') + routed_layers.clear() if len(fallback_layers) > 0: if enabled(): log.warning(f'Network load: type=LoRA quant=sdnq layers={len(fallback_layers)} non-factorable networks requantized in place (reduced fidelity on quantized weights)') diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index 53872f021..da9ffa12b 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -753,6 +753,129 @@ def test_hosted_transitions_and_rng_isolation(): return True +@contextmanager +def requant_rule(ratio, energy): + old_r, old_e = lora_sdnq.REQUANT_RATIO, lora_sdnq.REQUANT_ENERGY + lora_sdnq.REQUANT_RATIO, lora_sdnq.REQUANT_ENERGY = ratio, energy + try: + yield + finally: + lora_sdnq.REQUANT_RATIO, lora_sdnq.REQUANT_ENERGY = old_r, old_e + + +def test_route_fat_dense_delta_requantizes(): + layer = build_layer('uint4') + torch.manual_seed(21) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 # full-rank and well above the grid step: the grid retains it, truncation would cut it + net = make_dense_net('fatnet', layer, D) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'a fat full-rank delta must route to requantize' + assert isinstance(getattr(layer, 'network_weights_backup', None), torch.Tensor), 'the routed layer takes the requantize backup' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.7, f'the grid must retain the routed delta: rho={rho:.3f}' + activate() + assert torch.equal(dq(layer), Wdq0), 'restore from backup must be bit-exact' + return True + + +def test_route_rule_terms_gate_both_ways(): + layer = build_layer('uint4') + torch.manual_seed(23) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 # sr about 0.8, capture about 0.8 at cap 256: each term alone can hold it hosted + net = make_dense_net('gatenet', layer, D) + with host_rank(256), mock_model(lin=layer): + with requant_rule(ratio=10.0, energy=0.90): + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'sr below the ratio must host regardless of capture' + activate() + with requant_rule(ratio=0.30, energy=0.0): + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'capture above the energy floor must host regardless of sr' + activate() + with requant_rule(ratio=0.30, energy=0.90): + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'both terms crossed must requantize' + activate() + return True + + +def test_route_low_rank_fat_delta_stays_hosted(): + layer = build_layer('uint4') + _A, _B, D = make_delta(seed=22, sigma=3e-3) # rank-8: fat against the grid, exact under the cap + net = make_dense_net('fatlow', layer, D) + with host_rank(64), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'a low-rank delta hosts exactly at any magnitude' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.95, f'rho={rho:.4f}' + activate() + return True + + +def test_route_mixed_set_keeps_hosting(): + layer = build_layer('uint4') + A, B, _D1 = make_delta(seed=24) + torch.manual_seed(25) + D2 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net1 = make_net('mixp', layer, A, B) + net2 = make_dense_net('mixf', layer, D2) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net1, net2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'a set with factorable members keeps the side-channel' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_route_svd_checkpoint_keeps_hosting(): + layer = build_layer('uint4', use_svd=True) + torch.manual_seed(26) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net = make_dense_net('svdfat', layer, D) + with host_rank(256), mock_model(lin=layer): + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'svd checkpoints keep hosting; the rule is not grounded there' + activate() + return True + + +def test_route_dense_stack_keeps_hosting(): + layer = build_layer('uint4') + torch.manual_seed(27) + D1 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + D2 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 + net1, net2 = make_dense_net('df1', layer, D1), make_dense_net('df2', layer, D2) + with host_rank(256), stack_mode('ties'), mock_model(lin=layer): + activate(net1, net2) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'dense-combined deltas host at any magnitude' + activate() + return True + + +def test_route_replay_from_cache(): + import tempfile + layer = build_layer('uint4') + with tempfile.TemporaryDirectory() as tmp: + with host_rank(256), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer): + net, _D = cache_fixture(tmp, layer, name='fatcache', sigma=1e-2, seed=28) + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'fat delta must route on the fresh-sketch path' + activate() + real_svd = torch.svd_lowrank + torch.svd_lowrank = raise_no_svd + try: + activate(net) # the stored entry memoizes the routing: same decision, no sketch + finally: + torch.svd_lowrank = real_svd + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'cache replay must route the same way' + activate() + return True + + CAT_CALIB = category('calibration') @@ -1294,7 +1417,10 @@ def run_tests(): run_test(CAT_TRANS, fn) log.warning('=== Hosting ===') for fn in [test_hosted_low_rank_delta_is_kept, test_hosted_dense_delta_beats_requant, test_hosted_skips_int8, - test_hosted_disabled_by_option, test_hosted_transitions_and_rng_isolation]: + 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_dense_stack_keeps_hosting, + test_route_replay_from_cache]: 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, From 0c1b087e20a0ab63e523771dd6ba60a1c6e2a584 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 26 Jul 2026 03:51:01 +0100 Subject: [PATCH 7/8] perf(lora): serve hosted layers from the factor cache before delta assembly A cache hit still paid full calc_updown per layer, the dominant cost of a replayed apply. The walk now probes the pass's cache entry first through apply_cached, which evaluates the requantize routing rule from the delta rms stored in the entry and attaches the cached factors without assembling the delta; layers the rule declines fall through with the assembled delta reused for the requantize path instead of recomputing it. The entry format gains the rms and bumps to fmt 5, so older files recompute once and rewrite. Hit accounting stays single-count through a plain lookup plus an explicit hit note. - replayed krea2 LoKR apply drops calc from 9.1s to 0.2s and activate from 39.5s fresh to 4.8s replayed, hosted/routed split identical - three cache tests pin the calc skip on plain, mixed and dense-pair sets --- modules/lora/lora_factor_cache.py | 41 ++++++++--- modules/lora/lora_sdnq.py | 52 ++++++++++++-- modules/lora/networks.py | 32 +++++---- test/test-sdnq-lora-factors.py | 111 ++++++++++++++++++++++++++---- 4 files changed, 193 insertions(+), 43 deletions(-) diff --git a/modules/lora/lora_factor_cache.py b/modules/lora/lora_factor_cache.py index 6e7e46a0b..cda0dca39 100644 --- a/modules/lora/lora_factor_cache.py +++ b/modules/lora/lora_factor_cache.py @@ -34,6 +34,7 @@ from modules.logger import log cache_root = os.path.join(paths.models_path, 'lora-factor-cache') state = {'wn': None, 'sig': None, 'path': None, 'store': {}, 'dirty': False, 'hits': 0, 'misses': 0} +FMT = '5' # bump on entry-layout changes so older files recompute instead of replaying short def budget_gb(): @@ -53,7 +54,7 @@ def signature(wanted_names): parts = { 'model': model_name, 'rank': int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0), - 'calib': int(os.path.getmtime(calib_path)) if os.path.isfile(calib_path) else None, + 'calib': int(os.path.getmtime(calib_path)) if lora_calib.enabled() and os.path.isfile(calib_path) else None, # the toggle is part of the identity: factors computed under the other setting must not replay 'nets': [], } for name, te, unet, dyn in wanted_names: @@ -96,7 +97,7 @@ def begin_pass(wanted_names): from safetensors import safe_open with safe_open(path, framework='pt', device='cpu') as f: meta = f.metadata() or {} - if meta.get('sig') == sig and meta.get('fmt') == '3': + if meta.get('sig') == sig and meta.get('fmt') == FMT: for k in f.keys(): entries[k] = f.get_tensor(k) os.utime(path, None) # freshness for LRU eviction @@ -120,8 +121,12 @@ def dequantize_rowwise(q, scale): return q.to(torch.float32) * scale -def fetch(network_layer_name): - """Cached (up, down, energy, calibrated) for a layer, or None; factors return as fp32.""" +def lookup(network_layer_name): + """Cached (up, down, energy, calibrated, rms) for a layer, or None; factors return as fp32. + + Pure lookup with no hit/miss accounting: the fast-path probe uses it so a + layer is only counted once, by whichever caller consumes the answer. + """ if state['sig'] is None: return None st = state['store'] @@ -129,20 +134,35 @@ def fetch(network_layer_name): down_q, down_s = st.get(f'{network_layer_name}.down_q'), st.get(f'{network_layer_name}.down_s') energy = st.get(f'{network_layer_name}.energy') calib = st.get(f'{network_layer_name}.calib') - if up_q is None or up_s is None or down_q is None or down_s is None or energy is None or calib is None: - state['misses'] += 1 + rms = st.get(f'{network_layer_name}.rms') + if up_q is None or up_s is None or down_q is None or down_s is None or energy is None or calib is None or rms is None: + return None + return dequantize_rowwise(up_q, up_s), dequantize_rowwise(down_q, down_s), float(energy), bool(calib), float(rms) + + +def note_hit(): + state['hits'] += 1 + + +def fetch(network_layer_name): + """``lookup`` with accounting: a usable entry counts a hit, anything else a miss.""" + entry = lookup(network_layer_name) + if entry is None: + if state['sig'] is not None: + state['misses'] += 1 return None state['hits'] += 1 - return dequantize_rowwise(up_q, up_s), dequantize_rowwise(down_q, down_s), float(energy), bool(calib) + return entry -def store(network_layer_name, up, down, energy, calibrated): +def store(network_layer_name, up, down, energy, calibrated, rms): """Quantize-before-use: returns the pair 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. + the same tensors. ``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 @@ -155,6 +175,7 @@ def store(network_layer_name, up, down, energy, calibrated): 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) @@ -190,7 +211,7 @@ def flush(): from safetensors.torch import save_file os.makedirs(cache_root, exist_ok=True) tmp = state['path'] + '.tmp' - save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': '3'}) + save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': FMT}) os.replace(tmp, state['path']) evict() except Exception as e: diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py index 5a341c137..6aead3b12 100644 --- a/modules/lora/lora_sdnq.py +++ b/modules/lora/lora_sdnq.py @@ -242,6 +242,50 @@ def host_candidate(self, network_layer_name, wanted_names): return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks) +def apply_cached(self, network_layer_name, wanted_names): + """Attach a hosted set straight from the factor cache, before the delta exists. + + Probed by the walk ahead of delta assembly: on a usable entry the routing + rule is evaluated from the stored delta rms and the cached factors attach + exactly as a fetch inside ``apply_hosted`` would, so the pass skips + ``calc_updown`` for the layer entirely. Returns True when the layer was + served; None sends the caller down the assemble-and-host path (no entry, + or the rule wants the grid). + """ + from sdnq.quant_utils import rotate_hadamard + + lora_factor_cache.begin_pass(wanted_names) + entry = lora_factor_cache.lookup(network_layer_name) + if entry is None: + return None + up_h, down_h, energy, calibrated, rms = entry + 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 + members = [] + for net in l.loaded_networks: + module = net.modules.get(network_layer_name, None) + if module is None: + continue + factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape) + if factors is not None: + members.append(factors) + if len(members) == 0 and self.svd_up is None: + step = float(self.scale.detach().float().mean()) + if step > 0 and rms / step > REQUANT_RATIO and energy < REQUANT_ENERGY: + return None # routed to the grid: the caller assembles the delta and requantizes + ups, downs = [], [] + for up_eff, down in members: + 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) + 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)) + return True + + def apply_hosted(self, network_layer_name, updown, wanted_names): """Host a set's delta on the svd channel: exact factors for factorable members, the top-k singular directions of the remainder for the rest. @@ -281,11 +325,11 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): # low its capture, and a low-rank delta hosts exactly however fat it is. Scoped to # sets the side-channel would otherwise carry whole: factorable members ride # exactly. + delta_rms = float(updown.detach().float().square().mean().sqrt()) maybe_requant = len(members) == 0 and self.svd_up is None if maybe_requant: step = float(self.scale.detach().float().mean()) - rms = float(updown.detach().float().square().mean().sqrt()) - maybe_requant = step > 0 and rms / step > REQUANT_RATIO + maybe_requant = step > 0 and delta_rms / step > REQUANT_RATIO lora_factor_cache.begin_pass(wanted_names) cached = lora_factor_cache.fetch(network_layer_name) @@ -301,7 +345,7 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): downs.append(down) if cached is not None: - up_h, down_h, energy, calibrated = cached + up_h, down_h, energy, calibrated, _cached_rms = cached if maybe_requant and energy < REQUANT_ENERGY: routed_layers.append(network_layer_name) return None @@ -310,7 +354,7 @@ def apply_hosted(self, network_layer_name, updown, wanted_names): return True up_h, down_h, energy, calibrated = truncate_delta(self, D, dtype) - up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, calibrated) + up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, calibrated, delta_rms) 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 diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 95c7675a5..9e3b2e7a0 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -106,6 +106,7 @@ 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) + calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing 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): @@ -124,22 +125,24 @@ 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) # the hosted delta is measured against the pristine base - batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit) - if batch_ex_bias is None: # bias deltas need the plain path; weight-only sets ride the side-channel without a weight backup - hosted = lora_sdnq.apply_hosted(module, network_layer_name, batch_updown, component_wanted) + hosted = lora_sdnq.apply_cached(module, network_layer_name, component_wanted) # a stored entry serves the layer before the delta is assembled + if hosted is None: + batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit) + calced = True + if batch_ex_bias is None: # bias deltas need the plain path; weight-only sets ride the side-channel without a weight backup + hosted = lora_sdnq.apply_hosted(module, network_layer_name, batch_updown, component_wanted) if hosted is not None: - if hosted and component_wanted: - applied_layers.append(network_layer_name) - applied_weight += 1 - module.network_current_names = component_wanted - module.network_current_stack = stack_sig batch_updown, batch_ex_bias = None, None del batch_updown, batch_ex_bias - if task is not None: - pbar.update(task, advance=1) - continue - batch_updown, batch_ex_bias = None, None - del batch_updown, batch_ex_bias + if hosted is not None: + if hosted 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 stripped = lora_sdnq.remove_factors(module) # the mechanism gate can decline a layer still carrying attached factors; the weight path must start from the pristine channel if stripped and not component_wanted: # factor-mode layers have no tensor backup, dropping the factors is the whole restore module.network_current_names = () @@ -156,7 +159,8 @@ def network_activate(include=None, exclude=None): continue 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 not calced: # the host branch may have assembled the delta already; a declined layer reuses it + batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit) 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: diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index da9ffa12b..d350b633c 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -843,19 +843,6 @@ def test_route_svd_checkpoint_keeps_hosting(): return True -def test_route_dense_stack_keeps_hosting(): - layer = build_layer('uint4') - torch.manual_seed(27) - D1 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 - D2 = torch.randn(OUT_F, IN_F, device=DEVICE) * 1e-2 - net1, net2 = make_dense_net('df1', layer, D1), make_dense_net('df2', layer, D2) - with host_rank(256), stack_mode('ties'), mock_model(lin=layer): - activate(net1, net2) - assert hasattr(layer, 'sdnq_lora_svd_stash'), 'dense-combined deltas host at any magnitude' - activate() - return True - - def test_route_replay_from_cache(): import tempfile layer = build_layer('uint4') @@ -1254,6 +1241,99 @@ def test_factor_cache_disabled_at_zero(): return True +def test_factor_cache_invalidates_on_calib_toggle(): + import tempfile + from modules.lora import lora_calib + 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) + old_root = lora_calib.calib_root + lora_calib.calib_root = os.path.join(tmp, 'calib') + os.makedirs(lora_calib.calib_root, exist_ok=True) + with open(lora_calib.calib_file('test/cache-model'), 'wb') as f: + f.write(b'0' * 64) # the signature stats this file; its content is never read here + torch.manual_seed(77) + layer.sdnq_calib_rms = torch.rand(IN_F) * 4 + 0.1 + try: + with host_calib(True): + activate(net) + up_cal = layer.svd_up.detach().clone() + activate() + with host_calib(False): + activate(net) # same set with calibration off: the entry keyed under the other setting must miss + up_plain = layer.svd_up.detach().clone() + activate() + assert not torch.equal(up_cal, up_plain), 'toggling calibration must not replay factors computed under the other setting' + assert len(os.listdir(os.path.join(tmp, 'cache'))) == 2, 'the two settings must key separate cache entries' + finally: + del layer.sdnq_calib_rms + lora_calib.calib_root = old_root + return True + + +@contextmanager +def counting_calc(): + """Count NetworkModuleFull.calc_updown calls: zero on a pass proves the walk skipped delta assembly.""" + from modules.lora import network_full + calls = {'n': 0} + real = network_full.NetworkModuleFull.calc_updown + def wrapper(self, *args, **kwargs): + calls['n'] += 1 + return real(self, *args, **kwargs) + network_full.NetworkModuleFull.calc_updown = wrapper + try: + yield calls + finally: + network_full.NetworkModuleFull.calc_updown = real + + +def test_cache_fastpath_skips_calc(): + 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) + Wdq0 = dq(layer) + with counting_calc() as calls: + activate(net) + assert calls['n'] > 0, 'a fresh apply must assemble the delta' + first = dq(layer) + activate() + calls['n'] = 0 + activate(net) + assert calls['n'] == 0, f'a cache replay must not assemble the delta: calc_updown ran {calls["n"]} times' + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'the fast path must attach the cached factors' + assert torch.equal(dq(layer), first), 'fast-path replay must be bit-identical to the fresh apply' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_cache_fastpath_serves_mixed_set(): + import tempfile + layer = build_layer('uint4') + A, B, _D1 = make_delta(seed=63) + with tempfile.TemporaryDirectory() as tmp: + with host_rank(64), host_cache(10, os.path.join(tmp, 'cache')), mock_model(lin=layer): + net_full, _D2 = cache_fixture(tmp, layer, name='mixfull', seed=64) + net_plain = make_net('mixlora', layer, A, B) + lora_file = os.path.join(tmp, 'mixlora.safetensors') + with open(lora_file, 'wb') as f: + f.write(b'0' * 64) + net_plain.network_on_disk.filename = lora_file + activate(net_plain, net_full) + first = dq(layer) + activate() + with counting_calc() as calls: + activate(net_plain, net_full) # the factorable member re-extracts from its own weights; the hosted remainder replays + assert calls['n'] == 0, 'a mixed-set replay must not assemble the delta' + assert hasattr(layer, 'sdnq_lora_svd_stash') + assert torch.equal(dq(layer), first), 'mixed-set replay must be bit-identical to the fresh apply' + activate() + return True + + CAT_COMPILE = category('compile') @@ -1419,7 +1499,7 @@ def run_tests(): for fn in [test_hosted_low_rank_delta_is_kept, test_hosted_dense_delta_beats_requant, test_hosted_skips_int8, 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_dense_stack_keeps_hosting, + test_route_mixed_set_keeps_hosting, test_route_svd_checkpoint_keeps_hosting, test_route_replay_from_cache]: run_test(CAT_HOST, fn) log.warning('=== Calibration ===') @@ -1429,7 +1509,8 @@ def run_tests(): run_test(CAT_CALIB, fn) 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_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]: run_test(CAT_FCACHE, fn) log.warning('=== Compile ===') for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]: From 2306867b63e38b2b66184cdcebf1606c36f02cb0 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 27 Jul 2026 21:49:49 +0100 Subject: [PATCH 8/8] 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]: