fix(lora): scale the dora diff before the decompose norm

finalize_updown ran apply_weight_decompose on the unscaled delta and
multiplied the result by alpha/rank afterward. LyCORIS and ComfyUI both
bake alpha/rank into the diff before computing the row norms, so any
DoRA with alpha != rank renormalized against the wrong merged weight
(64% relative delta error for kohya-style alpha=1 rank=8; exact only
when alpha == rank, which full-matrix LoKR forces).

- scale updown by calc_scale() before apply_weight_decompose; apply
  only the multiplier afterward
- multiplier lerps the full merged delta (0 disables, 1 equals the
  trainer output); LyCORIS weight-mode ratio interpolation leaves the
  diff applied at multiplier 0 and is not used
- add a numeric regression test mirroring the LyCORIS forward reference
This commit is contained in:
CalamitousFelicitousness
2026-07-12 23:38:30 +01:00
parent ca729a01ca
commit 387a349dfd
2 changed files with 47 additions and 1 deletions
+7 -1
View File
@@ -268,7 +268,13 @@ class NetworkModule:
if ex_bias is not None:
ex_bias = ex_bias * self.multiplier()
if self.dora_scale is not None:
updown = self.apply_weight_decompose(updown, orig_weight)
# LyCORIS/ComfyUI convention: alpha/rank is baked into the diff
# before the decompose norm. The multiplier then lerps the full
# merged delta (ComfyUI semantics: 0 disables, 1 equals the
# trainer's output; LyCORIS weight-mode ratio interpolation is
# not used since it leaves the diff applied at multiplier 0).
updown = self.apply_weight_decompose(updown * self.calc_scale(), orig_weight)
return updown * self.multiplier(), ex_bias
return updown * self.calc_scale() * self.multiplier(), ex_bias
def calc_updown(self, target):
+40
View File
@@ -1355,6 +1355,45 @@ def test_full_calc_updown_shape():
return True
def test_dora_ordering_matches_lycoris():
"""DoRA merge equals the LyCORIS forward reference at multiplier 1 and
lerps the full merged delta at other multipliers (ComfyUI semantics).
Reference mirrors lycoris locon forward + apply_weight_decompose with
wd_on_out=True: alpha/rank scales the diff BEFORE the row norms. Uses
alpha != rank (the kohya-style case the old decompose-before-scale
ordering got ~64% wrong).
"""
torch.manual_seed(0)
alpha = 1.0
w_base = torch.randn(QKV_OUT, HIDDEN)
down = torch.randn(RANK_LORA, HIDDEN)
up = torch.randn(QKV_OUT, RANK_LORA)
dora = w_base.reshape(QKV_OUT, -1).norm(dim=1, keepdim=True)
def reference_delta(mult):
diff = (up @ down) * (alpha / RANK_LORA)
merged = w_base + diff
norm = merged.reshape(QKV_OUT, -1).norm(dim=1).reshape(QKV_OUT, 1) + torch.finfo(w_base.dtype).eps
return (merged * (dora / norm) - w_base) * mult
sd = {
'transformer.transformer_blocks.0.attn.to_q.lora_A.weight': down,
'transformer.transformer_blocks.0.attn.to_q.lora_B.weight': up,
'transformer.transformer_blocks.0.attn.to_q.alpha': torch.tensor(alpha),
'transformer.transformer_blocks.0.attn.to_q.dora_scale': dora,
}
for mult in (1.0, 0.5):
net = _load_via(F.try_load_lora, sd)
assert net is not None and len(net.modules) == 1
mod = make_network_for_module(next(iter(net.modules.values())), te_mul=mult, unet_mul=mult)
updown, _ex_bias = mod.calc_updown(w_base.clone())
ref = reference_delta(mult)
rel = ((updown - ref).norm() / (ref.norm() + 1e-12)).item()
assert torch.allclose(updown, ref, rtol=1e-4, atol=1e-5), f'mult={mult}: rel err {rel:.4f}'
return True
# ============================================================
# Tests — apply path / regressions in shared infra
# ============================================================
@@ -1479,6 +1518,7 @@ def run_tests():
test_ia3_calc_updown_shape,
test_glora_calc_updown_shape,
test_full_calc_updown_shape,
test_dora_ordering_matches_lycoris,
]:
run_test(CAT_MATH, fn)