diff --git a/cli/lora-quant-fidelity.py b/cli/lora-quant-fidelity.py new file mode 100644 index 000000000..74bd214cd --- /dev/null +++ b/cli/lora-quant-fidelity.py @@ -0,0 +1,483 @@ +#!/usr/bin/env python +"""LoRA fidelity analyzer for quantized base models. + +Measures, in weight space, how faithfully a LoRA lands on an SDNQ-quantized +model. Every targeted module is rebuilt with the loader's own module class and +its delta taken from the production ``calc_updown``, so all adapter families +(LoRA, LoKR, LoHA, OFT, full, IA3, GLoRA, norm, plus DoRA and bias variants) +are measured as they would actually apply: + +- factor path (plain additive LoRA riding the svd side-channel): storage is + lossless; the reported figure is the delta realized through the result-dtype + materialize, the same bf16 rounding an unquantized model applies. + Eligibility is decided by the loader's own predicate. +- hosted path (non-factorable families on sub-8-bit formats): the seeded svd + truncation at ``--host-rank``, realized the same way. +- requantize path (all other fallbacks): retention ``rho`` of the intended + delta. On-grid rounding erases sub-step deltas down to a ``2/group_size`` + floor, so low-bit formats (<=6 bits) typically show rho ~= 0.02-0.03. +- unquantized modules: the LoRA applies exactly regardless. + +Reported fidelity is per-module ``applied_rho`` (the measured figure for +whichever path the loader would take), summarized as a median and an +energy-weighted mean over the file's modules; ``requant_rho`` always carries +the if-merged figure. + +Works offline against a pre-quantized SDNQ repo (stored tensors + config) or +a bf16 repo with simulated quantization settings, so a combination can be +assessed before committing to a quantized checkpoint. + +Examples: + python cli/lora-quant-fidelity.py --model vladmandic/Krea-2-Base-sdnq-hadamard-uint4 --arch krea2 --lora "~/models/Lora/Krea 2/krea2_turbo_distill_r256.safetensors" + python cli/lora-quant-fidelity.py --model CalamitousFelicitousness/Krea-2-Base-Diffusers --arch krea2 --dtype uint4 --lora lora.safetensors --json report.json +""" + +import os +import sys +import json +import argparse + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +os.environ.setdefault('SD_INSTALL_QUIET', '1') + + +def parse_cli(): + parser = argparse.ArgumentParser(description='lora-quant-fidelity') + parser.add_argument('--model', required=True, help='model dir, transformer dir, or org/name repo id') + parser.add_argument('--arch', default='generic', help='lora key resolver: a native arch (e.g. krea2, zimage, f2) or generic') + parser.add_argument('--lora', required=True, nargs='+', help='lora safetensors file(s)') + parser.add_argument('--dtype', default=None, help='simulate quantization of a bf16 repo at this sdnq dtype (e.g. uint4, int8); bf16 measures the unquantized reference') + parser.add_argument('--group', type=int, default=0, help='sdnq group_size for simulation') + parser.add_argument('--hadamard-group', type=int, default=256, help='sdnq hadamard group for simulation') + parser.add_argument('--sample', type=int, default=40, help='max modules analyzed per lora (evenly sampled)') + parser.add_argument('--full', action='store_true', help='analyze every matched module') + parser.add_argument('--json', default=None, help='write full report to this json file') + parser.add_argument('--host-rank', type=int, default=256, help='svd hosting cap for non-factorable modules on sub-8-bit formats, mirroring lora_sdnq_host_rank; 0 scores the requantize path instead') + parser.add_argument('--calib', default=None, help='activation statistics file (models/calibration/*.safetensors): hosting truncation is then channel-weighted as with lora_sdnq_host_calib, and hosted rho is measured in the activation-weighted norm (the output-error proxy)') + parser.add_argument('--fail-under', type=float, default=None, help='exit 2 when median applied fidelity of any lora is below this') + return parser.parse_args() + + +cli_args = parse_cli() +sys.argv = [sys.argv[0]] # sdnext arg parsing during imports must not see tool args (prefix matching eats --model/--lora) + +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +modules.cmd_args.parse_args() +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _unknown = modules.cmd_args.parser.parse_known_args([]) + +import torch # pylint: disable=wrong-import-position +from safetensors import safe_open # pylint: disable=wrong-import-position +from rich import print as rprint # pylint: disable=wrong-import-position + +from modules.lora import native_adapter, network, network_lora, network_lokr, network_hada, network_oft, network_full, network_ia3, network_glora, network_norm, lora_sdnq # pylint: disable=wrong-import-position +from modules.lora.lora_load import NATIVE_DISPATCH # pylint: disable=wrong-import-position +from sdnq.quantizer import sdnq_quantize_layer_weight # pylint: disable=wrong-import-position +from sdnq.quant_utils import rotate_hadamard # pylint: disable=wrong-import-position + + +MODEL_ROOTS = [ + os.path.expanduser('~/database/models/huggingface'), + os.path.expanduser('~/database/models/Diffusers'), +] +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + +# every adapter family the native loader can build, with the module class that owns its +# apply-time math. deltas are taken from the production calc_updown so the tool cannot +# drift from the loader, and eligibility is decided by the production predicate itself. +FAMILY_SPECS = ( + ('lora', network_lora.NetworkModuleLora, native_adapter.LORA_SUFFIXES, native_adapter.LORA_MARKERS), + ('lokr', network_lokr.NetworkModuleLokr, native_adapter.LOKR_SUFFIXES, native_adapter.LOKR_MARKERS), + ('loha', network_hada.NetworkModuleHada, native_adapter.LOHA_SUFFIXES, native_adapter.LOHA_MARKERS), + ('oft', network_oft.NetworkModuleOFT, native_adapter.OFT_SUFFIXES, native_adapter.OFT_MARKERS), + ('full', network_full.NetworkModuleFull, native_adapter.FULL_SUFFIXES, native_adapter.FULL_MARKERS), + ('ia3', network_ia3.NetworkModuleIa3, native_adapter.IA3_SUFFIXES, native_adapter.IA3_MARKERS), + ('glora', network_glora.NetworkModuleGLora, native_adapter.GLORA_SUFFIXES, native_adapter.GLORA_MARKERS), + ('norm', network_norm.NetworkModuleNorm, native_adapter.NORM_SUFFIXES, native_adapter.NORM_MARKERS), +) + + +class StubOnDisk: + def __init__(self, path): + self.filename = path + self.name = os.path.splitext(os.path.basename(path))[0] + self.shorthash = '' + self.sd_version = 'unknown' + + +def resolve_model_dir(spec): + """Return the transformer directory for a local path or org/name repo id.""" + candidates = [spec, os.path.join(spec, 'transformer')] + cache_name = 'models--' + spec.replace('/', '--') + for root in MODEL_ROOTS: + snap_root = os.path.join(root, cache_name, 'snapshots') + if os.path.isdir(snap_root): + for snap in sorted(os.listdir(snap_root), reverse=True): + candidates.append(os.path.join(snap_root, snap, 'transformer')) + candidates.append(os.path.join(snap_root, snap)) + for c in candidates: + if os.path.isfile(os.path.join(c, 'config.json')): + return c + raise SystemExit(f'model not found: {spec}') + + +def resolve_arch(name): + """Return the arch lora module for key resolution, or None for generic matching.""" + if name == 'generic': + return None + path = NATIVE_DISPATCH.get({'flux2': 'f2', 'ernie': 'ernieimage'}.get(name, name)) + if path is None: + raise SystemExit(f'unknown arch {name}; choices: {sorted(NATIVE_DISPATCH)} or generic') + import importlib + return importlib.import_module(path) + + +def map_lora_modules(lora_path, arch_mod): + """Return {model_module_path: (family, weights)} across every adapter family, plus a census. + + Grouping mirrors the native loader: a family is only considered when its + marker is present, and groups resolve to model paths through the arch's own + resolver. Fused-split chunks are counted but not analyzed (their apply-time + math is arch-owned). + """ + with safe_open(lora_path, framework='pt', device='cpu') as f: + state_dict = {k: f.get_tensor(k) for k in f.keys()} + prefixes = getattr(arch_mod, 'KNOWN_PREFIXES', native_adapter.KNOWN_PREFIXES_DEFAULT) + bare = getattr(arch_mod, 'BARE_DIFFUSERS_PREFIXES', ()) + resolve = getattr(arch_mod, 'resolve_targets', None) or (lambda prefix, base: [(base, None)]) + mapped, census, chunked = {}, {}, 0 + for fam, _cls, suffixes, markers in FAMILY_SPECS: + if not native_adapter.has_marker(state_dict, markers): + continue + groups = native_adapter.group_by_suffixes(state_dict, suffixes, prefixes=prefixes, bare_diffusers_prefixes=bare) + if fam == 'lora': + groups = {k: w for k, w in groups.items() if 'lora_down.weight' in w and 'lora_up.weight' in w} + else: + groups = {k: w for k, w in groups.items() if native_adapter.has_marker({f'x.{s}': None for s in w}, markers)} + if not groups: + continue + census[fam] = len(groups) + for (prefix, base), w in groups.items(): + for path, chunk in native_adapter.resolve_group_targets(resolve, prefix, base): + if chunk is not None: + chunked += 1 + continue + mapped.setdefault(path, []).append((fam, w)) # a module can carry several families; the loader applies each + return mapped, census, chunked + + +def stamp_index(paths): + """Map each module path to its stamped form, the way the loader matches. + + The loader compares ``network_prefix + path.replace('.', '_')`` against each + module's stamped ``network_layer_name``, so kohya-style ``lora_unet_`` keys + (whose base arrives already underscored) resolve fine there. Matching on the + stamped form reproduces that and keeps dotted bases working unchanged. + """ + return {p.replace('.', '_'): p for p in paths} + + +def make_stub(shape, dtype=torch.bfloat16): + """Minimal sd_module standing in for a bf16 repo weight: the module classes key off its type and shape.""" + if len(shape) == 2: + return torch.nn.Linear(shape[1], shape[0], bias=False, dtype=dtype, device='meta') + return torch.nn.Conv2d(shape[1], shape[0], shape[2:], bias=False, dtype=dtype, device='meta') + + +def build_module(fam, path, w, net, sd_module): + """Instantiate the family's production NetworkModule for one target.""" + cls = next(c for f, c, _s, _m in FAMILY_SPECS if f == fam) + weights = network.NetworkWeights(network_key=path, sd_key=path, w=w, sd_module=sd_module) + return cls(net, weights) + + +def resolve_transformer_cls(arch, class_name): + """Prefer an sdnext-owned transformer class over the upstream diffusers one. + + Arches like krea2 keep checkpoint-style module names in their own class; + the diffusers class of the same name expects diffusers-style keys and + cannot load these state dicts. + """ + if arch and class_name: + try: + import importlib + pkg = importlib.import_module(f'pipelines.{ {"zimage": "z_image", "f2": "flux"}.get(arch, arch) }') + for attr in dir(pkg): + if attr.endswith('_SPEC'): + cls = getattr(getattr(pkg, attr), 'cls', None) + if cls is not None and cls.__name__ == class_name: + return cls + except Exception: + pass + return None + + +def load_quantized_model(model_dir, arch=None, class_name=None): + from sdnq.loader import load_sdnq_model + model = load_sdnq_model(model_dir, model_cls=resolve_transformer_cls(arch, class_name), dtype=torch.bfloat16, device='cpu') + layers = {} + for name, module in model.named_modules(): + if getattr(module, 'sdnq_dequantizer', None) is not None: + layers[name] = module + elif module.__class__.__name__ == 'Linear' and getattr(module, 'weight', None) is not None: + layers[name] = module + del model # layer modules own their tensors; the dict keeps them alive + return layers + + +class Bf16Repo: + """Lazy per-module weight access for a sharded bf16 transformer repo.""" + + def __init__(self, model_dir): + self.model_dir = model_dir + self.handles = {} # reopening a multi-gb shard per module dominates runtime over many loras + index = os.path.join(model_dir, 'diffusion_pytorch_model.safetensors.index.json') + if os.path.isfile(index): + with open(index, encoding='utf-8') as f: + self.weight_map = json.load(f)['weight_map'] + else: + single = os.path.join(model_dir, 'diffusion_pytorch_model.safetensors') + with safe_open(single, framework='pt', device='cpu') as f: + self.weight_map = dict.fromkeys(f.keys(), 'diffusion_pytorch_model.safetensors') + + def get(self, key): + shard = self.weight_map.get(key) + if shard is None: + return None + f = self.handles.get(shard) + if f is None: + f = safe_open(os.path.join(self.model_dir, shard), framework='pt', device='cpu') + self.handles[shard] = f + return f.get_tensor(key) + + +def analyze_module(W_dq, deq_params, mods, calib_rms=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 + loader stacks them, so every family (and dora / dense-bias / diff_b variant) + 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. + """ + D = None + for mod in mods: + d = mod.calc_updown(W_dq)[0].to(device, torch.float32).reshape(W_dq.shape) + D = d if D is None else D + d + nD = D.norm() + control = deq_params['weights_dtype'] == 'bf16' # unquantized reference: the delta just rounds into bf16 + factor_eligible = (not control) and all(lora_sdnq.get_module_factors(m, device, torch.bfloat16) is not None for m in mods) + if float(nD) == 0.0: # an all-zero delta (some full-rank extractions carry empty .diff): retention is undefined, not erased + return dict(rank=getattr(mods[0], 'dim', None), rms_delta=0.0, rms_weight=float(W_dq.pow(2).mean().sqrt()), + step_ratio=None, crossers=None, requant_rho=None, requant_resid=None, + factor_eligible=factor_eligible, hosted=False, applied_rho=None, delta_energy=0.0) + step_ratio, crossers = None, None + if control: + W2 = (W_dq + D).to(torch.bfloat16).float() + else: + # mirror network_add_weights: it requantizes with the layer's own svd setting and rank, + # and an svd checkpoint's dequantized weight is not on the plain integer grid + use_svd = deq_params.get('use_svd', False) + kw = dict(layer_class_name='Linear', torch_dtype=torch.bfloat16, group_size=deq_params['group_size'], + hadamard_group_size=deq_params['hadamard_group_size'], use_hadamard=deq_params['use_hadamard'], + weights_dtype=deq_params['weights_dtype'], use_svd=use_svd, svd_rank=deq_params.get('svd_rank', 32), + svd_steps=deq_params.get('svd_steps', 8), use_quantized_matmul=False, dequantize_fp32=False) + deq2, data2 = sdnq_quantize_layer_weight(W_dq + D, **kw) + W2 = deq2(data2['weight'], data2['scale'], zero_point=data2['zero_point'], + svd_up=data2['svd_up'], svd_down=data2['svd_down'], dtype=torch.float32, skip_compile=True) + Dh = rotate_hadamard(D, group_size=deq_params['hadamard_group_size']) if deq_params['use_hadamard'] else D + step = data2['scale'].float() + Dg = Dh.unflatten(-1, (step.shape[1], -1)) if step.ndim == 3 else Dh + step_ratio = float((Dg.abs() / step).mean()) + crossers = float((Dg.abs() > step / 2).float().mean()) + E = W2 - W_dq + rho = float(E.flatten() @ D.flatten() / nD.square()) + resid = float((E - D).norm() / nD) + hosted = False + if factor_eligible: + # the factor path stores the delta losslessly, but the dequantizer materializes + # base + factors in the result dtype (bf16 here), so realized fidelity floors at + # the same ULP rounding an unquantized bf16 model applies to a merged delta + base16 = W_dq.to(torch.bfloat16).float() + realized = (W_dq.to(torch.bfloat16) + D.to(torch.bfloat16)).float() - base16 + applied_rho = float(realized.flatten() @ D.flatten() / nD.square()) + else: + applied_rho = rho + if (not control) and cli_args.host_rank > 0: + from sdnq.common import dtype_dict + if dtype_dict[deq_params['weights_dtype']]['num_bits'] < 8: + # mirror lora_sdnq.apply_hosted: seeded svd truncation, realized through the bf16 materialize + q = min(cli_args.host_rank, *D.shape) + rms = None + if calib_rms is not None and calib_rms.shape[-1] == D.shape[-1]: + rms = calib_rms.to(D.device, torch.float32).clamp(min=1e-8) + 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) + Dk = (U * S) @ V.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, + delta_energy=float(nD.square())) + + +def main(): + args = cli_args + model_dir = resolve_model_dir(args.model) + arch_mod = resolve_arch(args.arch) + with open(os.path.join(model_dir, 'config.json'), encoding='utf-8') as f: + model_config = json.load(f) + pre_quantized = model_config.get('quantization_config') is not None + + quant_layers, bf16_repo = {}, None + quant_stamps, bf16_stamps = {}, {} + if pre_quantized: + rprint(f'model: "{model_dir}" pre-quantized={pre_quantized}') + quant_layers = load_quantized_model(model_dir, arch=args.arch, class_name=model_config.get('_class_name')) + quant_stamps = stamp_index(quant_layers) + else: + bf16_repo = Bf16Repo(model_dir) + bf16_stamps = stamp_index(k[:-len('.weight')] for k in bf16_repo.weight_map if k.endswith('.weight')) + if args.dtype is None: + rprint('model is not quantized and no --dtype given: loras apply exactly, nothing to analyze') + return 0 + rprint(f'model: "{model_dir}" simulating dtype={args.dtype} group={args.group} hadamard={args.hadamard_group}') + + calib_stats = {} + if args.calib: + with safe_open(os.path.expanduser(args.calib), framework='pt', device='cpu') as f: + calib_stats = {k: f.get_tensor(k) for k in f.keys()} + rprint(f'calib: "{args.calib}" layers={len(calib_stats)}') + + report = {'model': model_dir, 'pre_quantized': pre_quantized, 'loras': []} + worst_effective = 1.0 + def write_report(): + if args.json: + with open(args.json, 'w', encoding='utf-8') as f: + json.dump(report, f, indent=2) + + for lora_path in args.lora: + lora_path = os.path.expanduser(lora_path) + try: + mapped, census, chunked = map_lora_modules(lora_path, arch_mod) + net = network.Network(os.path.basename(lora_path), StubOnDisk(lora_path)) + rows, unquantized, unmatched, failed, non_matrix = [], [], [], [], [] + keys = sorted(mapped) + if not args.full and len(keys) > args.sample: + keys = keys[::max(1, len(keys) // args.sample)][:args.sample] + for path in keys: + entries = mapped[path] + lname = path + if pre_quantized: + if path not in quant_layers: + lname = quant_stamps.get(path.replace('.', '_'), '') + layer = quant_layers.get(lname) + if layer is None: + unmatched.append(path) + continue + deq = getattr(layer, 'sdnq_dequantizer', None) + if deq is None: + unquantized.append(path) + continue + if len(deq.original_shape) != 2: + non_matrix.append(path) + continue + W_dq = deq(layer.weight, layer.scale, zero_point=layer.zero_point, svd_up=layer.svd_up, svd_down=layer.svd_down, + 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) + sd_module = layer + else: + W = bf16_repo.get(f'{path}.weight') + if W is None: + W = bf16_repo.get(f'{bf16_stamps.get(path.replace(".", "_"), "")}.weight') + if W is None: + unmatched.append(path) + continue + if W.ndim != 2: # norm/scale targets (e.g. adaLN_modulation) are 1-D; the quantizer and the stub both expect a matrix + non_matrix.append(path) + continue + 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) + 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) + 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)) + 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 + continue + row.update(module=path, dtype=params['weights_dtype'], family='+'.join(f for f, _w in entries)) + rows.append(row) + del W_dq # the caching allocator reuses these; emptying it per module costs more than it saves + + scored = [r for r in rows if r['applied_rho'] is not None] # zero-delta modules have no retention to report + applied = sorted(r['applied_rho'] for r in scored) + median_applied = applied[len(applied) // 2] if applied else None + energy = sum(r['delta_energy'] for r in scored) + weighted = (sum(r['applied_rho'] * r['delta_energy'] for r in scored) / energy) if energy > 0 else None + n_exact = sum(1 for r in scored if r['factor_eligible']) + fb = [r['requant_rho'] for r in scored if not r['factor_eligible']] + fb_median = sorted(fb)[len(fb) // 2] if fb else None + if median_applied is not None: + worst_effective = min(worst_effective, median_applied) + report['loras'].append({'file': lora_path, 'families': census, 'targets': len(mapped), 'unquantized': unquantized, + 'unmatched': unmatched, 'non_matrix': non_matrix, 'chunked': chunked, 'failed': failed, + 'exact_modules': n_exact, 'fallback_modules': len(fb), 'fallback_median_rho': fb_median, + 'median_applied_rho': median_applied, 'weighted_applied_rho': weighted, 'modules': rows}) + write_report() # rewrite per file so a crash keeps completed work + rprint(f'\nlora: "{os.path.basename(lora_path)}" families={census or "none"} targets={len(mapped)} analyzed={len(rows)} scored={len(scored)} exact={n_exact} fallback={len(fb)} unquantized={len(unquantized)} unmatched={len(unmatched)} non_matrix={len(non_matrix)} chunked={chunked} failed={len(failed)}') + if median_applied is None: + rprint(' no analyzable modules: nothing measured') + else: + rprint(f' applied fidelity: median={median_applied:.3f} energy-weighted={weighted:.3f}' + (f' (fallback modules land at median rho={fb_median:.3f})' if fb_median is not None else '')) + for f in failed[:3]: + rprint(f' [red]could not rebuild[/red]: {f}') + if fb: + worst = sorted((r for r in scored if not r['factor_eligible']), key=lambda r: r['requant_rho'])[:5] + rprint(' lowest-retention modules:') + for r in worst: + grid = f'step-ratio={r["step_ratio"]:.3f} crossers={r["crossers"]*100:5.1f}%' if r['step_ratio'] is not None else 'unquantized reference' + rprint(f' {r["module"]:48s} fam={r["family"]:5s} dtype={r["dtype"]} {grid} rho={r["requant_rho"]:.3f}') + del mapped, net + except KeyboardInterrupt: + raise + except Exception as e: # one broken file must not cost the rest of the batch + rprint(f'\n[red]lora failed[/red]: "{os.path.basename(lora_path)}" {type(e).__name__}: {e}') + report['loras'].append({'file': lora_path, 'error': f'{type(e).__name__}: {e}'}) + write_report() + if device.type == 'cuda': + torch.cuda.empty_cache() # once per file, after its modules are done + + report['complete'] = True + write_report() + if args.json: + rprint(f'\nreport: "{args.json}"') + if args.fail_under is not None and worst_effective < args.fail_under: + rprint(f'FAIL: effective fidelity {worst_effective:.3f} < {args.fail_under}') + return 2 + return 0 + + +if __name__ == '__main__': + with torch.inference_mode(): + sys.exit(main()) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index 35605e95e..d0a6ffe85 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -178,6 +178,8 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers, strict=False)] def changed(self, requested: list[str], include: list[str] | None = None, exclude: list[str] | None = None) -> tuple[bool, str]: + from modules.lora import lora_sdnq + requested = requested + [f'stack={lora_sdnq.signature()}'] # settings-only mechanism changes must re-trigger activation if shared.opts.lora_force_reload: debug_log(f'Network check: type=LoRA requested={requested} status="forced"') return True, "forced" @@ -254,7 +256,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): if has_changed: jobid = shared.state.begin('LoRA') if len(l.previously_loaded_networks) > 0: - log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_native else "backup"}') + log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if lora_overrides.fuse_native() else "backup"}') networks.network_deactivate(include, exclude) networks.network_activate(include, exclude) debug_log(f'Network change: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]}') @@ -267,7 +269,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): prompt(p) if has_changed and len(include) == 0: # print only once actual_method = 'native' if any(len(n.modules) > 0 for n in l.loaded_networks) else load_method - log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if shared.opts.lora_fuse_native else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') + log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if lora_overrides.fuse_native() else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"') def deactivate(self, p, force=False): if len(lora_diffusers.diffuser_loaded) > 0 and (shared.opts.lora_force_reload or force): diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index c1b46c4ab..3cb8688c0 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)") -def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, wanted_names: tuple): +def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, wanted_names: tuple, fuse: bool): backup_size = 0 if len(l.loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in l.loaded_networks]): # noqa: C419 # pylint: disable=R1729 t0 = time.time() @@ -24,7 +24,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr weights_backup = getattr(self, "network_weights_backup", None) bias_backup = getattr(self, "network_bias_backup", None) if weights_backup is not None or bias_backup is not None: - if (shared.opts.lora_fuse_native and not isinstance(weights_backup, bool)) or (not shared.opts.lora_fuse_native and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly + if (fuse and not isinstance(weights_backup, bool)) or (not fuse and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly weights_backup = None bias_backup = None self.network_weights_backup = weights_backup @@ -33,7 +33,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr if weights_backup is None and wanted_names != (): # pylint: disable=C1803 weight = getattr(self, 'weight', None) self.network_weights_backup = None - if shared.opts.lora_fuse_native: + if fuse: self.network_weights_backup = True else: self.network_weights_backup = weight.clone().to(devices.cpu) @@ -53,7 +53,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr if bias_backup is None: if getattr(self, 'bias', None) is not None: - if shared.opts.lora_fuse_native: + if fuse: self.network_bias_backup = True else: bias_backup = self.bias.clone() diff --git a/modules/lora/lora_calib.py b/modules/lora/lora_calib.py new file mode 100644 index 000000000..d854016d9 --- /dev/null +++ b/modules/lora/lora_calib.py @@ -0,0 +1,175 @@ +"""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 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. +""" + +import os + +import torch + +from modules import paths, shared, script_callbacks +from modules.logger import log + + +TOKENS_DONE = 65536 +calib_root = os.path.join(paths.models_path, 'calibration') +capture = {'model': None, 'recs': {}, 'handles': [], '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 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: + return [] + from sdnq.common import dtype_dict + out = [] + for name, m in transformer.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['complete'] = False + + +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 completed 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 + 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(): + rms = (rec['ss'] / max(rec['n'], 1)).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']) + 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}"') + 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 + 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) diff --git a/modules/lora/lora_diffusers.py b/modules/lora/lora_diffusers.py index 94ad716c7..e8321f8a0 100644 --- a/modules/lora/lora_diffusers.py +++ b/modules/lora/lora_diffusers.py @@ -4,6 +4,7 @@ import diffusers from modules import shared, errors from modules.logger import log from modules.lora import network +from modules.lora import lora_overrides from modules.lora import lora_common as l @@ -54,7 +55,7 @@ def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale t0 = time.time() name = name.replace(".", "_") sd_model: diffusers.DiffusionPipeline = getattr(shared.sd_model, "pipe", shared.sd_model) - log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason={reason or "unknown"} scale={lora_scale} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}') + log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason={reason or "unknown"} scale={lora_scale} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}') if not hasattr(sd_model, 'load_lora_weights'): log.error(f'Network load: type=LoRA class={sd_model.__class__} does not implement load lora') return None diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index 7993ffe9a..b297108bd 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -149,7 +149,7 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne if l.debug: log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}') else: - log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}') + log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}') if len(matched_networks) == 0: return None lora_cache[name] = net @@ -350,7 +350,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non networks.network_activate() if len(l.loaded_networks) > 0 and l.debug: - log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}') + log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}') if recompile_model: log.info("Network load: type=LoRA recompiling model") diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index 2d27f7b88..eb629896e 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -77,13 +77,56 @@ def get_method(shorthash=''): return 'native', 'default' +# Roles a LoRA is fused into; a quantized component in any of them makes fusing unsafe. +fuse_roots = ('transformer', 'unet', 'text_encoder', 'llm_adapter') + + +def fuse_components(sd_model): + """Component names a network fuses into, matched by role prefix so numbered and reference siblings are covered.""" + names = getattr(sd_model, 'components', None) + if not isinstance(names, dict): + names = vars(sd_model) + return [name for name in names if name.startswith(fuse_roots)] + + +def is_quantized(module): + """Return True when ``module`` carries a quantization config. + + ``config.quantization_config`` is read first: SDNQ sets both it and the plain + attribute when it quantizes in place, but a checkpoint that ships pre-quantized + only reaches the plain attribute through the diffusers ConfigMixin name proxy, + which is deprecated for removal. + """ + if module is None: + return False + config = getattr(module, 'config', None) + if config is not None and getattr(config, 'quantization_config', None) is not None: + return True + return getattr(module, 'quantization_config', None) is not None + + def disable_fuse(): - if hasattr(shared.sd_model, 'quantization_config'): + """Return True when fusing a network into model weights is unsafe. + + Fusing keeps no pristine copy of the weight, so each apply and restore + round-trips it through its storage format. On quantized weights that is a + dequantize-add-requantize cycle per network swap whose error compounds. + """ + sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model) + if is_quantized(sd_model): return True - if hasattr(shared.sd_model, 'transformer') and hasattr(shared.sd_model.transformer, 'quantization_config'): + if any(is_quantized(getattr(sd_model, name, None)) for name in fuse_components(sd_model)): return True - if hasattr(shared.sd_model, 'transformer_2') and hasattr(shared.sd_model.transformer_2, 'quantization_config'): - return True - if hasattr(shared.sd_model, '_lora_partial'): + if hasattr(sd_model, '_lora_partial'): return True return shared.sd_model_type in fuse_ignore + + +def fuse_native(): + """Return True when the native apply path may fuse into model weights. + + The single source of truth for the native fuse decision: it must agree across + the backup, activate and deactivate passes, since backup mode restores from a + stored tensor while fuse mode restores by subtracting the delta. + """ + return shared.opts.lora_fuse_native and not disable_fuse() diff --git a/modules/lora/lora_sdnq.py b/modules/lora/lora_sdnq.py new file mode 100644 index 000000000..63597111b --- /dev/null +++ b/modules/lora/lora_sdnq.py @@ -0,0 +1,299 @@ +"""Exact LoRA application for SDNQ-quantized layers. + +Baking a LoRA into a quantized weight requantizes it: dequantize, add the +delta, re-round onto the integer grid. When the per-element delta is smaller +than half a quantization step (a rank-decomposed delta on a uint4 layer sits +at a few percent of a step), rounding erases it; what survives is the two +grid-extrema elements per quantization group (2/group_size of the signal) +plus grid-shift noise of the same norm as the delta. The optimal in-grid +representation provably retains ~0%, so no rewrite of the stored integers +can fix this. + +The exact path instead rides the SDNQ svd side-channel: the dequantizer +computes ``W = dq(q) + svd_up @ svd_down`` in the rotated domain at full +precision, in every forward mode. A LoRA delta ``B @ A`` is appended as +extra columns of ``svd_up`` and rows of ``svd_down``; because the Hadamard +rotation is block-diagonal, symmetric and self-inverse, storing ``A·H`` for +the down factor makes the round trip exact: ``(B @ (A·H)) · H = B @ A``. +Quantized weights are never touched, so apply and remove are exact and no +weight backup is needed. The side-channel storage is lossless; realized +fidelity floors at the compute dtype, because the dequantizer materializes +``base + factors`` in the result dtype and a delta below its ULP of the +base rounds exactly as it would on an unquantized model of that dtype. + +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 +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 +``lora_calib``), the truncation is channel-weighted to minimize output +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. +""" + +import torch + +from modules import devices, shared +from modules.lora import lora_calib +from modules.lora import lora_common as l +from modules.logger import log + + +fallback_layers: list[str] = [] +hosted_layers: list[tuple[str, float, bool]] = [] + + +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.""" + return getattr(shared.opts, 'lora_sdnq_apply', 'exact') != 'requantize' + + +def signature(): + """Identity suffix for the per-module apply stamp; empty on the default exact mechanism.""" + return '' if enabled() else '|quant=requantize' + + +def get_module_factors(module, device, dtype, original_shape=None): + """Return ``(up_eff, down)`` reproducing ``calc_updown`` exactly, or None. + + ``updown = up @ down * calc_scale() * multiplier()`` for a plain linear + LoRA; the scalars fold into the up factor. ``dyn_dim`` slices ranks the + same way ``lyco_helpers.rebuild_conventional`` does. + """ + if module.__class__.__name__ != 'NetworkModuleLora': + return None + if module.dora_scale is not None or module.bias is not None or module.ex_bias is not None: + return None + if getattr(module, 'mid_model', None) is not None: + return None + up = module.up_model.weight + down = module.down_model.weight + if up.ndim != 2 or down.ndim != 2: + return None + if original_shape is not None and (up.shape[0] != original_shape[0] or down.shape[1] != original_shape[-1]): + return None # factor_candidate skips shape checks for layers already in factor mode; recheck here so a malformed stack falls back instead of raising in cat + dyn_dim = module.network.dyn_dim + if dyn_dim is not None and up.shape[1] != dyn_dim: + up = up[:, :dyn_dim] + down = down[:dyn_dim] + scalar = module.calc_scale() * module.multiplier() + up_eff = up.to(device=device, dtype=torch.float32) * scalar + return up_eff.to(dtype=dtype), down.to(device=device, dtype=dtype) + + +def factor_candidate(self, network_layer_name, wanted_names, use_previous=False): + """True when this layer should take the exact svd-append path. + + Requires an SDNQ linear layer whose active networks all contribute plain + factorable LoRA modules for this layer. An empty ``wanted_names`` is a + removal request and qualifies whenever factors are currently attached. + """ + if not enabled(): + return False # declined layers with factors still attached are stripped by the activate fallthrough + if getattr(self, 'sdnq_dequantizer', None) is None or self.__class__.__name__ != 'SDNQLinear': + return False + if hasattr(self, 'sdnq_lora_svd_stash'): + 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: + module = net.modules.get(network_layer_name, None) + if module is None: + continue + seen = True + if module.__class__.__name__ != 'NetworkModuleLora': + return False + if module.dora_scale is not None or module.bias is not None or module.ex_bias is not None or getattr(module, 'mid_model', None) is not None: + return False + if module.up_model.weight.ndim != 2 or module.down_model.weight.ndim != 2: + return False + if module.up_model.weight.shape[0] != self.sdnq_dequantizer.original_shape[0] or module.down_model.weight.shape[1] != self.sdnq_dequantizer.original_shape[-1]: + return False + return seen + + +def remove_factors(self): + """Restore the layer's original svd factors; True when factors were attached.""" + stash = getattr(self, 'sdnq_lora_svd_stash', None) + if stash is None: + return False + svd_up, svd_down = stash + device = self.scale.device # the stash tuple does not follow module device moves; restore onto wherever the layer lives now + if svd_up is not None and svd_up.device != device: + svd_up = torch.nn.Parameter(svd_up.to(device=device), requires_grad=False) + svd_down = torch.nn.Parameter(svd_down.to(device=device), requires_grad=False) + self.svd_up = svd_up + self.svd_down = svd_down + del self.sdnq_lora_svd_stash + return True + + +def apply_factors(self, network_layer_name, wanted_names, use_previous=False): + """Attach the active networks' LoRA factors to this layer's svd side-channel. + + Replaces any previously attached factors (multiplier changes re-enter + here with a new ``wanted_names`` signature). Returns True when the layer + changed. Falls back to the caller's requantize path by returning None + when factor extraction fails at this stage. + """ + from sdnq.quant_utils import rotate_hadamard + + changed = remove_factors(self) + if wanted_names == (): + return changed + + 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: + 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: + return None + up_eff, down = factors + 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 not ups: + return changed + append_factors(self, ups, downs) + return True + + +def append_factors(self, ups, downs): + """Concatenate ``[out, r]`` / ``[r, in]`` factor pairs onto the layer's svd channel and stash the originals.""" + deq = self.sdnq_dequantizer + device = self.scale.device + dtype = deq.result_dtype + orig_up, orig_down = self.svd_up, self.svd_down + if deq.use_quantized_matmul: + # matmul layout stores factors transposed: svd_up [r, out], svd_down [in, r] + parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + [u.t() for u in ups] + parts_down = ([orig_down.to(device=devices.device, dtype=dtype)] if orig_down is not None else []) + [d.t() for d in downs] + new_up = torch.cat(parts_up, dim=0).contiguous() + new_down = torch.cat(parts_down, dim=1).contiguous() + else: + parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + ups + 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() + 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) + + +def host_candidate(self, network_layer_name, wanted_names, use_previous=False): + """True when a non-factorable set on this layer should be hosted as a truncated svd.""" + if not enabled(): + return False + if int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0) <= 0: + return False + if getattr(self, 'sdnq_dequantizer', None) is None or self.__class__.__name__ != 'SDNQLinear': + return False + if wanted_names == (): + return 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) + + +def apply_hosted(self, network_layer_name, updown, wanted_names, use_previous=False): + """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. + + The delta comes from the families' own ``calc_updown``, so every family + and scaling quirk is included; factorable members are subtracted out and + 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. + """ + from sdnq.quant_utils import rotate_hadamard + + deq = self.sdnq_dequantizer + changed = remove_factors(self) + if wanted_names == (): + return changed + 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) + + ups, downs = [], [] + loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks + for net in loaded: + 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 + 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) + + cap = int(shared.opts.lora_sdnq_host_rank) + q = min(cap, *D.shape) + rms = lora_calib.rms_for(self) + if rms is not None and rms.shape[-1] == D.shape[-1]: + # scale input channels by their activation RMS so truncation minimizes output error rather than weight error + rms = rms.to(device=D.device, dtype=torch.float32).clamp(min=1e-8) + D = D.mul_(rms) + else: + rms = None + # 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) + 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() + if rms is not None: + 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)]) + hosted_layers.append((network_layer_name, energy, rms is not None)) + return True + + +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: + fallback_layers.append(network_layer_name) + + +def report_fallbacks(): + if len(hosted_layers) > 0: + 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') + 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(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)') + else: + log.info(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(fallback_layers)} reason=setting') + if l.debug: + log.debug(f'Network load: type=LoRA quant=sdnq requantized={fallback_layers[:8]}{"..." if len(fallback_layers) > 8 else ""}') + fallback_layers.clear() diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 31945c27d..95c7675a5 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -3,6 +3,8 @@ import time import rich.progress as rp from modules.errorlimiter import limit_errors from modules.lora import lora_common as l +from modules.lora import lora_overrides +from modules.lora import lora_sdnq from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights from modules import shared, devices, sd_models from modules.logger import log, console @@ -53,6 +55,7 @@ def network_activate(include=None, exclude=None): net.unet_multiplier = pending['unet'] net.dyn_dim = pending['dyn'] t0 = time.time() + fuse = lora_overrides.fuse_native() # resolve once: backup and apply passes must agree with limit_errors("network_activate") as elimit: sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) if shared.opts.diffusers_offload_mode == "sequential": @@ -86,7 +89,10 @@ def network_activate(include=None, exclude=None): refused = 0 with devices.inference_context(), pbar: wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () + stack_sig = lora_sdnq.signature() # apply-mechanism token tracked beside network_current_names so a settings-only flip re-applies applied_layers.clear() + lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind + lora_sdnq.hosted_layers.clear() backup_size = 0 for component in modules.keys(): component_wanted = wanted_names if component in components else () @@ -94,13 +100,54 @@ def network_activate(include=None, exclude=None): for _, module in modules[component]: network_layer_name = getattr(module, 'network_layer_name', None) current_names = getattr(module, "network_current_names", ()) - if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == component_wanted): + if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == component_wanted and getattr(module, 'network_current_stack', '') == stack_sig): if task is not None: pbar.update(task, advance=1) 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) - backup_size += network_backup_weights(module, network_layer_name, component_wanted) + 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): + network_apply_weights(module, None, None, device=device) # an earlier non-factorable set requantized this layer, restore the pristine base before attaching factors + applied = lora_sdnq.apply_factors(module, network_layer_name, component_wanted) + if applied is not None: # exact path took the layer; None falls through to hosting or requantize + if applied 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 + if lora_sdnq.host_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): + 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) + 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 + 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 = () + module.network_current_stack = stack_sig + if task is not None: + pbar.update(task, advance=1) + continue + backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse) if not component_wanted: weights_backup = getattr(module, "network_weights_backup", None) if weights_backup is None or isinstance(weights_backup, bool): # fuse mode has no tensor backup, restore stays with network_deactivate @@ -110,7 +157,9 @@ def network_activate(include=None, exclude=None): 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 shared.opts.lora_fuse_native: + 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: weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device) else: weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device) @@ -122,6 +171,7 @@ def network_activate(include=None, exclude=None): batch_updown, batch_ex_bias = None, None del batch_updown, batch_ex_bias module.network_current_names = component_wanted + module.network_current_stack = stack_sig if task is not None: bs = round(backup_size/1024/1024/1024, 2) if backup_size > 0 else None pbar.update(task, advance=1, description=f'networks={len(l.loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={bs} device={device}') @@ -129,13 +179,14 @@ def network_activate(include=None, exclude=None): if task is not None and len(applied_layers) == 0: pbar.remove_task(task) # hide progress bar for no action global native_active, refused_writes # pylint: disable=global-statement + lora_sdnq.report_fallbacks() native_active = len(l.loaded_networks) > 0 refused_writes = refused l.timer.activate += time.time() - t0 if refused > 0: log.error(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} weights={applied_weight} bias={applied_bias} refused={refused} network partially applied') if l.debug and len(l.loaded_networks) > 0: - log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} refused={refused} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') + log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} refused={refused} backup={round(backup_size/1024/1024/1024, 2)} fuse={fuse}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') modules.clear() if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") @@ -146,7 +197,8 @@ def network_deactivate(include=None, exclude=None): exclude = [] if include is None: include = [] - if not shared.opts.lora_fuse_native or shared.opts.lora_force_diffusers: + fuse = lora_overrides.fuse_native() # must match network_activate: backup mode restores in its restore-only pass instead + if not fuse or shared.opts.lora_force_diffusers: return if len(l.previously_loaded_networks) == 0: return @@ -190,8 +242,14 @@ def network_deactivate(include=None, exclude=None): continue if group_offload and component not in group_stripped and group_will_mutate(module, network_layer_name, l.previously_loaded_networks): device = group_offload_strip(sd_model, component, group_stripped) + if lora_sdnq.remove_factors(module): # exact inverse for factor-mode layers, weights were never touched + applied_layers.append(network_layer_name) + module.network_current_names = () + if task is not None: + pbar.update(task, advance=1) + continue batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit) - if shared.opts.lora_fuse_native: + if fuse: weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) else: weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True) @@ -206,7 +264,7 @@ def network_deactivate(include=None, exclude=None): if refused > 0: log.error(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} unapply={len(applied_layers)} refused={refused} network partially removed') if l.debug and len(l.previously_loaded_networks) > 0: - log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} refused={refused} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') + log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} refused={refused} fuse={fuse}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') modules.clear() if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential" or len(group_stripped) > 0: sd_models.set_diffuser_offload(sd_model, op="model") diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 8f8a6be97..89e13de37 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -676,18 +676,6 @@ def create_settings(cmd_opts): "extra_network_reference_enable": OptionInfo(True, "Enable use of reference models", gr.Checkbox), "extra_network_reference_values": OptionInfo(False, "Use reference values when available", gr.Checkbox), - "extra_networks_lora_sep": OptionInfo("

