diff --git a/src/soup_cli/utils/adapter_arithmetic.py b/src/soup_cli/utils/adapter_arithmetic.py index 8d943cb..a4e8b39 100644 --- a/src/soup_cli/utils/adapter_arithmetic.py +++ b/src/soup_cli/utils/adapter_arithmetic.py @@ -165,15 +165,39 @@ def parse_expression(expr: str, known_names: set[str]) -> list[TaskTerm]: 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, un-normalized element-wise ``out[k] = Σ cᵢ·tensorᵢ[k]``. + """Signed task-vector combine over the intersection of tensor names. - Operates over the intersection of tensor names; 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). + 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 @@ -202,7 +226,7 @@ def merge_task_arithmetic( ) acc = np.zeros_like(tensors[0]) for c, t in zip(coeffs, tensors): - acc += float(c) * t + acc += _factor_coeff(name, float(c)) * t merged[name] = acc.astype(np.float32) skipped = tuple(sorted(all_keys - shared)) diff --git a/src/soup_cli/utils/lisa.py b/src/soup_cli/utils/lisa.py index af56386..e70b257 100644 --- a/src/soup_cli/utils/lisa.py +++ b/src/soup_cli/utils/lisa.py @@ -13,9 +13,13 @@ groups. This callback then toggles ``requires_grad`` — frozen parameters produ re-freeze — so peak optimizer memory ≈ (embed + head + ``num_layers`` active) rather than the whole model. -No top-level torch/transformers — the callback is duck-typed (mirrors -``ReLoRACallback``); torch is never imported here at all (pure ``requires_grad`` -toggling + optimizer ``.state`` dict operations). +No top-level torch/transformers — torch is never imported (pure +``requires_grad`` toggling + optimizer ``.state`` dict operations), and the HF +``TrainerCallback`` base is resolved lazily via ``_try_import_callback_base`` +(mirrors ``monitoring/curriculum_callback.py``) so the module stays import-cheap +while ``LisaCallback`` still inherits the no-op defaults for every Trainer event +(HF dispatches every event with ``getattr(cb, event)`` and no ``hasattr`` guard, +so a bare duck-typed callback would crash on ``on_epoch_begin``). """ from __future__ import annotations @@ -68,6 +72,21 @@ class LisaPolicy: raise TypeError("LisaPolicy.reset_optimizer must be bool") +def _try_import_callback_base(): + """Return HF ``TrainerCallback`` (or ``object`` when transformers is absent). + + Imported inside the function so the module has no top-level transformers + dependency; the class below still inherits every no-op event stub the HF + dispatch loop requires. Mirrors ``monitoring/curriculum_callback.py``. + """ + try: + from transformers import TrainerCallback # noqa: PLC0415 + + return TrainerCallback + except Exception: # noqa: BLE001 — transformers optional in slim test envs. + return object + + def _is_always_on(name: str) -> bool: return any(sub in name for sub in _ALWAYS_ON) @@ -82,11 +101,12 @@ def locate_decoder_layer_indices(model: Any) -> list[int]: return sorted(seen) -class LisaCallback: - """Duck-typed HF ``TrainerCallback`` implementing LISA layer sampling. +class LisaCallback(_try_import_callback_base()): # type: ignore[misc] + """HF ``TrainerCallback`` implementing LISA layer sampling. - Not a ``transformers.TrainerCallback`` subclass so importing this module - never loads transformers — the Trainer's callback dispatch is structural. + Subclasses the lazily-resolved ``TrainerCallback`` so it inherits the no-op + default for every Trainer event; only ``on_train_begin`` / + ``on_step_end`` are overridden. """ def __init__(self, policy: LisaPolicy, console: Any = None) -> None: @@ -128,11 +148,15 @@ class LisaCallback: def _resample(self, model: Any, optimizer: Any) -> None: indices = locate_decoder_layer_indices(model) if not indices: - logger.warning( + # Fail loud rather than silently full-fine-tune with none of LISA's + # memory savings: the callback left the whole model trainable and + # cannot select a decoder subset for this architecture. + raise RuntimeError( "LISA: could not detect numbered decoder layers " - "('layers.N.' / 'h.N.') — no layer sampling applied." + "('layers.N.' / 'h.N.') in the model — LISA layer sampling " + "cannot be applied to this architecture. Disable lisa_enabled " + "or use a supported decoder LM." ) - return k = min(self.policy.num_layers, len(indices)) chosen = set(self._rng.sample(indices, k)) diff --git a/tests/test_v07134.py b/tests/test_v07134.py index 9d38c8c..b2cb76c 100644 --- a/tests/test_v07134.py +++ b/tests/test_v07134.py @@ -123,22 +123,49 @@ class TestParseExpression: # Task A2 — signed merge + base reader # --------------------------------------------------------------------------- class TestMergeTaskArithmetic: - def test_subtract(self): + def test_linear_on_non_lora_tensor(self): from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic - a = {"lora_A": np.ones((2, 3), dtype=np.float32)} - b = {"lora_A": np.full((2, 3), 4.0, dtype=np.float32)} + # A tensor that is neither lora_A nor lora_B combines linearly by c. + a = {"modules_to_save.weight": np.ones((2, 3), dtype=np.float32)} + b = {"modules_to_save.weight": np.full((2, 3), 4.0, dtype=np.float32)} merged, skipped = merge_task_arithmetic([a, b], [1.0, -1.0]) - assert np.allclose(merged["lora_A"], -3.0) + assert np.allclose(merged["modules_to_save.weight"], -3.0) assert skipped == () - def test_scale(self): + def test_scale_non_lora(self): from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic a = {"w": np.ones((2, 2), dtype=np.float32)} merged, _ = merge_task_arithmetic([a], [2.5]) assert np.allclose(merged["w"], 2.5) + def test_reconstructed_delta_negates(self): + # For a real LoRA, negating the task vector must negate ΔW = B @ A. + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + rng = np.random.default_rng(0) + a_mat = rng.standard_normal((4, 8)).astype(np.float32) + b_mat = rng.standard_normal((8, 4)).astype(np.float32) + ak = "base_model.model.layers.0.mlp.down_proj.lora_A.weight" + bk = "base_model.model.layers.0.mlp.down_proj.lora_B.weight" + merged, _ = merge_task_arithmetic([{ak: a_mat, bk: b_mat}], [-1.0]) + delta_orig = b_mat @ a_mat + delta_neg = merged[bk] @ merged[ak] + assert np.allclose(delta_neg, -delta_orig, atol=1e-4) + + def test_reconstructed_delta_scales_linearly(self): + from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic + + rng = np.random.default_rng(1) + a_mat = rng.standard_normal((4, 8)).astype(np.float32) + b_mat = rng.standard_normal((8, 4)).astype(np.float32) + ak = "x.lora_A.weight" + bk = "x.lora_B.weight" + merged, _ = merge_task_arithmetic([{ak: a_mat, bk: b_mat}], [0.5]) + delta = merged[bk] @ merged[ak] + assert np.allclose(delta, 0.5 * (b_mat @ a_mat), atol=1e-4) + def test_mixed_rank_rejected(self): from soup_cli.utils.adapter_arithmetic import merge_task_arithmetic @@ -606,6 +633,34 @@ class TestLisaCallback: for p in frozen_now: assert p not in opt.state or opt.state[p] == {} + def test_is_real_trainer_callback_subclass(self): + # CRITICAL: HF dispatches every event via getattr(cb, event) with no + # hasattr guard, so LisaCallback must inherit TrainerCallback's no-op + # stubs or training crashes on on_epoch_begin. + from transformers import TrainerCallback + + from soup_cli.utils.lisa import LisaCallback, LisaPolicy + + cb = LisaCallback(LisaPolicy(num_layers=1, interval_steps=5)) + assert isinstance(cb, TrainerCallback) + # a non-overridden event exists and is callable (inherited no-op) + assert callable(cb.on_epoch_begin) + + def test_no_decoder_layers_raises(self): + import torch.nn as nn + + from soup_cli.utils.lisa import LisaCallback, LisaPolicy + + class NoLayers(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = nn.Embedding(4, 4) + self.lm_head = nn.Linear(4, 4) + + cb = LisaCallback(LisaPolicy(num_layers=1, interval_steps=5)) + with pytest.raises(RuntimeError, match="decoder layer"): + cb.on_train_begin(None, _State(0), None, model=NoLayers()) + def test_no_top_level_torch(self): import soup_cli.utils.lisa as mod