Blame for adapter_runtime.py

"""Request-scoped eager LoRA contributions without changing base weights."""

from __future__ import annotations

from contextlib import contextmanager
import math
from numbers import Real
from threading import Lock, RLock
from weakref import WeakKeyDictionary

import torch
from torch import nn
from torch.nn import functional as F


def _scale(value):
    if not isinstance(value, Real) or isinstance(value, bool):
        raise TypeError("adapter scale must be a real number")
    value = float(value)
    if not math.isfinite(value):
        raise ValueError("adapter scale must be finite")
    return value


def _factors(a, b, name):
    for label, tensor in (("A", a), ("B", b)):
        if not isinstance(tensor, torch.Tensor):
            raise TypeError(f"{name}: {label} must be a tensor")
        if tensor.layout != torch.strided or not tensor.is_floating_point():
            raise ValueError(f"{name}: {label} must be a dense floating-point tensor")
        if tensor.ndim != 2 or min(tensor.shape) <= 0:
            raise ValueError(f"{name}: {label} must be a nonempty matrix")
        if tensor.is_meta or not torch.isfinite(tensor).all().item():
            raise ValueError(f"{name}: {label} must contain finite materialized data")
    if a.shape[0] != b.shape[1]:
        raise ValueError(f"{name}: A/B ranks do not match")


def _targets(transformer, entries):
    """Validate the entire set without installing any hooks."""
    if not isinstance(transformer, nn.Module):
        raise TypeError("transformer must be a torch.nn.Module")
    targets = []
    seen = set()
    for name, a, b in entries:
        if not isinstance(name, str) or not (name == "weight" or name.endswith(".weight")):
            raise ValueError(f"invalid Linear weight target: {name!r}")
        path = "" if name == "weight" else name[:-len(".weight")]
        if name != "weight" and (not path or any(not part for part in path.split("."))):
            raise ValueError(f"invalid Linear weight target: {name!r}")
        try:
            module = transformer.get_submodule(path)
        except AttributeError as exc:
            raise KeyError(f"adapter target {name!r} not found") from exc
        if not isinstance(module, nn.Linear) or not isinstance(module._parameters.get("weight"), nn.Parameter):
            raise ValueError(f"{name}: only unparametrized torch.nn.Linear weights are supported")
        if id(module) in seen:
            raise ValueError(f"{name}: duplicate adapter target")
        seen.add(id(module))
        weight = module.weight
        if weight.is_meta or not weight.is_floating_point() or weight.layout != torch.strided:
            raise ValueError(f"{name}: weight must be a materialized floating-point matrix")
        _factors(a, b, name)
        if tuple(weight.shape) != (b.shape[0], a.shape[1]):
            raise ValueError(f"{name}: adapter shape does not match weight {tuple(weight.shape)}")
        targets.append((module, a, b))
    return targets


