mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
Merge pull request #5066 from vladmandic/lora-exact
merge lora-exact into dev
This commit is contained in:
Vladimir Mandic
GitHub
@@ -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())
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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")
|
||||
|
||||
+25
-12
@@ -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("<h2>LoRA</h2>", "", 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("<h2>Styles</h2>", "", 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("<h2>Load options</h2>", "", 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("<h2>Prompt helpers</h2>", "", 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("<h2>Apply method</h2>", "", 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("<h2>Quantization options</h2>", "", 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("<h2>Metadata</h2>", "", 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"]}),
|
||||
|
||||
@@ -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)
|
||||
@@ -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.<br>Useful for debugging or when LoRA files are being modified externally.<br>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.<br><br><b style=\"color: #ef4444\">Warning:</b> 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.<br><br><b style=\"color: #ef4444\">Warning:</b> 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.<br>Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.<br>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.<br>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.<br>Useful for debugging or when LoRA files are being modified externally.<br>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.<br><br><b style=\"color: #ef4444\">Warning:</b> 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.<br><br><b style=\"color: #ef4444\">Warning:</b> 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:<br>- <b>exact</b>: 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<br>- <b>requantize</b>: 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<br><br>With <b>requantize</b> selected, the host rank and calibration options below have no effect.<br><br>Default is <b>exact</b>.","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.<br>Set to the number of tags to auto-apply, e.g., 3 = add top 3 trigger tags.<br>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.<br>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"},
|
||||
|
||||
Reference in New Issue
Block a user