diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index 910c166f1d..119f4e34b6 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -1283,7 +1283,8 @@ async def _generate_from_intrinsic( # Two load-bearing assumptions: (1) __getattr__ falls through for attributes # accessed by generate_with_transformers (model.device, model.vocab_size, - # model.generation_config). (2) chat_completion_request_to_transformers_inputs + # model.generation_config, and for aLoRA activation model.peft_config, + # model.active_adapters, model.modules()). (2) chat_completion_request_to_transformers_inputs # always sets return_dict_in_generate=True, so .generate() always returns a # GenerateDecoderOnlyOutput — if that ever changes, the cell stays None and # logits silently won't be populated. diff --git a/mellea/formatters/granite/base/util.py b/mellea/formatters/granite/base/util.py index 94a9243b84..d7e6264710 100644 --- a/mellea/formatters/granite/base/util.py +++ b/mellea/formatters/granite/base/util.py @@ -6,10 +6,12 @@ # Standard from __future__ import annotations +import contextlib import itertools import json import os import uuid +from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Any, cast # Third Party @@ -378,6 +380,160 @@ def chat_completion_request_to_transformers_inputs( return generate_input, other_input +# aLoRA adapter names already warned about for a missing invocation sequence, +# so a mispackaged adapter logs once per process rather than once per call. +_ALORA_MISSING_INVOCATION_WARNED: set[str] = set() + + +def _generation_setting( + model: PreTrainedModel, generate_kwargs: Mapping[str, Any], name: str +) -> int: + """Effective integer `generate()` setting, resolved in `generate()`'s order. + + A flat kwarg wins, then a `generation_config=` kwarg, then the model's own + `generation_config`; 1 if none sets it. Unset `GenerationConfig` fields + are `None`, so each level falls through to the next. + """ + value = generate_kwargs.get(name) + if value is None: + value = getattr(generate_kwargs.get("generation_config"), name, None) + if value is None: + value = getattr(getattr(model, "generation_config", None), name, None) + return value if isinstance(value, int) else 1 + + +@contextlib.contextmanager +def _alora_activation_context( + model: PreTrainedModel, + input_tokens: torch.Tensor | None, + generate_kwargs: Mapping[str, Any] | None = None, +) -> Iterator[None]: + """Context manager that makes aLoRA weights actually apply during `generate()`. + + PEFT computes the per-token aLoRA activation mask (`alora_offsets`) only + from inside its `PeftModel` wrapper's own `generate()`/`forward()` + overrides. A model that loads adapters via transformers' native + `PeftAdapterMixin.load_adapter()` (as `LocalHFBackend` does) is never + wrapped in `PeftModel`, so without this the aLoRA weights silently never + apply and generation degrades to base-model behaviour. This mirrors what + `PeftModel` does internally (`_enable_peft_forward_hooks` in + `peft/tuners/lora/model.py`): compute the offsets with PEFT's + `calculate_alora_offsets` helper (internal to + `peft.tuners.lora.variants`, not part of PEFT's exported API), then inject + them into every LoRA layer's forward via temporary pre-forward hooks. The + hook is required because transformers' model code does not generally + propagate `**kwargs` down to the projection `nn.Linear` calls, so passing + `alora_offsets` to `generate()` alone is not enough. + + Yields immediately (no-op) unless exactly one adapter is active and its + PEFT config declares a non-empty `alora_invocation_tokens`. When the + declared sequence does not occur in `input_tokens`, the adapter cannot + activate: no hooks are registered (so generation uses the base-model + weights, as PEFT would) and a WARNING naming the adapter is logged once + per process. Offsets are repeated per returned sequence, since + `generate()` expands the prompt batch for `num_return_sequences`. Hooks + are removed on exit, including on error. Safe to call on any model; the + PEFT imports happen lazily so this module keeps working without the `hf` + extra. + + Args: + model: Hugging Face model object (or a proxy that forwards attribute + access to one). + input_tokens: Prompt token-id tensor on the model's device, as + passed to `generate(input_ids=...)`. + generate_kwargs: The other keyword arguments the caller is about to + pass to `generate()`, read for `num_beams` and + `num_return_sequences`. + + Yields: + None. Active for the duration of the `with` block. + + Raises: + ValueError: Beam search (`num_beams > 1`) was requested while an aLoRA + adapter can activate; PEFT does not support that combination. + """ + peft_config = getattr(model, "peft_config", None) + if not peft_config: + yield + return + try: + active_adapters = model.active_adapters + # `active_adapters` is a method on transformers' native + # PeftAdapterMixin but a property on peft's PeftModel wrapper. + active = cast( + "list[str]", + active_adapters() if callable(active_adapters) else active_adapters, + ) + except ValueError as e: + # transformers raises + # `ValueError("No adapter loaded. Please load an adapter first.")` + # when nothing is loaded; anything else is a real fault. + if "No adapter loaded" not in str(e): + raise + active = [] + if len(active) != 1: + yield + return + adapter_name = active[0] + invocation = getattr(peft_config.get(adapter_name), "alora_invocation_tokens", None) + if not invocation: + yield + return + if input_tokens is None or not hasattr(input_tokens, "device"): + yield + return + + # Third Party (lazy: peft is an optional `hf`-extra dependency) + from peft.tuners.lora import ( + layer as _peft_lora_layer, + variants as _peft_lora_variants, + ) + + offsets = _peft_lora_variants.calculate_alora_offsets( + peft_config, adapter_name, input_tokens + ) + if all(offset is None for offset in offsets): + if adapter_name not in _ALORA_MISSING_INVOCATION_WARNED: + _ALORA_MISSING_INVOCATION_WARNED.add(adapter_name) + MelleaLogger.get_logger().warning( + f"aLoRA adapter {adapter_name!r} is active but its declared " + f"invocation sequence {list(invocation)} does not occur in the " + "prompt, so generation used base-model weights only. Check that " + "the adapter's io.yaml instruction or chat template produces " + "that token sequence." + ) + yield + return + + generate_kwargs = generate_kwargs or {} + if _generation_setting(model, generate_kwargs, "num_beams") > 1: + # Same refusal as PeftModel's own path (`_enable_peft_forward_hooks`). + raise ValueError( + f"Beam search is not supported for aLoRA adapters (adapter {adapter_name!r})." + ) + # generate() expands each prompt row `num_return_sequences` times via + # `repeat_interleave`, so repeat each row's offset the same way. + n_sequences = _generation_setting(model, generate_kwargs, "num_return_sequences") + offsets = [offset for offset in offsets for _ in range(n_sequences)] + + def _inject_offsets( + module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> tuple[tuple[Any, ...], dict[str, Any]]: + kwargs["alora_offsets"] = offsets + return args, kwargs + + handles = [ + module.register_forward_pre_hook(_inject_offsets, with_kwargs=True) + for module in model.modules() + if isinstance(module, _peft_lora_layer.LoraLayer) + ] + try: + yield + finally: + for handle in handles: + handle.remove() + + def generate_with_transformers( tokenizer: PreTrainedTokenizerBase, model: PreTrainedModel, @@ -403,6 +559,11 @@ def generate_with_transformers( Returns: A chat completion response in OpenAI format. + + Raises: + ValueError: Beam search (`num_beams > 1`) was requested while an aLoRA + adapter is active and its invocation sequence is in the prompt; + PEFT does not support that combination. """ with import_optional("torch"): # Third Party @@ -412,7 +573,8 @@ def generate_with_transformers( generate_input = generate_input.copy() del generate_input["input_tokens"] - generate_result = model.generate(input_ids=input_tokens, **generate_input) # type: ignore[operator] + with _alora_activation_context(model, input_tokens, generate_input): + generate_result = model.generate(input_ids=input_tokens, **generate_input) # type: ignore[operator] # Result is a a 2D tensor of shape (num responses, prompt + max generated tokens) # containing tokens, plus a tuple of tensors of shape diff --git a/test/formatters/granite/base/test_base_alora_activation.py b/test/formatters/granite/base/test_base_alora_activation.py new file mode 100644 index 0000000000..26822fbc6c --- /dev/null +++ b/test/formatters/granite/base/test_base_alora_activation.py @@ -0,0 +1,445 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for aLoRA activation through the bare-model (`load_adapter`) path. + +Background: PEFT only computes and injects aLoRA activation masks +(`alora_offsets`) from inside its `PeftModel` wrapper's `generate()`/ +`forward()` overrides. Models that load adapters via transformers' native +`PeftAdapterMixin.load_adapter()` — which is what `LocalHFBackend` does — +are never wrapped in `PeftModel`, so the aLoRA weights silently never apply +and generation degrades to base-model behaviour. +`util._alora_activation_context` fixes this by mirroring `PeftModel`'s own +mechanism (offset computation + temporary per-layer pre-forward hooks). + +These tests use a tiny randomly-initialised model (no download) and assert +on whether the aLoRA variant layer actually receives `alora_offsets` during +`generate()` — the exact property that was missing before the fix. +""" + +# Standard +import logging +from unittest.mock import MagicMock + +# Third Party +import pytest + +torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]") +pytest.importorskip( + "transformers", reason="transformers not installed — install mellea[hf]" +) +peft = pytest.importorskip("peft", reason="peft not installed — install mellea[hf]") +import peft.tuners.lora.variants as peft_variants +from peft import LoraConfig +from transformers import GenerationConfig, LlamaConfig, LlamaForCausalLM + +# First Party +import mellea.formatters.granite.base.util as granite_util +from mellea.backends.huggingface import LocalHFBackend +from mellea.formatters.granite.base.util import ( + _alora_activation_context, + generate_with_transformers, +) +from test.conftest import cleanup_gpu_backend, hf_skip +from test.predicates import require_gpu + +INVOCATION = [1, 2, 3] + + +def _n_pre_hooks(model: LlamaForCausalLM) -> int: + return sum(len(getattr(m, "_forward_pre_hooks", {})) for m in model.modules()) + + +@pytest.fixture(autouse=True) +def _reset_missing_invocation_warnings(monkeypatch: pytest.MonkeyPatch) -> None: + # The missing-invocation WARNING is once per adapter name per process; + # isolate each test from the others' warnings. + monkeypatch.setattr(granite_util, "_ALORA_MISSING_INVOCATION_WARNED", set()) + + +def _tiny_model(adapter_config: LoraConfig | None = None) -> LlamaForCausalLM: + config = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=128, + ) + model = LlamaForCausalLM(config).eval() + if adapter_config is not None: + model.add_adapter(adapter_config, "rc") + model.set_adapter("rc") + return model + + +def _alora_config(invocation: list[int] = INVOCATION) -> LoraConfig: + return LoraConfig( + r=4, + lora_alpha=8, + target_modules=["q_proj", "v_proj"], + task_type="CAUSAL_LM", + alora_invocation_tokens=invocation, + ) + + +def _lora_config() -> LoraConfig: + return LoraConfig( + r=4, lora_alpha=8, target_modules=["q_proj", "v_proj"], task_type="CAUSAL_LM" + ) + + +class _VariantSpy: + """Count ALoraLinearVariant.forward calls and record the alora_offsets each received.""" + + def __init__(self): + self.n_calls = 0 + self.offsets_seen: list = [] + self._orig = peft_variants.ALoraLinearVariant.forward + + def __enter__(self): + def spy(module, active_adapter, x, result, **kwargs): + self.n_calls += 1 + self.offsets_seen.append(kwargs.get("alora_offsets")) + return self._orig(module, active_adapter, x, result, **kwargs) + + peft_variants.ALoraLinearVariant.forward = staticmethod(spy) + return self + + def __exit__(self, *exc): + # Restore as a staticmethod: class-attribute access to a staticmethod + # returns the bare function, so a plain reassignment would turn it + # into an instance method and break every later aLoRA forward. + peft_variants.ALoraLinearVariant.forward = staticmethod(self._orig) + return False + + +class TestAloraActivationContext: + def test_no_offsets_reach_variant_without_context(self): + """Canary on upstream PEFT behaviour, not the regression guard: bare + model + aLoRA adapter, no context -> the aLoRA variant is called but + never receives alora_offsets, so it masks to base-model behaviour. + Passes with or without the fix; if it starts failing, PEFT has begun + supplying offsets on the bare-model path and the context may be + redundant. The guard is `test_generate_with_transformers_uses_context`.""" + model = _tiny_model(_alora_config()) + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with _VariantSpy() as spy: + with torch.no_grad(): + model.generate(input_ids=input_ids, max_new_tokens=2) + assert spy.n_calls > 0, "aLoRA variant layer should be exercised at all" + assert spy.offsets_seen == [None] * spy.n_calls + + def test_context_delivers_expected_offsets(self): + model = _tiny_model(_alora_config()) + # Invocation sequence starts at index 1 of a 6-token prompt; PEFT + # counts the offset back from the end, so it is 6 - 1 = 5. + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with _VariantSpy() as spy: + with torch.no_grad(), _alora_activation_context(model, input_ids): + model.generate(input_ids=input_ids, max_new_tokens=2) + assert spy.n_calls > 0 + assert spy.offsets_seen[0] == [5] + assert all(o == [5] for o in spy.offsets_seen) + + def test_missing_invocation_warns_once_and_registers_no_hooks(self, caplog): + """Invocation sequence absent -> the adapter cannot activate. The + context registers no hooks (the variant sees `alora_offsets=None`, + PEFT's base-model mask) and logs one WARNING naming the adapter, not + one per call.""" + model = _tiny_model(_alora_config()) + input_ids = torch.tensor([[10, 11, 12, 40, 50]]) + with caplog.at_level(logging.WARNING, logger="mellea"): + for _ in range(2): + with _VariantSpy() as spy: + with torch.no_grad(), _alora_activation_context(model, input_ids): + assert _n_pre_hooks(model) == 0 + model.generate(input_ids=input_ids, max_new_tokens=2) + assert spy.offsets_seen and all(o is None for o in spy.offsets_seen) + warned = [r for r in caplog.records if "invocation sequence" in r.message] + assert len(warned) == 1 + assert "'rc'" in warned[0].message + + @pytest.mark.parametrize( + "generate_kwargs", + [{"num_beams": 2}, {"generation_config": GenerationConfig(num_beams=2)}], + ids=["kwarg", "generation-config"], + ) + def test_beam_search_raises_clear_error(self, generate_kwargs): + """PEFT's own `PeftModel` path rejects beam search for aLoRA; so must + this one, however the setting is passed, rather than failing later + with an opaque IndexError.""" + model = _tiny_model(_alora_config()) + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with pytest.raises(ValueError, match="Beam search"): + with _alora_activation_context(model, input_ids, generate_kwargs): + pass + assert _n_pre_hooks(model) == 0 + + def test_hooks_live_only_inside_context(self): + model = _tiny_model(_alora_config()) + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with torch.no_grad(), _alora_activation_context(model, input_ids): + assert _n_pre_hooks(model) > 0 + assert _n_pre_hooks(model) == 0, ( + "pre-forward hooks must not outlive the context" + ) + with pytest.raises(RuntimeError, match="boom"): + with torch.no_grad(), _alora_activation_context(model, input_ids): + raise RuntimeError("boom") + assert _n_pre_hooks(model) == 0, "hooks must be removed on error too" + + @pytest.mark.parametrize( + "adapter_config", + [None, _lora_config(), _alora_config(invocation=[])], + ids=["no-adapter", "plain-lora", "empty-invocation"], + ) + def test_non_activating_model_is_noop(self, adapter_config): + model = _tiny_model(adapter_config) + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with torch.no_grad(), _alora_activation_context(model, input_ids): + assert _n_pre_hooks(model) == 0 + model.generate(input_ids=input_ids, max_new_tokens=2) + + @staticmethod + def _generate_via_production_path(model, input_ids, **generate_kwargs): + tokenizer = MagicMock() + tokenizer.eos_token_id = 9999 # outside vocab: never in generated tokens + tokenizer.decode.side_effect = lambda *a, **k: "x" + tokenizer.batch_decode.side_effect = lambda seqs: ["x"] * len(seqs) + generate_with_transformers( + tokenizer, + model, + generate_input={ + "input_tokens": input_ids, + "max_new_tokens": 2, + "return_dict_in_generate": True, + **generate_kwargs, + }, + other_input={}, + ) + + def test_generate_with_transformers_uses_context(self): + """The regression guard: the production call path must deliver + offsets to the variant. Fails when the context wiring is removed.""" + model = _tiny_model(_alora_config()) + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with _VariantSpy() as spy: + self._generate_via_production_path(model, input_ids, do_sample=False) + assert spy.n_calls > 0 + assert spy.offsets_seen[0] == [5] + + @pytest.mark.parametrize( + "generate_kwargs", + [ + {"do_sample": True, "num_return_sequences": 2}, + { + "generation_config": GenerationConfig( + do_sample=True, num_return_sequences=2 + ) + }, + ], + ids=["kwarg", "generation-config"], + ) + def test_generate_with_transformers_passes_generation_kwargs(self, generate_kwargs): + """Wiring check: the production path hands its generate kwargs to the + context, so a multi-sequence request (flat kwarg or `generation_config` + object) expands the offsets instead of crashing inside the variant.""" + model = _tiny_model(_alora_config()) + input_ids = torch.tensor([[10, 1, 2, 3, 40, 50]]) + with _VariantSpy() as spy: + self._generate_via_production_path(model, input_ids, **generate_kwargs) + assert spy.offsets_seen[0] == [5, 5] + + +@pytest.fixture +def alora_backend(): + """Factory for a granite-4.1-3b `LocalHFBackend` with one catalogue aLoRA + registered as a composed `Adapter`. Hub errors skip the test, and every + backend built is freed afterwards.""" + from mellea.backends.adapters._core import Adapter, Identity, LocalFileBinding + from mellea.backends.adapters.catalog import AdapterType, fetch_intrinsic_metadata + from mellea.backends.adapters.io_contracts import get_io_contract + + backends: list[LocalHFBackend] = [] + + def _make(name: str, capability: str) -> LocalHFBackend: + md = fetch_intrinsic_metadata(name) + with hf_skip(): + backend = LocalHFBackend(model_id="ibm-granite/granite-4.1-3b") + backends.append(backend) + backend.add_adapter( + Adapter( + identity=Identity( + name=name, adapter_type="alora", capability=capability + ), + io_contract=get_io_contract(name), + weights=LocalFileBinding( + name=name, + adapter_type=AdapterType.ALORA, + repo_id=md.repo_id, + revision=md.revision, + ), + ) + ) + return backend + + yield _make + for backend in backends: + cleanup_gpu_backend(backend, "test_base_alora_activation") + + +@pytest.mark.huggingface +@pytest.mark.e2e +@pytest.mark.qualitative +@pytest.mark.slow +@require_gpu(min_vram_gb=20) +class TestAloraDifferentialEndToEnd: + """The aLoRA weights must measurably change the score versus no adapter. + + Guards against the silent-degradation failure mode where aLoRA + generation falls back to base-model behaviour (issue #1679): before the + fix, adapter-on and adapter-off scores were byte-identical on every item. + The cardiff probe item is one where granite-4.1-3b's base model is + confidently wrong (says "yes" to a bullet-points check on a plain + sentence) and the published requirement-check aLoRA is correct ("no") — + scores measured 0.047 vs 0.999 in the #1679 diagnostic eval. + """ + + @pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason="the published requirement-check aLoRA io.yaml instruction does " + "not tokenise to its declared invocation sequence, so the adapter " + "cannot activate until the adapter is republished", + ) + def test_requirement_check_adapter_moves_score(self, alora_backend): + from mellea.stdlib.components import Message + from mellea.stdlib.components.intrinsic import core + from mellea.stdlib.context import ChatContext + + backend = alora_backend("requirement-check", "requirement_check") + ctx = ( + ChatContext() + .add(Message("user", "Write one sentence about Cardiff.")) + .add(Message("assistant", "Cardiff is the capital city of Wales.")) + ) + requirement = "The response uses bullet points." + + # The model-level PEFT adapter name as registered by load_adapter. + qualified = next(iter(backend._model.peft_config)) + + # Adapter ON: the intrinsic path activates the adapter itself + # around the generate call; this is the real production flow. + score_on = core.requirement_check(ctx, backend, requirement) + + # Adapter OFF: the intrinsic scope re-activates the adapter for + # every call, so the only way to get a genuine no-adapter score + # is to clear the model's active adapters at the generate level + # (the same set_adapter([]) override the diagnostic harness used; + # disable_adapters() is a documented no-op on this model class). + orig_generate = backend._model.generate + + def off_generate(*a, **kw): + backend._model.set_adapter([]) + try: + return orig_generate(*a, **kw) + finally: + backend._model.set_adapter(qualified) + + backend._model.generate = off_generate + try: + score_off = core.requirement_check(ctx, backend, requirement) + finally: + backend._model.generate = orig_generate + + # The adapter must say "no" (low score); the base model says "yes". + assert score_on < 0.5, ( + f"aLoRA active: expected 'no' (score < 0.5), got {score_on}" + ) + assert score_off > 0.5, ( + f"base model: expected 'yes' (score > 0.5), got {score_off}" + ) + assert score_off - score_on > 0.5, ( + f"adapter contribution too small to be a real effect: " + f"on={score_on}, off={score_off}" + ) + + +@pytest.mark.huggingface +@pytest.mark.e2e +@pytest.mark.qualitative +@pytest.mark.slow +@require_gpu(min_vram_gb=20) +class TestUncertaintyAloraDifferentialEndToEnd: + """Second aLoRA capability (uncertainty): confirms the activation fix is + general, not requirement-check-specific. + + The published uncertainty aLoRA io.yaml has no tokenisation mismatch, so + this test exercises the fix on the as-published adapter files. Measured on granite-4.1-3b in the #1679 + diagnostic: base model is underconfident (~0.06) on both a right and a + wrong last response, while the adapter gives ~0.95 (right) and ~0.73 + (wrong) — a wide differential either way. + """ + + def test_check_certainty_adapter_moves_score(self, alora_backend): + from mellea.stdlib.components import Message + from mellea.stdlib.components.intrinsic import core + from mellea.stdlib.context import ChatContext + + backend = alora_backend("uncertainty", "uncertainty") + qualified = next(iter(backend._model.peft_config)) + + def score_with_adapter_off(user: str, assistant: str) -> float: + ctx = ( + ChatContext() + .add(Message("user", user)) + .add(Message("assistant", assistant)) + ) + orig_generate = backend._model.generate + + def off_generate(*a, **kw): + backend._model.set_adapter([]) + try: + return orig_generate(*a, **kw) + finally: + backend._model.set_adapter(qualified) + + backend._model.generate = off_generate + try: + return core.check_certainty(ctx, backend) + finally: + backend._model.generate = orig_generate + + user = "What is the capital of France?" + score_on_right = core.check_certainty( + ChatContext() + .add(Message("user", user)) + .add(Message("assistant", "The capital of France is Paris.")), + backend, + ) + score_off_right = score_with_adapter_off( + user, "The capital of France is Paris." + ) + + score_on_wrong = core.check_certainty( + ChatContext() + .add(Message("user", user)) + .add(Message("assistant", "The capital of France is Madrid.")), + backend, + ) + score_off_wrong = score_with_adapter_off( + user, "The capital of France is Madrid." + ) + + # The adapter must move the certainty score by a wide margin on both + # cases; the base 3b model sits near 0.06 in both. + assert score_on_right - score_off_right > 0.5, ( + f"right answer: adapter contribution too small: " + f"on={score_on_right}, off={score_off_right}" + ) + assert score_on_wrong - score_off_wrong > 0.5, ( + f"wrong answer: adapter contribution too small: " + f"on={score_on_wrong}, off={score_off_wrong}" + )