Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(cli): stream quantized repos and measure the base error in lora-quant-fidelity

Browse Source 
QuantRepo builds the transformer as a meta skeleton through the sdnq
conversion and streams each layer from the shards while it is analyzed,
so a repo larger than host memory costs one module at a time.
Remote-code classes resolve from the modeling file beside the
checkpoint. LoRA groups go through the arch grouper with the file-level
alpha and the adapt_weights hook, and fused saves are sliced onto their
modules. --reference measures each module's quantization error against
the unquantized repo and reports the delta against it; without one a
uniform-rounding estimate stands in. The tool imports again:
group_by_suffixes takes no bare prefixes and shared initializes before
the lora modules.
This commit is contained in:
CalamitousFelicitousness
2026-09-15 03:52:04 +01:00
parent 57328dc615
commit 5a2870f345
2 changed files with 248 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/lora-quant-fidelity.py
+236 -46
View File
@@ -21,11 +21,15 @@ are measured as they would actually apply:
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.
the if-merged figure. ``snr`` sets the delta against the base weight's own
quantization error (uniform rounding from the grid step, or measured with
``--reference``), in the calibrated norm with ``--calib``: it does not depend
on the apply path and compares checkpoints of different widths.
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.
Works offline against a pre-quantized SDNQ repo (stored tensors + config,
streamed one module at a time so the repo need not fit in memory) 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"
@@ -33,9 +37,13 @@ Examples:
"""
import os
import re
import sys
import json
import types
import argparse
import importlib
import importlib.util
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault('SD_INSTALL_QUIET', '1')
@@ -54,6 +62,7 @@ def parse_cli():
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('--reference', default=None, help='unquantized repo of the same model: the base quantization error each delta competes with is then measured per module instead of estimated from the grid step')
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()
@@ -71,6 +80,7 @@ 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 import shared # pylint: disable=wrong-import-position,unused-import # shared must initialize before sd_models, which imports back into it
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
@@ -82,6 +92,7 @@ MODEL_ROOTS = [
os.path.expanduser('~/database/models/Diffusers'),
]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ARCH_PACKAGES = {'zimage': 'z_image', 'f2': 'flux', 'minimaxh3': 'minimax'} # arches whose pipelines package is not spelled like the arch
# 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
@@ -104,6 +115,8 @@ class StubOnDisk:
self.name = os.path.splitext(os.path.basename(path))[0]
self.shorthash = ''
self.sd_version = 'unknown'
with safe_open(path, framework='pt', device='cpu') as f:
self.metadata = f.metadata() or {}
def resolve_model_dir(spec):
@@ -129,7 +142,6 @@ def resolve_arch(name):
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)
@@ -138,21 +150,27 @@ def map_lora_modules(lora_path, arch_mod):
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).
resolver. A fused save is sliced onto its targets for the lora family; the
other families' 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()}
metadata = f.metadata() or {}
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)])
grouper = getattr(arch_mod, 'group_by_suffixes', None) or native_adapter.group_by_suffixes # an arch that rewrites keys before parsing groups them itself
file_alpha = getattr(arch_mod, 'file_alpha', 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)
groups = grouper(state_dict, suffixes, prefixes=prefixes)
if fam == 'lora':
groups = {k: w for k, w in groups.items() if 'lora_down.weight' in w and 'lora_up.weight' in w}
alpha = file_alpha(types.SimpleNamespace(filename=lora_path, metadata=metadata)) if file_alpha is not None else None
if alpha is not None and not any('alpha' in w for w in groups.values()): # a file-level alpha applies only to files without alpha tensors, as in try_load_lora
groups = {k: {**w, 'alpha': torch.tensor(float(alpha))} for k, w in groups.items()}
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:
@@ -160,10 +178,18 @@ def map_lora_modules(lora_path, arch_mod):
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
target = w
if chunk is not None: # a fused save spans several modules; only the lora family slices its up factor
target = None
if fam == 'lora':
fused_out = w['lora_up.weight'].shape[0]
target = native_adapter.slice_lora_chunk(w, chunk)
target = native_adapter.slice_dora_scale(target, chunk, fused_out)
target = native_adapter.slice_bias_delta(target, chunk, fused_out) if target is not None else None
if target is None:
chunked += 1
continue
mapped.setdefault(path, []).append((fam, target)) # a module can carry several families; the loader applies each
return mapped, census, chunked
@@ -192,17 +218,14 @@ def build_module(fam, path, w, net, 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:
def resolve_transformer_cls(arch, class_name, model_dir=None):
"""Resolve the transformer class: an sdnext-owned spec class, then diffusers, then a modeling file beside the weights.
Arches like krea2 keep checkpoint-style names in their own class; pruned MiniMax repos ship theirs as remote code."""
if not class_name:
return None
if arch:
try:
import importlib
pkg = importlib.import_module(f'pipelines.{ {"zimage": "z_image", "f2": "flux"}.get(arch, arch) }')
pkg = importlib.import_module(f'pipelines.{ARCH_PACKAGES.get(arch, arch)}')
for attr in dir(pkg):
if attr.endswith('_SPEC'):
cls = getattr(getattr(pkg, attr), 'cls', None)
@@ -210,20 +233,107 @@ def resolve_transformer_cls(arch, class_name):
return cls
except Exception:
pass
import diffusers
cls = getattr(diffusers, class_name, None)
if cls is not None or not model_dir:
return cls
for fname in sorted(os.listdir(model_dir)):
path = os.path.join(model_dir, fname)
if not fname.endswith('.py'):
continue
with open(path, encoding='utf-8') as f:
if re.search(rf'^class {re.escape(class_name)}\b', f.read(), re.MULTILINE) is None:
continue
spec = importlib.util.spec_from_file_location(fname[:-3], path)
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return getattr(module, class_name)
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 QuantRepo:
"""Per-module access to a pre-quantized SDNQ repo: a meta skeleton built through the sdnq conversion, each layer's
tensors streamed from the shards while it is analyzed; tensors outside the block stacks stay resident for the arch hooks."""
LAYER_KEYS = ('weight', 'bias', 'scale', 'zero_point', 'svd_up', 'svd_down')
def __init__(self, model_dir, model_config, arch=None):
from accelerate import init_empty_weights
from sdnq import SDNQConfig
from sdnq.quantizer import sdnq_post_load_quant
from sdnq.utils import get_quant_args_from_config
cls = resolve_transformer_cls(arch, model_config.get('_class_name'), model_dir)
if cls is None:
raise SystemExit(f'cannot resolve transformer class {model_config.get("_class_name")} for "{model_dir}"')
quant_config = SDNQConfig.from_dict(model_config['quantization_config'])
with init_empty_weights():
config = cls.load_config(model_dir)
if hasattr(config, 'pop'):
config.pop('quantization_config', None)
model = cls.from_config(config)
model = sdnq_post_load_quant(model, torch_dtype=torch.bfloat16, pre_quantized=True, **get_quant_args_from_config(quant_config))
self.model = model
self.model_dir = model_dir
self.handles = {}
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:
weight_map = json.load(f)['weight_map']
else:
weight_map = {}
for fname in sorted(os.listdir(model_dir)):
if fname.endswith('.safetensors'):
with safe_open(os.path.join(model_dir, fname), framework='pt', device='cpu') as f:
weight_map.update(dict.fromkeys(f.keys(), fname))
mapping = getattr(model, '_checkpoint_conversion_mapping', None) or {}
self.shards = {} # model key -> (stored key, shard)
for stored, shard in weight_map.items():
key = stored
for pattern, replacement in mapping.items():
key = re.sub(pattern, replacement, key)
self.shards[key] = (stored, shard)
self.layers = {}
for name, module in model.named_modules():
if getattr(module, 'sdnq_dequantizer', None) is not None or (module.__class__.__name__ == 'Linear' and getattr(module, 'weight', None) is not None):
self.layers[name] = module
stacks = tuple(f'{name}.' for name, module in model.named_modules() if isinstance(module, torch.nn.ModuleList))
resident = {key: self.get(key) for key in self.shards if not key.startswith(stacks)}
model.load_state_dict(resident, strict=False, assign=True)
def get(self, key):
entry = self.shards.get(key)
if entry is None:
return None
stored, shard = entry
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(stored)
def materialize(self, name):
"""Load one layer's stored tensors onto the analysis device and return the layer."""
from sdnq.quant_utils import prepare_weight_for_matmul, prepare_svd_for_matmul
layer = self.layers[name]
state = {}
for local in self.LAYER_KEYS:
t = self.get(f'{name}.{local}')
if t is not None:
state[local] = t.to(device)
layer.load_state_dict(state, strict=False, assign=True)
deq = getattr(layer, 'sdnq_dequantizer', None)
if deq is not None: # the loader's post-processing, so the dequantizer sees the layout it expects
if deq.use_quantized_matmul and not deq.re_quantize_for_matmul:
layer.weight.data = prepare_weight_for_matmul(layer.weight, matmul_dtype=deq.quantized_matmul_dtype)
if getattr(layer, 'svd_up', None) is not None:
layer.svd_up.data, layer.svd_down.data = prepare_svd_for_matmul(layer.svd_up, layer.svd_down, deq.use_quantized_matmul)
return layer
def release(self, name):
layer = self.layers[name]
state = {local: torch.empty_like(t, device='meta') for local, t in layer.state_dict().items() if t.device.type != 'meta'}
layer.load_state_dict(state, strict=False, assign=True)
class Bf16Repo:
@@ -252,7 +362,28 @@ class Bf16Repo:
return f.get_tensor(key)
def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None):
def quant_noise_energy(step, shape, hadamard_group, rms):
"""Expected energy of the base weight's quantization error, ``step^2/12`` per element under uniform rounding;
with ``rms`` per input channel it is the output-space energy per token."""
_out, n_in = shape
s = step.to(device, torch.float32)
if s.ndim == 3:
s = s.squeeze(-1)
if s.ndim == 1:
s = s[:, None]
groups = s.shape[1]
glen = n_in // groups
if rms is None:
power = torch.full((groups,), float(glen), device=s.device)
else:
p = rms.to(s.device, torch.float32).square()
if hadamard_group and n_in % hadamard_group == 0:
p = p.view(-1, hadamard_group).mean(1, keepdim=True).expand(-1, hadamard_group).reshape(-1)
power = p.view(groups, glen).sum(1)
return float(((s.square() / 12) @ power).sum())
def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None, noise=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
@@ -263,6 +394,14 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None):
With ``step_live`` (the layer's own pre-add scale), the production routing
rule applies: a delta fat against the grid whose truncation capture is low
reports the requantize path, the way the loader would route it.
``snr`` is the delta's energy over the energy of the base weight's own
quantization error, in the calibrated norm when ``calib_rms`` is given:
``noise`` carries that error's measured ``(plain, weighted)`` energies
against an unquantized reference, else the uniform-rounding estimate from
``step_live`` stands in. It is independent of the apply path and says
whether what the adapter adds stands above the error the checkpoint
already carries.
"""
D = None
for mod in mods:
@@ -274,7 +413,23 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None):
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)
factor_eligible=factor_eligible, hosted=False, applied_rho=None, delta_energy=0.0,
snr=None, snr_plain=None, noise_rms=None, noise_source=None)
snr, snr_plain, noise_rms, noise_source = None, None, None, None
if (noise is not None or step_live is not None) and not deq_params.get('use_codebook', False):
rms = calib_rms.to(device, torch.float32) if (calib_rms is not None and calib_rms.shape[-1] == D.shape[-1]) else None
if noise is not None:
plain, weighted = noise
noise_source = 'measured'
else:
hadamard = deq_params['hadamard_group_size'] if deq_params['use_hadamard'] else 0
plain = quant_noise_energy(step_live, D.shape, hadamard, None)
weighted = quant_noise_energy(step_live, D.shape, hadamard, rms) if rms is not None else plain
noise_source = 'uniform'
delta_w = float((D * rms).square().sum()) if rms is not None else float(nD.square())
noise_rms = (plain / D.numel()) ** 0.5
snr_plain = float(nD.square()) / plain if plain > 0 else None
snr = delta_w / weighted if weighted > 0 else None
step_ratio, crossers = None, None
if control:
W2 = (W_dq + D).to(torch.bfloat16).float()
@@ -339,7 +494,7 @@ def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None):
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()))
delta_energy=float(nD.square()), snr=snr, snr_plain=snr_plain, noise_rms=noise_rms, noise_source=noise_source)
def main():
@@ -350,12 +505,17 @@ def main():
model_config = json.load(f)
pre_quantized = model_config.get('quantization_config') is not None
quant_layers, bf16_repo = {}, None
quant_stamps, bf16_stamps = {}, {}
repo, bf16_repo, reference = None, None, None
quant_stamps, bf16_stamps, ref_stamps = {}, {}, {}
adapt = getattr(arch_mod, 'adapt_weights', None) # an arch refits deltas onto a live layout that differs from the trained one
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)
repo = QuantRepo(model_dir, model_config, arch=args.arch)
quant_stamps = stamp_index(repo.layers)
if args.reference:
reference = Bf16Repo(resolve_model_dir(args.reference))
ref_stamps = stamp_index(k[:-len('.weight')] for k in reference.weight_map if k.endswith('.weight'))
rprint(f'reference: "{reference.model_dir}" tensors={len(reference.weight_map)}')
else:
bf16_repo = Bf16Repo(model_dir)
bf16_stamps = stamp_index(k[:-len('.weight')] for k in bf16_repo.weight_map if k.endswith('.weight'))
@@ -370,6 +530,14 @@ def main():
calib_stats = {k: f.get_tensor(k) for k in f.keys()}
rprint(f'calib: "{args.calib}" layers={len(calib_stats)}')
noise_energies = {} # per module, shared by every lora: the base error does not depend on the adapter
def measure_noise(lname, err):
rms = calib_stats.get(lname)
plain = float(err.square().sum())
weighted = float((err * rms.to(device, torch.float32)).square().sum()) if rms is not None and rms.shape[-1] == err.shape[-1] else plain
noise_energies[lname] = (plain, weighted)
return noise_energies[lname]
report = {'model': model_dir, 'pre_quantized': pre_quantized, 'loras': []}
worst_effective = 1.0
def write_report():
@@ -389,10 +557,11 @@ def main():
for path in keys:
entries = mapped[path]
lname = path
noise = None
if pre_quantized:
if path not in quant_layers:
if path not in repo.layers:
lname = quant_stamps.get(path.replace('.', '_'), '')
layer = quant_layers.get(lname)
layer = repo.layers.get(lname)
if layer is None:
unmatched.append(path)
continue
@@ -403,12 +572,23 @@ def main():
if len(deq.original_shape) != 2:
non_matrix.append(path)
continue
layer = repo.materialize(lname)
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)
use_hadamard=deq.use_hadamard, use_svd=layer.svd_up is not None, svd_rank=deq.svd_rank, svd_steps=deq.svd_steps,
use_codebook=getattr(deq, 'use_codebook', False))
step_live = layer.scale.detach().to(device)
repo.release(lname)
sd_module = layer
if reference is not None:
noise = noise_energies.get(lname)
if noise is None:
W_ref = reference.get(f'{lname}.weight')
if W_ref is None:
W_ref = reference.get(f'{ref_stamps.get(lname.replace(".", "_"), "")}.weight')
if W_ref is not None and tuple(W_ref.shape) == tuple(W_dq.shape):
noise = measure_noise(lname, W_dq - W_ref.to(device, torch.float32))
else:
W = bf16_repo.get(f'{path}.weight')
if W is None:
@@ -430,10 +610,13 @@ def main():
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)
step_live = data0['scale'].detach()
noise = noise_energies.get(path) or measure_noise(path, W_dq - W.to(device, torch.float32)) # the simulated grid's own error, measured
sd_module = make_stub(W.shape)
try:
if adapt is not None and pre_quantized:
entries = [(fam, (adapt(sd_module, path, w, transformers=[repo.model]) or w) if fam == 'lora' else w) for fam, w in entries]
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), step_live=step_live)
row = analyze_module(W_dq, params, mods, calib_rms=calib_stats.get(lname), step_live=step_live, noise=noise)
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
@@ -452,16 +635,23 @@ def main():
fb_median = sorted(fb)[len(fb) // 2] if fb else None
if median_applied is not None:
worst_effective = min(worst_effective, median_applied)
snr_rows = [r for r in scored if r['snr'] is not None]
snr_median = sorted(r['snr'] for r in snr_rows)[len(snr_rows) // 2] if snr_rows else None
snr_energy = sum(r['delta_energy'] for r in snr_rows)
snr_weighted = (sum(r['snr'] * r['delta_energy'] for r in snr_rows) / snr_energy) if snr_energy > 0 else None
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})
'median_applied_rho': median_applied, 'weighted_applied_rho': weighted,
'snr_median': snr_median, 'snr_weighted': snr_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 ''))
if snr_median is not None:
rprint(f' delta over base quantization error: snr median={snr_median:.3g} energy-weighted={snr_weighted:.3g} ({snr_rows[0]["noise_source"]} noise, {"calibrated" if calib_stats else "weight-space"} norm)')
for f in failed[:3]:
rprint(f' [red]could not rebuild[/red]: {f}')
if fb:
pipelines/minimax/minimax_lora.py
+12 -8
View File
@@ -224,12 +224,13 @@ _BIND_KWARGS = dict(
)
def pruned_basis(sd_module, rank, width):
"""The AdaLN curve basis of the loaded transformer that owns ``sd_module``, or None on an unpruned model."""
from modules import shared
pipe = getattr(shared.sd_model, "pipe", shared.sd_model)
for component in ("transformer", "transformer_ref"):
transformer = getattr(pipe, component, None)
def pruned_basis(sd_module, rank, width, transformers=None):
"""The AdaLN curve basis of the transformer owning ``sd_module``, or None on an unpruned model."""
if transformers is None:
from modules import shared
pipe = getattr(shared.sd_model, "pipe", shared.sd_model)
transformers = [getattr(pipe, component, None) for component in ("transformer", "transformer_ref")]
for transformer in transformers:
basis = getattr(getattr(transformer, "time_embedder", None), "basis", None)
if basis is None or tuple(basis.shape) != (rank, width):
continue
@@ -238,14 +239,14 @@ def pruned_basis(sd_module, rank, width):
return None
def project_pruned_adaln(sd_module, network_key, w):
def project_pruned_adaln(sd_module, network_key, w, transformers=None):
"""Refit an AdaLN delta trained on the released time embedding onto the pruned curve basis: the pruned class
stores ``W @ P``, so ``up @ down`` lands exactly as ``up @ (down @ P)``."""
down = w.get("lora_down.weight")
shape = native_adapter.module_shape(sd_module)
if down is None or down.ndim != 2 or shape is None or len(shape) != 2 or down.shape[1] == shape[1]:
return None
basis = pruned_basis(sd_module, shape[1], down.shape[1])
basis = pruned_basis(sd_module, shape[1], down.shape[1], transformers)
if basis is None:
return None
projected = dict(w)
@@ -254,6 +255,9 @@ def project_pruned_adaln(sd_module, network_key, w):
return projected
adapt_weights = project_pruned_adaln # offline tools refit deltas through the same hook the loader binds
def try_load_lora(name, network_on_disk, lora_scale):
return native_adapter.try_load_lora(name, network_on_disk, lora_scale, network_alpha=file_alpha(network_on_disk), adapt_weights=project_pruned_adaln, **_BIND_KWARGS)