diff --git a/stable_audio_3/models/lora/model.py b/stable_audio_3/models/lora/model.py index ed330f8..2aacb13 100644 --- a/stable_audio_3/models/lora/model.py +++ b/stable_audio_3/models/lora/model.py @@ -6,13 +6,16 @@ import math import re +import types from functools import partial from itertools import product import torch +import torch.nn.functional as F import torch.nn.utils.parametrize as parametrize from torch import nn +from ...utils.device import disable_autocast from ...verbose import vprint @@ -296,6 +299,109 @@ def from_embedding(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bas ) +# --- fast LoRA/DoRA forward: reformulated, no W' = f(W0) materialization (see PR) --- + +_FAST_LORA_TYPES = ("lora", "dora", "dora-rows", "dora-cols") + + +def _fast_lora_param(module): + """The single fast-path-eligible LoRAParametrization on module.weight, else None.""" + pzs = getattr(module, "parametrizations", None) + if pzs is None or "weight" not in pzs: + return None + plist = pzs["weight"] + if len(plist) != 1: + return None # stacked adapters -> naive path + p = plist[0] + if type(p) is not LoRAParametrization or p.adapter_type not in _FAST_LORA_TYPES: + return None + if p.swap((0, 1)) != (0, 1): + return None # fan_in_fan_out (embedding) layout + if p.forward_fn is not p._adapter_forward_fn: + return None # disable_lora() active + if plist.original.requires_grad: + return None # trainable base weight + return p + + +@disable_autocast +def _dora_scale(module, p, W0, A, B, s): + """m / (norm(V) + 1e-12) along p._norm_dim, via the norm^2 expansion (fp32).""" + W0f = W0.view(W0.shape[0], -1).to(p.lora_A.dtype) + dim = p._norm_dim + # rowsum/colsum(W0^2) is constant for the frozen base -> cache it (never in state_dict) + ck = (W0.data_ptr(), W0._version) + cache = module.__dict__.get("_fast_lora_w0_sq") + if cache is None or cache[0] != ck: + with torch.no_grad(): + cache = (ck, (W0f * W0f).sum(dim)) + module._fast_lora_w0_sq = cache + if dim == 1: # dora-rows: per-output-row norm + cross = (F.linear(W0f, A) * B).sum(-1) + quad = (torch.matmul(B, torch.matmul(A, A.t())) * B).sum(-1) + else: # dora-cols: per-input-column norm + cross = (torch.matmul(W0f.t(), B) * A.t()).sum(-1) + quad = (torch.matmul(torch.matmul(B.t(), B), A) * A).sum(0) + norm = (cache[1] + (2.0 * s) * cross + (s * s) * quad).clamp_min(0).sqrt() + return p.magnitude / (norm + 1e-12) + + +def _fast_lora_linear_forward(self, x): + p = _fast_lora_param(self) + if p is None: + return nn.Linear.forward(self, x) # naive path (reads self.weight) + + W0 = self.parametrizations["weight"].original + bias = self.bias + is_dora = p.adapter_type != "lora" + strength = p.lora_strength + if is_dora and strength == 0: + return F.linear(x, W0, bias) # dora_forward's exact early-out + + A = p.dropout_fn(p.lora_A) # one dropout draw, shared by the data + norm terms + B = p.lora_B + s = p.scaling * strength # 0-dim tensor -> strength stays in the graph + + if is_dora and p._norm_dim == 1: + # dora-rows scales the OUTPUT by m/rownorm; bias added un-scaled AFTER + y = F.linear(x, W0) + F.linear(F.linear(x, A), B) * s + y = y * _dora_scale(self, p, W0, A, B, s).to(y.dtype) + return y if bias is None else y + bias.to(y.dtype) + + if is_dora: # dora-cols scales the INPUT by c = m/colnorm + x = x * _dora_scale(self, p, W0, A, B, s).to(x.dtype) + # lora / dora-cols: un-scaled bias rides the base GEMM (fused, bit-exact at delta=0) + return F.linear(x, W0, bias) + F.linear(F.linear(x, A), B) * s + + +def _install_fast_lora_forward(model): + """Instance-attribute swap of each LoRA-parametrized nn.Linear's forward (idempotent). + + Leaves state_dict layout and parametrization registration untouched; + custom Linear subclasses are skipped. + """ + for module in model.modules(): + if not isinstance(module, nn.Linear): + continue + if not parametrize.is_parametrized(module, "weight"): + continue + if type(module).forward is not nn.Linear.forward: + continue # custom Linear subclass: forward may do more than F.linear + if module.__dict__.get("_fast_lora_wrapped", False): + continue + module.forward = types.MethodType(_fast_lora_linear_forward, module) + module._fast_lora_wrapped = True + + +def _uninstall_fast_lora_forward(model): + """Restore the original class-level forward on every wrapped module.""" + for module in model.modules(): + if module.__dict__.get("_fast_lora_wrapped", False): + module.__dict__.pop("forward", None) + module.__dict__.pop("_fast_lora_w0_sq", None) + module._fast_lora_wrapped = False + + default_lora_config = { # specify which layers to add lora to, by default only add to linear layers nn.Linear: { "weight": partial(LoRAParametrization.from_linear, rank=4), @@ -413,6 +519,8 @@ def add_lora(model, lora_config=default_lora_config, include=None, exclude=None, # Fast path: original behavior, no name inspection needed model.apply(partial(apply_lora, lora_config=lora_config)) + _install_fast_lora_forward(model) + def add_lora_by_name(model, target_module_names, lora_config=default_lora_config, svd_bases=None): """Add LoRA to specific layers by name. Convenience wrapper around add_lora().""" add_lora(model, lora_config=lora_config, include=target_module_names, svd_bases=svd_bases) @@ -421,10 +529,12 @@ def add_lora_by_name(model, target_module_names, lora_config=default_lora_config def merge_lora(model): """merge lora parametrization to all layers in a model. This will remove all parametrization""" model.apply(partial(apply_lora, register=False, merge=True)) + _uninstall_fast_lora_forward(model) def remove_lora(model): """remove lora parametrization to all layers in a model. This will remove all parametrization""" model.apply(partial(apply_lora, register=False, merge=False)) + _uninstall_fast_lora_forward(model) def remove_lora_by_index(model, lora_index): """Remove all LoRA parametrizations with a specific lora_index. diff --git a/tests/test_fast_lora_forward.py b/tests/test_fast_lora_forward.py new file mode 100644 index 0000000..1de0029 --- /dev/null +++ b/tests/test_fast_lora_forward.py @@ -0,0 +1,282 @@ +"""Fast (reformulated) LoRA/DoRA forward vs the naive materialized parametrization. + +The fast path (the fast-lora forward (models/lora/model.py)) must +be a drop-in for torch.nn.utils.parametrize's materialize-W'-per-access +forward: same outputs, same gradients, same state_dict layout, and graceful +fallback everywhere the reformulation doesn't apply. +""" + +import copy +from functools import partial + +import pytest +import torch +import torch.nn.utils.parametrize as parametrize +from torch import nn + +from stable_audio_3.models.lora.model import ( + LoRAParametrization, + add_lora, + _uninstall_fast_lora_forward, + _install_fast_lora_forward, + merge_lora, + remove_lora, + set_lora_strength, +) + +ADAPTERS = ["lora", "dora-rows", "dora-cols"] +DEVICES = ["cpu"] + (["mps"] if torch.backends.mps.is_available() else []) + + +def make_model(bias=True, device="cpu", seed=0): + torch.manual_seed(seed) + model = nn.Sequential( + nn.Linear(64, 96, bias=bias), + nn.GELU(), + nn.Linear(96, 48, bias=bias), + ).to(device) + model.requires_grad_(False) + return model + + +def lora_config(adapter_type, rank=8, alpha=16): + return { + nn.Linear: { + "weight": partial( + LoRAParametrization.from_linear, + rank=rank, + lora_alpha=alpha, + adapter_type=adapter_type, + ), + }, + } + + +def randomize_lora(model, seed=1, magnitude_jitter=0.3): + """add_lora inits lora_B to zeros (delta == 0) — make the test non-trivial.""" + g = torch.Generator().manual_seed(seed) + for mod in model.modules(): + if not isinstance(mod, LoRAParametrization): + continue + with torch.no_grad(): + for name in ("lora_A", "lora_B"): + p = getattr(mod, name) + p.copy_(torch.randn(p.shape, generator=g).to(p) * 0.05) + for name in ("magnitude",): + if hasattr(mod, name): + p = getattr(mod, name) + p.mul_( + 1.0 + magnitude_jitter * torch.rand(p.shape, generator=g).to(p) + ) + + +def lora_trainables(model): + out = {} + for name, p in model.named_parameters(): + if name.split(".")[-1] in ("lora_A", "lora_B", "magnitude"): + out[name] = p + return out + + +def run_fwd_bwd(model, x, proj, grads=True): + for p in lora_trainables(model).values(): + p.requires_grad_(True) + if p.grad is not None: + p.grad = None + y = model(x) + if grads: + (y * proj).sum().backward() + gs = {k: p.grad.detach().clone() for k, p in lora_trainables(model).items()} + return y.detach(), gs + return y.detach(), None + + +def rel_err(a, b): + denom = b.norm().item() + if denom == 0: + return a.norm().item() + return (a - b).norm().item() / denom + + +def build_pair(adapter, bias, device, strength=None): + """Same model twice: fast-enabled and naive (fast disabled).""" + model = make_model(bias=bias, device=device) + add_lora(model, lora_config(adapter)) # add_lora auto-enables the fast path + randomize_lora(model) + if strength is not None: + set_lora_strength(model, strength) + naive = copy.deepcopy(model) + _uninstall_fast_lora_forward(naive) + _install_fast_lora_forward(model) # idempotent; explicit for clarity + return model, naive + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_output_parity_fp32(adapter, bias, device): + fast, naive = build_pair(adapter, bias, device) + x = torch.randn(4, 7, 64, device=device) + y_fast = fast(x) + y_naive = naive(x) + assert fast[0].__dict__.get("_fast_lora_wrapped"), "fast wrapper not installed" + assert not naive[0].__dict__.get("_fast_lora_wrapped", False) + torch.testing.assert_close(y_fast, y_naive, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_gradient_parity_fp32(adapter, device): + fast, naive = build_pair(adapter, bias=True, device=device) + x = torch.randn(4, 7, 64, device=device) + proj = torch.randn(4, 7, 48, device=device) + _, g_fast = run_fwd_bwd(fast, x, proj) + _, g_naive = run_fwd_bwd(naive, x, proj) + assert set(g_fast) == set(g_naive) and len(g_fast) > 0 + worst = max(rel_err(g_fast[k], g_naive[k]) for k in g_naive) + print(f"[grad-parity fp32 {adapter} {device}] worst rel err = {worst:.3e}") + assert worst <= 1e-4, f"{adapter}/{device}: worst grad rel err {worst:.3e} > 1e-4" + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="needs MPS") +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_gradient_parity_mps_fp16_autocast(adapter): + device = "mps" + fast, naive = build_pair(adapter, bias=True, device=device) + x = torch.randn(4, 7, 64, device=device) + proj = torch.randn(4, 7, 48, device=device) + with torch.amp.autocast(device, dtype=torch.float16): + y_fast, g_fast = run_fwd_bwd(fast, x, proj) + y_naive, g_naive = run_fwd_bwd(naive, x, proj) + out_err = rel_err(y_fast.float(), y_naive.float()) + worst = max(rel_err(g_fast[k].float(), g_naive[k].float()) for k in g_naive) + print( + f"[autocast fp16 {adapter}] output rel err = {out_err:.3e}, worst grad rel err = {worst:.3e}" + ) + # fp16 GEMM noise dominates; ~1e-2-class is the expected scale, this bound is a regression guard + assert out_err <= 3e-2 + assert worst <= 5e-2 + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_strength_buffer(adapter, device): + # strength 0.5: fast == naive at the same strength + fast, naive = build_pair(adapter, bias=True, device=device, strength=0.5) + x = torch.randn(3, 5, 64, device=device) + torch.testing.assert_close(fast(x), naive(x), atol=1e-5, rtol=1e-5) + + # strength 0: exactly the frozen base + set_lora_strength(fast, 0.0) + base = make_model(bias=True, device=device) # same seed -> identical base weights + y_fast = fast(x) + y_base = base(x) + assert torch.equal(y_fast, y_base), ( + f"strength=0 must be bit-exact to the base ({adapter}/{device})" + ) + + +@pytest.mark.parametrize("device", DEVICES) +def test_stacked_parametrization_falls_back(device): + fast, _ = build_pair("dora-rows", bias=True, device=device) + naive = copy.deepcopy(fast) + _uninstall_fast_lora_forward(naive) + + # stack a second adapter at lora_index=1 on BOTH copies (same init seed) + for m in (fast, naive): + torch.manual_seed(7) + for mod in m.modules(): + if isinstance(mod, nn.Linear) and parametrize.is_parametrized( + mod, "weight" + ): + p2 = LoRAParametrization.from_linear( + mod, rank=4, lora_alpha=8, adapter_type="dora-rows", lora_index=1 + ) + with torch.no_grad(): + p2.lora_B.copy_(torch.randn_like(p2.lora_B) * 0.05) + parametrize.register_parametrization(mod, "weight", p2, unsafe=True) + _install_fast_lora_forward(fast) + + for mod in fast.modules(): + if isinstance(mod, nn.Linear): + assert len(mod.parametrizations["weight"]) == 2 + + x = torch.randn(2, 9, 64, device=device) + y_fast = fast(x) # wrapper must detect the stack and fall back per-forward + y_naive = naive(x) + assert torch.equal(y_fast, y_naive), "stacked modules must use the exact naive path" + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_merge_and_remove_with_wrapper(adapter, device): + # NOTE: fresh models per case — deepcopying a parametrized module shares the + # injected ParametrizedLinear class, so un-parametrizing one copy breaks the + # other (pre-existing torch.nn.utils.parametrize behavior, wrapper-unrelated). + x = torch.randn(2, 6, 64, device=device) + + merged, _ = build_pair(adapter, bias=True, device=device) + y_before = merged(x) + merge_lora(merged) + assert not parametrize.is_parametrized(merged[0], "weight") + assert not merged[0].__dict__.get("_fast_lora_wrapped", False) + torch.testing.assert_close(merged(x), y_before, atol=1e-5, rtol=1e-5) + + removed, _ = build_pair(adapter, bias=True, device=device) + remove_lora(removed) + assert not parametrize.is_parametrized(removed[0], "weight") + base = make_model(bias=True, device=device) + torch.testing.assert_close(removed(x), base(x), atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize("adapter", ADAPTERS) +def test_state_dict_layout_unchanged(adapter): + fast = make_model() + add_lora(fast, lora_config(adapter)) + randomize_lora(fast) + + naive = make_model() + add_lora(naive, lora_config(adapter)) + _uninstall_fast_lora_forward(naive) # naive reference: class-level forward + assert not naive[0].__dict__.get("_fast_lora_wrapped", False) + + # populate the norm-constant cache before snapshotting + fast(torch.randn(2, 4, 64)) + assert fast[0].__dict__.get("_fast_lora_w0_sq") is not None or adapter == "lora" + + k_fast, k_naive = set(fast.state_dict()), set(naive.state_dict()) + assert k_fast == k_naive + assert not any("fast" in k or "w0_sq" in k for k in k_fast) + # canonical parametrize key layout intact + assert any(".parametrizations.weight.original" in k for k in k_fast) + assert any(".parametrizations.weight.0.lora_A" in k for k in k_fast) + + +@pytest.mark.parametrize("device", DEVICES) +def test_dropout_shared_between_direction_and_norm(device): + """With lora_dropout_p > 0 in train mode, fast and naive draw dropout once per + forward; under a fixed RNG seed both paths must consume the same draws.""" + cfg = { + nn.Linear: { + "weight": partial( + LoRAParametrization.from_linear, + rank=8, + lora_alpha=16, + adapter_type="dora-rows", + lora_dropout_p=0.5, + ), + }, + } + model = make_model(device=device) + add_lora(model, cfg) + randomize_lora(model) + naive = copy.deepcopy(model) + _uninstall_fast_lora_forward(naive) + model.train(), naive.train() + + x = torch.randn(2, 5, 64, device=device) + torch.manual_seed(123) + y_fast = model(x) + torch.manual_seed(123) + y_naive = naive(x) + torch.testing.assert_close(y_fast, y_naive, atol=1e-5, rtol=1e-5)