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

@@ -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

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

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,