mirror of https://github.com/razor-ai/soup.git
273 lines
9.4 KiB
Python
273 lines
9.4 KiB
Python
"""LoRA adapter task-vector arithmetic (v0.71.34).
|
|
|
|
Task arithmetic (arXiv:2212.04089) applied to LoRA adapters: add / scale /
|
|
NEGATE task vectors via an expression such as ``"coder + 0.5*math - toxic"``.
|
|
|
|
The engine is **signed, un-normalized, element-wise** over the intersection of
|
|
``lora_A`` / ``lora_B`` tensor names (mirrors PEFT ``combination_type="linear"``):
|
|
``out[k] = Σ cᵢ·tensorᵢ[k]``. Same-rank inputs only — a shape mismatch on a shared
|
|
tensor is a rank mismatch and is rejected loudly (harmonize rank first). Exact
|
|
concatenation+SVD arithmetic for mixed-rank adapters is a future enhancement.
|
|
|
|
No top-level torch/transformers/peft — the parser + numpy merge stay light.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
import re
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Mapping, Sequence
|
|
|
|
_MAX_EXPR_LEN = 4096
|
|
_MAX_TERMS = 64
|
|
_MAX_ADAPTER_CONFIG_BYTES = 256 * 1024
|
|
|
|
_NAME_RE = re.compile(r"[A-Za-z0-9_.\-]+")
|
|
_FLOAT_RE = re.compile(r"[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TaskTerm:
|
|
"""A single ``coeff * adapter`` term of a task-arithmetic expression."""
|
|
|
|
name: str
|
|
coeff: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArithmeticReport:
|
|
"""Result of a ``soup adapters arithmetic`` run."""
|
|
|
|
expression: str
|
|
terms: tuple[TaskTerm, ...]
|
|
output_dir: str
|
|
merged_layers: int
|
|
skipped_layers: tuple[str, ...]
|
|
base_model: str | None
|
|
|
|
|
|
def parse_expression(expr: str, known_names: set[str]) -> list[TaskTerm]:
|
|
"""Parse a task-arithmetic expression into signed ``TaskTerm``s.
|
|
|
|
Grammar (NO ``eval``): ``expr := term (('+'|'-') term)*`` where
|
|
``term := [sign] [coeff '*'] name`` (also ``name ['*' coeff]``). Signs fold
|
|
into the coefficient; an omitted coefficient is ``1.0``. Duplicate adapter
|
|
names sum their coefficients; a term summing to ``0.0`` is dropped.
|
|
|
|
Raises:
|
|
ValueError: empty / over-length expression, unknown token, adapter name
|
|
not in ``known_names``, non-finite coefficient, all terms cancel,
|
|
or more than ``_MAX_TERMS`` distinct adapters.
|
|
"""
|
|
if not isinstance(expr, str):
|
|
raise TypeError("expression must be a string")
|
|
s = expr.strip()
|
|
if not s:
|
|
raise ValueError("empty expression")
|
|
if len(s) > _MAX_EXPR_LEN:
|
|
raise ValueError(f"expression too long (> {_MAX_EXPR_LEN} chars)")
|
|
|
|
n = len(s)
|
|
i = 0
|
|
coeffs: dict[str, float] = {}
|
|
order: list[str] = []
|
|
seen_term = False
|
|
|
|
def _skip_ws() -> None:
|
|
nonlocal i
|
|
while i < n and s[i] in " \t":
|
|
i += 1
|
|
|
|
_skip_ws()
|
|
while i < n:
|
|
_skip_ws()
|
|
if i >= n:
|
|
break
|
|
# Leading sign(s) — only valid before the first term or between terms.
|
|
sign = 1.0
|
|
if s[i] in "+-":
|
|
while i < n and s[i] in "+- \t":
|
|
if s[i] == "-":
|
|
sign = -sign
|
|
i += 1
|
|
elif seen_term:
|
|
# Two terms with no operator between them → malformed.
|
|
raise ValueError(
|
|
f"expected '+' or '-' between terms at pos {i}: {s[i:i + 8]!r}"
|
|
)
|
|
_skip_ws()
|
|
if i >= n:
|
|
raise ValueError("expression ends with a dangling operator")
|
|
|
|
coeff = 1.0
|
|
name: str | None = None
|
|
m = _FLOAT_RE.match(s, i)
|
|
if m:
|
|
coeff = float(m.group())
|
|
i = m.end()
|
|
_skip_ws()
|
|
if i < n and s[i] == "*":
|
|
i += 1
|
|
_skip_ws()
|
|
nm = _NAME_RE.match(s, i)
|
|
if not nm:
|
|
raise ValueError(
|
|
f"expected adapter name after coefficient at pos {i}"
|
|
)
|
|
name = nm.group()
|
|
i = nm.end()
|
|
else:
|
|
nm = _NAME_RE.match(s, i)
|
|
if not nm:
|
|
raise ValueError(
|
|
f"unexpected token at pos {i}: {s[i:i + 8]!r}"
|
|
)
|
|
name = nm.group()
|
|
i = nm.end()
|
|
_skip_ws()
|
|
if i < n and s[i] == "*":
|
|
i += 1
|
|
_skip_ws()
|
|
cm = _FLOAT_RE.match(s, i)
|
|
if not cm:
|
|
raise ValueError(
|
|
f"expected coefficient after '*' at pos {i}"
|
|
)
|
|
coeff = float(cm.group())
|
|
i = cm.end()
|
|
|
|
if not math.isfinite(coeff):
|
|
raise ValueError("coefficient must be finite")
|
|
if name not in known_names:
|
|
raise ValueError(
|
|
f"unknown adapter name {name!r} "
|
|
f"(declare it with --adapter {name}=<path>)"
|
|
)
|
|
signed = sign * coeff
|
|
if name not in coeffs:
|
|
coeffs[name] = 0.0
|
|
order.append(name)
|
|
if len(order) > _MAX_TERMS:
|
|
raise ValueError(f"too many distinct terms (> {_MAX_TERMS})")
|
|
coeffs[name] += signed
|
|
seen_term = True
|
|
_skip_ws()
|
|
|
|
if not seen_term:
|
|
raise ValueError("expression has no terms")
|
|
terms = [TaskTerm(nm, coeffs[nm]) for nm in order if coeffs[nm] != 0.0]
|
|
if not terms:
|
|
raise ValueError("all terms cancelled to zero — nothing to merge")
|
|
return terms
|
|
|
|
|
|
def _factor_coeff(name: str, c: float) -> float:
|
|
"""Per-factor coefficient so the reconstructed delta scales *linearly*.
|
|
|
|
A LoRA contributes ``ΔW = B @ A``. Applying the raw coefficient ``c`` to
|
|
both ``lora_A`` and ``lora_B`` would scale ``ΔW`` by ``c²`` (so negation is
|
|
a no-op and 0.5·adapter halves twice). Instead split the magnitude as
|
|
``√|c|`` across both factors and carry the sign on ``lora_B`` only, so the
|
|
self (diagonal) term of the reconstruction is exactly ``c·B@A``
|
|
(``√|c| · sign(c)√|c| = c``). Mirrors PEFT's ``combination_type='linear'``
|
|
task-arithmetic. Non-LoRA tensors (biases / direct deltas) scale linearly
|
|
by ``c``.
|
|
"""
|
|
lname = name.lower()
|
|
if "lora_a" in lname or "lora_embedding_a" in lname:
|
|
return math.sqrt(abs(c))
|
|
if "lora_b" in lname or "lora_embedding_b" in lname:
|
|
return math.copysign(math.sqrt(abs(c)), c)
|
|
return float(c)
|
|
|
|
|
|
def merge_task_arithmetic(
|
|
weights_list: Sequence[Mapping[str, Any]],
|
|
coeffs: Sequence[float],
|
|
) -> tuple[dict[str, Any], tuple[str, ...]]:
|
|
"""Signed task-vector combine over the intersection of tensor names.
|
|
|
|
Per adapter ``i`` and tensor ``k`` the effective factor coefficient is
|
|
:func:`_factor_coeff` of ``coeffs[i]`` so that the reconstructed LoRA delta
|
|
``B_out @ A_out`` scales linearly with each coefficient (negation flips the
|
|
delta, ``0.5·`` halves it — not the ``c²`` a naive element-wise sum gives).
|
|
Names present in only some adapters are reported in ``skipped``. A shape
|
|
mismatch on a *shared* name is a rank mismatch and raises (same-rank
|
|
contract).
|
|
"""
|
|
import numpy as np
|
|
|
|
if len(weights_list) != len(coeffs):
|
|
raise ValueError(
|
|
f"weights_list ({len(weights_list)}) and coeffs "
|
|
f"({len(coeffs)}) length mismatch"
|
|
)
|
|
if not weights_list:
|
|
raise ValueError("need at least one adapter")
|
|
|
|
shared = set(weights_list[0].keys())
|
|
all_keys = set(weights_list[0].keys())
|
|
for w in weights_list[1:]:
|
|
shared &= set(w.keys())
|
|
all_keys |= set(w.keys())
|
|
|
|
merged: dict[str, Any] = {}
|
|
for name in sorted(shared):
|
|
tensors = [np.asarray(w[name], dtype=np.float64) for w in weights_list]
|
|
if len({t.shape for t in tensors}) > 1:
|
|
raise ValueError(
|
|
f"rank/shape mismatch on {name!r} across adapters — task "
|
|
f"arithmetic requires same-rank adapters (harmonize the LoRA "
|
|
f"rank first, or merge with `soup adapters merge --strategy svd`)"
|
|
)
|
|
acc = np.zeros_like(tensors[0])
|
|
for c, t in zip(coeffs, tensors):
|
|
acc += _factor_coeff(name, float(c)) * t
|
|
merged[name] = acc.astype(np.float32)
|
|
|
|
skipped = tuple(sorted(all_keys - shared))
|
|
return merged, skipped
|
|
|
|
|
|
def read_adapter_base(adapter_dir: str) -> str | None:
|
|
"""Read ``base_model_name_or_path`` from an adapter's ``adapter_config.json``.
|
|
|
|
Returns ``None`` when the config is absent or the field is missing. The read
|
|
is symlink-rejecting (O_NOFOLLOW) and size-capped at 256 KiB (mirrors
|
|
``adapter_merge.write_merged_adapter``'s config-read guards).
|
|
"""
|
|
cfg_path = Path(adapter_dir) / "adapter_config.json"
|
|
if not os.path.lexists(str(cfg_path)):
|
|
return None
|
|
try:
|
|
fd = os.open(str(cfg_path), os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
|
|
except OSError as exc:
|
|
raise ValueError(
|
|
f"adapter_config.json unreadable: {type(exc).__name__}"
|
|
) from exc
|
|
fh = None
|
|
try:
|
|
if os.fstat(fd).st_size > _MAX_ADAPTER_CONFIG_BYTES:
|
|
raise ValueError(
|
|
f"adapter_config.json exceeds {_MAX_ADAPTER_CONFIG_BYTES} byte cap"
|
|
)
|
|
fh = os.fdopen(fd, "r", encoding="utf-8")
|
|
raw = fh.read()
|
|
finally:
|
|
if fh is not None:
|
|
fh.close()
|
|
else:
|
|
os.close(fd)
|
|
try:
|
|
data = json.loads(raw)
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise ValueError(f"adapter_config.json is not valid JSON: {exc}") from exc
|
|
if not isinstance(data, dict):
|
|
return None
|
|
base = data.get("base_model_name_or_path")
|
|
return base if isinstance(base, str) else None
|