403 lines
18 KiB
Python
403 lines
18 KiB
Python
"""
|
||
Triple LoRA Merging Methods for EasyLoRAMerger - Version 1.3.0
|
||
Specialized algorithms for merging three LoRAs simultaneously.
|
||
|
||
This module is now a thin wrapper around engine.triple_merge_core,
|
||
preserving the same public API (merge_triple_method returning 4 values)
|
||
while eliminating ~700 lines of duplicated merge function bodies.
|
||
"""
|
||
|
||
import math
|
||
from typing import List, Dict, Optional, Tuple
|
||
|
||
import torch
|
||
|
||
from ..utils import (
|
||
safe_get_rank,
|
||
silent_pad_or_truncate,
|
||
categorize_key,
|
||
compute_component_energy_ratios,
|
||
ProgressTracker,
|
||
DeviceManager,
|
||
)
|
||
from .methods import apply_magnitude_scaling
|
||
from .scale_utils import find_alpha_value, apply_alpha_correction
|
||
|
||
# ── Import shared merge functions and helpers from unified core ──────────
|
||
from .triple_merge_core import (
|
||
merge_linear,
|
||
merge_cross,
|
||
merge_ties_strict,
|
||
merge_ties_gentle,
|
||
merge_ties_contrast,
|
||
merge_slerp,
|
||
merge_feature_mix,
|
||
merge_magnitude,
|
||
merge_subtract,
|
||
merge_dare_rescale,
|
||
merge_dare_lite,
|
||
merge_svd_preserve,
|
||
merge_block_swap,
|
||
merge_noise_aware,
|
||
merge_gradient_alignment,
|
||
ensure_energy_preservation_triple,
|
||
MERGE_METHOD_REGISTRY as _CORE_MERGE_REGISTRY,
|
||
)
|
||
|
||
# ── Backward-compatible aliases (so internal references still work) ─────
|
||
# These allow `triple_merge_svd_preserve` to call `triple_merge_linear`,
|
||
# and `merge_triple_method` to use `triple_merge_linear` as fallback.
|
||
triple_merge_linear = merge_linear
|
||
triple_merge_cross = merge_cross
|
||
triple_merge_ties_strict = merge_ties_strict
|
||
triple_merge_ties_gentle = merge_ties_gentle
|
||
triple_merge_ties_contrast = merge_ties_contrast
|
||
triple_merge_slerp = merge_slerp
|
||
triple_merge_feature_mix = merge_feature_mix
|
||
triple_merge_magnitude = merge_magnitude
|
||
triple_merge_subtract = merge_subtract
|
||
triple_merge_dare_rescale = merge_dare_rescale
|
||
triple_merge_dare_lite = merge_dare_lite
|
||
triple_merge_svd_preserve = merge_svd_preserve
|
||
triple_merge_block_swap = merge_block_swap
|
||
triple_merge_noise_aware = merge_noise_aware
|
||
triple_merge_gradient_alignment = merge_gradient_alignment
|
||
|
||
|
||
# ==================== HELPER FUNCTIONS ====================
|
||
|
||
|
||
def adjust_ranks(tensors: List[torch.Tensor], key: str) -> List[torch.Tensor]:
|
||
"""
|
||
Pad/truncate all tensors to the same rank (max rank among them).
|
||
Returns adjusted tensors.
|
||
"""
|
||
ranks = [safe_get_rank(t, key) for t in tensors]
|
||
target_rank = max(ranks)
|
||
|
||
adjusted_tensors = []
|
||
for i, t in enumerate(tensors):
|
||
if ranks[i] != target_rank:
|
||
t = silent_pad_or_truncate(t, target_rank, f"{key}_{i}")
|
||
adjusted_tensors.append(t)
|
||
|
||
return adjusted_tensors
|
||
|
||
|
||
def apply_lora_scaling_triple(tensor: torch.Tensor, original_sd: Dict[str, torch.Tensor],
|
||
key: str, mapping: Optional[Dict[str, str]] = None,
|
||
is_converted: bool = False) -> torch.Tensor:
|
||
"""
|
||
Apply alpha/rank scaling correction to LoRA tensors (up‑weight only).
|
||
Delegates to :func:`scale_utils.find_alpha_value` and
|
||
:func:`scale_utils.apply_alpha_correction`.
|
||
|
||
Preserves exact behavior:
|
||
- Down weights (lora_A / lora_down) are never scaled
|
||
- Converted LoRAs skip scaling
|
||
- No alpha found → no scaling (assume already scaled or alpha == rank)
|
||
"""
|
||
# Down weights should not be scaled; scaling is applied only to up weights
|
||
if any(suffix in key for suffix in ['.lora_A.weight', '.lora_down.weight']):
|
||
return tensor
|
||
|
||
# Determine rank
|
||
rank = safe_get_rank(tensor, key)
|
||
rank = max(1, rank)
|
||
|
||
# Find alpha value (uses mapping internally for candidate generation)
|
||
alpha_value = find_alpha_value(original_sd, key, mapping=mapping, rank=rank)
|
||
|
||
# Skip if converted or no alpha found
|
||
if is_converted or alpha_value is None:
|
||
return tensor
|
||
|
||
# Apply linear (alpha / rank) scaling to up-weight
|
||
return apply_alpha_correction(tensor, alpha_value, rank, mode="linear")
|
||
|
||
|
||
# ==================== REGISTRY ====================
|
||
|
||
# Re-export the core registry unchanged — all 15 method names map to the same
|
||
# unified merge functions that both wrappers now share.
|
||
TRIPLE_METHOD_REGISTRY = dict(_CORE_MERGE_REGISTRY)
|
||
|
||
|
||
# ── Triple-capability warnings ────────────────────────────────────────
|
||
|
||
|
||
METHOD_3WAY_WARNINGS: Dict[str, str] = {
|
||
"slerp": (
|
||
"slerp only works with exactly 2 active LoRAs. "
|
||
"When merging 3+ LoRAs, only the first 2 are used."
|
||
),
|
||
}
|
||
|
||
|
||
def _warn_triple_capability(method: str, num_tensors: int) -> None:
|
||
"""Emit warnings for methods with limited 3-way support."""
|
||
if num_tensors > 2 and method in METHOD_3WAY_WARNINGS:
|
||
print(f"⚠️ WARNING [{method}]: {METHOD_3WAY_WARNINGS[method]}")
|
||
|
||
|
||
# ==================== MAIN TRIPLE MERGE FUNCTION ====================
|
||
|
||
|
||
def merge_triple_method(sds: List[Dict[str, torch.Tensor]],
|
||
weights: List[float],
|
||
method: str = "linear",
|
||
density: float = 1.0,
|
||
uniqueness: float = 0.7,
|
||
threshold: float = 0.0,
|
||
blend: float = 0.5,
|
||
blend_mode: str = "auto",
|
||
device: str = "auto",
|
||
magnitude_scaling: str = "none",
|
||
max_scaling_factor: float = 10.0,
|
||
batch_size: int = 32,
|
||
streaming: bool = True,
|
||
energy_preservation: bool = True,
|
||
balancing_mode: str = "disabled",
|
||
mappings: Optional[List[Dict[str, str]]] = None,
|
||
original_sds: Optional[List[Dict[str, torch.Tensor]]] = None,
|
||
metas: Optional[List[Dict[str, str]]] = None) -> Tuple[Dict[str, torch.Tensor], Dict[str, str], Optional[Dict], Optional[List[float]]]:
|
||
"""
|
||
Core triple merge logic with rank adjustment and magnitude scaling.
|
||
|
||
Parameters:
|
||
sds: List of state dictionaries (normalized).
|
||
weights: List of weights for each LoRA.
|
||
method: Merge method name.
|
||
density: Density parameter (for DARE and sparsity).
|
||
uniqueness: Uniqueness parameter for feature_mix.
|
||
threshold: Threshold parameter for subtract.
|
||
blend: Blend parameter for magnitude.
|
||
blend_mode: Blend mode ("active" or "dense").
|
||
Note: "slerp" is now a first-class method in the method dropdown.
|
||
magnitude_scaling: Magnitude scaling mode ("none", "rms", "top_5%", etc.).
|
||
max_scaling_factor: Maximum scaling factor for magnitude equalization.
|
||
batch_size: Number of keys to process per batch (streaming).
|
||
streaming: Whether to process keys in batches to save VRAM.
|
||
energy_preservation: Whether to apply energy preservation safety check.
|
||
balancing_mode: Weight balancing mode ("disabled", "safe", "creative").
|
||
mappings: Optional list of key mappings from normalized to original keys.
|
||
original_sds: Optional list of original state dicts (for alpha scaling).
|
||
metas: Optional list of metadata dicts.
|
||
|
||
Returns:
|
||
Tuple of (merged state dictionary, master_key_map, energy_by_component, adjusted_weights).
|
||
Energy data is None if balancing is disabled or only 2 LoRAs.
|
||
"""
|
||
# Forensic data collectors (populated inside the balancing block if active)
|
||
_forensic_energy_data: Optional[Dict] = None
|
||
_forensic_adjusted_weights: Optional[List[float]] = None
|
||
|
||
# Emit triple-capability warning (e.g., slerp only works with 2 LoRAs)
|
||
_warn_triple_capability(method, len(sds))
|
||
|
||
all_keys = set()
|
||
for sd in sds:
|
||
all_keys.update(sd.keys())
|
||
|
||
# Auto-weight balancing based on energy ratios
|
||
if balancing_mode != "disabled" and len(sds) == 3:
|
||
# Keys present in all three LoRAs
|
||
common_keys = [key for key in all_keys if all(key in sd for sd in sds)]
|
||
if common_keys:
|
||
print(f" ⚖️ Auto-weight balancing ({balancing_mode}) analyzing {len(common_keys)} common keys...")
|
||
print(f" Energy computed via per-element mean (rank-independent) on shared layers only")
|
||
# Use shared utility for rank-independent mean energy computation
|
||
energy_by_component = compute_component_energy_ratios(
|
||
norm_sds=sds,
|
||
common_keys=common_keys,
|
||
original_sds=original_sds,
|
||
mappings=mappings,
|
||
converted_flags=[False, False, False], # triple uses original SDs directly
|
||
key_categorizer=categorize_key,
|
||
)
|
||
# Accumulate global energy across all components
|
||
energy_per_lora = [0.0, 0.0, 0.0]
|
||
for component, energies in energy_by_component.items():
|
||
for i in range(3):
|
||
energy_per_lora[i] += energies[i]
|
||
|
||
epsilon = 1e-12
|
||
# Compute average energy
|
||
avg_energy = sum(energy_per_lora) / 3.0
|
||
# Compute scaling factors per LoRA based on balancing_mode
|
||
factors = [1.0, 1.0, 1.0]
|
||
if balancing_mode == "safe":
|
||
# Safe mode: reduce weight of louder LoRAs to match average
|
||
for i in range(3):
|
||
if energy_per_lora[i] > avg_energy and energy_per_lora[i] > epsilon:
|
||
ratio = math.sqrt(avg_energy / energy_per_lora[i])
|
||
factors[i] = max(0.1, min(10.0, ratio))
|
||
print(f" LoRA {i} energy {energy_per_lora[i]:.2e} is {energy_per_lora[i]/avg_energy:.2f}x average, weight scaled by {factors[i]:.2f}")
|
||
elif balancing_mode == "creative":
|
||
# Creative mode: compromise scaling for all LoRAs
|
||
for i in range(3):
|
||
if energy_per_lora[i] > epsilon:
|
||
ratio = math.sqrt(avg_energy / energy_per_lora[i])
|
||
creative_factor = (1.0 + ratio) / 2.0 # 50% compromise
|
||
factors[i] = max(0.1, min(10.0, creative_factor))
|
||
print(f" LoRA {i} energy {energy_per_lora[i]:.2e} → compromise factor {factors[i]:.2f}")
|
||
# Apply factors to weights
|
||
adjusted_weights = [weights[i] * factors[i] for i in range(3)]
|
||
print(f" ⚖️ Original weights: {weights}")
|
||
print(f" ⚖️ Adjusted weights: {adjusted_weights}")
|
||
if method == "slerp":
|
||
print(f" ℹ️ SLERP uses adjusted weight ratios for the interpolation factor t = |w| / Σ|w|")
|
||
# Replace weights for the rest of the merge
|
||
weights = adjusted_weights
|
||
# Capture forensic data for upstream report
|
||
_forensic_energy_data = energy_by_component
|
||
_forensic_adjusted_weights = adjusted_weights
|
||
|
||
# Use DeviceManager instead of hardcoded inline resolution
|
||
device = DeviceManager.get_device(device)
|
||
merged = {}
|
||
|
||
# Convert to list for deterministic ordering
|
||
key_list = list(all_keys)
|
||
total_keys = len(key_list)
|
||
|
||
# Memory guard helper — checks VRAM pressure, empties cache if >80%
|
||
def memory_guard():
|
||
if device.type == 'cuda':
|
||
allocated = torch.cuda.memory_allocated()
|
||
total = torch.cuda.get_device_properties(0).total_memory
|
||
ratio = allocated / total if total > 0 else 0
|
||
if ratio > 0.8:
|
||
print(f"⚠️ High GPU memory usage: {allocated / (1024**3):.2f} GB / {total / (1024**3):.2f} GB ({ratio:.1%}). Clearing cache.")
|
||
torch.cuda.empty_cache()
|
||
|
||
# Process keys in batches
|
||
with ProgressTracker(total=total_keys, desc="Merging triple LoRA keys") as merge_progress:
|
||
for batch_start in range(0, total_keys, batch_size):
|
||
memory_guard()
|
||
batch_keys = key_list[batch_start:batch_start + batch_size]
|
||
for key in batch_keys:
|
||
tensors = []
|
||
valid_weights = []
|
||
|
||
# Collect tensors that exist in each SD, along with original state dicts and mappings
|
||
original_sd_list = []
|
||
mapping_list = []
|
||
for i, sd in enumerate(sds):
|
||
if key in sd:
|
||
# Async H2D transfer to overlap with computation
|
||
t = sd[key]
|
||
if t.device.type == 'cpu' and device.type == 'cuda':
|
||
t = t.to(device, non_blocking=True)
|
||
t = t.to(torch.bfloat16)
|
||
tensors.append(t)
|
||
valid_weights.append(weights[i])
|
||
# Store corresponding original_sd and mapping if available
|
||
if original_sds is not None and i < len(original_sds):
|
||
original_sd_list.append(original_sds[i])
|
||
else:
|
||
original_sd_list.append(None)
|
||
if mappings is not None and i < len(mappings):
|
||
mapping_list.append(mappings[i])
|
||
else:
|
||
mapping_list.append(None)
|
||
|
||
# Apply alpha/rank scaling (up‑weight only) if original_sds provided
|
||
if original_sds is not None:
|
||
for idx, (tensor, orig_sd, mapping) in enumerate(zip(tensors, original_sd_list, mapping_list)):
|
||
if orig_sd is not None:
|
||
tensors[idx] = apply_lora_scaling_triple(tensor, orig_sd, key, mapping=mapping, is_converted=False)
|
||
|
||
if len(tensors) < 2:
|
||
if tensors:
|
||
# Unique key (present in only one LoRA) – still apply scaling already done
|
||
result = tensors[0] * valid_weights[0]
|
||
merged[key] = result
|
||
continue
|
||
|
||
# Signal magnitude equalization (RMS/percentile scaling) if enabled
|
||
if magnitude_scaling != "none" and len(tensors) >= 2:
|
||
# Determine reference tensor (LoRA A) - index 0 in original sds order
|
||
orig_indices = [i for i, sd in enumerate(sds) if key in sd]
|
||
if 0 in orig_indices:
|
||
ref_idx = orig_indices.index(0)
|
||
tensors = apply_magnitude_scaling(tensors, valid_weights,
|
||
magnitude_scaling, max_scaling_factor,
|
||
ref_idx=ref_idx)
|
||
|
||
# Rank adjustment (pad/truncate to same rank)
|
||
tensors = adjust_ranks(tensors, key)
|
||
|
||
# Apply method with adjusted tensors
|
||
merge_fn = TRIPLE_METHOD_REGISTRY.get(method)
|
||
if merge_fn is None:
|
||
# Fallback to linear
|
||
merge_fn = triple_merge_linear
|
||
# Prepare kwargs for the method
|
||
METHOD_KWARGS = {
|
||
"feature_mix": {"uniqueness": uniqueness},
|
||
"magnitude": {"blend": blend},
|
||
"subtract": {"threshold": threshold},
|
||
"dare_rescale": {"density": density},
|
||
"dare_lite": {"density": density},
|
||
}
|
||
kwargs = METHOD_KWARGS.get(method, {})
|
||
|
||
|
||
result = merge_fn(tensors, valid_weights, blend_mode=blend_mode, **kwargs)
|
||
|
||
# Energy safety check
|
||
if energy_preservation:
|
||
result = ensure_energy_preservation_triple(result, tensors, valid_weights, threshold=0.8, gain_min=1.0, target='avg')
|
||
|
||
# Apply density if needed
|
||
# Note: dare_rescale and dare_lite handle density internally
|
||
if density < 1.0 and method not in ("dare_rescale", "dare_lite"):
|
||
flat = result.abs().flatten()
|
||
k = max(1, int(flat.numel() * density))
|
||
threshold_val = torch.topk(flat, k).values.min()
|
||
mask = result.abs() >= threshold_val
|
||
result = result * mask
|
||
|
||
# Keep result on GPU within batch; batch transfer to CPU below
|
||
merged[key] = result
|
||
|
||
# Cleanup
|
||
del tensors
|
||
|
||
merge_progress += 1
|
||
|
||
# Transfer entire batch from GPU → CPU in one shot, then clear cache
|
||
if device.type == 'cuda':
|
||
for key in batch_keys:
|
||
if key in merged and merged[key].device.type == 'cuda':
|
||
merged[key] = merged[key].cpu()
|
||
if streaming:
|
||
torch.cuda.empty_cache()
|
||
|
||
# Build master_map from input mappings (identity map for key restoration)
|
||
master_map: Dict[str, str] = {}
|
||
mapped_count = 0
|
||
if mappings:
|
||
for key in merged:
|
||
found = False
|
||
# Priority order: LoRA A > B > C
|
||
for i in range(len(mappings)):
|
||
if i < len(mappings) and mappings[i] is not None and key in mappings[i]:
|
||
master_map[key] = mappings[i][key]
|
||
found = True
|
||
break
|
||
if not found:
|
||
# Key not in any mapping (e.g., generated alpha key) — identity fallback
|
||
master_map[key] = key
|
||
else:
|
||
mapped_count += 1
|
||
else:
|
||
# No mappings provided — identity fallback for all keys
|
||
master_map = {key: key for key in merged}
|
||
|
||
if mappings:
|
||
print(f" 🗺️ Built master_map: {mapped_count}/{len(merged)} keys mapped, {len(merged) - mapped_count} identity fallback")
|
||
|
||
return merged, master_map, _forensic_energy_data, _forensic_adjusted_weights
|