feat(lora): per-layer select stack modes klora and estlora

Two-network subject+style sets select a winner per layer instead of
summing: scores are top-K magnitude sums (klora) or Frobenius energies
(estlora), and a timestep ramp shifts layers from the subject network
toward the style network across sampling, reduced to at most one
precomputed flip per layer per pass. On sub-8-bit SDNQ the pair rides
the side-channel as separate segments flipped in place; other layers
recompute the winner from the pristine backup, so select modes force
backup mode. Selection resets per pass from the callback setup and is
gated off under model compile. estlora's measured style-discrepancy
term is exposed as an option. Adds XYZ axes for the stack settings.
This commit is contained in:
CalamitousFelicitousness
2026-07-18 02:58:11 +01:00
parent 2396185393
commit 259e15fafe
9 changed files with 426 additions and 19 deletions
+3
View File
@@ -225,6 +225,9 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
reason = ''
load_method, load_reason = lora_overrides.get_method()
from modules.lora import lora_stack
if load_method != 'native' and lora_stack.mode() != 'sum':
lora_stack.warn_once(f'method-{load_method}', f'Network stack: mode={lora_stack.mode()} method={load_method} unsupported, using sum')
if debug:
import sys
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+5 -3
View File
@@ -68,7 +68,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr
return backup_size
def network_calc_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, use_previous: bool = False, *, elimit: Callable[[], None] | None = None):
def network_calc_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, use_previous: bool = False, *, elimit: Callable[[], None] | None = None, per_net: bool = False):
if shared.opts.diffusers_offload_mode == "none":
try:
self.to(devices.device)
@@ -77,8 +77,8 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
batch_updown = None
batch_ex_bias = None
stack_deltas = None
if lora_stack.mode() in lora_stack.DENSE_MODES and network_layer_name is not None and not network_layer_name.startswith('lora_te'):
stack_deltas = [] # collect per-net deltas; combined after the loop (bias deltas stay summed)
if per_net or (lora_stack.mode() in lora_stack.DENSE_MODES and network_layer_name is not None and not network_layer_name.startswith('lora_te')):
stack_deltas = [] # collect per-net deltas; combined after the loop unless the caller wants them separate (bias deltas stay summed)
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)
@@ -142,6 +142,8 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
if elimit is not None:
elimit()
continue
if per_net:
return stack_deltas, batch_ex_bias
if stack_deltas is not None and stack_deltas:
if len(stack_deltas) >= 2:
t0 = time.time()
+3
View File
@@ -112,6 +112,9 @@ def disable_fuse():
round-trips it through its storage format. On quantized weights that is a
dequantize-add-requantize cycle per network swap whose error compounds.
"""
from modules.lora import lora_stack
if lora_stack.mode() in lora_stack.SELECT_MODES:
return True # select stack modes flip per-layer winners against the pristine backup
sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model)
if is_quantized(sd_model):
return True
+66 -1
View File
@@ -174,6 +174,7 @@ def remove_factors(self):
self.svd_up = svd_up
self.svd_down = svd_down
del self.sdnq_lora_svd_stash
lora_stack.drop(getattr(self, 'network_layer_name', None)) # a selection schedule must not outlive the segments it points into
return True
@@ -214,11 +215,23 @@ def apply_factors(self, network_layer_name, wanted_names):
def append_factors(self, ups, downs):
"""Concatenate ``[out, r]`` / ``[r, in]`` factor pairs onto the layer's svd channel and stash the originals."""
"""Concatenate ``[out, r]`` / ``[r, in]`` factor pairs onto the layer's svd channel and stash the originals.
Returns the appended parts' rank ranges plus the transposed-layout flag; the
checkpoint's own factors occupy the range before the first entry and bucket
padding lands after the last, so the ranges stay valid on the live buffers.
"""
deq = self.sdnq_dequantizer
device = self.scale.device
dtype = deq.result_dtype
orig_up, orig_down = self.svd_up, self.svd_down
orig_rank = 0
if orig_up is not None:
orig_rank = orig_up.shape[0] if deq.use_quantized_matmul else orig_up.shape[1]
segments, offset = [], orig_rank
for u in ups:
segments.append((offset, offset + u.shape[1]))
offset += u.shape[1]
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]
@@ -240,6 +253,7 @@ def append_factors(self, ups, downs):
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)
return segments, deq.use_quantized_matmul
def host_candidate(self, network_layer_name, wanted_names):
@@ -434,6 +448,57 @@ def truncate_delta(self, D, dtype):
return up_h, down_h, energy, rms is not None
def apply_select(self, network_layer_name, per_net, wanted_names):
"""Attach two networks' contributions as separate side-channel segments for per-layer selection.
Factorable members ride exactly; the rest host as their own truncated svd
with per-net cache entries. Segment ranges and selection scores register
with ``lora_stack``; the flip schedule executes from the step callback.
Returns None when the pair cannot ride the channel; the caller falls back.
"""
from sdnq.quant_utils import rotate_hadamard
deq = self.sdnq_dequantizer
changed = remove_factors(self)
if wanted_names == ():
return changed
if per_net is None or len(per_net) != 2:
return None
dtype = deq.result_dtype
lora_factor_cache.begin_pass(wanted_names)
pairs, ranks = [], []
for i, (net_name, D) in enumerate(per_net):
if D is None or D.ndim != 2 or tuple(D.shape) != tuple(deq.original_shape):
return None
net = next((n for n in l.loaded_networks if n.name == net_name), None)
module = net.modules.get(network_layer_name, None) if net is not None else None
if module is None:
return None
ranks.append(int(getattr(module, 'dim', 0) or 0) or min(int(shared.opts.lora_sdnq_host_rank), *deq.original_shape))
factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
if factors is not None:
up_i, down_i = factors
if deq.use_hadamard:
down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
else:
key = f'{network_layer_name}#{i}'
cached = lora_factor_cache.fetch(key)
if cached is not None:
up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype)
hosted_layers.append((key, cached[2], cached[3]))
else:
up_i, down_i, energy, calibrated = truncate_delta(self, D.detach().to(devices.device, torch.float32), dtype)
up_i, down_i = lora_factor_cache.store(key, up_i, down_i, energy, calibrated, float(D.detach().float().square().mean().sqrt()))
hosted_layers.append((key, energy, calibrated))
pairs.append((up_i, down_i))
d0 = per_net[0][1].detach().to(devices.device, torch.float32)
d1 = per_net[1][1].detach().to(devices.device, torch.float32)
scores, abs_sums = lora_stack.score_pair(d0, d1, ranks[0], ranks[1])
del d0, d1
segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]])
lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums)
return True
def note_fallback(self, network_layer_name):
"""Record a quantized layer taking the requantize path (summary-logged per pass); layers the routing rule sent there are counted apart."""
if getattr(self, 'sdnq_dequantizer', None) is not None and network_layer_name not in routed_layers:
+66 -15
View File
@@ -30,7 +30,7 @@ KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable
ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence
SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits)
state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'total_steps': 0, 'finalized': False}
state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_num': 0.0, 'gamma_den': 0.0, 'total_steps': 0, 'finalized': False}
warned: set = set()
@@ -134,6 +134,45 @@ def combine(named_deltas, layer_name):
return result.to(out_dtype)
def score_pair(d0, d1, rank0, rank1):
"""Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance."""
abs_sums = (float(d0.abs().sum()), float(d1.abs().sum()))
if mode() == 'klora':
k = max(1, int(rank0) * int(rank1))
s0 = float(torch.topk(d0.abs().flatten(), min(k, d0.numel()), sorted=False).values.sum())
s1 = float(torch.topk(d1.abs().flatten(), min(k, d1.numel()), sorted=False).values.sum())
else:
s0 = float(d0.float().square().sum())
s1 = float(d1.float().square().sum())
return (s0, s1), abs_sums
def register_weight_pair(layer_name, module, per_net):
"""Score and register a weight-kind selection pair; True when the layer is scheduled."""
from modules.lora import lora_common as l
if per_net is None or len(per_net) != 2:
return False
ranks, names = [], []
for net_name, d in per_net:
if d is None:
return False
net = next((n for n in l.loaded_networks if n.name == net_name), None)
net_module = net.modules.get(layer_name, None) if net is not None else None
if net_module is None:
return False
names.append(net_name)
ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64)
scores, abs_sums = score_pair(per_net[0][1].float(), per_net[1][1].float(), ranks[0], ranks[1])
register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums)
return True
def drop(layer_name):
"""Forget a layer's selection entry (its factors were removed or restored)."""
if layer_name is not None and state['entries'].pop(layer_name, None) is not None:
state['finalized'] = False
def score_topk(up, down, k):
"""K-LoRA layer score: sum of the top-K absolute delta entries (one dense materialization)."""
d = (up.to(torch.float32) @ down.to(torch.float32)).abs().flatten()
@@ -152,22 +191,28 @@ def clear():
state['entries'] = {}
state['flips'] = {}
state['gamma'] = 1.0
state['gamma_num'] = 0.0
state['gamma_den'] = 0.0
state['total_steps'] = 0
state['finalized'] = False
def register(layer_name, module, kind, segments, scores, factors=None):
"""Record a select-mode layer: its two segments (or bf16 factor pairs) and static scores.
def register(layer_name, module, kind, scores, segments=None, nets=None, abs_sums=None):
"""Record a select-mode layer for schedule finalization.
kind 'factor': segments = [(start, stop), (start, stop)] column ranges in svd_up/svd_down
with the transposed-layout flag appended; stashes both segments' values for flips.
kind 'weight': factors = [(up0, down0), (up1, down1)] kept for recompute-from-backup.
kind 'factor': segments = ((s0, s1), (t0, t1), transposed) column ranges on the svd
channel; both segments' pristine values are stashed for flips. kind 'weight': nets =
the two network names; the winner delta is recomputed from the layer backup at
selection time. abs_sums feeds the global magnitude balance (klora gamma).
"""
entry = {'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'factors': factors, 'stash': None}
entry = {'layer': layer_name, 'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'nets': nets, 'stash': None}
if kind == 'factor':
(s0, s1), (t0, t1), transposed = segments
up = module.svd_up.data
entry['stash'] = (segment_view(up, s0, s1, transposed).clone(), segment_view(up, t0, t1, transposed).clone())
if abs_sums is not None:
state['gamma_num'] += abs_sums[0]
state['gamma_den'] += abs_sums[1]
state['entries'][layer_name] = entry
state['finalized'] = False
@@ -196,6 +241,7 @@ def layer_flip_step(scores, total_steps):
def finalize(total_steps):
"""Build the inverted flip map for the pass; select-mode layers start at their step-0 winner."""
state['total_steps'] = int(total_steps)
state['gamma'] = (state['gamma_num'] / state['gamma_den']) if state['gamma_den'] > 0 else 1.0
state['flips'] = {}
for layer_name, entry in state['entries'].items():
flip_at = layer_flip_step(entry['scores'], state['total_steps'])
@@ -208,7 +254,7 @@ def finalize(total_steps):
def reset(total_steps):
"""Per-pass reset from set_callbacks_p: restore initial selections and reschedule for this pass's step count."""
if mode() not in SELECT_MODES or not state['entries']:
if mode() not in SELECT_MODES or not state['entries'] or int(total_steps) <= 0:
return
finalize(total_steps)
@@ -231,20 +277,25 @@ def apply_selection(layer_name, entry, winner):
if entry['kind'] == 'factor':
(s0, s1), (t0, t1), transposed = entry['segments']
up = module.svd_up.data
keep, drop = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1))
keep_seg, drop_seg = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1))
stash = entry['stash'][winner]
segment_view(up, keep[0], keep[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype))
segment_view(up, drop[0], drop[1], transposed).zero_()
segment_view(up, keep_seg[0], keep_seg[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype))
segment_view(up, drop_seg[0], drop_seg[1], transposed).zero_()
else:
weight_selection(module, entry, winner)
def weight_selection(module, entry, winner):
from modules.lora import lora_common as l
from modules.lora.lora_apply import network_apply_weights
backup = getattr(module, 'network_weights_backup', None)
if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy
warn_once('select-nobackup', 'Network stack: select flip skipped, no weight backup')
return
up, down = entry['factors'][winner]
weight = backup.to(device=module.weight.device, dtype=torch.float32)
delta = up.to(device=module.weight.device, dtype=torch.float32) @ down.to(device=module.weight.device, dtype=torch.float32)
module.weight.data.copy_((weight + delta.reshape(weight.shape)).to(module.weight.dtype))
net = next((n for n in l.loaded_networks if n.name == entry['nets'][winner]), None)
net_module = net.modules.get(entry['layer'], None) if net is not None else None
if net_module is None:
return
device = module.weight.device
updown = net_module.calc_updown(backup.to(device))[0]
network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it
+33
View File
@@ -91,6 +91,7 @@ def network_activate(include=None, exclude=None):
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_stack.signature() + lora_sdnq.signature() # tracked beside network_current_names so settings-only stack or mechanism changes re-apply
select_active = lora_stack.active_select(len(l.loaded_networks))
applied_layers.clear()
lora_sdnq.fallback_layers.clear() # a raise mid-pass leaves stale entries behind
lora_sdnq.hosted_layers.clear()
@@ -108,6 +109,37 @@ def network_activate(include=None, exclude=None):
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)
calced = False # tracks whether this iteration assembled the delta, so the fallthrough reuses it instead of recomputing
if select_active and component_wanted and not network_layer_name.startswith('lora_te'):
if lora_sdnq.host_candidate(module, network_layer_name, component_wanted): # sub-8-bit SDNQ pairs ride the channel as separate segments
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)
per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True)
if sel_bias is None:
applied = lora_sdnq.apply_select(module, network_layer_name, per_net, component_wanted)
if applied is not None:
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
else: # other layers select by recomputing the winner from the pristine backup at schedule time
backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse)
weights_backup = getattr(module, "network_weights_backup", None)
if weights_backup is not None and not isinstance(weights_backup, bool):
per_net, sel_bias = network_calc_weights(module, network_layer_name, elimit=elimit, per_net=True)
if sel_bias is None and lora_stack.register_weight_pair(network_layer_name, module, per_net):
network_apply_weights(module, None, None, device=device) # pristine until the schedule applies the winner
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.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):
@@ -153,6 +185,7 @@ def network_activate(include=None, exclude=None):
continue
backup_size += network_backup_weights(module, network_layer_name, component_wanted, fuse)
if not component_wanted:
lora_stack.drop(network_layer_name) # a restored layer must leave the selection schedule
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
if task is not None:
+6
View File
@@ -17,6 +17,8 @@ def set_callbacks_p(processing):
global p, warned # pylint: disable=global-statement
p = processing
warned = False
from modules.lora import lora_stack
lora_stack.reset(int(getattr(processing, 'steps', 0) or 0)) # per-pass: restore initial selections and reschedule flips before any step runs
def prompt_callback(step, kwargs):
@@ -37,6 +39,8 @@ def prompt_callback(step, kwargs):
def diffusers_callback_legacy(step: int, timestep: int, latents: torch.FloatTensor | np.ndarray):
if p is None:
return
from modules.lora import lora_stack
lora_stack.on_step(step)
if isinstance(latents, np.ndarray): # latents from Onnx pipelines is ndarray.
latents = torch.from_numpy(latents)
shared.state.sampling_step = step
@@ -56,6 +60,8 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
if kwargs is None:
kwargs = {}
t0 = time.time()
from modules.lora import lora_stack
lora_stack.on_step(step)
if shared.opts.torch_sync:
if devices.backend == "ipex":
+16
View File
@@ -104,6 +104,10 @@ class SharedSettingsStackHelper():
todo_ratio = None
teacache_thresh = None
extra_networks_default_multiplier = None
lora_stack_mode = None
lora_stack_density = None
lora_stack_alpha = None
lora_stack_discrepancy = None
disable_apply_metadata = None
disable_apply_params = None
sdnq_quant_mode = None
@@ -136,6 +140,10 @@ class SharedSettingsStackHelper():
self.sd_unet = shared.opts.sd_unet
self.sd_text_encoder = shared.opts.sd_text_encoder
self.extra_networks_default_multiplier = shared.opts.extra_networks_default_multiplier
self.lora_stack_mode = shared.opts.lora_stack_mode
self.lora_stack_density = shared.opts.lora_stack_density
self.lora_stack_alpha = shared.opts.lora_stack_alpha
self.lora_stack_discrepancy = shared.opts.lora_stack_discrepancy
self.teacache_thresh = shared.opts.teacache_thresh
self.disable_apply_metadata = shared.opts.disable_apply_metadata
self.disable_apply_params = shared.opts.disable_apply_params
@@ -148,6 +156,10 @@ class SharedSettingsStackHelper():
shared.opts.data["disable_apply_metadata"] = self.disable_apply_metadata
shared.opts.data["disable_apply_params"] = self.disable_apply_params
shared.opts.data["extra_networks_default_multiplier"] = self.extra_networks_default_multiplier
shared.opts.data["lora_stack_mode"] = self.lora_stack_mode
shared.opts.data["lora_stack_density"] = self.lora_stack_density
shared.opts.data["lora_stack_alpha"] = self.lora_stack_alpha
shared.opts.data["lora_stack_discrepancy"] = self.lora_stack_discrepancy
shared.opts.data["prompt_attention"] = self.prompt_attention
shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order
shared.opts.data["schedulers_sigma_adjust"] = self.schedulers_sigma_adjust
@@ -205,6 +217,10 @@ axis_options = [
AxisOption("[Prompt] Prompt parser", str, apply_setting("prompt_attention"), choices=lambda: ["native", "compel", "xhinker", "a1111", "fixed"]),
AxisOption("[Network] LoRA", str, apply_lora, cost=0.5, choices=list_lora),
AxisOption("[Network] LoRA strength", float, apply_lora_strength, cost=0.6),
AxisOption("[Network] LoRA stack mode", str, apply_setting("lora_stack_mode"), cost=0.6, choices=lambda: ["sum", "ties", "dare_ties", "dare_linear", "magnitude_prune", "klora", "estlora"]),
AxisOption("[Network] LoRA stack density", float, apply_setting("lora_stack_density"), cost=0.6),
AxisOption("[Network] LoRA stack ramp", float, apply_setting("lora_stack_alpha"), cost=0.6),
AxisOption("[Network] LoRA stack discrepancy", float, apply_setting("lora_stack_discrepancy"), cost=0.6),
AxisOption("[Network] Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]),
AxisOption("[Param] Width", int, apply_field("width")),
AxisOption("[Param] Height", int, apply_field("height")),
+228
View File
@@ -1562,6 +1562,228 @@ def test_sum_mode_keeps_exact_stacking():
return True
CAT_SELECT = category('stack-select')
@contextmanager
def select_mode(name, alpha=None, disc=None):
old = {k: getattr(shared.opts, k, None) for k in ('lora_stack_mode', 'lora_stack_alpha', 'lora_stack_discrepancy')}
shared.opts.lora_stack_mode = name
if alpha is not None:
shared.opts.lora_stack_alpha = alpha
if disc is not None:
shared.opts.lora_stack_discrepancy = disc
lora_stack.clear()
lora_stack.warned.clear()
try:
yield
finally:
for k, v in old.items():
setattr(shared.opts, k, v)
lora_stack.clear()
def select_pair(layer, seed0=41, seed1=42, scale1=1.0):
A1, B1, D1 = make_delta(seed=seed0, sigma=1e-2)
A2, B2, D2 = make_delta(seed=seed1, sigma=1e-2)
if scale1 != 1.0:
A2, D2 = A2 * scale1, D2 * scale1
n1 = make_net('subject', layer, A1, B1)
n2 = make_net('style', layer, A2, B2)
return n1, n2, D1, D2
def test_select_flip_schedule_end_to_end():
layer = build_layer('uint4')
n1, n2, D1, D2 = select_pair(layer)
with mock_model(lin=layer), select_mode('klora', alpha=1.5):
Wdq0 = dq(layer)
activate(n1, n2)
entry = lora_stack.state['entries'].get('lora_transformer_test')
assert entry is not None and entry['kind'] == 'factor', 'a factorable pair must register factor segments'
assert entry['segments'][0] == (0, 8) and entry['segments'][1] == (8, 16), f'segments {entry["segments"]}'
total = 20
lora_stack.reset(total)
flips = [s for s, layers in lora_stack.state['flips'].items() for _ in layers]
assert len(flips) <= 1, 'a monotone ramp allows at most one flip per layer'
eff0 = dq(layer) - Wdq0
winner0 = 0 if rho_of(eff0, D1) > rho_of(eff0, D2) else 1
for s in range(total):
lora_stack.on_step(s)
eff1 = dq(layer) - Wdq0
if flips:
assert rho_of(eff1, D2) > 0.99, 'after the flip the style delta must be selected'
assert rho_of(eff0, D1) > 0.99, 'before the flip the subject delta must be selected'
else:
assert rho_of(eff1, [D1, D2][winner0]) > 0.99
activate()
assert torch.equal(dq(layer), Wdq0), 'removal from an end-of-schedule state must restore bit-exact'
return True
def test_select_initial_style_when_ramp_starts_won():
layer = build_layer('uint4')
n1, n2, _D1, D2 = select_pair(seed0=43, seed1=44, scale1=8.0, layer=layer) # style delta dominates
with mock_model(lin=layer), select_mode('estlora', alpha=1.0, disc=0.5):
Wdq0 = dq(layer)
activate(n1, n2)
lora_stack.reset(20)
eff = dq(layer) - Wdq0
assert rho_of(eff, D2) > 0.99, 'a layer whose style side wins at step 0 must start style-selected'
return True
def test_select_flip_is_inplace_and_shape_stable():
layer = build_layer('uint4')
n1, n2, _D1, _D2 = select_pair(layer, seed0=45, seed1=46)
with mock_model(lin=layer), select_mode('klora'):
activate(n1, n2)
param_id = id(layer.svd_up)
shape = tuple(layer.svd_up.shape)
lora_stack.reset(20)
entry = lora_stack.state['entries']['lora_transformer_test']
(s0, s1), (t0, t1), transposed = entry['segments']
zeroed = lora_stack.segment_view(layer.svd_up.data, t0, t1, transposed)
kept = lora_stack.segment_view(layer.svd_up.data, s0, s1, transposed)
assert float(zeroed.abs().sum()) == 0.0 or float(kept.abs().sum()) == 0.0, 'exactly one segment must be zeroed initially'
for s in range(20):
lora_stack.on_step(s)
assert id(layer.svd_up) == param_id and tuple(layer.svd_up.shape) == shape, 'flips must mutate in place, never reassign'
return True
def test_select_matmul_transposed_layout():
layer = build_layer('uint4', use_quantized_matmul=True)
n1, n2, D1, D2 = select_pair(layer, seed0=47, seed1=48)
with mock_model(lin=layer), select_mode('klora'):
Wdq0 = dq(layer)
activate(n1, n2)
entry = lora_stack.state['entries']['lora_transformer_test']
assert entry['segments'][2] is True, 'quantized-matmul layout must register as transposed'
lora_stack.reset(20)
eff = dq(layer) - Wdq0
assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.99, 'initial selection must realize one delta exactly'
activate()
assert torch.equal(dq(layer), Wdq0)
return True
def test_select_per_net_hosted_pair():
layer = build_layer('uint4')
torch.manual_seed(49)
Dd1 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3 # rank inside the host cap so truncation is near-lossless
Dd2 = (torch.randn(OUT_F, 24, device=DEVICE) @ torch.randn(24, IN_F, device=DEVICE)) * 1e-3
n1 = make_dense_net('lk1', layer, Dd1)
n2 = make_dense_net('lk2', layer, Dd2)
with host_rank(32), mock_model(lin=layer), select_mode('klora'):
Wdq0 = dq(layer)
activate(n1, n2)
entry = lora_stack.state['entries'].get('lora_transformer_test')
assert entry is not None, 'non-factorable pairs must register through per-net hosting'
assert entry['segments'][0] == (0, 24) and entry['segments'][1] == (24, 48), f'segments {entry["segments"]}' # hosting stores the effective rank (24), not the cap
lora_stack.reset(20)
eff = dq(layer) - Wdq0
best = max(rho_of(eff, Dd1), rho_of(eff, Dd2))
assert best > 0.9, f'initial selection must realize one hosted delta, rho={best:.3f}'
activate()
assert torch.equal(dq(layer), Wdq0)
return True
def test_select_reset_restores_initial_state():
layer = build_layer('uint4')
n1, n2, _D1, _D2 = select_pair(layer, seed0=51, seed1=52)
with mock_model(lin=layer), select_mode('klora'):
activate(n1, n2)
lora_stack.reset(20)
initial = dq(layer)
for s in range(20):
lora_stack.on_step(s)
lora_stack.reset(20)
assert torch.equal(dq(layer), initial), 'a fresh pass must restore the initial selection without re-activation'
return True
def test_select_deactivate_from_midflip():
layer = build_layer('uint4')
n1, n2, _D1, _D2 = select_pair(layer, seed0=53, seed1=54)
with mock_model(lin=layer), select_mode('klora'):
Wdq0 = dq(layer)
activate(n1, n2)
lora_stack.reset(20)
for s in range(10):
lora_stack.on_step(s)
activate()
assert torch.equal(dq(layer), Wdq0), 'removal mid-schedule must restore bit-exact'
assert not lora_stack.state['entries'], 'removal must drop the selection entry'
return True
def test_select_requires_exactly_two_nets():
layer = build_layer('uint4')
A3, B3, _D3 = make_delta(seed=55)
n1, n2, _D1, _D2 = select_pair(layer, seed0=56, seed1=57)
n3 = make_net('third', layer, A3, B3)
with mock_model(lin=layer), select_mode('klora'):
activate(n1, n2, n3)
assert hasattr(layer, 'sdnq_lora_svd_stash'), 'three nets must fall back to the exact concat path'
assert not lora_stack.state['entries'], 'no selection entries outside the two-net case'
activate()
return True
def test_select_gated_off_when_compiled():
layer = build_layer('uint4')
n1, n2, _D1, _D2 = select_pair(layer, seed0=58, seed1=59)
old_compile = getattr(shared.opts, 'cuda_compile', None)
try:
shared.opts.cuda_compile = ['Model']
with mock_model(lin=layer), select_mode('klora'):
activate(n1, n2)
assert not lora_stack.state['entries'], 'select must gate off under model compile'
assert hasattr(layer, 'sdnq_lora_svd_stash'), 'gated select behaves as sum'
activate()
finally:
shared.opts.cuda_compile = old_compile
return True
def test_est_energy_matches_full_frobenius():
torch.manual_seed(60)
up = torch.randn(64, 8, device=DEVICE)
down = torch.randn(8, 96, device=DEVICE)
gram = lora_stack.score_energy(up, down)
full = float((up @ down).square().sum())
assert abs(gram - full) / full < 1e-5, f'{gram} vs {full}'
return True
def test_select_weight_kind_plain_layer():
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.02)
lin.network_layer_name = 'lora_transformer_plain'
lin.network_current_names = ()
A1, B1, D1 = make_delta(seed=61, sigma=1e-2)
A2, B2, D2 = make_delta(seed=62, sigma=1e-2)
n1 = make_net('w1', lin, A1, B1)
n2 = make_net('w2', lin, A2, B2)
W0 = lin.weight.detach().float().clone()
with mock_model(lin=lin), select_mode('klora'):
activate(n1, n2)
entry = lora_stack.state['entries'].get('lora_transformer_plain')
assert entry is not None and entry['kind'] == 'weight', 'plain layers must register weight-kind selection'
assert torch.equal(lin.weight.detach().float(), W0), 'weights stay pristine until the schedule applies a winner'
lora_stack.reset(20)
eff = lin.weight.detach().float() - W0
assert max(rho_of(eff, D1), rho_of(eff, D2)) > 0.95, 'initial selection must apply one delta from backup'
for s in range(20):
lora_stack.on_step(s)
activate()
assert torch.equal(lin.weight.detach().float(), W0), 'restore-only pass must return the pristine weight'
return True
CAT_COMPILE = category('compile')
@@ -1746,6 +1968,12 @@ def run_tests():
test_magnitude_prune_keeps_top_density, test_dense_two_plain_loras_hosted_not_summed, test_single_net_ignores_dense_mode,
test_te_layer_stays_plain_sum, test_sum_mode_keeps_exact_stacking]:
run_test(CAT_STACK, fn)
log.warning('=== Stack modes: select ===')
for fn in [test_select_flip_schedule_end_to_end, test_select_initial_style_when_ramp_starts_won, test_select_flip_is_inplace_and_shape_stable,
test_select_matmul_transposed_layout, test_select_per_net_hosted_pair, test_select_reset_restores_initial_state,
test_select_deactivate_from_midflip, test_select_requires_exactly_two_nets, test_select_gated_off_when_compiled,
test_est_energy_matches_full_frobenius, test_select_weight_kind_plain_layer]:
run_test(CAT_SELECT, fn)
log.warning('=== Compile ===')
for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse]:
run_test(CAT_COMPILE, fn)