592 lines
20 KiB
Python
592 lines
20 KiB
Python
"""
|
||
Delta reconstruction, alpha/rank scaling, and shape alignment utilities.
|
||
|
||
Extracted from baking_processor.py to reduce file size and isolate
|
||
tensor-level operations from key matching logic.
|
||
"""
|
||
|
||
import torch
|
||
from typing import (
|
||
Dict, Iterator, Iterable, Mapping, MutableMapping,
|
||
Optional, Tuple, Union,
|
||
)
|
||
|
||
from ..utils import ProgressTracker
|
||
from .scale_utils import apply_alpha_correction
|
||
|
||
|
||
# ===================================================================
|
||
# LoRA Pair Grouping
|
||
# ===================================================================
|
||
|
||
def _group_lora_pairs(
|
||
lora_sd: Dict[str, torch.Tensor],
|
||
) -> Tuple[Dict[str, Dict[str, torch.Tensor]], Dict[str, torch.Tensor]]:
|
||
"""
|
||
Group LoRA state dict entries by their base key and component type.
|
||
|
||
For each key in lora_sd, determines if it's a component of a decomposed
|
||
LoRA pair (``lora_A``, ``lora_B``, ``lora_up``, ``lora_down``, ``.alpha``,
|
||
``.diff``) or a standalone tensor (pre-baked delta, bias).
|
||
|
||
Returns:
|
||
pairs: dict of ``base_key → {component_name: tensor, ...}``
|
||
standalone: dict of ``base_key → tensor`` (non-decomposed entries)
|
||
"""
|
||
pairs: Dict[str, Dict[str, torch.Tensor]] = {}
|
||
standalone: Dict[str, torch.Tensor] = {}
|
||
|
||
for key, tensor in lora_sd.items():
|
||
base_key = None
|
||
component = None
|
||
|
||
if key.endswith('.lora_A.weight'):
|
||
base_key = key[:-14]
|
||
component = 'A'
|
||
elif key.endswith('.lora_B.weight'):
|
||
base_key = key[:-14]
|
||
component = 'B'
|
||
elif key.endswith('.lora_up.weight'):
|
||
base_key = key[:-15]
|
||
component = 'A'
|
||
elif key.endswith('.lora_down.weight'):
|
||
base_key = key[:-17]
|
||
component = 'B'
|
||
elif key.endswith('.lora_A'):
|
||
base_key = key[:-7]
|
||
component = 'A'
|
||
elif key.endswith('.lora_B'):
|
||
base_key = key[:-7]
|
||
component = 'B'
|
||
elif key.endswith('.lora_up'):
|
||
base_key = key[:-8]
|
||
component = 'A'
|
||
elif key.endswith('.lora_down'):
|
||
base_key = key[:-10]
|
||
component = 'B'
|
||
elif key.endswith('.alpha'):
|
||
base_key = key[:-6]
|
||
component = 'alpha'
|
||
elif key.endswith('.diff'):
|
||
base_key = key[:-5]
|
||
component = 'diff'
|
||
else:
|
||
standalone[key] = tensor
|
||
continue
|
||
|
||
if base_key not in pairs:
|
||
pairs[base_key] = {}
|
||
pairs[base_key][component] = tensor
|
||
|
||
return pairs, standalone
|
||
|
||
|
||
# ===================================================================
|
||
# Single Delta Reconstruction
|
||
# ===================================================================
|
||
|
||
def _reconstruct_one(
|
||
base_key: str,
|
||
components: Dict[str, torch.Tensor],
|
||
device: Optional[Union[str, torch.device]] = None,
|
||
) -> Optional[torch.Tensor]:
|
||
"""
|
||
Reconstruct a single LoRA weight delta from its A/B components.
|
||
|
||
``delta = B @ A`` (Kohya convention) or ``delta = A @ B`` (transposed),
|
||
followed by ``(alpha / rank)`` scaling.
|
||
|
||
Handles 4D Conv2d squeeze, GPU acceleration, and convention-agnostic
|
||
rank detection.
|
||
|
||
Returns:
|
||
The reconstructed delta on CPU, or **None** if the components are
|
||
incompatible (non-2D after squeeze, ambiguous rank).
|
||
"""
|
||
if 'A' in components and 'B' in components:
|
||
B = components['B']
|
||
A = components['A']
|
||
|
||
# Handle 4D Conv2d tensors (proj_in/proj_out) — squeeze spatial dims
|
||
if B.dim() == 4 and A.dim() == 4:
|
||
if B.shape[2] == 1 and B.shape[3] == 1 and A.shape[2] == 1 and A.shape[3] == 1:
|
||
B = B.squeeze(-1).squeeze(-1)
|
||
A = A.squeeze(-1).squeeze(-1)
|
||
else:
|
||
return None
|
||
|
||
if B.dim() != 2 or A.dim() != 2:
|
||
return None
|
||
|
||
# GPU acceleration for matmul
|
||
if device is not None:
|
||
A = A.to(device=device)
|
||
B = B.to(device=device)
|
||
|
||
# Convention-agnostic delta reconstruction
|
||
all_dims = [A.shape[0], A.shape[1], B.shape[0], B.shape[1]]
|
||
rank = min(all_dims)
|
||
|
||
if A.shape[1] == rank:
|
||
delta = A @ B # Kohya convention
|
||
elif B.shape[1] == rank:
|
||
delta = B @ A # transposed convention
|
||
else:
|
||
return None
|
||
|
||
# Alpha scaling
|
||
alpha = components.get('alpha')
|
||
if alpha is not None:
|
||
alpha_val = alpha.item() if hasattr(alpha, 'item') else float(alpha)
|
||
delta = apply_alpha_correction(delta, alpha_val, rank, mode="linear")
|
||
|
||
return delta.cpu()
|
||
|
||
elif 'diff' in components:
|
||
# Pre-baked delta — pass through unchanged
|
||
return components['diff']
|
||
|
||
return None
|
||
|
||
|
||
# ===================================================================
|
||
# Shape Inference (no matmul needed)
|
||
# ===================================================================
|
||
|
||
def _infer_delta_shape(
|
||
base_key: str,
|
||
components: Dict[str, torch.Tensor],
|
||
) -> Optional[Tuple[int, ...]]:
|
||
"""
|
||
Infer the shape of a delta tensor from its A/B component shapes.
|
||
|
||
This avoids the cost of a full matmul (O(n³)) just to determine
|
||
shape compatibility during key matching. Only the dimensions are
|
||
needed, which can be derived from the component tensor metadata.
|
||
|
||
For Kohya convention: ``A=(out_dim,rank), B=(rank,in_dim)``
|
||
→ ``delta = A @ B = (out_dim, in_dim)``
|
||
|
||
For transposed convention: ``A=(rank,in_dim), B=(out_dim,rank)``
|
||
→ ``delta = B @ A = (out_dim, in_dim)``
|
||
"""
|
||
if 'A' not in components or 'B' not in components:
|
||
return None
|
||
|
||
B, A = components['B'], components['A']
|
||
|
||
# Handle 4D Conv2d tensors — extract spatial dims from metadata
|
||
if B.dim() == 4 and A.dim() == 4:
|
||
if B.shape[2] == 1 and B.shape[3] == 1 and A.shape[2] == 1 and A.shape[3] == 1:
|
||
B_dims = (B.shape[0], B.shape[1])
|
||
A_dims = (A.shape[0], A.shape[1])
|
||
else:
|
||
return None
|
||
elif B.dim() == 2 and A.dim() == 2:
|
||
B_dims = (B.shape[0], B.shape[1])
|
||
A_dims = (A.shape[0], A.shape[1])
|
||
else:
|
||
return None
|
||
|
||
all_dims = [A_dims[0], A_dims[1], B_dims[0], B_dims[1]]
|
||
rank = min(all_dims)
|
||
|
||
if A_dims[1] == rank:
|
||
return (A_dims[0], B_dims[1]) # Kohya: delta = A @ B
|
||
elif B_dims[1] == rank:
|
||
return (B_dims[0], A_dims[1]) # Transposed: delta = B @ A
|
||
else:
|
||
return None
|
||
|
||
|
||
# ===================================================================
|
||
# Lightweight Shape Proxy
|
||
# ===================================================================
|
||
|
||
class _ShapeInfo:
|
||
"""
|
||
Lightweight shape metadata — quacks like a ``torch.Tensor`` for shape checks.
|
||
|
||
Exposes ``.shape``, ``.dim()``, and ``.numel()`` so that
|
||
:func:`_check_shape_compatible` and related functions can operate without
|
||
a materialised tensor. Used by :class:`_LazyDeltaDict.items()` to avoid
|
||
reconstructing deltas during key matching.
|
||
"""
|
||
__slots__ = ('shape',)
|
||
|
||
def __init__(self, shape: Tuple[int, ...]):
|
||
self.shape = shape
|
||
|
||
def dim(self) -> int:
|
||
return len(self.shape)
|
||
|
||
def numel(self) -> int:
|
||
n = 1
|
||
for s in self.shape:
|
||
n *= s
|
||
return n
|
||
|
||
def __repr__(self):
|
||
return f"_ShapeInfo({self.shape})"
|
||
|
||
|
||
# ===================================================================
|
||
# Lazy Delta Dictionary
|
||
# ===================================================================
|
||
|
||
class _LazyDeltaDict(Mapping[str, torch.Tensor]):
|
||
"""
|
||
Dict-like wrapper storing LoRA A/B components for lazy delta reconstruction.
|
||
|
||
Stores only the tiny component tensors (A, B, alpha — ≈14 MB total for
|
||
112 LoRA pairs) instead of reconstructing all full deltas (≈9.6 GB).
|
||
Each delta is reconstructed on demand via :meth:`pop` or :meth:`__getitem__`.
|
||
|
||
During matching, :meth:`items` yields :class:`_ShapeInfo` objects (shape
|
||
metadata only) to avoid unnecessary matmul computation — the matcher only
|
||
needs shapes for ``_is_shape_compatible`` checks.
|
||
|
||
Typical workflow::
|
||
|
||
pairs, standalone = _group_lora_pairs(lora_sd)
|
||
lazy = _LazyDeltaDict(pairs, standalone, device="cuda:0")
|
||
|
||
# Matching — shapes only, no matmul
|
||
for key, shape_info in lazy.items():
|
||
if is_shape_compatible(ckpt_key, shape_info):
|
||
matched[ckpt_key] = (lora_base, shape_info)
|
||
|
||
# Baking — reconstruct one delta at a time
|
||
for ckpt_key in bake_order:
|
||
delta = lazy.reconstruct(lora_base) # actual matmul
|
||
bake(delta)
|
||
# delta freed after loop iteration
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
pairs: Dict[str, Dict[str, torch.Tensor]],
|
||
standalone: Dict[str, torch.Tensor],
|
||
device: Optional[Union[str, torch.device]] = None,
|
||
):
|
||
self._pairs = pairs
|
||
self._standalone = standalone
|
||
self._device = device
|
||
self._shape_cache: Dict[str, Optional[Tuple[int, ...]]] = {}
|
||
|
||
# ── Shape Inference ───────────────────────────────────────────────
|
||
|
||
def _infer_shape(self, base_key: str) -> Optional[Tuple[int, ...]]:
|
||
"""Infer delta shape from component shapes (cached)."""
|
||
if base_key in self._shape_cache:
|
||
return self._shape_cache[base_key]
|
||
components = self._pairs.get(base_key)
|
||
if components is not None:
|
||
shape = _infer_delta_shape(base_key, components)
|
||
else:
|
||
shape = None
|
||
self._shape_cache[base_key] = shape
|
||
return shape
|
||
|
||
def get_shape_info(
|
||
self, base_key: str
|
||
) -> Union['_ShapeInfo', torch.Tensor, None]:
|
||
"""
|
||
Get shape metadata without reconstructing the delta.
|
||
|
||
Returns:
|
||
``_ShapeInfo`` for A/B pairs (shape only),
|
||
the actual tensor for standalone entries (already materialised),
|
||
or **None** if the key doesn't exist or shape can't be inferred.
|
||
"""
|
||
if base_key in self._standalone:
|
||
return self._standalone[base_key]
|
||
shape = self._infer_shape(base_key)
|
||
if shape is not None:
|
||
return _ShapeInfo(shape)
|
||
return None
|
||
|
||
# ── Reconstruction ────────────────────────────────────────────────
|
||
|
||
def reconstruct(self, base_key: str) -> torch.Tensor:
|
||
"""
|
||
Fully reconstruct a single delta (matmul + alpha scaling).
|
||
|
||
Raises ``KeyError`` if *base_key* is not found.
|
||
"""
|
||
components = self._pairs.get(base_key)
|
||
if components is not None:
|
||
result = _reconstruct_one(base_key, components, self._device)
|
||
if result is not None:
|
||
return result
|
||
if base_key in self._standalone:
|
||
return self._standalone[base_key]
|
||
raise KeyError(f"Key {base_key!r} not found in _LazyDeltaDict")
|
||
|
||
# ── Dict Protocol ────────────────────────────────────────────────
|
||
|
||
def __getitem__(self, key: str) -> torch.Tensor:
|
||
return self.reconstruct(key)
|
||
|
||
def __iter__(self) -> Iterator[str]:
|
||
seen: set = set()
|
||
for k in self._pairs:
|
||
if k not in seen:
|
||
yield k
|
||
seen.add(k)
|
||
for k in self._standalone:
|
||
if k not in seen:
|
||
yield k
|
||
seen.add(k)
|
||
|
||
def __len__(self) -> int:
|
||
return len(self._pairs) + len(self._standalone)
|
||
|
||
def __contains__(self, key: object) -> bool:
|
||
return key in self._pairs or key in self._standalone
|
||
|
||
def keys(self):
|
||
return self._pairs.keys() | self._standalone.keys()
|
||
|
||
def items(self):
|
||
"""
|
||
Yield ``(key, value)`` pairs.
|
||
|
||
For A/B pairs, yields ``_ShapeInfo`` instead of the full delta tensor
|
||
to avoid matmul during key matching. Standalone tensors are yielded
|
||
as-is since they are already materialised.
|
||
|
||
Use :meth:`pop` or :meth:`__getitem__` to get the actual reconstructed
|
||
tensor.
|
||
"""
|
||
for key in self._pairs:
|
||
shape_info = self.get_shape_info(key)
|
||
if shape_info is not None:
|
||
yield key, shape_info
|
||
for key, tensor in self._standalone.items():
|
||
yield key, tensor
|
||
|
||
def values(self):
|
||
for key in self:
|
||
yield self[key]
|
||
|
||
def get(self, key: str, default=None):
|
||
try:
|
||
return self[key]
|
||
except KeyError:
|
||
return default
|
||
|
||
def pop(self, key: str, *args):
|
||
"""
|
||
Pop and reconstruct a delta on demand.
|
||
|
||
Accepts an optional default as the second argument, matching
|
||
``dict.pop(key, default)``.
|
||
"""
|
||
if len(args) > 1:
|
||
raise TypeError(
|
||
f"pop expected at most 2 arguments, got {1 + len(args)}"
|
||
)
|
||
if key in self._pairs:
|
||
delta = self.reconstruct(key)
|
||
del self._pairs[key]
|
||
if key in self._shape_cache:
|
||
del self._shape_cache[key]
|
||
return delta
|
||
if key in self._standalone:
|
||
return self._standalone.pop(key)
|
||
if args:
|
||
return args[0]
|
||
raise KeyError(key)
|
||
|
||
def clear(self):
|
||
"""Clear all stored data and shape caches."""
|
||
self._pairs.clear()
|
||
self._standalone.clear()
|
||
self._shape_cache.clear()
|
||
|
||
|
||
# ===================================================================
|
||
# Materialised Delta Wrapper for Baking
|
||
# ===================================================================
|
||
|
||
class _MatchedDeltas(MutableMapping[str, torch.Tensor]):
|
||
"""
|
||
Dict-like wrapper that reconstructs LoRA deltas on demand during baking.
|
||
|
||
Wraps a :class:`_LazyDeltaDict` together with a ``ckpt_key → lora_base``
|
||
mapping produced by :func:`_find_matching_keys`. Each :meth:`pop` or
|
||
:meth:`__getitem__` call triggers a single delta reconstruction — no
|
||
materialised tensor is held between calls.
|
||
|
||
This eliminates the ≈9.6 GB ``matched_deltas`` dict that previously
|
||
persisted during the entire bake loop. Only one delta (≈86 MB for
|
||
768×768 float32) is in memory at any time.
|
||
|
||
Supports the full ``MutableMapping`` protocol needed by the three bake
|
||
methods (``bake_linear``, ``bake_impact_weighted``, ``bake_orthogonal``),
|
||
including :meth:`items` (which reconstructs incrementally during energy
|
||
analysis), :meth:`pop`, :meth:`clear`, :meth:`keys`, and :meth:`values`.
|
||
|
||
Some matched keys (fused QKV, ``to_out.0`` fallback) produce tensors
|
||
that cannot be reconstructed from a single ``_LazyDeltaDict`` entry
|
||
(they are concatenations or already-materialised tensors). These are
|
||
stored in *precomputed* and returned directly.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
lazy_deltas: '_LazyDeltaDict',
|
||
ckpt_to_lora: Dict[str, str],
|
||
precomputed: Optional[Dict[str, torch.Tensor]] = None,
|
||
):
|
||
"""
|
||
Args:
|
||
lazy_deltas: Component-based lazy delta store.
|
||
ckpt_to_lora: Mapping from checkpoint key → LoRA base key,
|
||
produced during matching.
|
||
precomputed: Optional dict of checkpoint keys whose tensors
|
||
are already materialised (fused QKV, to_out.0 fallback)
|
||
and cannot be reconstructed from ``_LazyDeltaDict`` alone.
|
||
"""
|
||
self._lazy = lazy_deltas
|
||
self._ckpt_to_lora = ckpt_to_lora
|
||
self._precomputed = precomputed if precomputed is not None else {}
|
||
|
||
def _get_lora_key(self, ckpt_key: str) -> str:
|
||
return self._ckpt_to_lora[ckpt_key]
|
||
|
||
def pop(self, key: str, *args):
|
||
"""
|
||
Pop a checkpoint key and reconstruct its delta on demand.
|
||
|
||
Checks *precomputed* dict first (fused QKV, ``to_out.0``), then
|
||
falls back to :meth:`_LazyDeltaDict.reconstruct` via the key map.
|
||
|
||
Accepts optional default, matching ``dict.pop(key, default)``.
|
||
"""
|
||
if len(args) > 1:
|
||
raise TypeError(
|
||
f"pop expected at most 2 arguments, got {1 + len(args)}"
|
||
)
|
||
# 1. Check precomputed dict (fused QKV, to_out.0 fallback)
|
||
if key in self._precomputed:
|
||
return self._precomputed.pop(key)
|
||
# 2. Reconstruct from lazy components on demand
|
||
lora_base = self._ckpt_to_lora.pop(key, None)
|
||
if lora_base is not None:
|
||
return self._lazy.reconstruct(lora_base)
|
||
if args:
|
||
return args[0]
|
||
raise KeyError(key)
|
||
|
||
def __getitem__(self, key: str) -> torch.Tensor:
|
||
# Check precomputed first
|
||
if key in self._precomputed:
|
||
return self._precomputed[key]
|
||
return self._lazy.reconstruct(self._get_lora_key(key))
|
||
|
||
def __setitem__(self, key: str, value: torch.Tensor) -> None:
|
||
# Not used during baking — raise for safety if it ever is.
|
||
raise NotImplementedError("_MatchedDeltas does not support __setitem__")
|
||
|
||
def __delitem__(self, key: str) -> None:
|
||
if key in self._precomputed:
|
||
del self._precomputed[key]
|
||
else:
|
||
del self._ckpt_to_lora[key]
|
||
|
||
def __iter__(self) -> Iterator[str]:
|
||
seen: set = set()
|
||
for k in self._precomputed:
|
||
if k not in seen:
|
||
yield k
|
||
seen.add(k)
|
||
for k in self._ckpt_to_lora:
|
||
if k not in seen:
|
||
yield k
|
||
seen.add(k)
|
||
|
||
def __len__(self) -> int:
|
||
return len(self._ckpt_to_lora) + len(self._precomputed)
|
||
|
||
def __contains__(self, key: object) -> bool:
|
||
return key in self._ckpt_to_lora or key in self._precomputed
|
||
|
||
def keys(self):
|
||
seen: set = set()
|
||
for k in self._precomputed:
|
||
if k not in seen:
|
||
yield k
|
||
seen.add(k)
|
||
for k in self._ckpt_to_lora:
|
||
if k not in seen:
|
||
yield k
|
||
seen.add(k)
|
||
|
||
def items(self):
|
||
"""Yield ``(ckpt_key, delta)`` pairs, reconstructing each in turn."""
|
||
# Yield precomputed entries first (no reconstruction needed)
|
||
for ckpt_key, tensor in self._precomputed.items():
|
||
yield ckpt_key, tensor
|
||
for ckpt_key, lora_base in self._ckpt_to_lora.items():
|
||
yield ckpt_key, self._lazy.reconstruct(lora_base)
|
||
|
||
def values(self):
|
||
for ckpt_key in list(self._precomputed):
|
||
yield self._precomputed[ckpt_key]
|
||
for ckpt_key, lora_base in self._ckpt_to_lora.items():
|
||
yield self._lazy.reconstruct(lora_base)
|
||
|
||
def clear(self):
|
||
self._ckpt_to_lora.clear()
|
||
self._precomputed.clear()
|
||
|
||
|
||
# ===================================================================
|
||
# Full Delta Reconstruction (Original API — backward compatible)
|
||
# ===================================================================
|
||
|
||
def _reconstruct_lora_delta(
|
||
lora_sd: Dict[str, torch.Tensor],
|
||
device: Optional[Union[str, torch.device]] = None,
|
||
) -> Dict[str, torch.Tensor]:
|
||
"""
|
||
Reconstruct LoRA weight deltas from decomposed A/B pairs.
|
||
|
||
For each pair ``(lora_B, lora_A)`` sharing a base key::
|
||
|
||
delta = B @ A
|
||
|
||
Then applies ``(alpha / rank)`` scaling to produce the final delta
|
||
that represents the weight change for this LoRA key.
|
||
|
||
Also handles standalone tensors (pre-baked deltas, biases) — passed
|
||
through unchanged.
|
||
|
||
This is the **original eager API** — use :class:`_LazyDeltaDict` for
|
||
on-demand reconstruction.
|
||
|
||
Args:
|
||
lora_sd: LoRA state dict with A/B decomposed tensors.
|
||
device: If provided (e.g. ``"cuda:0"``), moves A/B tensors to this
|
||
device before the matmul for GPU-accelerated reconstruction.
|
||
|
||
Returns:
|
||
``dict`` of ``base_key → delta tensor`` (on CPU).
|
||
"""
|
||
pairs, standalone = _group_lora_pairs(lora_sd)
|
||
deltas: Dict[str, torch.Tensor] = {}
|
||
|
||
with ProgressTracker(total=len(pairs), desc="Reconstructing deltas") as delta_progress:
|
||
for base_key, components in pairs.items():
|
||
delta = _reconstruct_one(base_key, components, device)
|
||
if delta is not None:
|
||
deltas[base_key] = delta
|
||
delta_progress += 1
|
||
|
||
# Add standalone tensors (pre-baked deltas, biases)
|
||
deltas.update(standalone)
|
||
|
||
return deltas
|