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("