315 lines
11 KiB
Python
315 lines
11 KiB
Python
"""
|
||
Unified alpha/rank scaling utilities for EasyLoRAMerger.
|
||
|
||
Consolidates 4 separate implementations of alpha/rank scaling logic into a
|
||
single shared module, eliminating duplicated code while preserving exact
|
||
behavior for all consumers.
|
||
|
||
=== Formula Reference ===
|
||
|
||
Runtime scaling (merge_engine_v2, triple_methods):
|
||
scale = alpha / rank # linear mode
|
||
result = up_tensor * scale # only up-weight scaled
|
||
|
||
Baking (bake_alphas in utils.py):
|
||
scale = sqrt(alpha / rank) # sqrt mode
|
||
down_tensor *= scale # both weights scaled
|
||
up_tensor *= scale # because B * A * sqrt(a/r) * sqrt(a/r) = B * A * (a/r)
|
||
|
||
Delta reconstruction (baking_processor_delta):
|
||
scale = alpha / rank # linear mode
|
||
delta = B @ A # reconstructed weight delta
|
||
delta *= scale # scale the combined result
|
||
|
||
=== Consumer Mapping ===
|
||
|
||
| Consumer | Mode | Target | Data Source |
|
||
|-----------------------------------|-----------|--------|-----------------------|
|
||
| merge_engine_v2._apply_lora_scaling | linear | up | original_sd + mapping |
|
||
| triple_methods.apply_lora_scaling_triple | linear | up | original_sd + mapping |
|
||
| baking_processor_delta._reconstruct_lora_delta | linear | delta | lora_sd (grouped) |
|
||
| utils.bake_alphas | sqrt | both | state_dict (in-place) |
|
||
"""
|
||
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
import torch
|
||
|
||
|
||
# ===================================================================
|
||
# Alpha Key Candidate Generation
|
||
# ===================================================================
|
||
|
||
|
||
def generate_alpha_candidates(key: str) -> List[str]:
|
||
"""
|
||
Generate all possible alpha key candidates for a given weight key.
|
||
|
||
Merges candidate generation logic from:
|
||
- merge_engine_v2._apply_lora_scaling (7 patterns)
|
||
- triple_methods.apply_lora_scaling_triple (7 patterns, identical)
|
||
- utils._lookup_alpha_value (7 patterns, adds lora_A/B separately)
|
||
|
||
Returns a list of candidate alpha keys ordered by specificity.
|
||
The first match in ``find_alpha_value`` will be used.
|
||
"""
|
||
candidates: List[str] = []
|
||
|
||
# Pattern 1: replace .weight with .alpha
|
||
if key.endswith(".weight"):
|
||
candidates.append(key.replace(".weight", ".alpha"))
|
||
|
||
# Pattern 2: replace specific LoRA suffix with 'alpha'
|
||
if "lora_A.weight" in key:
|
||
candidates.append(key.replace("lora_A.weight", "alpha"))
|
||
if "lora_B.weight" in key:
|
||
candidates.append(key.replace("lora_B.weight", "alpha"))
|
||
if "lora_down.weight" in key:
|
||
candidates.append(key.replace("lora_down.weight", "alpha"))
|
||
if "lora_up.weight" in key:
|
||
candidates.append(key.replace("lora_up.weight", "alpha"))
|
||
|
||
# Pattern 3: strip suffix after last .lora_ or _lora_
|
||
if ".lora_" in key:
|
||
base = key.split(".lora_")[0]
|
||
candidates.append(f"{base}.alpha")
|
||
if "_lora_" in key:
|
||
base = key.split("_lora_")[0]
|
||
candidates.append(f"{base}.alpha")
|
||
|
||
return candidates
|
||
|
||
|
||
# ===================================================================
|
||
# Alpha Value Lookup
|
||
# ===================================================================
|
||
|
||
|
||
def find_alpha_value(
|
||
original_sd: Dict[str, torch.Tensor],
|
||
key: str,
|
||
mapping: Optional[Dict[str, str]] = None,
|
||
alpha_one_is_rank: bool = False,
|
||
rank: int = 1,
|
||
metadata: Optional[Dict[str, Any]] = None,
|
||
default: Optional[float] = None,
|
||
) -> Optional[float]:
|
||
"""
|
||
Find the alpha value for a LoRA weight key in a state dict.
|
||
|
||
Parameters
|
||
----------
|
||
original_sd : Dict[str, torch.Tensor]
|
||
The original LoRA state dict to search for alpha keys.
|
||
key : str
|
||
The weight key (e.g. ``some_block.lora_A.weight``).
|
||
mapping : Optional[Dict[str, str]]
|
||
Optional mapping from normalized key back to original weight key.
|
||
If provided, ``mapping.get(key, key)`` is used as the base for
|
||
candidate generation.
|
||
alpha_one_is_rank : bool
|
||
If True and the found alpha value is exactly 1.0, treat it as
|
||
``rank`` instead. Used by ``IdentityMergeEngine`` for converted
|
||
LoRAs where alpha=1 is meaningless but stored as metadata.
|
||
rank : int
|
||
The rank of the tensor. Only used when ``alpha_one_is_rank=True``.
|
||
metadata : Optional[Dict[str, Any]]
|
||
Optional metadata dict to check for alpha values injected by
|
||
EasyLoRAMerger during previous merges (e.g. ``ss_network_alpha``,
|
||
``lora_a_ss_network_alpha``, ``lora_b_ss_network_alpha``).
|
||
default : Optional[float]
|
||
Default value to return if no alpha key or metadata is found.
|
||
If None, returns None on miss (original behavior).
|
||
|
||
Returns
|
||
-------
|
||
Optional[float]
|
||
The alpha value, or *default* (or None) if no matching alpha
|
||
key was found.
|
||
"""
|
||
# Resolve the original weight key via mapping
|
||
orig_key = mapping.get(key, key) if mapping is not None else key
|
||
|
||
# Generate candidates and search
|
||
candidates = generate_alpha_candidates(orig_key)
|
||
alpha_value: Optional[float] = None
|
||
|
||
for ak in candidates:
|
||
if ak in original_sd:
|
||
alpha_tensor = original_sd[ak]
|
||
if isinstance(alpha_tensor, torch.Tensor):
|
||
if alpha_tensor.numel() == 1:
|
||
alpha_value = alpha_tensor.item()
|
||
else:
|
||
alpha_value = alpha_tensor.mean().item()
|
||
break
|
||
|
||
# Apply alpha_one_is_rank correction
|
||
if alpha_one_is_rank and alpha_value is not None and alpha_value == 1.0:
|
||
alpha_value = float(rank)
|
||
|
||
# Metadata check (EasyLoRAMerger metadata injection from previous merges)
|
||
if alpha_value is None and metadata is not None:
|
||
for mk in ['ss_network_alpha', 'lora_a_ss_network_alpha', 'lora_b_ss_network_alpha']:
|
||
if mk in metadata:
|
||
try:
|
||
alpha_value = float(metadata[mk])
|
||
break
|
||
except (ValueError, TypeError):
|
||
continue
|
||
|
||
# Default fallback
|
||
if alpha_value is None and default is not None:
|
||
alpha_value = default
|
||
|
||
return alpha_value
|
||
|
||
|
||
# ===================================================================
|
||
# Scale Factor Resolution
|
||
# ===================================================================
|
||
|
||
|
||
def resolve_scale(
|
||
alpha: float,
|
||
rank: int,
|
||
mode: str = "linear",
|
||
) -> float:
|
||
"""
|
||
Resolve the alpha/rank scaling factor.
|
||
|
||
Parameters
|
||
----------
|
||
alpha : float
|
||
The alpha value (from LoRA metadata).
|
||
rank : int
|
||
The rank of the LoRA tensor (>= 1).
|
||
mode : str
|
||
``"linear"`` → returns ``alpha / rank``
|
||
``"sqrt"`` → returns ``sqrt(alpha / rank)``
|
||
|
||
Returns
|
||
-------
|
||
float
|
||
The scaling factor. Returns 1.0 if ``rank <= 0``.
|
||
"""
|
||
if rank <= 0:
|
||
return 1.0
|
||
|
||
if mode == "sqrt":
|
||
return (alpha / rank) ** 0.5
|
||
return alpha / rank
|
||
|
||
|
||
# ===================================================================
|
||
# Tensor-Level Alpha Correction
|
||
# ===================================================================
|
||
|
||
|
||
def apply_alpha_correction(
|
||
tensor: torch.Tensor,
|
||
alpha: float,
|
||
rank: int,
|
||
mode: str = "linear",
|
||
) -> torch.Tensor:
|
||
"""
|
||
Apply alpha/rank scaling correction to a single tensor.
|
||
|
||
Equivalent to ``tensor * resolve_scale(alpha, rank, mode)``
|
||
but skips multiplication when the scale factor is essentially 1.0.
|
||
|
||
Parameters
|
||
----------
|
||
tensor : torch.Tensor
|
||
The LoRA weight tensor to scale.
|
||
alpha : float
|
||
The alpha value.
|
||
rank : int
|
||
The rank of the tensor (>= 1).
|
||
mode : str
|
||
``"linear"`` or ``"sqrt"``.
|
||
|
||
Returns
|
||
-------
|
||
torch.Tensor
|
||
The scaled tensor (or the original tensor if scaling ∼= 1.0).
|
||
"""
|
||
scale = resolve_scale(alpha, rank, mode)
|
||
if abs(scale - 1.0) > 1e-6:
|
||
return tensor * scale
|
||
return tensor
|
||
|
||
|
||
# ===================================================================
|
||
# Baking: Embed Alpha into Both Down and Up Weights
|
||
# ===================================================================
|
||
|
||
|
||
def bake_alpha_key(
|
||
down_tensor: torch.Tensor,
|
||
up_tensor: torch.Tensor,
|
||
alpha: float,
|
||
rank: Optional[int] = None,
|
||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""
|
||
Bake an alpha value into both down and up weight tensors.
|
||
|
||
Uses ``sqrt(alpha / rank)`` scaling so that the overall delta
|
||
``B @ A`` is scaled by ``alpha / rank`` (since
|
||
``(B * s) @ (A * s) = B @ A * s^2 = B @ A * (alpha / rank)``).
|
||
|
||
Parameters
|
||
----------
|
||
down_tensor : torch.Tensor
|
||
The down (A) weight tensor (e.g. ``lora_A.weight``).
|
||
up_tensor : torch.Tensor
|
||
The up (B) weight tensor (e.g. ``lora_B.weight``).
|
||
alpha : float
|
||
The alpha value to bake in.
|
||
rank : Optional[int]
|
||
The rank of the LoRA (inferred from ``down_tensor`` shape if None).
|
||
For a 2D down tensor ``(R, in_dim)``, rank is ``R``.
|
||
|
||
Returns
|
||
-------
|
||
Tuple[torch.Tensor, torch.Tensor]
|
||
``(scaled_down, scaled_up)`` — both tensors scaled by
|
||
``sqrt(alpha / rank)``.
|
||
"""
|
||
if rank is None:
|
||
# Infer rank: down tensor shape is typically (rank, in_dim)
|
||
rank = down_tensor.shape[0] if down_tensor.dim() >= 2 else 1
|
||
|
||
rank = max(1, rank)
|
||
scale = resolve_scale(alpha, rank, mode="sqrt")
|
||
|
||
if abs(scale - 1.0) > 1e-6:
|
||
return down_tensor * scale, up_tensor * scale
|
||
return down_tensor, up_tensor
|
||
|
||
|
||
def build_alpha_mapping(state_dict: Dict[str, torch.Tensor], key_map: Dict[str, str]) -> Dict[str, str]:
|
||
"""
|
||
Build mapping from normalized alpha keys to original alpha keys.
|
||
|
||
For each weight key pair (norm→orig), generates the corresponding
|
||
alpha key pair using the ``.lora_A.weight`` → ``.alpha`` convention.
|
||
|
||
This is a **map-building** operation (not a value lookup). It returns
|
||
a dict mapping ``normalized_alpha_key -> original_alpha_key`` for use
|
||
in key-mapping pipelines such as :func:`identity_normalize`.
|
||
|
||
Returns:
|
||
Dict mapping normalized_alpha_key -> original_alpha_key
|
||
"""
|
||
alpha_map: Dict[str, str] = {}
|
||
for norm_key, orig_key in key_map.items():
|
||
if "lora_A.weight" in norm_key or "lora_down.weight" in norm_key:
|
||
base = norm_key.replace(".lora_A.weight", "").replace(".lora_down.weight", "")
|
||
alpha_key = f"{base}.alpha"
|
||
orig_base = orig_key.replace(".lora_A.weight", "").replace(".lora_down.weight", "")
|
||
target_alpha_key = f"{orig_base}.alpha"
|
||
if target_alpha_key in state_dict:
|
||
alpha_map[alpha_key] = target_alpha_key
|
||
return alpha_map
|