LoRA

", "", gr.HTML), - "extra_networks_default_multiplier": OptionInfo(1.0, "Default strength", gr.Slider, {"minimum": 0.0, "maximum": 2.0, "step": 0.01}), - "lora_force_reload": OptionInfo(False, "LoRA force reload always"), - "lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA load using Diffusers method"), - - "lora_apply_te": OptionInfo(False, "LoRA native apply to text encoder"), - "lora_fuse_native": OptionInfo(True, "LoRA native fuse with model"), - "lora_fuse_diffusers": OptionInfo(False, "LoRA diffusers fuse with model"), - "lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}), - "lora_in_memory_limit": OptionInfo(1, "LoRA memory cache", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), - "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info to metadata"), - "extra_networks_styles_sep": OptionInfo("

Styles

", "", gr.HTML), "extra_networks_styles": OptionInfo(True, "Show reference styles"), "extra_networks_apply_unparsed": OptionInfo(True, "Restore unparsed prompt"), @@ -700,6 +688,31 @@ def create_settings(cmd_opts): "wildcards_enabled": OptionInfo(True, "Enable file wildcards support"), })) + # --- LoRA --- + options_templates.update(options_section(('lora', "LoRA"), { + "lora_load_sep": OptionInfo("

Load options

", "", gr.HTML), + "lora_force_reload": OptionInfo(False, "LoRA force reload always"), + "lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA load using Diffusers method"), + "lora_in_memory_limit": OptionInfo(1, "LoRA memory cache", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), + + "lora_prompt_sep": OptionInfo("

Prompt helpers

", "", gr.HTML), + "extra_networks_default_multiplier": OptionInfo(1.0, "Default strength", gr.Slider, {"minimum": 0.0, "maximum": 2.0, "step": 0.01}), + "lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}), + + "lora_apply_sep": OptionInfo("

Apply method

", "", gr.HTML), + "lora_apply_te": OptionInfo(False, "LoRA native apply to text encoder"), + "lora_fuse_native": OptionInfo(True, "LoRA native fuse with model"), + "lora_fuse_diffusers": OptionInfo(False, "LoRA diffusers fuse with model"), + + "lora_quant_sep": OptionInfo("

Quantization options

", "", gr.HTML), + "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_meta_sep": OptionInfo("

Metadata

", "", gr.HTML), + "lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info to metadata"), + })) + # --- Extensions --- options_templates.update(options_section(('extensions', "Extensions"), { "disable_all_extensions": OptionInfo("none", "Disable all extensions", gr.Radio, {"choices": ["none", "user", "all"]}), diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py new file mode 100644 index 000000000..1c0ccb4f6 --- /dev/null +++ b/test/test-sdnq-lora-factors.py @@ -0,0 +1,992 @@ +#!/usr/bin/env python +""" +Offline unit tests for LoRA application on SDNQ-quantized layers. + +Pins two facts established on real checkpoints (see cli/lora-quant-fidelity.py +for the per-model analyzer): + +- The requantize path (dequantize + add + requantize) erases sub-step deltas + on low-bit formats: retention collapses to the ~2/group_size grid-extrema + floor on uint4, while int8 retains most of the delta. Guards against the + erasure law silently changing. +- The factor path (modules/lora/lora_sdnq.py) applies plain LoRA deltas + through the svd side-channel exactly, in both svd layouts and across + quantization configs (hadamard on/off, checkpoint svd correction present + or absent), with exact stacking, multiplier scaling and bit-exact + removal, wired through the real networks.network_activate / + network_deactivate control flow. +- Multi-LoRA set transitions keep the base pristine: a layer that fell back + to requantize (mixed factorable/non-factorable set) restores from backup + before re-entering the factor path, layers targeted by only some of the + loaded networks stay independent, and untargeted quantized layers are not + flagged as requantized. +- Robustness: factor removal restores onto the layer's current device after + an offload-style move, and a shape-mismatched network stacked onto a + factor-mode layer downgrades to the legacy path instead of raising. +- Hosting: on sub-8-bit layers, non-factorable sets ride the side-channel as + a truncated svd of their calc_updown delta: low-rank content survives + whole, dense content beats the requantize floor by a wide margin, int8 + and rank 0 keep the requantize path, removal stays bit-exact, and the + svd's random projections never touch the generation rng stream. +- Calibration: per-channel activation statistics weight the hosted + truncation toward loud input channels for better output-space retention; + low-rank content still survives whole, disabling the option reproduces + plain truncation bit-exact, and the capture hooks accumulate, persist + and reload statistics correctly, gated by option, format width and + model compile. + +All tensors are synthetic; no model files or running server required. + +Usage: + python test/test-sdnq-lora-factors.py +""" + +import os +import sys +import time +from contextlib import contextmanager + +import torch + +script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, script_dir) +os.chdir(script_dir) + +os.environ['SD_INSTALL_QUIET'] = '1' + +# Bootstrap cmd_args before any module that pulls in shared.py. +import modules.cmd_args # pylint: disable=wrong-import-position +import installer # pylint: disable=wrong-import-position +_orig_argv = sys.argv +sys.argv = [sys.argv[0]] +try: + modules.cmd_args.parse_args() +finally: + sys.argv = _orig_argv +installer.add_args(modules.cmd_args.parser) +modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + +from modules.errors import log # pylint: disable=wrong-import-position +from modules import shared, sd_models # pylint: disable=wrong-import-position +from modules.lora import network, network_lora, lora_sdnq, networks # pylint: disable=wrong-import-position +from modules.lora import lora_common as l_common # pylint: disable=wrong-import-position +from sdnq.quantizer import sdnq_quantize_layer, SDNQConfig # pylint: disable=wrong-import-position + +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +OUT_F, IN_F, RANK = 512, 512, 8 + +results: dict[str, dict] = {} + + +def category(name: str): + if name not in results: + results[name] = {'passed': 0, 'failed': 0, 'tests': []} + return name + + +def record(cat: str, passed: bool, name: str, detail: str = ''): + status = 'PASS' if passed else 'FAIL' + results[cat]['passed' if passed else 'failed'] += 1 + results[cat]['tests'].append((status, name)) + msg = f' {status}: {name}' + if detail: + msg += f' ({detail})' + if passed: + log.info(msg) + else: + log.error(msg) + + +def run_test(cat: str, fn): + name = fn.__name__ + try: + ok = fn() + record(cat, ok is not False, name) + except AssertionError as e: + record(cat, False, name, str(e)) + except Exception as e: # pylint: disable=broad-except + record(cat, False, name, f'exception: {e}') + import traceback + traceback.print_exc() + + +def build_layer(weights_dtype='uint4', use_quantized_matmul=False, seed=0, use_hadamard=True, use_svd=False): + torch.manual_seed(seed) + lin = torch.nn.Linear(IN_F, OUT_F, bias=False, dtype=torch.bfloat16, device=DEVICE) + with torch.no_grad(): + lin.weight.copy_(torch.randn(OUT_F, IN_F, device=DEVICE) * 0.04) + cfg = SDNQConfig(weights_dtype=weights_dtype, group_size=0, hadamard_group_size=256, use_hadamard=use_hadamard, + use_svd=use_svd, svd_rank=32, use_quantized_matmul=use_quantized_matmul, dequantize_fp32=False, + quantization_device=str(DEVICE), return_device=str(DEVICE)) + layer, _ = sdnq_quantize_layer(lin, cfg, torch_dtype=torch.bfloat16, param_name='test.weight') + layer.network_layer_name = 'lora_transformer_test' + layer.network_current_names = () + return layer + + +def dq(layer): + 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, skip_compile=True) + + +def make_delta(seed=1, sigma=3e-4): + torch.manual_seed(seed) + A = torch.randn(RANK, IN_F, device=DEVICE) * (sigma ** 0.5) + B = torch.randn(OUT_F, RANK, device=DEVICE) * (sigma ** 0.5) + return A, B, B @ A + + +class MockNOD: + def __init__(self, name): + self.filename = f'/tmp/{name}.safetensors' + self.name = name + self.shorthash = '' + self.sd_version = 'unknown' + + +def make_net(name, layer, A, B, te_mult=1.0, alpha=None, dora=False): + net = network.Network(name, MockNOD(name)) + net.te_multiplier = te_mult + net.unet_multiplier = [te_mult] * 3 + w = {'lora_up.weight': B.cpu(), 'lora_down.weight': A.cpu()} + if alpha is not None: + w['alpha'] = torch.tensor(float(alpha)) + if dora: + w['dora_scale'] = torch.ones(B.shape[0], 1) + nw = network.NetworkWeights(network_key=layer.network_layer_name, sd_key=layer.network_layer_name, w=w, sd_module=layer) + mod = network_lora.NetworkModuleLora(net, nw) + net.modules[layer.network_layer_name] = mod + return net + + +def rho_of(E, D): + return float(E.flatten() @ D.flatten() / D.flatten().square().sum()) + + +def requant_effective(layer, D): + """The lossy fallback path: quantize(W_dq + D) fresh with the layer's own params.""" + from sdnq.quantizer import sdnq_quantize_layer_weight + deq = layer.sdnq_dequantizer + Wdq = dq(layer) + deq2, data2 = sdnq_quantize_layer_weight(Wdq + D, layer_class_name='Linear', weights_dtype=deq.weights_dtype, + group_size=deq.group_size, hadamard_group_size=deq.hadamard_group_size, + use_hadamard=deq.use_hadamard, use_svd=False, use_quantized_matmul=False, + dequantize_fp32=False, torch_dtype=torch.bfloat16) + W2 = deq2(data2['weight'], data2['scale'], zero_point=data2['zero_point'], svd_up=None, svd_down=None, dtype=torch.float32, skip_compile=True) + return W2 - Wdq + + +# ============================================================ +# Tests - the erasure law (why the factor path exists) +# ============================================================ + +CAT_LAW = category('erasure-law') + + +def test_uint4_erases_substep_delta(): + layer = build_layer('uint4') + _A, _B, D = make_delta(sigma=2e-4) + rho = rho_of(requant_effective(layer, D), D) + group = layer.sdnq_dequantizer.group_size + floor = 2.0 / group + assert rho < 4 * floor, f'rho={rho:.4f} expected near extrema floor {floor:.4f}' + return True + + +def test_int8_retains_delta(): + layer = build_layer('int8') + _A, _B, D = make_delta(sigma=2e-4) + rho = rho_of(requant_effective(layer, D), D) + assert rho > 0.5, f'rho={rho:.4f} expected int8 to retain most of the delta' + return True + + +# ============================================================ +# Tests - factor path exactness +# ============================================================ + +CAT_FACTOR = category('factor-path') + + +def test_apply_exact_and_remove_bitexact(): + layer = build_layer('uint4') + A, B, D = make_delta() + net = make_net('one', layer, A, B) + l_common.loaded_networks.clear() + l_common.loaded_networks.append(net) + wanted = (('one', 1.0, 1.0, None),) + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) is True + Wdq0 = dq(layer) + assert lora_sdnq.apply_factors(layer, layer.network_layer_name, wanted) is True + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + assert lora_sdnq.remove_factors(layer) is True + assert torch.equal(dq(layer), Wdq0), 'remove must be bit-exact' + assert layer.svd_up is None and not hasattr(layer, 'sdnq_lora_svd_stash') + l_common.loaded_networks.clear() + return True + + +def test_multiplier_and_alpha_scaling(): + layer = build_layer('uint4') + A, B, D = make_delta() + net = make_net('one', layer, A, B, te_mult=0.5, alpha=RANK // 2) # alpha/rank = 0.5 + l_common.loaded_networks.clear() + l_common.loaded_networks.append(net) + Wdq0 = dq(layer) + lora_sdnq.apply_factors(layer, layer.network_layer_name, (('one', 0.5, 0.5, None),)) + rho = rho_of(dq(layer) - Wdq0, D) + assert abs(rho - 0.25) < 0.01, f'expected 0.5*0.5 scaling, rho={rho:.4f}' + lora_sdnq.remove_factors(layer) + l_common.loaded_networks.clear() + return True + + +def test_stacking_two_networks(): + layer = build_layer('uint4') + A1, B1, D1 = make_delta(seed=1) + A2, B2, D2 = make_delta(seed=2) + l_common.loaded_networks.clear() + l_common.loaded_networks.extend([make_net('a', layer, A1, B1), make_net('b', layer, A2, B2)]) + Wdq0 = dq(layer) + lora_sdnq.apply_factors(layer, layer.network_layer_name, (('a', 1.0, 1.0, None), ('b', 1.0, 1.0, None))) + rho = rho_of(dq(layer) - Wdq0, D1 + D2) + assert rho > 0.99, f'rho={rho:.4f}' + lora_sdnq.remove_factors(layer) + l_common.loaded_networks.clear() + return True + + +def test_matmul_layout_transposed(): + layer = build_layer('uint4', use_quantized_matmul=True) + A, B, D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + Wdq0 = dq(layer) + res = lora_sdnq.apply_factors(layer, layer.network_layer_name, (('one', 1.0, 1.0, None),)) + rho = rho_of(dq(layer) - Wdq0, D) + assert res is True and rho > 0.99, f'rho={rho:.4f}' + lora_sdnq.remove_factors(layer) + assert torch.equal(dq(layer), Wdq0) + l_common.loaded_networks.clear() + return True + + +def test_dora_falls_back(): + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('dora', layer, A, B, dora=True)) + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, (('dora', 1.0, 1.0, None),)) is False + l_common.loaded_networks.clear() + return True + + +def assert_factor_roundtrip(layer, tag): + """Apply-exact plus bit-exact removal on the given layer, whatever its quantization config.""" + A, B, D = make_delta() + Wdq0 = dq(layer) + orig_up = layer.svd_up + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + wanted = (('one', 1.0, 1.0, None),) + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) is True, f'{tag}: not a factor candidate' + assert lora_sdnq.apply_factors(layer, layer.network_layer_name, wanted) is True, f'{tag}: apply failed' + E = dq(layer) - Wdq0 + rho = rho_of(E, D) + resid = float((E - D).norm() / D.norm()) + assert rho > 0.99 and resid < 0.2, f'{tag}: rho={rho:.4f} resid={resid:.4f}' + assert lora_sdnq.remove_factors(layer) and torch.equal(dq(layer), Wdq0), f'{tag}: remove not bit-exact' + assert layer.svd_up is orig_up, f'{tag}: original svd factors not restored' + l_common.loaded_networks.clear() + + +def test_no_hadamard_checkpoint(): + """Checkpoints quantized without hadamard: factors attach unrotated.""" + assert_factor_roundtrip(build_layer('uint4', use_hadamard=False), 'plain') + assert_factor_roundtrip(build_layer('uint4', use_hadamard=False, use_quantized_matmul=True), 'matmul') + return True + + +def test_checkpoint_svd_factors_preserved(): + """Checkpoints quantized with their own svd correction keep it under apply/remove.""" + layer = build_layer('uint4', use_svd=True) + assert layer.svd_up is not None, 'quantizer produced no svd correction' + assert_factor_roundtrip(layer, 'plain') + assert_factor_roundtrip(build_layer('uint4', use_svd=True, use_quantized_matmul=True), 'matmul') + return True + + +# ============================================================ +# Tests - memory accounting across apply modes +# ============================================================ + +CAT_MEM = category('memory') + + +def tensor_bytes(t): + return t.numel() * t.element_size() if isinstance(t, torch.Tensor) else 0 + + +def test_factor_path_memory_is_factors_only(): + """Factor path: no weight/quant-state backups; added memory = the factor tensors.""" + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + lora_sdnq.apply_factors(layer, layer.network_layer_name, (('one', 1.0, 1.0, None),)) + assert getattr(layer, 'network_weights_backup', None) is None + assert not hasattr(layer, 'sdnq_dequantizer_backup') and not hasattr(layer, 'sdnq_scale_backup') + added = tensor_bytes(layer.svd_up) + tensor_bytes(layer.svd_down) + expected = RANK * (OUT_F + IN_F) * 2 # bf16 factors + assert added == expected, f'factor bytes {added} != expected {expected}' + would_be_backup = tensor_bytes(layer.weight) + tensor_bytes(layer.scale) + tensor_bytes(layer.zero_point) + assert added < would_be_backup / 4, f'factors {added}B should undercut the {would_be_backup}B backup this layer would otherwise clone' + lora_sdnq.remove_factors(layer) + assert layer.svd_up is None and layer.svd_down is None + l_common.loaded_networks.clear() + return True + + +def test_backup_mode_clones_full_quant_state(): + """Fallback in backup mode: packed weight + scale + zero_point are cloned to cpu.""" + from modules.lora.lora_apply import network_backup_weights + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('dora', layer, A, B, dora=True)) # non-factorable + reported = network_backup_weights(layer, layer.network_layer_name, (('dora', 1.0, 1.0, None),), fuse=False) + assert isinstance(layer.network_weights_backup, torch.Tensor) and layer.network_weights_backup.device.type == 'cpu' + assert hasattr(layer, 'sdnq_dequantizer_backup') and isinstance(layer.sdnq_scale_backup, torch.Tensor) + assert reported == tensor_bytes(layer.weight), f'reported {reported} != packed weight bytes {tensor_bytes(layer.weight)}' + total = reported + tensor_bytes(layer.sdnq_scale_backup) + tensor_bytes(layer.sdnq_zero_point_backup) + expected_min = OUT_F * IN_F // 2 # uint4 packs two weights per byte + assert total >= expected_min, f'backup {total}B below packed-weight floor {expected_min}B' + l_common.loaded_networks.clear() + return True + + +def test_fuse_mode_marker_takes_no_memory(): + """Fuse mode stores a boolean marker instead of tensors; guard forces backup on quantized models.""" + from modules.lora.lora_apply import network_backup_weights + from modules.lora import lora_overrides + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('dora', layer, A, B, dora=True)) + reported = network_backup_weights(layer, layer.network_layer_name, (('dora', 1.0, 1.0, None),), fuse=True) + assert layer.network_weights_backup is True and reported == 0 + assert not hasattr(layer, 'sdnq_dequantizer_backup') + + # the guard: a quantized component forces fuse off model-wide regardless of the option + class MockCfg: + quantization_config = {'quant_method': 'sdnq'} + class MockSd: + pass + sd = MockSd() + sd.transformer = torch.nn.Linear(4, 4) + sd.transformer.config = MockCfg() + from modules.modeldata import model_data + prev_model = model_data.sd_model + old_fuse = shared.opts.lora_fuse_native + try: + model_data.sd_model = sd + shared.opts.lora_fuse_native = True + assert lora_overrides.disable_fuse() is True + assert lora_overrides.fuse_native() is False + finally: + shared.opts.lora_fuse_native = old_fuse + model_data.sd_model = prev_model + l_common.loaded_networks.clear() + return True + + +# ============================================================ +# Tests - integration through networks.network_activate +# ============================================================ + +CAT_E2E = category('activate-e2e') + + +class MockHolder(torch.nn.Module): + @property + def device(self): + return DEVICE + + +@contextmanager +def mock_model(**layers): + """Install a one-component mock pipeline holding the given layers as shared.sd_model.""" + class MockPipe: + pass + class MockSd: + pass + holder = MockHolder() + for attr, lyr in layers.items(): + setattr(holder, attr, lyr) + pipe = MockPipe() + pipe.transformer = holder + sd = MockSd() + sd.pipe = pipe + from modules.modeldata import model_data + model_data.sd_model = sd + real_offload = sd_models.set_diffuser_offload + sd_models.set_diffuser_offload = lambda *a, **k: None + old_fuse = shared.opts.lora_fuse_native + shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config, so pin it instead of inheriting the running config + try: + yield + finally: + shared.opts.lora_fuse_native = old_fuse + sd_models.set_diffuser_offload = real_offload + l_common.loaded_networks.clear() + l_common.previously_loaded_networks.clear() + + +def activate(*nets): + l_common.loaded_networks.clear() + l_common.loaded_networks.extend(nets) + networks.network_activate() + + +def test_network_activate_roundtrip(): + layer = build_layer('uint4') + A, B, D = make_delta() + net = make_net('one', layer, A, B) + + with mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + assert getattr(layer, 'network_weights_backup', None) is None, 'factor path must not take weight backups' + + activate() # restore pass + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + + # fuse-mode deactivate route + layer.network_current_names = () + activate(net) + l_common.previously_loaded_networks[:] = l_common.loaded_networks + shared.opts.lora_fuse_native = True + networks.network_deactivate() + assert torch.equal(dq(layer), Wdq0), 'fuse-mode deactivate must restore bit-exact' + return True + + +# ============================================================ +# Tests - multi-LoRA set transitions between the two paths +# ============================================================ + +CAT_TRANS = category('transitions') + + +def test_mixed_family_transition_restores_base(): + layer = build_layer('uint4') + bystander = build_layer('uint4', seed=7) + bystander.network_layer_name = 'lora_transformer_bystander' + A, B, D = make_delta() + net_plain = make_net('plain', layer, A, B) + A2, B2, _ = make_delta(seed=5, sigma=3e-3) + net_dora = make_net('doranet', layer, A2, B2, dora=True) + + noted = [] + real_report = lora_sdnq.report_fallbacks + def capture_report(): + noted.append(len(lora_sdnq.fallback_layers)) + real_report() + lora_sdnq.report_fallbacks = capture_report + try: + with host_rank(0), mock_model(lin=layer, bystander=bystander): # pins the requantize fallback; hosted transitions are covered in the hosting category + Wdq0 = dq(layer) + activate(net_plain) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'plain set must take the factor path' + assert noted[-1] == 0, f'untargeted quantized layers must not be flagged as requantized: noted={noted[-1]}' + + activate(net_plain, net_dora) + assert not hasattr(layer, 'sdnq_lora_svd_stash') and isinstance(layer.network_weights_backup, torch.Tensor), 'mixed set must fall back with a tensor backup' + assert not torch.equal(dq(layer), Wdq0), 'fallback must have requantized the weights' + assert noted[-1] == 1, f'exactly the requantized layer must be flagged: noted={noted[-1]}' + + activate(net_plain) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'plain-only set must re-enter the factor path' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + stash, up, down = layer.sdnq_lora_svd_stash, layer.svd_up, layer.svd_down + lora_sdnq.remove_factors(layer) + base_clean = torch.equal(dq(layer), Wdq0) + layer.sdnq_lora_svd_stash, layer.svd_up, layer.svd_down = stash, up, down + assert base_clean, 'base under factors must be restored from backup on mixed-set exit' + + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must return bit-exact pristine' + finally: + lora_sdnq.report_fallbacks = real_report + return True + + +def test_partial_coverage_layers_stay_independent(): + layer_plain = build_layer('uint4') + layer_dora = build_layer('uint4', seed=7) + layer_dora.network_layer_name = 'lora_transformer_other' + A, B, D = make_delta() + net_plain = make_net('plain', layer_plain, A, B) + A2, B2, _ = make_delta(seed=5, sigma=3e-3) + net_dora = make_net('dorafar', layer_dora, A2, B2, dora=True) + + with host_rank(0), mock_model(lin=layer_plain, other=layer_dora): # pins the requantize fallback for the non-factorable layer + Wdq0, Wdq0_dora = dq(layer_plain), dq(layer_dora) + activate(net_plain, net_dora) + assert hasattr(layer_plain, 'sdnq_lora_svd_stash') and getattr(layer_plain, 'network_weights_backup', None) is None, 'plain layer must stay on the factor path' + assert isinstance(getattr(layer_dora, 'network_weights_backup', None), torch.Tensor), 'dora layer must take the backup fallback' + rho = rho_of(dq(layer_plain) - Wdq0, D) + assert rho > 0.99, f'rho={rho:.4f}' + activate() + assert torch.equal(dq(layer_plain), Wdq0), 'factor layer must restore bit-exact' + assert torch.equal(dq(layer_dora), Wdq0_dora), 'fallback layer must restore bit-exact' + return True + + +@contextmanager +def apply_method(value): + old = getattr(shared.opts, 'lora_sdnq_apply', 'exact') + shared.opts.lora_sdnq_apply = value + try: + yield + finally: + shared.opts.lora_sdnq_apply = old + + +def test_mechanism_gate_declines_candidates(): + """The requantize option must gate every svd-channel entry point and flip the apply-stamp token.""" + layer = build_layer('uint4') + A, B, _D = make_delta() + l_common.loaded_networks.clear() + l_common.loaded_networks.append(make_net('one', layer, A, B)) + wanted = (('one', 1.0, 1.0, None),) + try: + assert lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) + with host_rank(64): + assert lora_sdnq.host_candidate(layer, layer.network_layer_name, wanted) + assert lora_sdnq.signature() == '' + with apply_method('requantize'): + assert not lora_sdnq.factor_candidate(layer, layer.network_layer_name, wanted) + with host_rank(64): + assert not lora_sdnq.host_candidate(layer, layer.network_layer_name, wanted) + assert lora_sdnq.signature() == '|quant=requantize' + finally: + l_common.loaded_networks.clear() + return True + + +def test_requantize_option_routes_to_legacy_path(): + """With the option set, a factorable set must take the classic backup-and-requantize path end to end.""" + layer = build_layer('uint4') + A, B, _D = make_delta(sigma=3e-3) + net = make_net('one', layer, A, B) + with apply_method('requantize'), mock_model(lin=layer): + shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config + Wdq0 = dq(layer) + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'legacy path must not touch the svd channel' + assert layer.svd_up is None, 'legacy path must leave the channel empty' + assert isinstance(layer.network_weights_backup, torch.Tensor), 'legacy path must take a tensor backup' + assert not torch.equal(dq(layer), Wdq0), 'legacy path must requantize the weights' + activate() + assert torch.equal(dq(layer), Wdq0), 'legacy restore must be bit-exact from backup' + return True + + +def test_mechanism_flip_strips_attached_factors(): + """Flipping to requantize with factors attached must strip them before the weight path takes the layer; flipping back must re-enter the factor path.""" + layer = build_layer('uint4') + A, B, _D = make_delta(sigma=3e-3) + net = make_net('one', layer, A, B) + with mock_model(lin=layer): + shared.opts.lora_fuse_native = False # a real quantized model forces backup mode; the mock carries no quantization config + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'default mechanism must take the factor path' + E_exact = dq(layer) - Wdq0 + with apply_method('requantize'): + activate(net) # same set; the mechanism token in the apply stamp must force re-processing + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'flip must strip the attached factors' + assert layer.svd_up is None, 'stripped channel must be empty, or the requantized delta double-applies' + assert isinstance(layer.network_weights_backup, torch.Tensor), 'flipped layer must continue on the backup path' + activate(net) # flip back within the same loaded set + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'flip back must re-enter the factor path' + assert torch.equal(dq(layer) - Wdq0, E_exact), 'exact re-apply must restore the base from backup before attaching' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must return bit-exact pristine' + return True + + +def test_mechanism_flip_restore_pass_strips(): + """A restore-only pass under the requantize option must still drop attached factors.""" + layer = build_layer('uint4') + A, B, _D = make_delta(sigma=3e-3) + net = make_net('one', layer, A, B) + with mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash') + with apply_method('requantize'): + activate() # unload with the gate closed: the fallthrough strip is the only removal route + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'restore pass must strip the factors' + assert torch.equal(dq(layer), Wdq0), 'strip must restore bit-exact' + assert layer.network_current_names == (), 'stripped layer must be stamped restored' + return True + + +CAT_HOST = category('hosting') + + +@contextmanager +def host_rank(rank): + old = getattr(shared.opts, 'lora_sdnq_host_rank', 0) + shared.opts.lora_sdnq_host_rank = rank + try: + yield + finally: + shared.opts.lora_sdnq_host_rank = old + + +def make_dense_net(name, layer, D): + """A full-family (dense diff) network module: non-factorable by construction.""" + from modules.lora import network_full + net = network.Network(name, MockNOD(name)) + net.te_multiplier = 1.0 + net.unet_multiplier = [1.0] * 3 + nw = network.NetworkWeights(network_key=layer.network_layer_name, sd_key=layer.network_layer_name, + w={'diff': D.cpu()}, sd_module=layer) + net.modules[layer.network_layer_name] = network_full.NetworkModuleFull(net, nw) + return net + + +def test_hosted_low_rank_delta_is_kept(): + layer = build_layer('uint4') + _A, _B, D = make_delta(sigma=3e-3) + net = make_dense_net('densenet', layer, D) # low-rank content in a non-factorable container + with host_rank(64), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash'), 'hosted set must ride the side-channel' + assert getattr(layer, 'network_weights_backup', None) is None, 'hosted layers must not take a weight backup' + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.95, f'rank-8 delta under cap 64 must be kept nearly whole: rho={rho:.4f}' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + return True + + +def test_hosted_dense_delta_beats_requant(): + layer = build_layer('uint4') + torch.manual_seed(3) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 3e-4 # full-rank, sub-step: requant erases it + requant_rho = rho_of(requant_effective(layer, D), D) + net = make_dense_net('densefull', layer, D) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + hosted_rho = rho_of(dq(layer) - Wdq0, D) + assert hosted_rho > 0.4, f'hosted rho={hosted_rho:.3f}' + assert hosted_rho > requant_rho + 0.3, f'hosting must beat requant by a wide margin: {hosted_rho:.3f} vs {requant_rho:.3f}' + activate() + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_hosted_skips_int8(): + layer = build_layer('int8') + _A, _B, D = make_delta(sigma=3e-3) + net = make_dense_net('int8net', layer, D) + with host_rank(256), mock_model(lin=layer): + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'int8 must keep the requantize path' + assert isinstance(getattr(layer, 'network_weights_backup', None), torch.Tensor), 'int8 fallback must take the backup' + activate() + return True + + +def test_hosted_disabled_by_option(): + layer = build_layer('uint4') + _A, _B, D = make_delta(sigma=3e-3) + net = make_dense_net('offnet', layer, D) + with host_rank(0), mock_model(lin=layer): + activate(net) + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'rank 0 must disable hosting' + activate() + return True + + +def test_hosted_transitions_and_rng_isolation(): + layer = build_layer('uint4') + A, B, D = make_delta() + net_plain = make_net('plainh', layer, A, B) + _A2, _B2, D2 = make_delta(seed=9, sigma=3e-3) + net_dense = make_dense_net('denseh', layer, D2) + with host_rank(256), mock_model(lin=layer): + Wdq0 = dq(layer) + rng0 = torch.cuda.get_rng_state() if DEVICE.type == 'cuda' else torch.get_rng_state() + activate(net_dense) # hosted + rng1 = torch.cuda.get_rng_state() if DEVICE.type == 'cuda' else torch.get_rng_state() + assert torch.equal(rng0, rng1), 'hosting must not consume the generation rng stream' + assert hasattr(layer, 'sdnq_lora_svd_stash') + activate(net_plain) # exact replaces hosted + rho = rho_of(dq(layer) - Wdq0, D) + assert rho > 0.99, f'exact set after hosted set: rho={rho:.4f}' + activate(net_plain, net_dense) # mixed set hosts the combined delta + rho_mix = rho_of(dq(layer) - Wdq0, D + D2) + assert rho_mix > 0.9, f'mixed hosted rho={rho_mix:.4f}' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + return True + + +CAT_CALIB = category('calibration') + + +@contextmanager +def host_calib(value): + old = getattr(shared.opts, 'lora_sdnq_host_calib', False) + shared.opts.lora_sdnq_host_calib = value + try: + yield + finally: + shared.opts.lora_sdnq_host_calib = old + + +def test_calibrated_hosting_beats_plain(): + layer = build_layer('uint4') + torch.manual_seed(11) + scale = torch.ones(IN_F, device=DEVICE) + scale[:32] = 40.0 # a few loud input channels, the shape real activations have + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 3e-4 + X = torch.randn(1024, IN_F, device=DEVICE) * scale + Y = X @ D.t() + net = make_dense_net('calnet', layer, D) + + def out_rho(E): + return float((X @ E.t()).flatten() @ Y.flatten() / Y.square().sum()) + + with host_rank(32), host_calib(True), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net) + plain = out_rho(dq(layer) - Wdq0) + activate() + layer.sdnq_calib_rms = scale.cpu() # statistics as the capture leaves them + activate(net) + weighted = out_rho(dq(layer) - Wdq0) + activate() + del layer.sdnq_calib_rms + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact' + assert weighted > plain + 0.2, f'calibrated hosting must beat plain in output space: {weighted:.3f} vs {plain:.3f}' + return True + + +def test_calibrated_low_rank_delta_survives(): + layer = build_layer('uint4') + _A, _B, D = make_delta(sigma=3e-3) + net = make_dense_net('calfull', layer, D) + with host_rank(64), host_calib(True), mock_model(lin=layer): + Wdq0 = dq(layer) + torch.manual_seed(21) + layer.sdnq_calib_rms = torch.rand(IN_F) * 10 + 0.1 # arbitrary positive statistics: unscale must round-trip + activate(net) + rho = rho_of(dq(layer) - Wdq0, D) + activate() + del layer.sdnq_calib_rms + assert rho > 0.95, f'rank-8 delta under weighted cap 64 must be kept nearly whole: rho={rho:.4f}' + assert torch.equal(dq(layer), Wdq0) + return True + + +def test_calib_option_off_matches_plain(): + layer = build_layer('uint4') + torch.manual_seed(31) + D = torch.randn(OUT_F, IN_F, device=DEVICE) * 3e-4 + net = make_dense_net('caloff', layer, D) + with host_rank(64), mock_model(lin=layer): + Wdq0 = dq(layer) + with host_calib(False): + layer.sdnq_calib_rms = torch.rand(IN_F) + 0.5 + activate(net) + off = dq(layer) + activate() + del layer.sdnq_calib_rms + with host_calib(True): + activate(net) # no statistics attribute: plain truncation + plain = dq(layer) + activate() + assert torch.equal(off, plain), 'option off must reproduce the uncalibrated truncation bit-exact' + assert torch.equal(dq(layer), Wdq0) + return True + + +class MockCheckpointInfo: + def __init__(self, name): + self.name = name + + +class MockCalibSd: + def __init__(self, name, **layers): + self.transformer = MockHolder() + for attr, lyr in layers.items(): + setattr(self.transformer, attr, lyr) + self.sd_checkpoint_info = MockCheckpointInfo(name) + + +def test_calib_capture_persist_roundtrip(): + import tempfile + from modules.lora import lora_calib + layer_a = build_layer('uint4', seed=41) + layer_b = build_layer('uint4', seed=42) + sd = MockCalibSd('test/calib-model', la=layer_a, lb=layer_b) + old_root, old_tokens = lora_calib.calib_root, lora_calib.TOKENS_DONE + with tempfile.TemporaryDirectory() as tmp, host_calib(True): + try: + 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' + torch.manual_seed(51) + scale = torch.linspace(0.1, 4.0, IN_F, device=DEVICE) + xs = [] + for _ in range(2): # exactly the completion threshold, so statistics cover every forward + x = (torch.randn(1024, IN_F, device=DEVICE) * scale).to(torch.bfloat16) + xs.append(x.float()) + layer_a(x) + layer_b(x) + assert lora_calib.capture['complete'], 'capture must complete once enough tokens are seen' + path = lora_calib.calib_file('test/calib-model') + assert os.path.isfile(path), f'statistics must persist to {path}' + expected = torch.cat(xs).square().mean(dim=0).sqrt().cpu() + assert torch.allclose(layer_a.sdnq_calib_rms, expected, rtol=1e-3, atol=1e-5), 'streamed rms must match the seen activations' + del layer_a.sdnq_calib_rms, layer_b.sdnq_calib_rms + lora_calib.on_model_loaded(sd) # second load takes the cached path + assert len(lora_calib.capture['handles']) == 0, 'cached statistics must not re-attach capture hooks' + assert torch.allclose(layer_a.sdnq_calib_rms, expected, rtol=1e-3, atol=1e-5), 'reload must restore the persisted rms' + del layer_a.sdnq_calib_rms, layer_b.sdnq_calib_rms + finally: + lora_calib.calib_root, lora_calib.TOKENS_DONE = old_root, old_tokens + lora_calib.detach_capture() + return True + + +def test_calib_capture_gates(): + from modules.lora import lora_calib + sd_int8 = MockCalibSd('test/calib-int8', lin=build_layer('int8', seed=43)) + with host_calib(True): + lora_calib.on_model_loaded(sd_int8) + assert len(lora_calib.capture['handles']) == 0, 'int8-only models have nothing to calibrate' + sd_u4 = MockCalibSd('test/calib-gates', lin=build_layer('uint4', seed=44)) + with host_calib(False): + lora_calib.on_model_loaded(sd_u4) + assert len(lora_calib.capture['handles']) == 0, 'option off must disable capture' + old_compile = getattr(shared.opts, 'cuda_compile', None) + with host_calib(True): + shared.opts.cuda_compile = ['Model'] + try: + lora_calib.on_model_loaded(sd_u4) + assert len(lora_calib.capture['handles']) == 0, 'model compile must disable capture' + finally: + shared.opts.cuda_compile = old_compile + lora_calib.detach_capture() + return True + + +CAT_ROBUST = category('robustness') + + +def test_remove_factors_after_device_move(): + layer = build_layer('uint4', use_svd=True) # checkpoint svd correction so the stash holds real tensors + A, B, _D = make_delta() + net = make_net('mover', layer, A, B) + with mock_model(lin=layer): + Wdq0 = dq(layer) + orig_up = layer.svd_up.detach().clone() + activate(net) + assert hasattr(layer, 'sdnq_lora_svd_stash') + layer.to('cpu') # offload moves registered params, never the stash tuple + activate() + assert layer.svd_up.device == layer.scale.device, f'restored svd must live on the layer device, got {layer.svd_up.device} vs {layer.scale.device}' + assert torch.equal(layer.svd_up, orig_up.to('cpu')), 'restored svd values must match the original factors' + layer.to(DEVICE) + assert torch.equal(dq(layer), Wdq0), 'round trip must restore bit-exact' + return True + + +def test_stacked_shape_mismatch_falls_back(): + from types import SimpleNamespace + layer = build_layer('uint4') + A, B, _D = make_delta() + net_good = make_net('good', layer, A, B) + torch.manual_seed(9) + A_bad = torch.randn(RANK, IN_F, device=DEVICE) * 0.01 + B_bad = torch.randn(OUT_F // 2, RANK, device=DEVICE) * 0.01 # wrong out_features for this layer + net_bad = make_net('badshape', layer, A_bad, B_bad) + prev_enl = l_common.extra_network_lora + l_common.extra_network_lora = SimpleNamespace(errors={}) # the error path reports through the extra-networks registry + try: + with host_rank(0), mock_model(lin=layer): + Wdq0 = dq(layer) + activate(net_good) + assert hasattr(layer, 'sdnq_lora_svd_stash') + activate(net_good, net_bad) # must not raise: a malformed stack downgrades the layer to the legacy path + assert not hasattr(layer, 'sdnq_lora_svd_stash'), 'shape-mismatched stack must leave factor mode' + activate() + assert torch.equal(dq(layer), Wdq0), 'unload must restore bit-exact pristine' + finally: + l_common.extra_network_lora = prev_enl + return True + + +def run_tests(): + t0 = time.time() + log.warning('=== Erasure law ===') + for fn in [test_uint4_erases_substep_delta, test_int8_retains_delta]: + run_test(CAT_LAW, fn) + log.warning('=== Factor path ===') + for fn in [test_apply_exact_and_remove_bitexact, test_multiplier_and_alpha_scaling, test_stacking_two_networks, test_matmul_layout_transposed, test_dora_falls_back, test_no_hadamard_checkpoint, test_checkpoint_svd_factors_preserved]: + run_test(CAT_FACTOR, fn) + log.warning('=== Memory accounting ===') + for fn in [test_factor_path_memory_is_factors_only, test_backup_mode_clones_full_quant_state, test_fuse_mode_marker_takes_no_memory]: + run_test(CAT_MEM, fn) + log.warning('=== Activate integration ===') + for fn in [test_network_activate_roundtrip]: + run_test(CAT_E2E, fn) + log.warning('=== Set transitions ===') + for fn in [test_mixed_family_transition_restores_base, test_partial_coverage_layers_stay_independent, + test_mechanism_gate_declines_candidates, test_requantize_option_routes_to_legacy_path, + test_mechanism_flip_strips_attached_factors, test_mechanism_flip_restore_pass_strips]: + 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]: + 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]: + run_test(CAT_CALIB, fn) + log.warning('=== Robustness ===') + for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: + run_test(CAT_ROBUST, fn) + + elapsed = time.time() - t0 + log.warning('=== Results ===') + total_pass = total_fail = 0 + for cat, info in results.items(): + status = 'PASS' if info['failed'] == 0 else 'FAIL' + log.info(f' {cat}: {info["passed"]} passed, {info["failed"]} failed [{status}]') + total_pass += info['passed'] + total_fail += info['failed'] + log.warning(f'Total: {total_pass} passed, {total_fail} failed in {elapsed:.2f}s') + return total_fail == 0 + + +if __name__ == '__main__': + with torch.inference_mode(): + ok = run_tests() + sys.exit(0 if ok else 1) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 089b239b6..cdd66f749 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -414,7 +414,7 @@ {"id":"","label":"Desktop","localized":"","hint":"","ui":"settings_ui"}, {"id":"","label":"Downscale high resolution live previews","localized":"","hint":"","ui":"settings_live-preview"}, {"id":"","label":"Detailer use model augment","localized":"","hint":"Run detailer detection models at extra precision","ui":"settings_postprocessing"}, - {"id":"","label":"Default strength","localized":"","hint":"When adding extra network such as Lora to prompt, use this multiplier for it","ui":"settings_extra_networks"}, + {"id":"","label":"Default strength","localized":"","hint":"When adding extra network such as Lora to prompt, use this multiplier for it","ui":"settings_lora"}, {"id":"","label":"Do not change selected model when reading generation parameters","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"Do conditional and unconditional denoising in one batch","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"Disable NaN check","localized":"","hint":"","ui":"settings_legacy_options"}, @@ -857,14 +857,17 @@ {"id":"","label":"Log view update period","localized":"","hint":"Log view update period, in milliseconds","ui":"settings_ui"}, {"id":"","label":"Live preview display period","localized":"","hint":"Request preview image every n steps, set to 0 to disable","ui":"settings_live-preview"}, {"id":"","label":"Load custom Diffusers pipeline","localized":"","hint":"","ui":"settings_huggingface"}, - {"id":"","label":"LoRA force reload always","localized":"","hint":"Forces LoRA networks to reload from storage on every generation, even if already cached.
Useful for debugging or when LoRA files are being modified externally.
Disable for normal use to benefit from caching.","ui":"settings_extra_networks"}, - {"id":"","label":"LoRA load using Diffusers method","localized":"","hint":"Alternative method uses diffusers built-in LoRA capabilities instead of native SD.Next implementation (may reduce LoRA compatibility)","ui":"settings_extra_networks"}, - {"id":"","label":"LoRA native apply to text encoder","localized":"","hint":"","ui":"settings_extra_networks"}, - {"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_extra_networks"}, - {"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_extra_networks"}, - {"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_extra_networks"}, - {"id":"","label":"LoRA memory cache","localized":"","hint":"How many LoRAs to keep in network for future use before requiring reloading from storage","ui":"settings_extra_networks"}, - {"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_extra_networks"}, + {"id":"","label":"LoRA force reload always","localized":"","hint":"Forces LoRA networks to reload from storage on every generation, even if already cached.
Useful for debugging or when LoRA files are being modified externally.
Disable for normal use to benefit from caching.","ui":"settings_lora"}, + {"id":"","label":"LoRA load using Diffusers method","localized":"","hint":"Alternative method uses diffusers built-in LoRA capabilities instead of native SD.Next implementation (may reduce LoRA compatibility)","ui":"settings_lora"}, + {"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 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 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"}, {"id":"","label":"LDSR Path","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"LoRA load using legacy method","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"Loaded LoRA","localized":"","hint":"","ui":"component-5851"},