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,