feat(assets): SA3 音效补齐 radio 11+1 与 dungeon 13源+15pan（token 走 hf-mirror 实测可用）
- 模型 17/17 文件经镜像全量缓存（gated+token，直连本机不通；sfx.md/setup.md/asset-pipeline.md 镜像结论同步更新） - 修复迁移遗留：stable_audio_3/models/ 源码子包被 .gitignore models/ 规则误伤丢失，从上游恢复并加白名单 - radio：arg_muffled/brake_screech 等叙事件落地（接线原有 L111/L154 零改动）+ phone_ring SYN→SA3 升级；dungeon 13 源 A/B 对齐（RMS 偏差 ≤0.3dB）全量同名替换，SYN 原版备份 tools/asset_fix_backup/sa3_ab_20260902/ - 关键玩法语义达标：echo 距离档 0.83<1.43<2.18<2.98s 严格递增；one-shot 真峰值全 ≤-3dBFS（解码回读法） - check_all_assets 8 游戏全绿（radio 254/254、dungeon 303/303）；playtest radio 221/dungeon 1610 全绿
This commit is contained in:
@@ -0,0 +1,28 @@
|
||||
from .model import LoRAParametrization, add_lora, default_lora_config, merge_lora, remove_lora, remove_lora_by_index, set_lora_strength
|
||||
from .utils import (
|
||||
apply_to_lora,
|
||||
convert_lora_ckpt_to_safetensors,
|
||||
disable_lora,
|
||||
enable_lora,
|
||||
filter_lora_layers,
|
||||
get_lora_count,
|
||||
has_lora,
|
||||
get_bias_params,
|
||||
get_lora_params,
|
||||
get_lora_layers,
|
||||
get_lora_state_dict,
|
||||
load_lora_checkpoint,
|
||||
load_multiple_lora,
|
||||
name_is_lora,
|
||||
remap_lora_state_dict,
|
||||
save_lora_safetensors,
|
||||
select_lora,
|
||||
tie_weights,
|
||||
untie_weights,
|
||||
infer_global_rank,
|
||||
cast_base_to_precision,
|
||||
detect_dora_variant,
|
||||
prepare_dora_state_dict,
|
||||
resolve_adapter_type
|
||||
)
|
||||
from .loader import load_and_apply_loras
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,112 @@
|
||||
"""Helpers for loading pre-trained LoRA checkpoints onto a live model."""
|
||||
import os
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
|
||||
from .model import LoRAParametrization, add_lora
|
||||
from .utils import (
|
||||
get_lora_layers,
|
||||
infer_global_rank,
|
||||
load_lora_checkpoint,
|
||||
prepare_dora_state_dict,
|
||||
remap_lora_state_dict,
|
||||
resolve_adapter_type,
|
||||
)
|
||||
from ...verbose import vprint
|
||||
|
||||
|
||||
def load_and_apply_loras(model, lora_ckpt_paths, model_type, svd_bases_path=None):
|
||||
"""Load LoRA checkpoints from disk and attach them to `model`.
|
||||
|
||||
- Uses a two-pass approach: first resolves each LoRA's adapter type, then
|
||||
loads SVD bases (only if any LoRA is -XS), then applies each LoRA.
|
||||
- Handles legacy "dora" adapter_type via `resolve_adapter_type`.
|
||||
- Only passes svd_bases to -XS LoRAs to avoid spurious "no -XS layers"
|
||||
messages for non-XS LoRAs.
|
||||
- Sets `model.use_lora = True` and `model.lora_names = [...]` as a
|
||||
convenience for downstream UI code.
|
||||
|
||||
Args:
|
||||
model: The loaded model. For `diffusion_cond` / `diffusion_cond_inpaint`
|
||||
types, LoRAs are applied to `model.model` and `model.conditioner`;
|
||||
otherwise they are applied to `model` directly.
|
||||
lora_ckpt_paths: List of paths to `.ckpt` or `.safetensors` files.
|
||||
model_type: String from the model config, e.g. "diffusion_cond".
|
||||
svd_bases_path: Optional path to a precomputed SVD bases .pt file
|
||||
(used only by -XS adapter types).
|
||||
|
||||
Returns:
|
||||
List of display names (file stems) in the same order as the input paths.
|
||||
"""
|
||||
if not lora_ckpt_paths:
|
||||
model.use_lora = False
|
||||
model.lora_names = []
|
||||
return []
|
||||
|
||||
is_cond = model_type in ("diffusion_cond", "diffusion_cond_inpaint")
|
||||
|
||||
# Pass 1: load each checkpoint and resolve adapter type.
|
||||
lora_entries = [] # (path, state_dict, config_dict, adapter_type)
|
||||
for i, lora_path in enumerate(lora_ckpt_paths):
|
||||
print(f"Loading LoRA {i} from {lora_path}")
|
||||
state_dict, config_dict = load_lora_checkpoint(lora_path)
|
||||
adapter_type_raw = config_dict.get("adapter_type", "lora")
|
||||
adapter_type = resolve_adapter_type(adapter_type_raw, state_dict)
|
||||
if adapter_type != adapter_type_raw:
|
||||
vprint(f"Resolved legacy '{adapter_type_raw}' -> {adapter_type}")
|
||||
lora_entries.append((lora_path, state_dict, config_dict, adapter_type))
|
||||
|
||||
# Load SVD bases only if at least one LoRA is -XS.
|
||||
model_svd_bases = None
|
||||
cond_svd_bases = None
|
||||
any_xs = any(t.endswith("-xs") for _, _, _, t in lora_entries)
|
||||
if any_xs:
|
||||
if svd_bases_path is not None:
|
||||
vprint(f"Loading SVD bases from {svd_bases_path}")
|
||||
svd_bases_all = torch.load(svd_bases_path, map_location="cpu", weights_only=True)
|
||||
model_svd_bases = {k[len("model."):]: v for k, v in svd_bases_all.items() if k.startswith("model.")}
|
||||
cond_svd_bases = {k[len("conditioner."):]: v for k, v in svd_bases_all.items() if k.startswith("conditioner.")}
|
||||
vprint(f" model bases: {len(model_svd_bases)}, conditioner bases: {len(cond_svd_bases)}")
|
||||
else:
|
||||
print("WARNING: -XS adapter type present without svd_bases_path -- SVD will be computed on current device")
|
||||
|
||||
# Pass 2: apply each LoRA.
|
||||
lora_names = []
|
||||
for i, (lora_path, state_dict, config_dict, adapter_type) in enumerate(lora_entries):
|
||||
rank = config_dict.get("rank", infer_global_rank(state_dict))
|
||||
alpha = config_dict.get("alpha", rank)
|
||||
include = config_dict.get("include", None)
|
||||
exclude = config_dict.get("exclude", None)
|
||||
is_xs = adapter_type.endswith("-xs")
|
||||
|
||||
lora_config = {
|
||||
torch.nn.Linear: {
|
||||
"weight": partial(LoRAParametrization.from_linear, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
|
||||
},
|
||||
torch.nn.Conv1d: {
|
||||
"weight": partial(LoRAParametrization.from_conv1d, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
|
||||
},
|
||||
}
|
||||
if is_cond:
|
||||
add_lora(model.model, lora_config, include=include, exclude=exclude,
|
||||
svd_bases=model_svd_bases if is_xs else None)
|
||||
add_lora(model.conditioner, lora_config, include=include, exclude=exclude,
|
||||
svd_bases=cond_svd_bases if is_xs else None)
|
||||
else:
|
||||
add_lora(model, lora_config, include=include, exclude=exclude,
|
||||
svd_bases=model_svd_bases if is_xs else None)
|
||||
|
||||
prepare_dora_state_dict(state_dict)
|
||||
remapped_sd = remap_lora_state_dict(state_dict, i)
|
||||
if is_cond:
|
||||
model.model.load_state_dict(remapped_sd, strict=False)
|
||||
model.conditioner.load_state_dict(remapped_sd, strict=False)
|
||||
else:
|
||||
model.load_state_dict(remapped_sd, strict=False)
|
||||
lora_names.append(os.path.splitext(os.path.basename(lora_path))[0])
|
||||
|
||||
vprint("lora layers:", len(get_lora_layers(model)))
|
||||
model.use_lora = True
|
||||
model.lora_names = lora_names
|
||||
return lora_names
|
||||
@@ -0,0 +1,471 @@
|
||||
"""
|
||||
References:
|
||||
1) the official LoRA implementation released by Microsoft:
|
||||
https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
|
||||
"""
|
||||
|
||||
import math
|
||||
import re
|
||||
from functools import partial
|
||||
from itertools import product
|
||||
|
||||
import torch
|
||||
import torch.nn.utils.parametrize as parametrize
|
||||
from torch import nn
|
||||
|
||||
from ...verbose import vprint
|
||||
|
||||
|
||||
def _canonicalize_svd_signs(U, Vh):
|
||||
"""Enforce deterministic sign convention: largest-magnitude element of each U column is positive."""
|
||||
max_abs_idx = U.abs().argmax(dim=0)
|
||||
signs = U[max_abs_idx, torch.arange(U.shape[1])].sign()
|
||||
signs[signs == 0] = 1
|
||||
return U * signs.unsqueeze(0), Vh * signs.unsqueeze(1)
|
||||
|
||||
|
||||
class LoRAParametrization(nn.Module):
|
||||
def __init__(self, fan_in, fan_out, fan_in_fan_out=False, rank=4, lora_dropout_p=0.0, lora_alpha=1,
|
||||
adapter_type="lora", W0=None, lora_index=0, svd_bases=None
|
||||
):
|
||||
super().__init__()
|
||||
dtype = torch.get_default_dtype()
|
||||
device = W0.device if W0 is not None else None
|
||||
# if weight is stored as (fan_out, fan_in), the memory layout of A & B follows (W + BA)x
|
||||
# otherwise, it's x(W + AB). This allows us to tie the weights between linear layers and embeddings
|
||||
self.swap = (lambda x: (x[1], x[0])) if fan_in_fan_out else (lambda x: x)
|
||||
self.lora_alpha, self.rank = lora_alpha, rank
|
||||
self.scaling = lora_alpha / rank
|
||||
self.lora_dropout = nn.Dropout(p=lora_dropout_p) if lora_dropout_p > 0 else lambda x: x
|
||||
self.dropout_fn = self._dropout if lora_dropout_p > 0 else lambda x: x
|
||||
self.register_buffer("lora_dropout_mask", torch.ones(self.swap((1, fan_in)), dtype=dtype, device=device))
|
||||
self.register_buffer("lora_strength", torch.tensor(1.0, dtype=dtype, device=device), persistent=False)
|
||||
self.adapter_type = adapter_type
|
||||
self.lora_index = lora_index
|
||||
|
||||
if self.adapter_type == "lora":
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
self.forward_fn = self.lora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type in ("dora-rows", "dora"):
|
||||
# DoRA-rows: magnitude per row (per-output-neuron, dim=1) - paper-correct variant
|
||||
# "dora" maps here as the default for backward compat
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
self._norm_dim = 1
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "dora-cols":
|
||||
# DoRA-cols: magnitude per column (per-input-feature, dim=0) - legacy variant
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
self._norm_dim = 0
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "bora":
|
||||
# BoRA: bi-dimensional DoRA - magnitude on both row and column axes
|
||||
self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
|
||||
self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
|
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
self.magnitude_r = nn.Parameter(W0_2d.norm(dim=1).to(dtype))
|
||||
self.magnitude_c = nn.Parameter(W0_2d.norm(dim=0).to(dtype))
|
||||
else:
|
||||
self.magnitude_r = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.magnitude_c = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.bora_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "lora-xs":
|
||||
# frozen bases U:(fan_out,r), V:(fan_in,r); tiny trainable core M:(r,r)
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
elif W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1).cpu().float()
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d, full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.forward_fn = self.lora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "dora-rows-xs":
|
||||
# DoRA-rows + LoRA-XS: frozen SVD bases + trainable M_xs core + row magnitude
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
self._norm_dim = 1
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
if svd_bases is None:
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "dora-cols-xs":
|
||||
# DoRA-cols + LoRA-XS: frozen SVD bases + trainable M_xs core + column magnitude
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
self._norm_dim = 0
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
if svd_bases is None:
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
|
||||
else:
|
||||
self.magnitude = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.dora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
elif self.adapter_type == "bora-xs":
|
||||
# BoRA + LoRA-XS: frozen SVD bases + trainable M_xs core + row & column magnitudes
|
||||
self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
|
||||
self.register_buffer("V", torch.empty(self.swap((fan_in, self.rank)), device=device), persistent=True)
|
||||
self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
|
||||
if svd_bases is not None:
|
||||
self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
|
||||
self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
|
||||
if W0 is not None:
|
||||
W0_2d = W0.view(W0.shape[0], -1)
|
||||
if svd_bases is None:
|
||||
U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
|
||||
U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
|
||||
self.U.copy_(U_full[:, :self.rank].to(dtype))
|
||||
self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
|
||||
self.magnitude_r = nn.Parameter(W0_2d.norm(dim=1).to(dtype))
|
||||
self.magnitude_c = nn.Parameter(W0_2d.norm(dim=0).to(dtype))
|
||||
else:
|
||||
self.magnitude_r = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
|
||||
self.magnitude_c = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
|
||||
self.forward_fn = self.bora_xs_forward
|
||||
self._adapter_forward_fn = self.forward_fn
|
||||
|
||||
|
||||
def lora_forward(self, W):
|
||||
delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A)))).view(W.shape)
|
||||
delta = self.scaling * self.lora_strength * delta
|
||||
return (W + delta.to(W.dtype))
|
||||
|
||||
def dora_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
# work in 2D so norm/magnitude ops don't broadcast badly for Conv1d/Conv2d
|
||||
# Upcast W to LoRA param dtype for stable norm computation, cast result back.
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.lora_A.dtype)
|
||||
# low-rank update on the *direction*
|
||||
delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A))))
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
# normalize along _norm_dim, then scale by per-element magnitude
|
||||
V_hat = V / (V.norm(dim=self._norm_dim, keepdim=True) + 1e-12)
|
||||
return (V_hat * self.magnitude.unsqueeze(self._norm_dim)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def bora_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.lora_A.dtype)
|
||||
delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A))))
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
# Row-normalize, then scale by row magnitudes
|
||||
V_r = V / (V.norm(dim=1, keepdim=True) + 1e-12)
|
||||
intermediate = self.magnitude_r.unsqueeze(1) * V_r
|
||||
# Column-normalize, then scale by column magnitudes
|
||||
H_c = intermediate / (intermediate.norm(dim=0, keepdim=True) + 1e-12)
|
||||
return (H_c * self.magnitude_c.unsqueeze(0)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def lora_xs_forward(self, W):
|
||||
delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
|
||||
delta = delta.view(W.shape)
|
||||
return (W + (self.scaling * self.lora_strength * delta).to(W.dtype))
|
||||
|
||||
def dora_xs_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.U.dtype)
|
||||
delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
V_hat = V / (V.norm(dim=self._norm_dim, keepdim=True) + 1e-12)
|
||||
return (V_hat * self.magnitude.unsqueeze(self._norm_dim)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def bora_xs_forward(self, W):
|
||||
if self.lora_strength == 0:
|
||||
return W
|
||||
orig_shape = W.shape
|
||||
W_2d = W.view(W.shape[0], -1).to(self.U.dtype)
|
||||
delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
|
||||
V = W_2d + self.scaling * self.lora_strength * delta
|
||||
V_r = V / (V.norm(dim=1, keepdim=True) + 1e-12)
|
||||
intermediate = self.magnitude_r.unsqueeze(1) * V_r
|
||||
H_c = intermediate / (intermediate.norm(dim=0, keepdim=True) + 1e-12)
|
||||
return (H_c * self.magnitude_c.unsqueeze(0)).view(orig_shape).to(W.dtype)
|
||||
|
||||
def forward(self, X):
|
||||
return self.forward_fn(X)
|
||||
|
||||
def _dropout(self, A):
|
||||
# to mimic the original implementation: A @ dropout(x), we do (A * dropout(ones)) @ x
|
||||
return A * self.lora_dropout(self.lora_dropout_mask)
|
||||
|
||||
def disable_lora(self):
|
||||
self.forward_fn = lambda x: x
|
||||
|
||||
def enable_lora(self):
|
||||
self.forward_fn = self._adapter_forward_fn
|
||||
|
||||
@classmethod
|
||||
def _get_original_weight(cls, layer, attr_name="weight"):
|
||||
"""Get the true original weight, even when the layer is already parametrized."""
|
||||
if hasattr(layer, 'parametrizations') and hasattr(layer.parametrizations, attr_name):
|
||||
return getattr(layer.parametrizations, attr_name).original.detach()
|
||||
return getattr(layer, attr_name).detach()
|
||||
|
||||
@classmethod
|
||||
def from_linear(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_out, fan_in = layer.weight.shape
|
||||
W0 = cls._get_original_weight(layer)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_conv2d(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_out, fan_in = layer.weight.view(layer.weight.shape[0], -1).shape
|
||||
W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_conv1d(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_out, fan_in = layer.weight.view(layer.weight.shape[0], -1).shape
|
||||
W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_embedding(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
|
||||
fan_in, fan_out = layer.weight.shape
|
||||
W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
|
||||
return cls(
|
||||
fan_in, fan_out, fan_in_fan_out=True, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
|
||||
W0=W0, svd_bases=svd_bases, **kwargs
|
||||
)
|
||||
|
||||
|
||||
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),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# --- layer name filtering helpers ---
|
||||
|
||||
def _expand(s):
|
||||
"""Expand bracket notation in layer filter strings, e.g. 'layers[0-5]' -> ['layers0', ..., 'layers5']."""
|
||||
t = re.split(r'\[(\d+)-(\d+)\]', s)
|
||||
if len(t) == 1:
|
||||
return [s]
|
||||
lits, starts, ends = t[0::3], t[1::3], t[2::3]
|
||||
pools = []
|
||||
for a, b in zip(starts, ends):
|
||||
ai, bi = int(a), int(b)
|
||||
step = 1 if bi >= ai else -1
|
||||
pools.append([str(n) for n in range(ai, bi + step, step)])
|
||||
out = []
|
||||
for combo in product(*pools):
|
||||
pieces = []
|
||||
for i, lit in enumerate(lits):
|
||||
pieces.append(lit)
|
||||
if i < len(combo):
|
||||
pieces.append(combo[i])
|
||||
out.append(''.join(pieces))
|
||||
return out
|
||||
|
||||
|
||||
def _matches_any(name, patterns):
|
||||
"""Check if name contains any pattern substring (with bracket expansion)."""
|
||||
if not patterns:
|
||||
return False
|
||||
for pattern in patterns:
|
||||
for expanded in _expand(pattern.strip()):
|
||||
if expanded in name:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# --- core LoRA application ---
|
||||
|
||||
def _match_layer_type(layer, lora_config):
|
||||
"""Find the matching lora_config key for a layer, using isinstance to handle ParametrizedLinear etc."""
|
||||
for layer_type in lora_config:
|
||||
if isinstance(layer, layer_type):
|
||||
return layer_type
|
||||
return None
|
||||
|
||||
|
||||
def apply_lora(layer, register=True, merge=False, lora_config=default_lora_config):
|
||||
"""add lora parametrization to a layer, designed to be used with model.apply"""
|
||||
if register:
|
||||
matched_type = _match_layer_type(layer, lora_config)
|
||||
if matched_type is not None:
|
||||
for attr_name, parametrization in lora_config[matched_type].items():
|
||||
parametrize.register_parametrization(layer, attr_name, parametrization(layer), unsafe=True)
|
||||
else: # this will remove all parametrizations, use with caution
|
||||
if hasattr(layer, "parametrizations"):
|
||||
for attr_name in list(layer.parametrizations.keys()):
|
||||
parametrize.remove_parametrizations(layer, attr_name, leave_parametrized=merge)
|
||||
|
||||
|
||||
def add_lora(model, lora_config=default_lora_config, include=None, exclude=None, svd_bases=None):
|
||||
"""Add LoRA parametrization to layers in a model.
|
||||
|
||||
Args:
|
||||
model: The model to add LoRA to.
|
||||
lora_config: Dict mapping nn.Module types to parametrization configs.
|
||||
include: Optional list of substrings. If provided, only modules whose
|
||||
name contains at least one pattern get LoRA. Supports bracket
|
||||
expansion (e.g. "layers[0-11]").
|
||||
exclude: Optional list of substrings. Modules matching any pattern are
|
||||
skipped even if they match an include pattern.
|
||||
svd_bases: Optional dict of pre-computed SVD bases keyed by
|
||||
"layer_name.weight". When provided, -XS adapter types use
|
||||
these instead of computing SVD.
|
||||
"""
|
||||
if svd_bases is not None or include is not None or exclude is not None:
|
||||
# Named-module path: iterate by name so we can look up SVD bases per layer
|
||||
n_cached = 0
|
||||
n_recomputed = 0
|
||||
for name, module in model.named_modules():
|
||||
matched_type = _match_layer_type(module, lora_config)
|
||||
if matched_type is None:
|
||||
continue
|
||||
|
||||
if include is not None and not _matches_any(name, include):
|
||||
continue
|
||||
|
||||
if exclude is not None and _matches_any(name, exclude):
|
||||
continue
|
||||
|
||||
for attr_name, parametrization_fn in lora_config[matched_type].items():
|
||||
layer_bases = None
|
||||
if svd_bases is not None:
|
||||
bases_key = f"{name}.{attr_name}"
|
||||
layer_bases = svd_bases.get(bases_key)
|
||||
param = parametrization_fn(module, svd_bases=layer_bases)
|
||||
parametrize.register_parametrization(module, attr_name, param, unsafe=True)
|
||||
if hasattr(param, 'U'): # -XS layer that needs SVD bases
|
||||
if layer_bases is not None:
|
||||
n_cached += 1
|
||||
else:
|
||||
n_recomputed += 1
|
||||
if n_cached > 0:
|
||||
vprint(f"SVD bases: {n_cached} layers loaded from cache")
|
||||
if n_recomputed > 0:
|
||||
vprint(f"SVD bases: {n_recomputed} layers RECOMPUTED (missing from cache)")
|
||||
if svd_bases is not None and n_cached == 0 and n_recomputed == 0:
|
||||
vprint("SVD bases: provided but no -XS layers found")
|
||||
else:
|
||||
# Fast path: original behavior, no name inspection needed
|
||||
model.apply(partial(apply_lora, lora_config=lora_config))
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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))
|
||||
|
||||
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))
|
||||
|
||||
def remove_lora_by_index(model, lora_index):
|
||||
"""Remove all LoRA parametrizations with a specific lora_index.
|
||||
|
||||
For modules where other LoRA entries remain, only the matching entry is deleted.
|
||||
For modules where the removed entry was the only one, the parametrization is
|
||||
removed entirely and the weight reverts to its original value.
|
||||
"""
|
||||
for name, mod in model.named_modules():
|
||||
if not hasattr(mod, 'parametrizations'):
|
||||
continue
|
||||
for attr in list(mod.parametrizations.keys()):
|
||||
params = mod.parametrizations[attr]
|
||||
to_remove = [i for i, p in enumerate(params)
|
||||
if isinstance(p, LoRAParametrization) and p.lora_index == lora_index]
|
||||
if not to_remove:
|
||||
continue
|
||||
if len(params) == len(to_remove):
|
||||
# All entries being removed — remove parametrization entirely
|
||||
parametrize.remove_parametrizations(mod, attr, leave_parametrized=False)
|
||||
else:
|
||||
for i in reversed(to_remove):
|
||||
del params[i]
|
||||
|
||||
def set_lora_strength(model, strength: float, lora_index=None):
|
||||
"""Set lora strength. If lora_index is None, sets all LoRAs. If specified, sets only that index."""
|
||||
strength = float(strength)
|
||||
if lora_index is None:
|
||||
for name, buf in model.named_buffers(recurse=True):
|
||||
if name.endswith("lora_strength"):
|
||||
buf.fill_(strength)
|
||||
else:
|
||||
for p in _iter_lora_params(model):
|
||||
if p.lora_index == lora_index:
|
||||
p.lora_strength.fill_(strength)
|
||||
|
||||
def _iter_lora_params(model):
|
||||
for _, mod in model.named_modules():
|
||||
plist = getattr(getattr(mod, "parametrizations", None), "weight", None)
|
||||
if plist is None:
|
||||
continue
|
||||
for p in plist:
|
||||
if isinstance(p, LoRAParametrization):
|
||||
yield p
|
||||
@@ -0,0 +1,420 @@
|
||||
import json
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file as _st_save_file
|
||||
from torch import nn
|
||||
|
||||
from .model import LoRAParametrization, _expand
|
||||
from ...verbose import vprint
|
||||
|
||||
def apply_to_lora(fn):
|
||||
"""apply a function to LoRAParametrization layers, designed to be used with model.apply"""
|
||||
|
||||
def apply_fn(layer):
|
||||
if isinstance(layer, LoRAParametrization):
|
||||
fn(layer)
|
||||
|
||||
return apply_fn
|
||||
|
||||
def enable_lora(model, lora_index=None):
|
||||
"""Enable LoRA layers. If lora_index is None, enables all. If specified, enables only that index."""
|
||||
def _enable(layer):
|
||||
if isinstance(layer, LoRAParametrization):
|
||||
if lora_index is None or layer.lora_index == lora_index:
|
||||
layer.enable_lora()
|
||||
model.apply(_enable)
|
||||
|
||||
def disable_lora(model, lora_index=None):
|
||||
"""Disable LoRA layers. If lora_index is None, disables all. If specified, disables only that index."""
|
||||
def _disable(layer):
|
||||
if isinstance(layer, LoRAParametrization):
|
||||
if lora_index is None or layer.lora_index == lora_index:
|
||||
layer.disable_lora()
|
||||
model.apply(_disable)
|
||||
|
||||
def get_lora_layers(model):
|
||||
layers = []
|
||||
for name, m in model.named_modules():
|
||||
plist = getattr(getattr(m, "parametrizations", None), "weight", None)
|
||||
if plist is None:
|
||||
continue
|
||||
for p in plist:
|
||||
if isinstance(p, LoRAParametrization):
|
||||
layers.append((name, p))
|
||||
return layers
|
||||
|
||||
def has_lora(model):
|
||||
"""Return True if the model has at least one LoRAParametrization on any weight."""
|
||||
return len(get_lora_layers(model))>=1
|
||||
|
||||
def filter_lora_layers(model, lora_layer_filter, lora_index=None):
|
||||
"""Enable/disable LoRA layers by name filter. If lora_index is specified, only affects that index."""
|
||||
lora_layer_filter = (lora_layer_filter or "").strip()
|
||||
# lora_layer_filter: logical OR on comma separated substrings
|
||||
# e.g. ".transformer.layers, .to_global_embed."
|
||||
def is_filtered(layer_name):
|
||||
filters = [x.lower().strip() for x in lora_layer_filter.split(",") if len(x)]
|
||||
for f in filters:
|
||||
for _f in _expand(f):
|
||||
if _f in layer_name:
|
||||
return True
|
||||
return False
|
||||
for name, p in get_lora_layers(model):
|
||||
if lora_index is not None and p.lora_index != lora_index:
|
||||
continue
|
||||
if is_filtered(name):
|
||||
p.disable_lora()
|
||||
else:
|
||||
p.enable_lora()
|
||||
|
||||
|
||||
|
||||
# ------------------- helper function for collecting parameters for training/saving -------------------
|
||||
|
||||
|
||||
def name_is_lora(name):
|
||||
return (
|
||||
len(name.split(".")) >= 4
|
||||
and (name.split(".")[-4]) == "parametrizations"
|
||||
and name.split(".")[-1] in ["lora_A", "lora_B", "M_xs", "magnitude", "magnitude_r", "magnitude_c"]
|
||||
)
|
||||
|
||||
|
||||
def name_is_bias(name):
|
||||
return name.split(".")[-1] == "bias"
|
||||
|
||||
|
||||
def get_params_by_name(model, print_shapes=False, name_filter=None):
|
||||
for n, p in model.named_parameters():
|
||||
if name_filter is None or name_filter(n):
|
||||
if print_shapes:
|
||||
print(n, p.shape)
|
||||
yield p
|
||||
|
||||
|
||||
def get_lora_params(model, print_shapes=False):
|
||||
return get_params_by_name(model, print_shapes=print_shapes, name_filter=name_is_lora)
|
||||
|
||||
|
||||
def get_bias_params(model, print_shapes=False):
|
||||
return get_params_by_name(model, print_shapes=print_shapes, name_filter=name_is_bias)
|
||||
|
||||
|
||||
def get_lora_state_dict(model):
|
||||
return {k: v for k, v in model.state_dict().items() if name_is_lora(k)}
|
||||
|
||||
|
||||
def cast_base_to_precision(model, precision):
|
||||
"""Cast frozen base weights to lower precision, keeping LoRA params in fp32.
|
||||
|
||||
Args:
|
||||
model: Model with LoRA parametrizations applied.
|
||||
precision: One of "bf16", "bfloat16", "fp16", "float16".
|
||||
"""
|
||||
import torch
|
||||
|
||||
dtype_map = {"bf16": torch.bfloat16, "bfloat16": torch.bfloat16,
|
||||
"fp16": torch.float16, "float16": torch.float16}
|
||||
target_dtype = dtype_map.get(precision)
|
||||
if target_dtype is None:
|
||||
print(f"Unknown base_precision={precision!r}, keeping fp32")
|
||||
return
|
||||
vprint(f"Casting frozen base weights to {target_dtype}")
|
||||
model.to(target_dtype)
|
||||
# Restore LoRA params to fp32 — optimizer needs full precision
|
||||
n_restored = 0
|
||||
for p in get_lora_params(model):
|
||||
p.data = p.data.to(torch.float32)
|
||||
n_restored += 1
|
||||
vprint(f" Restored {n_restored} LoRA params to fp32")
|
||||
|
||||
|
||||
# ------------------- helper function for loading LoRa at inference -------------------
|
||||
|
||||
def infer_global_rank(lora_sd: dict) -> int:
|
||||
candidates = []
|
||||
for k, v in lora_sd.items():
|
||||
if k.endswith(".lora_A"):
|
||||
candidates.append(v.shape[0])
|
||||
elif k.endswith(".lora_B"):
|
||||
candidates.append(v.shape[1])
|
||||
elif k.endswith(".M_xs"):
|
||||
candidates.append(v.shape[0])
|
||||
if not candidates:
|
||||
raise ValueError("No LoRA/LoRA-XS tensors found to infer rank")
|
||||
r = candidates[0]
|
||||
# sanity check they all match (most checkpoints use one global r)
|
||||
if any(c != r for c in candidates):
|
||||
print(f"Multiple ranks detected in ckpt: {sorted(set(candidates))}. "
|
||||
"You’ll need per-layer ranks.")
|
||||
return r
|
||||
|
||||
|
||||
def detect_dora_variant(state_dict):
|
||||
"""Detect DoRA variant from a state dict.
|
||||
|
||||
Returns "dora-cols" (dim=0, per-input-feature) or "dora-rows" (dim=1,
|
||||
per-output-neuron) based on the shape of the first .magnitude tensor
|
||||
found. Returns None if no magnitude tensors exist or axis is ambiguous.
|
||||
"""
|
||||
for key, val in state_dict.items():
|
||||
if not key.endswith(".magnitude"):
|
||||
continue
|
||||
if val.dim() == 2:
|
||||
if val.shape[0] == 1:
|
||||
return "dora-cols" # old format: (1, fan_in)
|
||||
elif val.shape[1] == 1:
|
||||
return "dora-rows" # (fan_out, 1)
|
||||
# 1D magnitude — can't determine axis without model context
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def prepare_dora_state_dict(state_dict):
|
||||
"""Squeeze 2D DoRA magnitude tensors to 1D for loading. Modifies in-place."""
|
||||
for key in list(state_dict.keys()):
|
||||
if key.endswith(".magnitude") and state_dict[key].dim() == 2:
|
||||
state_dict[key] = state_dict[key].squeeze()
|
||||
|
||||
|
||||
def resolve_adapter_type(adapter_type, state_dict=None):
|
||||
"""Resolve legacy "dora" adapter_type to the correct variant.
|
||||
|
||||
If adapter_type is "dora" and a checkpoint state_dict is provided,
|
||||
attempts to detect dora-rows vs dora-cols from magnitude tensor shapes.
|
||||
Falls back to "dora-rows" (paper-correct default).
|
||||
All other adapter types are returned unchanged.
|
||||
"""
|
||||
if adapter_type != "dora":
|
||||
return adapter_type
|
||||
if state_dict is not None:
|
||||
detected = detect_dora_variant(state_dict)
|
||||
if detected:
|
||||
return detected
|
||||
return "dora-rows"
|
||||
|
||||
|
||||
# ------------------- helper functions for safetensors LoRA checkpoints -------------------
|
||||
|
||||
|
||||
def save_lora_safetensors(state_dict, lora_config, path):
|
||||
"""Save a LoRA checkpoint in safetensors format with config as metadata.
|
||||
|
||||
The lora_config dict is JSON-serialized and stored under the metadata key
|
||||
``"lora_config"``, which safetensors stores in the file header as a string.
|
||||
|
||||
Args:
|
||||
state_dict: Dict of LoRA tensors (from get_lora_state_dict).
|
||||
lora_config: Dict with keys like rank, alpha, adapter_type, include, exclude.
|
||||
path: Output file path (should end in .safetensors).
|
||||
"""
|
||||
metadata = {"lora_config": json.dumps(lora_config)}
|
||||
fp16_dict = {k: v.half() if v.is_floating_point() else v for k, v in state_dict.items()}
|
||||
_st_save_file(fp16_dict, str(path), metadata=metadata)
|
||||
|
||||
|
||||
def load_lora_checkpoint(path):
|
||||
"""Load a LoRA checkpoint from either .safetensors or .ckpt format.
|
||||
|
||||
Returns:
|
||||
Tuple of (state_dict, lora_config) where lora_config may be an empty
|
||||
dict if the checkpoint has no embedded config.
|
||||
"""
|
||||
path = str(path)
|
||||
if path.endswith(".safetensors"):
|
||||
with safe_open(path, framework="pt", device="cpu") as f:
|
||||
state_dict = {key: f.get_tensor(key) for key in f.keys()}
|
||||
metadata = f.metadata()
|
||||
lora_config = {}
|
||||
if metadata and "lora_config" in metadata:
|
||||
lora_config = json.loads(metadata["lora_config"])
|
||||
return state_dict, lora_config
|
||||
else:
|
||||
ckpt = torch.load(path, map_location="cpu", weights_only=False)
|
||||
state_dict = ckpt["state_dict"]
|
||||
lora_config = ckpt.get("lora_config", {})
|
||||
return state_dict, lora_config
|
||||
|
||||
|
||||
def convert_lora_ckpt_to_safetensors(ckpt_path, output_path=None):
|
||||
"""Convert a .ckpt LoRA checkpoint to .safetensors format.
|
||||
|
||||
Args:
|
||||
ckpt_path: Path to the input .ckpt file.
|
||||
output_path: Path for the output .safetensors file.
|
||||
If None, replaces .ckpt extension with .safetensors.
|
||||
|
||||
Returns:
|
||||
The output path.
|
||||
"""
|
||||
ckpt_path = str(ckpt_path)
|
||||
if output_path is None:
|
||||
if ckpt_path.endswith(".ckpt"):
|
||||
output_path = ckpt_path[:-5] + ".safetensors"
|
||||
else:
|
||||
output_path = ckpt_path + ".safetensors"
|
||||
state_dict, lora_config = load_lora_checkpoint(ckpt_path)
|
||||
save_lora_safetensors(state_dict, lora_config, output_path)
|
||||
return output_path
|
||||
|
||||
|
||||
# ------------------- helper functions for multi-LoRA stacking -------------------
|
||||
|
||||
|
||||
def get_lora_count(model):
|
||||
"""Return the number of stacked LoRAs (distinct lora_index values)."""
|
||||
indices = set()
|
||||
for _, p in get_lora_layers(model):
|
||||
indices.add(p.lora_index)
|
||||
return len(indices)
|
||||
|
||||
|
||||
def remap_lora_state_dict(state_dict, target_index):
|
||||
"""Remap LoRA state dict keys from parametrization index 0 to target_index.
|
||||
|
||||
Checkpoints always save keys with index 0 (e.g. *.parametrizations.weight.0.lora_A).
|
||||
When loading the Nth stacked LoRA, we need to remap to index N.
|
||||
"""
|
||||
if target_index == 0:
|
||||
return state_dict
|
||||
remapped = {}
|
||||
for k, v in state_dict.items():
|
||||
new_key = k.replace(
|
||||
".parametrizations.weight.0.",
|
||||
".parametrizations.weight.{}.".format(target_index)
|
||||
)
|
||||
remapped[new_key] = v
|
||||
return remapped
|
||||
|
||||
|
||||
# ------------------- helper function for inferencing with multiple lora -------------------
|
||||
|
||||
|
||||
def _get_adapter_param_names(lora_layer):
|
||||
"""Return the list of trainable parameter names for this adapter type."""
|
||||
if lora_layer.adapter_type == "lora":
|
||||
return ["lora_A", "lora_B"]
|
||||
elif lora_layer.adapter_type in ("dora", "dora-rows", "dora-cols"):
|
||||
return ["lora_A", "lora_B", "magnitude"]
|
||||
elif lora_layer.adapter_type == "bora":
|
||||
return ["lora_A", "lora_B", "magnitude_r", "magnitude_c"]
|
||||
elif lora_layer.adapter_type == "lora-xs":
|
||||
return ["M_xs"]
|
||||
elif lora_layer.adapter_type in ("dora-rows-xs", "dora-cols-xs"):
|
||||
return ["M_xs", "magnitude"]
|
||||
elif lora_layer.adapter_type == "bora-xs":
|
||||
return ["M_xs", "magnitude_r", "magnitude_c"]
|
||||
return []
|
||||
|
||||
|
||||
def _prepare_for_multiple_lora(lora_layer):
|
||||
lora_layer._multi_lora_slots = []
|
||||
|
||||
|
||||
def _append_lora(lora_layer):
|
||||
slot = {}
|
||||
for name in _get_adapter_param_names(lora_layer):
|
||||
slot[name] = nn.Parameter(getattr(lora_layer, name).clone())
|
||||
lora_layer._multi_lora_slots.append(slot)
|
||||
|
||||
|
||||
def load_multiple_lora(model, lora_state_dicts):
|
||||
model.apply(apply_to_lora(_prepare_for_multiple_lora))
|
||||
for state_dict in lora_state_dicts:
|
||||
_ = model.load_state_dict(state_dict, strict=False)
|
||||
model.apply(apply_to_lora(_append_lora))
|
||||
return model
|
||||
|
||||
|
||||
def _select_lora(lora_layer, index):
|
||||
slot = lora_layer._multi_lora_slots[index]
|
||||
for name, param in slot.items():
|
||||
setattr(lora_layer, name, param)
|
||||
|
||||
|
||||
def select_lora(model, index):
|
||||
model.apply(apply_to_lora(lambda x: _select_lora(x, index)))
|
||||
return model
|
||||
|
||||
# ------------------- helper function for merging multiple LoRAs -------------------
|
||||
|
||||
def merge_loras_into_base_model(model, lora_configurations):
|
||||
"""
|
||||
Merges multiple LoRAs with specified application weights directly into the base model's weights.
|
||||
Works with all adapter types (LoRA, DoRA, LoRA-XS) by using the parametrization's own forward.
|
||||
After merging, the LoRA parametrization on each layer is disabled.
|
||||
|
||||
Args:
|
||||
model: The PyTorch model with LoRAParametrization already applied (via add_lora).
|
||||
lora_configurations: A list of dicts, each specifying a LoRA to merge:
|
||||
[
|
||||
{
|
||||
'name': 'style1',
|
||||
'state_dict': loaded_torch_state_dict,
|
||||
'application_weight': 0.7
|
||||
},
|
||||
...
|
||||
]
|
||||
"""
|
||||
import torch
|
||||
|
||||
if not lora_configurations:
|
||||
return
|
||||
|
||||
# Collect all parametrized layers with their LoRA instances
|
||||
parametrized_layers = []
|
||||
for name, module in model.named_modules():
|
||||
if not hasattr(module, "parametrizations"):
|
||||
continue
|
||||
for attr_name in list(module.parametrizations.keys()):
|
||||
plist = module.parametrizations[attr_name]
|
||||
for p in plist:
|
||||
if isinstance(p, LoRAParametrization):
|
||||
parametrized_layers.append((name, module, attr_name, p, plist.original))
|
||||
break
|
||||
|
||||
# Initialize per-layer delta accumulators
|
||||
deltas = {id(orig): torch.zeros_like(orig.data) for _, _, _, _, orig in parametrized_layers}
|
||||
|
||||
for lora_cfg in lora_configurations:
|
||||
lora_sd = lora_cfg['state_dict']
|
||||
app_weight = lora_cfg['application_weight']
|
||||
|
||||
# Load this LoRA's adapter weights into the model's parametrizations
|
||||
model.load_state_dict(lora_sd, strict=False)
|
||||
|
||||
# Compute each layer's delta using the parametrization's own forward
|
||||
with torch.no_grad():
|
||||
for name, module, attr_name, lora_p, original in parametrized_layers:
|
||||
merged = lora_p(original)
|
||||
deltas[id(original)] += app_weight * (merged - original)
|
||||
|
||||
# Apply accumulated deltas to original weights and disable LoRA
|
||||
with torch.no_grad():
|
||||
for name, module, attr_name, lora_p, original in parametrized_layers:
|
||||
original.data.add_(deltas[id(original)])
|
||||
lora_p.disable_lora()
|
||||
|
||||
# ------------------- helper function for tying and untieing weights -------------------
|
||||
|
||||
|
||||
def tie_weights(linear: nn.Linear, embedding: nn.Embedding):
|
||||
"""tie the weights of the linear layer and the embedding layer both with the same lora"""
|
||||
lora_p = embedding.parametrizations.weight[0]
|
||||
if lora_p.adapter_type != "lora":
|
||||
raise NotImplementedError(f"tie_weights only supports standard LoRA, got adapter_type='{lora_p.adapter_type}'")
|
||||
# this line below is optional if the original is already tied
|
||||
embedding.parametrizations.weight.original = linear.parametrizations.weight.original
|
||||
embedding.parametrizations.weight[0].lora_A = linear.parametrizations.weight[0].lora_B
|
||||
embedding.parametrizations.weight[0].lora_B = linear.parametrizations.weight[0].lora_A
|
||||
|
||||
|
||||
def untie_weights(linear: nn.Linear, embedding: nn.Embedding):
|
||||
"""untie the weights of the linear layer and the embedding layer"""
|
||||
lora_p = embedding.parametrizations.weight[0]
|
||||
if lora_p.adapter_type != "lora":
|
||||
raise NotImplementedError(f"untie_weights only supports standard LoRA, got adapter_type='{lora_p.adapter_type}'")
|
||||
embedding.parametrizations.weight.original = nn.Parameter(embedding.weight.original.clone())
|
||||
embedding.parametrizations.weight[0].lora_A = nn.Parameter(embedding.parametrizations.weight[0].lora_A.clone())
|
||||
embedding.parametrizations.weight[0].lora_B = nn.Parameter(embedding.parametrizations.weight[0].lora_B.clone())
|
||||
Reference in New Issue
Block a user