Spaces:
Running on Zero
Running on Zero
| """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() | |
| 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 | |
| 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() | |