Part 6: Serving many adapters over one base model

"""Serve several LoRA adapters over ONE base model: load the base once, swap adapter weights per request."""
import time

import torch

from examples.m09_dpo import fixed_prompt
from examples.m09_lora import add_lora, with_adapter
from examples.m09_setup import MODELS, classify, prompt_for
from supportdesk.data import load_tickets
from supportdesk.tinylm import SamplingParams, generate, load


class AdapterServer:
    """One base model with LoRA slots of rank `max_rank`; each adapter is zero-padded to fit.

    Padding is exact: extra rows of A and columns of B are zero, and each adapter's alpha/r
    scaling is folded into its B, so every adapter computes exactly what it computed alone.
    """

    def __init__(self, max_rank: int = 8) -> None:
        self.model, self.tok = load()
        add_lora(self.model, r=max_rank, alpha=max_rank)  # slot scaling = 1
        self.max_rank = max_rank
        self.adapters: dict[str, dict[str, torch.Tensor]] = {}
        self.active = None

    def register(self, name: str, path) -> None:
        state = torch.load(path, weights_only=True)
        r, alpha = state["meta"]["r"], state["meta"]["alpha"]
        padded = {}
        for key, value in state["lora"].items():
            if key.endswith("lora_A"):
                padded[key] = torch.cat([value, torch.zeros(self.max_rank - r, value.shape[1])])
            else:
                padded[key] = torch.cat([value * (alpha / r), torch.zeros(value.shape[0], self.max_rank - r)], dim=1)
        self.adapters[name] = padded

    @torch.no_grad()
    def use(self, name: str) -> float:
        """Copy one adapter into the LoRA slots; returns milliseconds spent (0 if already active)."""
        if name == self.active:
            return 0.0
        started = time.perf_counter()
        params = dict(self.model.named_parameters())
        for key, value in self.adapters[name].items():
            params[key].copy_(value)
        self.active = name
        return (time.perf_counter() - started) * 1000


if __name__ == "__main__":
    started = time.perf_counter()
    server = AdapterServer(max_rank=8)
    base_ms = (time.perf_counter() - started) * 1000
    for name in ("m09-lora-triage", "m09-reply-sft", "m09-reply-dpo"):
        server.register(name, MODELS / name / "adapter.pt")
    print(f"base loaded once in {base_ms:.0f} ms; adapters registered: {list(server.adapters)}")

    # Exactness check: padded adapter in the server == the adapter loaded on its own.
    x = torch.tensor([server.tok.encode(fixed_prompt(0)).ids])
    server.use("m09-lora-triage")
    alone = with_adapter("m09-lora-triage")
    with torch.no_grad():
        print("max logit difference vs standalone adapter:", float((server.model(x) - alone(x)).abs().max()))

    tickets = load_tickets("test")[:3]
    requests = [("triage", t) for t in tickets] + [("reply", 10), ("triage", load_tickets("test")[3]), ("reply", 11)]
    swaps = []
    for kind, item in requests:
        if kind == "triage":
            ms = server.use("m09-lora-triage")
            out = classify(server.model, server.tok, prompt_for(server.tok, item))
            label = item.id
        else:
            ms = server.use("m09-reply-dpo")
            out = generate(server.model, server.tok, fixed_prompt(item),
                           SamplingParams(max_new_tokens=30, temperature=0, stop=["\n"])).text.strip()
            label = f"question {item + 1}"
        if ms:
            swaps.append(ms)
        print(f"  {kind:<6} {label:<11} swap {ms:5.2f} ms  -> {out[:70]!r}")
    timings = []
    for _ in range(50):
        server.active = None
        timings.append(server.use("m09-reply-dpo"))
    timings.sort()
    started = time.perf_counter()
    load(MODELS / "m09-full")
    full_ms = (time.perf_counter() - started) * 1000
    print(f"adapter swap: median {timings[25]:.2f} ms over 50 swaps; loading a full fine-tuned model from disk: {full_ms:.0f} ms")

base loaded once in 32 ms; adapters registered: ['m09-lora-triage', 'm09-reply-sft', 'm09-reply-dpo']
max logit difference vs standalone adapter: 0.0
  triage T-1003      swap  0.00 ms  -> 'cancellation'
  triage T-1006      swap  0.00 ms  -> 'how_to'
  triage T-1009      swap  0.00 ms  -> 'how_to'
  reply  question 11 swap  0.23 ms  -> 'Thank you. Team costs 12 USD per user per month. Business costs 24 USD'
  triage T-1012      swap  0.19 ms  -> 'billing'
  reply  question 12 swap  0.19 ms  -> 'Thank you. Please check status.brightlane.example for live updates.'
adapter swap: median 0.11 ms over 50 swaps; loading a full fine-tuned model from disk: 23 ms