"""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()