"""Per-checkpoint activation calibration for svd hosting on quantized layers. Plain svd truncation of a hosted delta is optimal in weight space but not in output space: transformer activations concentrate energy in a few input channels (per-channel RMS spreads by one to three orders of magnitude), so the directions that matter most for the output are not the largest in Frobenius norm. Scaling the delta by per-channel input RMS before the svd and folding the inverse scale into the down factor spends the same rank budget on output error instead; measured on real checkpoints this raises output-delta retention by ~0.05 at rank 256 and ~0.09 at rank 64, most on 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, 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 from typing import Optional, TypedDict import torch from modules import paths, shared, script_callbacks from modules.logger import log class CaptureRecord(TypedDict): m: torch.nn.Module ss: Optional[torch.Tensor] n: int done: bool 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': [], 'forwards': 0, 'complete': False} def enabled(): return bool(getattr(shared.opts, 'lora_sdnq_host_calib', False)) def calib_file(model_name): key = model_name.replace('/', '--').replace('\\', '--').replace(':', '-') return os.path.join(calib_root, f'{key}.safetensors') def checkpoint_name(sd_model): info = getattr(sd_model, 'sd_checkpoint_info', None) 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 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 root.named_modules(): deq = getattr(m, 'sdnq_dequantizer', None) if deq is None or len(deq.original_shape) != 2: continue if dtype_dict[deq.weights_dtype]['num_bits'] >= 8: continue out.append((name, m)) return out def detach_capture(): for h in capture['handles']: h.remove() 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']: return x = hook_args[0] if hook_args else None if not torch.is_tensor(x) or x.shape[-1] != in_features: return ss = x.detach().reshape(-1, in_features).float().square().sum(dim=0) if rec['ss'] is None: rec['ss'] = ss else: if rec['ss'].device != ss.device: # offload moves blocks between devices mid-run rec['ss'] = rec['ss'].to(ss.device) rec['ss'] += ss rec['n'] += x.numel() // in_features if rec['n'] >= TOKENS_DONE: rec['done'] = True if all(r['done'] for r in capture['recs'].values()): persist() return hook def persist(): """Write accumulated statistics and stamp them onto the layers. 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). """ if capture['complete']: return capture['complete'] = True from safetensors.torch import save_file tensors, min_n = {}, None for name, rec in capture['recs'].items(): 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)}/{len(capture["recs"])} tokens={min_n} saved="{path}"') except Exception as e: log.warning(f'Network calibration: save failed path="{path}" {e}') def maybe_detach(): """Remove inert hooks once capture finished; safe only outside a model forward.""" if capture['complete'] and capture['handles']: detach_capture() def load_stats(model_name, modules_list): from safetensors import safe_open path = calib_file(model_name) loaded = 0 with safe_open(path, framework='pt', device='cpu') as f: keys = set(f.keys()) for name, m in modules_list: if name in keys: m.sdnq_calib_rms = f.get_tensor(name) loaded += 1 log.info(f'Network calibration: model="{model_name}" layers={loaded} loaded="{path}"') def on_model_loaded(sd_model): detach_capture() if not enabled(): return name = checkpoint_name(sd_model) if name is None: return modules_list = eligible_modules(sd_model) if len(modules_list) == 0: return if os.path.isfile(calib_file(name)): load_stats(name, modules_list) return 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 capture['handles'].append(m.register_forward_pre_hook(hook_for(rec, int(m.sdnq_dequantizer.original_shape[-1])))) log.info(f'Network calibration: model="{name}" layers={len(modules_list)} collecting activation statistics') def rms_for(layer): """Per-channel input RMS for a layer, or None when absent or disabled.""" maybe_detach() if not enabled(): return None return getattr(layer, 'sdnq_calib_rms', None) script_callbacks.on_model_loaded(on_model_loaded)