def adapter_entries(
    profile, transformer: nn.Module, state: dict[str, torch.Tensor]
) -> tuple[list[tuple[str, torch.Tensor, torch.Tensor]], float]:
    """Map an already-loaded Larry/LightX state, without I/O or fusion.

    Larry's target transforms come from the unchanged vendored spec. LightX
    uses direct names and its upstream ``(alpha or rank) / rank`` convention.
    Returned factors stay on their source device; activation snapshots them
    onto the current target weights' device and dtype.
    """
    kind = getattr(profile, "kind", None)
    if kind not in ("larry", "lightx"):
        raise ValueError(f"unsupported adapter kind: {kind!r}")
    if not isinstance(state, dict):
        raise TypeError("adapter state must be a dict of tensors")
    if not state:
        raise ValueError("adapter state is empty")
    suffix = ".weight" if kind == "larry" else ".default.weight"
    suffix_a, suffix_b = ".lora_A" + suffix, ".lora_B" + suffix
    pairs = {}
    for key, tensor in state.items():
        if not isinstance(key, str):
            raise ValueError(f"malformed adapter key: {key!r}")
        if key.endswith(suffix_a):
            name, label = key[:-len(suffix_a)], "A"
        elif key.endswith(suffix_b):
            name, label = key[:-len(suffix_b)], "B"
        else:
            raise ValueError(f"malformed {kind} adapter key: {key!r}")
        if not name or any(not part for part in name.split(".")):
            raise ValueError(f"malformed adapter target: {name!r}")
        pairs.setdefault(name, {})[label] = tensor
    for name, pair in pairs.items():
        if set(pair) != {"A", "B"}:
            raise ValueError(f"{name}: unpaired LoRA A/B tensors")
        _factors(pair["A"], pair["B"], name)

    entries = []
    if kind == "larry":
        # Lazy import keeps generic activation independent of the H3 spec.
        from h3_adapter_spec import _larry_targets

        config = getattr(transformer, "config", None)
        heads = getattr(config, "num_attention_heads", None)
        head_dim = getattr(config, "attention_head_dim", None)
        if type(heads) is not int or type(head_dim) is not int or heads <= 0 or head_dim <= 0:
            raise ValueError("Larry mapping requires positive integer attention dimensions")
        inner_dim = heads * head_dim
        for name, pair in sorted(pairs.items()):
            a, b = pair["A"], pair["B"]
            # The upstream zip truncates excess QKV chunks; reject them before
            # invoking it. A gate/value projection must have two equal halves.
            if name.endswith(".attn.qkv_proj") and b.shape[0] != 3 * inner_dim:
                raise ValueError(f"{name}: QKV B rows must equal 3 * inner_dim")
            if name.endswith(".mlp.fc1") and b.shape[0] % 2:
                raise ValueError(f"{name}: gate/value B rows must be even")
            entries.extend((target, a, part) for target, part in _larry_targets(name, b, inner_dim))
        scale = 1.0
    else:
        ranks = {pair["A"].shape[0] for pair in pairs.values()}
        if len(ranks) != 1:
            raise ValueError(f"mixed LoRA ranks: {sorted(ranks)}")
        rank = ranks.pop()
        alpha = getattr(profile, "alpha", None)
        if alpha is not None:
            alpha = _scale(alpha)
        scale = _scale((alpha or rank) / rank)
        entries = [(name + ".weight", pair["A"], pair["B"]) for name, pair in sorted(pairs.items())]
    _targets(transformer, entries)
    return entries, scale


class _Access:
    def __init__(self):
        self.lock = RLock()
        self.active = False


_ACCESS = WeakKeyDictionary()
_ACCESS_LOCK = Lock()


@contextmanager
def adapter_context(transformer, entries, scale):
    """Serialize one eager inference context per transformer, including base.

    Every caller sharing a transformer must enter this context, even with no
    adapter. Direct calls outside it are not isolated. Reentrant activation is
    rejected. Do not move or replace modules/weights while a context is active.
    Compiled/AoTI execution that bypasses module forwards is unsupported.
    """
    if not isinstance(transformer, nn.Module):
        raise TypeError("transformer must be a torch.nn.Module")
    with _ACCESS_LOCK:
        access = _ACCESS.get(transformer)
        if access is None:
            access = _Access()
            _ACCESS[transformer] = access
    with access.lock:
        if access.active:
            raise RuntimeError("nested/reentrant adapter context for an active transformer")
        access.active = True
        try:
            with _activation(transformer, entries, scale) as model:
                yield model
        finally:
            access.active = False


@contextmanager
def _activation(transformer, entries, scale):
    """Validate and snapshot all factors before the first hook is installed."""
    scale = _scale(scale)
    targets = _targets(transformer, entries)
    active = []
    for module, a, b in targets:
        # Snapshot the caller's factors, including mixed dtypes, at the live
        # weight's device/dtype. Never mutate or reassign the base parameter.
        a = a.detach().to(device=module.weight.device, dtype=module.weight.dtype, copy=True)
        b = b.detach().to(device=module.weight.device, dtype=module.weight.dtype, copy=True)
        _factors(a, b, module.__class__.__name__)
        active.append((module, a, b))
    handles = []
    try:
        for module, a, b in active:
            def add_adapter(module, args, kwargs, output, a=a, b=b):
                x = args[0] if args else kwargs["input"]
                return output + scale * F.linear(F.linear(x, a), b)

            before = set(module._forward_hooks)
            try:
                handles.append(module.register_forward_hook(add_adapter, with_kwargs=True))
            except BaseException:
                # A wrapped registration method can insert a hook and then
                # raise without returning its handle. Roll back only its new
                # hook IDs, leaving pre-existing hooks and their flags intact.
                for hook_id in set(module._forward_hooks) - before:
                    module._forward_hooks.pop(hook_id, None)
                    module._forward_hooks_with_kwargs.pop(hook_id, None)
                    module._forward_hooks_always_called.pop(hook_id, None)
                raise
        yield transformer
    finally:
        for handle in reversed(handles):
            handle.remove()