| """Gated LoRA adapters of Vela 3.0. |
| |
| The adapters (vela3_config.json "adapters") sit on the base's Linear layers; each adapter's layers, rank and alpha come from its |
| *_lora.safetensors. At most one adapter is on in any forward; with no adapter active, a wrapped layer returns base(x) itself, so every |
| other pass is the untouched base computation. |
| |
| The adapter math is the training code's GatedLoRA in eval mode (dropout off): |
| y = base(x) + ((x.float() @ A.T) @ B.T * alpha / r).to(y.dtype) |
| A row mask applies the adapter to some rows of a batch only: torch.where keeps base(x) bit for bit on the other rows.""" |
| from __future__ import annotations |
|
|
| import contextlib |
| import json |
| import threading |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class _Gate(threading.local): |
| name = None |
| rows = None |
|
|
|
|
| GATE = _Gate() |
|
|
|
|
| @contextlib.contextmanager |
| def active(name, rows=None): |
| """Activate one adapter for the forwards inside the block (rows: optional bool mask over the batch).""" |
| prev = (GATE.name, GATE.rows) |
| GATE.name, GATE.rows = name, rows |
| try: |
| yield |
| finally: |
| GATE.name, GATE.rows = prev |
|
|
|
|
| class MultiLoRA(nn.Module): |
| def __init__(self, base: nn.Linear): |
| super().__init__() |
| self.base = base |
| self.ad = {} |
|
|
| def add(self, name, A, B, scale): |
| dev = self.base.weight.device |
| self.ad[name] = (A.to(dev, torch.float32), B.to(dev, torch.float32), float(scale)) |
|
|
| def forward(self, x): |
| y = self.base(x) |
| name = GATE.name |
| if name is None or name not in self.ad: |
| return y |
| A, B, s = self.ad[name] |
| d = ((x.float() @ A.T) @ B.T * s).to(y.dtype) |
| rows = GATE.rows |
| if rows is None: |
| return y + d |
| if rows.shape[0] != y.shape[0]: |
| raise RuntimeError("row gate does not match the batch") |
| return torch.where(rows.view(-1, *([1] * (y.dim() - 1))), y + d, y) |
|
|
|
|
| def attach(lm, name, path): |
| """Attach the adapter in `path` (safetensors: '<module path>.A' / '.B', metadata r / alpha / first / targets) to the text model.""" |
| from safetensors import safe_open |
| with safe_open(str(path), framework="pt") as f: |
| meta = f.metadata() or {} |
| sd = {k: f.get_tensor(k) for k in f.keys()} |
| r, alpha = int(meta["r"]), int(meta["alpha"]) |
| mods = sorted({k.rsplit(".", 1)[0] for k in sd}) |
| for mp in mods: |
| parent_path, attr = mp.rsplit(".", 1) |
| parent = lm.get_submodule(parent_path) |
| cur = getattr(parent, attr) |
| if isinstance(cur, nn.Linear): |
| cur = MultiLoRA(cur) |
| setattr(parent, attr, cur) |
| elif not isinstance(cur, MultiLoRA): |
| raise ValueError(f"{mp} is not a Linear layer") |
| A, B = sd[mp + ".A"], sd[mp + ".B"] |
| if A.shape != (r, cur.base.in_features) or B.shape != (cur.base.out_features, r): |
| raise ValueError(f"adapter shape mismatch at {mp}") |
| cur.add(name, A, B, alpha / r) |
| if len(sd) != 2 * len(mods): |
| raise ValueError(f"adapter {name}: unpaired A / B tensors") |
| return dict(modules=len(mods), r=r, alpha=alpha, first=int(meta.get("first", -1)), |
| targets=json.loads(meta.get("targets", "[]"))) |
|
|
