mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
fix(lora): apply bias deltas onto the bias, not the weight
network_add_weights defaulted its base tensor to self.weight for the bias delta as well, so in fuse mode a diff_b was added to the weight matrix and the result written into the bias. Layers where in and out differ threw a shape error and had the weight matrix installed as their bias, square layers broadcast silently, and either way the summary still counted the delta as applied. - pick the base tensor from the bias flag - name the layer, target and both shapes in the mismatch error - return which of (weight, bias) took a write, count the rest as refused - report refused= on partially applied and partially removed networks - cover both apply paths in test/test-lora-apply.py
This commit is contained in:
+25
-13
@@ -157,14 +157,16 @@ def assign_weight(self: torch.nn.Module, new_weight: torch.Tensor, device: torch
|
||||
self.weight = torch.nn.Parameter(new_weight, requires_grad=False)
|
||||
|
||||
|
||||
def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, model_weights: torch.Tensor | None = None, lora_weights: torch.Tensor = None, deactivate: bool = False, device: torch.device = None, bias: bool = False):
|
||||
def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, model_weights: torch.Tensor | None = None, lora_weights: torch.Tensor = None, deactivate: bool = False, device: torch.device = None, bias: bool = False) -> bool:
|
||||
"""Add a delta onto the module's weight or bias; False when nothing was written."""
|
||||
if lora_weights is None:
|
||||
return
|
||||
return False
|
||||
if deactivate:
|
||||
lora_weights *= -1
|
||||
if model_weights is None: # weights are used if provided-from-backup else use self.weight
|
||||
model_weights = self.weight
|
||||
if model_weights is None: # weights are used if provided-from-backup else use the live tensor the delta targets
|
||||
model_weights = self.bias if bias else self.weight
|
||||
weight, new_weight = None, None
|
||||
written = True
|
||||
if not bias and hasattr(self, "sdnq_dequantizer"):
|
||||
try:
|
||||
from sdnq import SDNQConfig, sdnq_quantize_layer
|
||||
@@ -226,21 +228,26 @@ def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Group
|
||||
del dequant_weight
|
||||
except Exception as e:
|
||||
log.error(f'Network load: type=LoRA quant=sdnq cls={self.__class__.__name__} weight={self.weight} lora_weights={lora_weights} {e}')
|
||||
written = False
|
||||
else:
|
||||
try:
|
||||
new_weight = model_weights.to(devices.device) + lora_weights.to(devices.device)
|
||||
except Exception as e:
|
||||
log.warning(f'Network load: {e}')
|
||||
if 'The size of tensor' in str(e):
|
||||
log.error(f'Network load: type=LoRA model={shared.sd_model.__class__.__name__} incompatible lora shape')
|
||||
target = 'bias' if bias else 'weight'
|
||||
log.error(f'Network load: type=LoRA model={shared.sd_model.__class__.__name__} layer="{getattr(self, "network_layer_name", None)}" target={target} shape={tuple(model_weights.shape)} lora={tuple(lora_weights.shape)} incompatible lora shape')
|
||||
new_weight = model_weights
|
||||
written = False
|
||||
else:
|
||||
new_weight = model_weights + lora_weights # try without device cast
|
||||
assign_weight(self, new_weight, device, bias=bias)
|
||||
del model_weights, lora_weights, new_weight, weight # required to avoid memory leak
|
||||
return written
|
||||
|
||||
|
||||
def network_apply_direct(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, updown: torch.Tensor, ex_bias: torch.Tensor, deactivate: bool = False, device: torch.device = devices.device):
|
||||
def network_apply_direct(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, updown: torch.Tensor, ex_bias: torch.Tensor, deactivate: bool = False, device: torch.device = devices.device) -> tuple[bool, bool]:
|
||||
"""Add the deltas onto the live tensors; returns which of (weight, bias) was written."""
|
||||
weights_backup = getattr(self, "network_weights_backup", False)
|
||||
bias_backup = getattr(self, "network_bias_backup", False)
|
||||
if not isinstance(weights_backup, bool): # remove previous backup if we switched settings
|
||||
@@ -248,37 +255,41 @@ def network_apply_direct(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
|
||||
if not isinstance(bias_backup, bool):
|
||||
bias_backup = True
|
||||
if not weights_backup and not bias_backup:
|
||||
return
|
||||
return False, False
|
||||
t0 = time.time()
|
||||
weight_written, bias_written = False, False
|
||||
|
||||
if weights_backup:
|
||||
if updown is not None and len(self.weight.shape) == 4 and self.weight.shape[1] == 9: # inpainting model so zero pad updown to make channel 4 to 9
|
||||
updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
|
||||
if updown is not None:
|
||||
network_add_weights(self, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
|
||||
weight_written = network_add_weights(self, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
|
||||
|
||||
if bias_backup:
|
||||
if ex_bias is not None:
|
||||
network_add_weights(self, lora_weights=ex_bias, deactivate=deactivate, device=device, bias=True)
|
||||
bias_written = network_add_weights(self, lora_weights=ex_bias, deactivate=deactivate, device=device, bias=True)
|
||||
|
||||
if hasattr(self, "qweight") and hasattr(self, "freeze"):
|
||||
self.freeze()
|
||||
|
||||
l.timer.apply += time.time() - t0
|
||||
return weight_written, bias_written
|
||||
|
||||
|
||||
def network_apply_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, updown: torch.Tensor, ex_bias: torch.Tensor, device: torch.device, deactivate: bool = False):
|
||||
def network_apply_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, updown: torch.Tensor, ex_bias: torch.Tensor, device: torch.device, deactivate: bool = False) -> tuple[bool, bool]:
|
||||
"""Add the deltas onto the backup copies; returns which of (weight, bias) was written."""
|
||||
weights_backup = getattr(self, "network_weights_backup", None)
|
||||
bias_backup = getattr(self, "network_bias_backup", None)
|
||||
if weights_backup is None and bias_backup is None:
|
||||
return
|
||||
return False, False
|
||||
t0 = time.time()
|
||||
weight_written, bias_written = False, False
|
||||
|
||||
if weights_backup is not None and not isinstance(weights_backup, bool):
|
||||
if updown is not None and len(weights_backup.shape) == 4 and weights_backup.shape[1] == 9: # inpainting model. zero pad updown to make channel[1] 4 to 9
|
||||
updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
|
||||
if updown is not None:
|
||||
network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
|
||||
weight_written = network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
|
||||
else:
|
||||
assign_weight(self, weights_backup, device)
|
||||
if hasattr(self, "sdnq_dequantizer_backup"):
|
||||
@@ -297,7 +308,7 @@ def network_apply_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gro
|
||||
|
||||
if bias_backup is not None and not isinstance(bias_backup, bool):
|
||||
if ex_bias is not None:
|
||||
network_add_weights(self, model_weights=bias_backup, lora_weights=ex_bias, deactivate=deactivate, device=device, bias=True)
|
||||
bias_written = network_add_weights(self, model_weights=bias_backup, lora_weights=ex_bias, deactivate=deactivate, device=device, bias=True)
|
||||
else:
|
||||
assign_weight(self, bias_backup, device, bias=True)
|
||||
|
||||
@@ -305,3 +316,4 @@ def network_apply_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gro
|
||||
self.freeze()
|
||||
|
||||
l.timer.apply += time.time() - t0
|
||||
return weight_written, bias_written
|
||||
|
||||
@@ -82,6 +82,7 @@ def network_activate(include=None, exclude=None):
|
||||
pbar = nullcontext()
|
||||
applied_weight = 0
|
||||
applied_bias = 0
|
||||
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 ()
|
||||
applied_layers.clear()
|
||||
@@ -109,13 +110,14 @@ def network_activate(include=None, exclude=None):
|
||||
else:
|
||||
batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit)
|
||||
if shared.opts.lora_fuse_native:
|
||||
network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
|
||||
weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
|
||||
else:
|
||||
network_apply_weights(module, batch_updown, batch_ex_bias, device=device)
|
||||
weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device)
|
||||
if batch_updown is not None or batch_ex_bias is not None:
|
||||
applied_layers.append(network_layer_name)
|
||||
applied_weight += 1 if batch_updown is not None else 0
|
||||
applied_bias += 1 if batch_ex_bias is not None else 0
|
||||
applied_weight += 1 if weight_written else 0
|
||||
applied_bias += 1 if bias_written else 0
|
||||
refused += (batch_updown is not None and not weight_written) + (batch_ex_bias is not None and not bias_written) # a delta the module would not take leaves that layer on its base value
|
||||
batch_updown, batch_ex_bias = None, None
|
||||
del batch_updown, batch_ex_bias
|
||||
module.network_current_names = component_wanted
|
||||
@@ -128,8 +130,10 @@ def network_activate(include=None, exclude=None):
|
||||
global native_active # pylint: disable=global-statement
|
||||
native_active = len(l.loaded_networks) > 0
|
||||
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} 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={shared.opts.lora_fuse_native}:{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")
|
||||
@@ -171,6 +175,7 @@ def network_deactivate(include=None, exclude=None):
|
||||
else:
|
||||
task = None
|
||||
pbar = nullcontext()
|
||||
refused = 0
|
||||
with devices.inference_context(), pbar:
|
||||
applied_layers.clear()
|
||||
for component in modules.keys():
|
||||
@@ -185,18 +190,21 @@ def network_deactivate(include=None, exclude=None):
|
||||
device = group_offload_strip(sd_model, component, group_stripped)
|
||||
batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit)
|
||||
if shared.opts.lora_fuse_native:
|
||||
network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
|
||||
weight_written, bias_written = network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
|
||||
else:
|
||||
network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
|
||||
weight_written, bias_written = network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
|
||||
if batch_updown is not None or batch_ex_bias is not None:
|
||||
applied_layers.append(network_layer_name)
|
||||
refused += (batch_updown is not None and not weight_written) + (batch_ex_bias is not None and not bias_written) # a delta the module would not take stays applied on that layer
|
||||
del batch_updown, batch_ex_bias
|
||||
module.network_current_names = ()
|
||||
if task is not None:
|
||||
pbar.update(task, advance=1, description=f'networks={len(l.previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}')
|
||||
l.timer.deactivate = time.time() - t0
|
||||
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)} 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={shared.opts.lora_fuse_native}:{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")
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Offline unit tests for the native LoRA apply paths.
|
||||
|
||||
Two paths write a delta onto a module: fuse mode (``network_apply_direct``)
|
||||
adds it straight onto the live tensors, backup mode (``network_apply_weights``)
|
||||
adds it onto a cloned copy. Both funnel through ``network_add_weights``, and
|
||||
only backup mode names the base tensor a bias delta targets, so the fuse-side
|
||||
bias case is the one with nothing pinning it.
|
||||
|
||||
Bias deltas ride the ``diff_b`` key that trainers pair with the weight LoRA on
|
||||
projection layers. A Linear's weight is ``[out, in]`` and its bias is ``[out]``,
|
||||
so reading the wrong base tensor for a bias delta throws whenever ``in != out``
|
||||
and broadcasts silently when they match: both shapes are covered here.
|
||||
|
||||
No running server required.
|
||||
|
||||
Usage:
|
||||
python test/test-lora-apply.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
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 devices # pylint: disable=wrong-import-position
|
||||
from modules.lora.lora_apply import network_apply_direct, network_apply_weights # pylint: disable=wrong-import-position
|
||||
|
||||
devices.device = torch.device('cpu') # apply moves operands to devices.device; keep the suite off the gpu
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test infrastructure
|
||||
# ============================================================
|
||||
|
||||
results: dict[str, dict] = {}
|
||||
|
||||
CAT_FUSE = 'fuse mode'
|
||||
CAT_BACKUP = 'backup mode'
|
||||
|
||||
|
||||
def record(cat: str, passed: bool, name: str, detail: str = ''):
|
||||
if cat not in results:
|
||||
results[cat] = {'passed': 0, 'failed': 0}
|
||||
status = 'PASS' if passed else 'FAIL'
|
||||
results[cat]['passed' if passed else 'failed'] += 1
|
||||
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()
|
||||
|
||||
|
||||
CPU = torch.device('cpu')
|
||||
|
||||
|
||||
def make_linear(out_features: int, in_features: int, seed: int = 0):
|
||||
"""Linear with deterministic weight and bias, plus copies of both as they started."""
|
||||
torch.manual_seed(seed)
|
||||
module = torch.nn.Linear(in_features, out_features, bias=True)
|
||||
with torch.no_grad():
|
||||
module.weight.copy_(torch.randn(out_features, in_features) * 0.02)
|
||||
module.bias.copy_(torch.randn(out_features) * 0.02)
|
||||
return module, module.weight.detach().clone(), module.bias.detach().clone()
|
||||
|
||||
|
||||
def stamp_fuse(module):
|
||||
"""Mark the module as network_backup_weights leaves it in fuse mode: no tensor backup."""
|
||||
module.network_weights_backup = True
|
||||
module.network_bias_backup = True
|
||||
|
||||
|
||||
def stamp_backup(module, weight, bias):
|
||||
"""Mark the module as network_backup_weights leaves it in backup mode: cloned tensors."""
|
||||
module.network_weights_backup = weight.clone().to(CPU)
|
||||
module.network_bias_backup = bias.clone().to(CPU)
|
||||
|
||||
|
||||
def assert_close(actual, expected, label):
|
||||
assert actual.shape == expected.shape, f'{label} shape {tuple(actual.shape)} != {tuple(expected.shape)}'
|
||||
assert torch.allclose(actual, expected, atol=1e-6), f'{label} values drifted'
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Fuse mode
|
||||
# ============================================================
|
||||
|
||||
|
||||
def test_fuse_bias_delta_asymmetric():
|
||||
"""A diff_b delta on a Linear whose in and out differ lands on the bias, not the weight."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_fuse(module)
|
||||
updown = torch.full_like(w0, 0.5)
|
||||
ex_bias = torch.full_like(b0, 0.25)
|
||||
written = network_apply_direct(module, updown, ex_bias, device=CPU)
|
||||
assert written == (True, True), f'reported {written}'
|
||||
assert_close(module.weight.detach(), w0 + 0.5, 'weight')
|
||||
assert_close(module.bias.detach(), b0 + 0.25, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_fuse_bias_delta_square():
|
||||
"""The same delta on a square Linear, where a wrong base tensor broadcasts instead of throwing."""
|
||||
module, w0, b0 = make_linear(16, 16)
|
||||
stamp_fuse(module)
|
||||
updown = torch.full_like(w0, 0.5)
|
||||
ex_bias = torch.full_like(b0, 0.25)
|
||||
network_apply_direct(module, updown, ex_bias, device=CPU)
|
||||
assert module.bias.dim() == 1, f'bias became {module.bias.dim()}d'
|
||||
assert_close(module.weight.detach(), w0 + 0.5, 'weight')
|
||||
assert_close(module.bias.detach(), b0 + 0.25, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_fuse_weight_only_leaves_bias():
|
||||
"""A LoRA with no bias delta leaves the bias untouched."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_fuse(module)
|
||||
written = network_apply_direct(module, torch.full_like(w0, 0.5), None, device=CPU)
|
||||
assert written == (True, False), f'reported {written}' # nothing to write is reported the same as refused; the caller knows which by whether it passed a delta
|
||||
assert_close(module.weight.detach(), w0 + 0.5, 'weight')
|
||||
assert_close(module.bias.detach(), b0, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_fuse_deactivate_restores():
|
||||
"""Deactivate subtracts the same deltas, returning both tensors to their loaded values."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_fuse(module)
|
||||
updown = torch.full_like(w0, 0.5)
|
||||
ex_bias = torch.full_like(b0, 0.25)
|
||||
network_apply_direct(module, updown.clone(), ex_bias.clone(), device=CPU)
|
||||
network_apply_direct(module, updown.clone(), ex_bias.clone(), device=CPU, deactivate=True)
|
||||
assert_close(module.weight.detach(), w0, 'weight')
|
||||
assert_close(module.bias.detach(), b0, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_fuse_mismatched_bias_delta_is_refused():
|
||||
"""A bias delta that genuinely does not fit is dropped, leaving the bias intact.
|
||||
|
||||
The refusal has to reach the caller: network_activate counts the layer as
|
||||
refused rather than applied, which is what keeps the summary line honest.
|
||||
"""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_fuse(module)
|
||||
written = network_apply_direct(module, torch.full_like(w0, 0.5), torch.full((7,), 0.25), device=CPU)
|
||||
assert written == (True, False), f'reported {written}'
|
||||
assert_close(module.weight.detach(), w0 + 0.5, 'weight')
|
||||
assert_close(module.bias.detach(), b0, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_fuse_mismatched_weight_delta_is_refused():
|
||||
"""A weight delta that does not fit is dropped while the bias delta still lands."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_fuse(module)
|
||||
written = network_apply_direct(module, torch.full((32, 5), 0.5), torch.full_like(b0, 0.25), device=CPU)
|
||||
assert written == (False, True), f'reported {written}'
|
||||
assert_close(module.weight.detach(), w0, 'weight')
|
||||
assert_close(module.bias.detach(), b0 + 0.25, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Backup mode
|
||||
# ============================================================
|
||||
|
||||
|
||||
def test_backup_bias_delta_asymmetric():
|
||||
"""Backup mode adds the deltas onto the cloned base rather than the live tensors."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_backup(module, w0, b0)
|
||||
updown = torch.full_like(w0, 0.5)
|
||||
ex_bias = torch.full_like(b0, 0.25)
|
||||
written = network_apply_weights(module, updown, ex_bias, device=CPU)
|
||||
assert written == (True, True), f'reported {written}'
|
||||
assert_close(module.weight.detach(), w0 + 0.5, 'weight')
|
||||
assert_close(module.bias.detach(), b0 + 0.25, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_backup_reapply_is_not_cumulative():
|
||||
"""Applying twice from the same backup yields one delta, not two."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_backup(module, w0, b0)
|
||||
for _ in range(2):
|
||||
network_apply_weights(module, torch.full_like(w0, 0.5), torch.full_like(b0, 0.25), device=CPU)
|
||||
assert_close(module.weight.detach(), w0 + 0.5, 'weight')
|
||||
assert_close(module.bias.detach(), b0 + 0.25, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def test_backup_restore_without_delta():
|
||||
"""Applying with no deltas restores the module to its backup."""
|
||||
module, w0, b0 = make_linear(32, 8)
|
||||
stamp_backup(module, w0, b0)
|
||||
network_apply_weights(module, torch.full_like(w0, 0.5), torch.full_like(b0, 0.25), device=CPU)
|
||||
network_apply_weights(module, None, None, device=CPU)
|
||||
assert_close(module.weight.detach(), w0, 'weight')
|
||||
assert_close(module.bias.detach(), b0, 'bias')
|
||||
return True
|
||||
|
||||
|
||||
def run_tests():
|
||||
t0 = time.time()
|
||||
log.warning('=== fuse mode ===')
|
||||
for fn in [
|
||||
test_fuse_bias_delta_asymmetric,
|
||||
test_fuse_bias_delta_square,
|
||||
test_fuse_weight_only_leaves_bias,
|
||||
test_fuse_deactivate_restores,
|
||||
test_fuse_mismatched_bias_delta_is_refused,
|
||||
test_fuse_mismatched_weight_delta_is_refused,
|
||||
]:
|
||||
run_test(CAT_FUSE, fn)
|
||||
|
||||
log.warning('=== backup mode ===')
|
||||
for fn in [
|
||||
test_backup_bias_delta_asymmetric,
|
||||
test_backup_reapply_is_not_cumulative,
|
||||
test_backup_restore_without_delta,
|
||||
]:
|
||||
run_test(CAT_BACKUP, fn)
|
||||
|
||||
elapsed = time.time() - t0
|
||||
log.warning('=== Results ===')
|
||||
total_pass = 0
|
||||
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__':
|
||||
ok = run_tests()
|
||||
sys.exit(0 if ok else 1)
|
||||
Reference in New Issue
Block a user