model_loader.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
from __future__ import annotations
 
import os
import threading
from collections import OrderedDict
from pathlib import Path
from typing import Any
 
# Model config registry — mirrors inspection-backend/model_config.py
MODEL_CONFIGS: dict[str, dict[str, Any]] = {
    "llama-3.2-1b": {
        "hf_name": "meta-llama/Llama-3.2-1B-Instruct",
        "d_model": 2048,
        "n_layers": 16,
        "sae_layer": 8,
        "sae_source": "native",
        "sae_layers": {i: f"sae_layer{i}.pt" for i in range(16)},
        "lora_target_modules": ["q_proj", "v_proj"],
    },
    "pythia-2.8b": {
        "hf_name": "EleutherAI/pythia-2.8b",
        "d_model": 2560,
        "n_layers": 32,
        "sae_layer": 16,
        "sae_source": "native",
        "sae_layers": {},
        "lora_target_modules": ["query_key_value", "dense"],
    },
    "gpt2-small": {
        "hf_name": "gpt2",
        "d_model": 768,
        "n_layers": 12,
        "sae_layer": 8,
        "sae_source": "native",
        "sae_layers": {8: "sae_layer8.pt"},
        "lora_target_modules": ["c_attn", "c_fc"],
    },
    "pythia-70m": {
        "hf_name": "EleutherAI/pythia-70m-deduped",
        "d_model": 512,
        "n_layers": 6,
        "sae_layer": 3,
        "sae_source": "native",
        "sae_layers": {3: "sae_layer3.pt"},
        "lora_target_modules": ["query_key_value", "dense"],
    },
    "sarvam-30b": {
        "hf_name": "sarvamai/sarvam-30b",
        "d_model": 4096,
        "n_layers": 19,
        "n_heads": 16,
        "sae_layer": 9,
        "sae_source": "native",
        "sae_layers": {9: "sarvam_l9.pt"},
        "trust_remote_code": True,
        "hf_only": True,
        "lora_target_modules": ["query_key_value"],
    },
    "lfm2.5-1.2b-instruct": {
        "hf_name": "LiquidAI/LFM2.5-1.2B-Instruct",
        "d_model": 2048,
        "n_layers": 16,
        "n_heads": 32,
        "sae_layer": 8,
        "sae_source": "native",
        "sae_layers": {8: "sae_layer8.pt"},
        "norm_layers": {8: "norm_layer8.pt"},
        "hf_only": True,
        "lora_target_modules": ["q_proj", "v_proj"],
    },
    "lfm2.5-1.2b-thinking": {
        "hf_name": "LiquidAI/LFM2.5-1.2B-Thinking",
        "d_model": 2048,
        "n_layers": 16,
        "n_heads": 32,
        "sae_layer": 8,
        "sae_source": "native",
        "sae_layers": {8: "sae_layer8.pt"},
        "norm_layers": {8: "norm_layer8.pt"},
        "hf_only": True,
        "lora_target_modules": ["q_proj", "v_proj"],
    },
    "lfm2.5-230m": {
        "hf_name": "LiquidAI/LFM2.5-230M",
        "d_model": 1024,
        "n_layers": 14,
        "n_heads": 16,
        "sae_layer": 7,
        "sae_source": "native",
        "sae_layers": {i: f"sae_layer{i}.pt" for i in range(14)},
        "norm_layers": {i: f"norm_layer{i}.pt" for i in range(14)},
        "hf_only": True,
        "lora_target_modules": ["q_proj", "v_proj"],
    },
}
 
_HF_TO_SHORT: dict[str, str] = {
    cfg["hf_name"].lower(): key for key, cfg in MODEL_CONFIGS.items()
}
 
MAX_LOADED_MODELS = 1
_models: OrderedDict = OrderedDict()
 
# One in-flight model build at a time. Concurrent from_pretrained / TL conversion
# on MPS (unified memory) doubles host RAM and looks like a multi-10GB "leak".
_model_io_lock = threading.RLock()
 
# Resident SAEs keyed by (short_model_id, layer). Persists across dispatches in
# the daemon process so repeated feature/inspect commands do not reload from disk.
_sae_cache: dict[tuple[str, int], Any] = {}
 
_ACTIVE_MODEL_PATH = Path.home() / ".aquin" / "active_model.txt"
 
 
def clear_sae_cache() -> None:
    _sae_cache.clear()
 
 
def evict_sae_cache(model_id: str, layer: int | None = None) -> None:
    """Drop a cached SAE so the next load re-reads from disk (after a rebind)."""
    short = resolve_model_id(model_id)
    if layer is None:
        for key in [k for k in _sae_cache if k[0] == short]:
            _sae_cache.pop(key, None)
    else:
        _sae_cache.pop((short, int(layer)), None)
 
 
def _save_active_model(model_id: str) -> None:
    _ACTIVE_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    _ACTIVE_MODEL_PATH.write_text(model_id)
 
 
def get_active_model_id(*, allow_daemon: bool = True) -> str | None:
    """Return the model slug written by the last successful `aquin load`."""
    if _ACTIVE_MODEL_PATH.exists():
        txt = _ACTIVE_MODEL_PATH.read_text().strip()
        if txt:
            return txt
    if not allow_daemon:
        return None
    # If the user suspended `aquin load --model ...` in their shell before the CLI
    # wrote active_model.txt, the background daemon may still finish loading. Treat
    # the daemon's resident model as active so follow-up commands still work.
    try:
        from aquin.engine import model_daemon
 
        resident = model_daemon.loaded_model()
        if resident:
            return resident
    except Exception:
        pass
    return None
 
 
def reload_vram_for_model(model_id: str) -> None:
    """Reload the given model into VRAM (used after heavy jobs release the GPU)."""
    short = resolve_model_id(model_id)
    if get_loaded_llm_id() != short:
        clear_llm_models()
        load_model(short)
 
 
def clear_active_model_file() -> None:
    if _ACTIVE_MODEL_PATH.exists():
        _ACTIVE_MODEL_PATH.unlink(missing_ok=True)
 
 
def clear_llm_models() -> None:
    from aquin.compute.device import empty_device_cache
 
    with _model_io_lock:
        while _models:
            _, evicted = _models.popitem(last=False)
            del evicted
        _sae_cache.clear()
        empty_device_cache()
 
 
def get_loaded_llm_id() -> str | None:
    if not _models:
        return None
    return next(reversed(_models))
 
 
class ComputeNotAvailableError(RuntimeError):
    pass
 
 
def resolve_model_id(model_id: str) -> str:
    if model_id in MODEL_CONFIGS:
        return model_id
    short = _HF_TO_SHORT.get(model_id.lower())
    if short:
        return short
    from aquin.compute.model_families import resolve_llm_family
 
    slug, _ = resolve_llm_family(model_id)
    return slug
 
 
def get_config(model_id: str) -> dict:
    short = resolve_model_id(model_id)
    if short in MODEL_CONFIGS:
        return MODEL_CONFIGS[short]
    from aquin.compute.model_families import get_runtime_llm_config
 
    runtime = get_runtime_llm_config(short)
    if runtime:
        return runtime
    raise ValueError(f"No config for model '{model_id}'")
 
 
def get_hf_name(model_id: str) -> str:
    """HuggingFace repo id for transformers.from_pretrained (resolves Aquin short slugs)."""
    return get_config(model_id)["hf_name"]
 
 
def get_sae_layer(model_id: str) -> int:
    return get_config(model_id)["sae_layer"]
 
 
def _layer_from_sae_filename(path: Path) -> int | None:
    stem = path.stem
    if stem.startswith("sae_layer"):
        try:
            return int(stem.replace("sae_layer", ""))
        except ValueError:
            return None
    return None
 
 
def _is_loadable_sae_file(path: Path) -> bool:
    from aquin.compute.torch_io import is_valid_sae_checkpoint_path
 
    return path.is_file() and is_valid_sae_checkpoint_path(path)
 
 
def corrupt_sae_checkpoint_hint(model_id: str, layer: int) -> str | None:
    """If a file exists on disk but cannot be loaded, return a re-download hint."""
    short = resolve_model_id(model_id)
    resolved_layer = int(layer)
    path = resolve_sae_path(short, resolved_layer)
    if not path.is_file():
        return None
    if _is_loadable_sae_file(path):
        return None
    return (
        f"SAE checkpoint at {path} is corrupted or incomplete. "
        f"Delete it and re-download:\n"
        f"  rm \"{path}\"\n"
        f"  aquin load sae {short}-l{resolved_layer}"
    )
 
 
def load_sae_from_disk(
    model_id: str,
    layer: int,
    *,
    device: str | None = None,
) -> Any:
    """Load a validated on-disk SAE or raise with a re-download hint."""
    import torch
    from aquin.compute.sae import SparseAutoencoder
 
    from aquin.compute.device import resolve_compute_device
 
    short = resolve_model_id(model_id)
    resolved_layer = int(layer)
    dev = device or resolve_compute_device()
 
    hint = corrupt_sae_checkpoint_hint(short, resolved_layer)
    if hint:
        raise ValueError(hint)
 
    path = resolve_sae_checkpoint_path(short, resolved_layer)
    if path is None:
        raise FileNotFoundError(
            format_sae_layer_choice_message(
                short,
                command="feature logit",
                requested_layer=resolved_layer,
            )
        )
 
    try:
        return SparseAutoencoder.load(str(path), device=dev)
    except Exception as exc:
        from aquin.compute.torch_io import is_corrupt_pytorch_zip
 
        cause = exc.__cause__ if exc.__cause__ is not None else exc
        if is_corrupt_pytorch_zip(exc) or is_corrupt_pytorch_zip(cause):
            raise ValueError(
                corrupt_sae_checkpoint_hint(short, resolved_layer)
                or (
                    f"SAE checkpoint at {path} is corrupted or incomplete. "
                    f"Re-download: aquin load sae {short}-l{resolved_layer}"
                )
            ) from exc
        raise
 
 
def resolve_sae_checkpoint_path(model_id: str, layer: int) -> Path | None:
    """Return a loadable SAE checkpoint for model+layer, or None if missing/ambiguous."""
    short = resolve_model_id(model_id)
    resolved_layer = int(layer)
 
    try:
        from aquin.compute.user_sae import get_active_user_sae_path, list_user_saes
 
        active = get_active_user_sae_path(short, resolved_layer)
        if active is not None and _is_loadable_sae_file(active):
            return active
 
        catalog = resolve_sae_path(short, resolved_layer)
        if _is_loadable_sae_file(catalog):
            return catalog
 
        user_matches = [
            Path(str(r["path"]))
            for r in list_user_saes(short)
            if r.get("layer") == resolved_layer and r.get("path")
        ]
        existing = [p for p in user_matches if _is_loadable_sae_file(p)]
        if len(existing) == 1:
            return existing[0]
    except Exception:
        pass
 
    return None
 
 
def get_available_sae_layers(model_id: str) -> list[int]:
    """SAE layer indices with loadable checkpoints on disk for this model."""
    short = resolve_model_id(model_id)
    found: set[int] = set()
 
    for layer in get_catalog_sae_layers(short):
        if resolve_sae_checkpoint_path(short, int(layer)) is not None:
            found.add(int(layer))
 
    try:
        from aquin.compute.user_sae import list_user_saes
 
        for row in list_user_saes(short):
            layer_raw = row.get("layer")
            if layer_raw is None:
                continue
            layer = int(layer_raw)
            if resolve_sae_checkpoint_path(short, layer) is not None:
                found.add(layer)
    except Exception:
        pass
 
    sae_dir = Path.home() / ".aquin" / "sae" / short
    if sae_dir.is_dir():
        for path in sae_dir.glob("sae_layer*.pt"):
            layer = _layer_from_sae_filename(path)
            if layer is not None and path.is_file():
                found.add(layer)
 
    return sorted(found)
 
 
def get_catalog_sae_layers(model_id: str) -> list[int]:
    """All SAE layer indices published for this model (may not be downloaded yet)."""
    short = resolve_model_id(model_id)
    cfg = get_config(short)
    layers: dict = cfg.get("sae_layers", {})
    if layers:
        return sorted(layers.keys())
    return [int(cfg["sae_layer"])]
 
 
def format_sae_layer_choice_message(
    model_id: str,
    *,
    command: str,
    layer_flag: str = "--layer",
    example_suffix: str = "",
    requested_layer: int | None = None,
) -> str:
    """Human-readable hint listing on-disk SAEs and aquin load sae pull commands."""
    short = resolve_model_id(model_id)
    available = get_available_sae_layers(short)
    catalog = get_catalog_sae_layers(short)
    not_downloaded = [layer for layer in catalog if layer not in available]
    sae_dir = Path.home() / ".aquin" / "sae" / short
 
    lines: list[str] = []
    if requested_layer is not None:
        corrupt = corrupt_sae_checkpoint_hint(short, requested_layer)
        if corrupt:
            lines.append(corrupt)
        else:
            lines.append(f"SAE layer {requested_layer} is not downloaded for {short}.")
    else:
        lines.append(f"No {layer_flag} specified for {command}.")
    lines.append(f"Model: {short}")
    lines.append("")
 
    if available:
        lines.append("Downloaded on this machine:")
        for layer in available:
            path = resolve_sae_checkpoint_path(short, layer) or resolve_sae_path(short, layer)
            lines.append(f"  {layer_flag} {layer}  ({path})")
    else:
        lines.append("No SAE checkpoints found on disk.")
        lines.append(f"  Directory: {sae_dir}/")
 
    lines.append("")
    if available:
        pick = requested_layer if requested_layer in available else available[0]
        example = f"  aquin {command}"
        if example_suffix:
            example += f" {example_suffix}"
        example += f" {layer_flag} {pick}"
        lines.append("Example:")
        lines.append(example)
 
    if not_downloaded:
        lines.append("")
        lines.append("Pull more layers:")
        pull_layers = (
            [requested_layer]
            if requested_layer is not None and requested_layer not in available
            else not_downloaded
        )
        for layer in pull_layers:
            lines.append(f"  aquin load sae {short}-l{layer}")
 
    lines.append("")
    lines.append(f"Try: aquin load sae {short}-l<n>  or  aquin load sae --path <weights.pt>")
    lines.append("Self-hosted: aquin load sae --path <weights.pt> --layer <n> [--model <id>]")
    return "\n".join(lines)
 
 
def require_sae_layer(
    model_id: str,
    layer: int | None,
    *,
    command: str = "trace",
    example_suffix: str = "",
) -> int:
    """
    Resolve an explicit SAE layer for tools that must not silently default.
 
    Raises ValueError with an actionable message when layer is omitted or missing on disk.
    """
    short = resolve_model_id(model_id)
    available = get_available_sae_layers(short)
 
    if layer is not None:
        resolved = int(layer)
        if resolve_sae_checkpoint_path(short, resolved) is not None:
            return resolved
        raise ValueError(
            format_sae_layer_choice_message(
                short,
                command=command,
                example_suffix=example_suffix,
                requested_layer=resolved,
            )
        )
 
    if not available:
        raise ValueError(
            format_sae_layer_choice_message(
                short,
                command=command,
                example_suffix=example_suffix,
            )
        )
 
    raise ValueError(
        format_sae_layer_choice_message(
            short,
            command=command,
            example_suffix=example_suffix,
        )
    )
 
 
def resolve_sae_path(model_id: str, layer: int | None = None) -> Path:
    """Canonical on-disk path for a model SAE (existing file, or expected location)."""
    short = resolve_model_id(model_id)
    cfg = get_config(short)
    sae_dir = Path.home() / ".aquin" / "sae" / short
    resolved_layer = int(layer if layer is not None else cfg["sae_layer"])
 
    candidates: list[Path] = []
    sae_layers: dict = cfg.get("sae_layers", {})
    if rel := sae_layers.get(resolved_layer):
        candidates.append(sae_dir / Path(rel).name)
    if layer is None and (fn := cfg.get("sae_filename")):
        candidates.append(sae_dir / Path(fn).name)
    candidates.append(sae_dir / f"sae_layer{resolved_layer}.pt")
 
    seen: set[Path] = set()
    unique: list[Path] = []
    for p in candidates:
        if p not in seen:
            seen.add(p)
            unique.append(p)
 
    for p in unique:
        if p.exists():
            return p
    return unique[0]
 
 
def get_lora_target_modules(model_id: str) -> list[str]:
    cfg = get_config(model_id)
    return list(cfg.get("lora_target_modules", ["q_proj", "v_proj"]))
 
 
def infer_lora_target_modules(model: Any) -> list[str]:
    """Pick LoRA targets from module names when config defaults do not match."""
    leaf_names = {name.split(".")[-1] for name, _ in model.named_modules()}
    presets: list[list[str]] = [
        ["q_proj", "v_proj"],
        ["query_key_value", "dense"],
        ["query_key_value"],
        ["c_attn", "c_fc"],
        ["Wqkv", "out_proj"],
    ]
    for candidates in presets:
        if all(c in leaf_names for c in candidates):
            return candidates
    for candidates in presets:
        hit = [c for c in candidates if c in leaf_names]
        if hit:
            return hit
    return ["q_proj", "v_proj"]
 
 
def get_sae_source(model_id: str) -> str:
    return get_config(model_id).get("sae_source", "native")
 
 
 
 
def get_loaded_model() -> Any | None:
    """Return the most recently used loaded model, or None."""
    if not _models:
        return None
    return next(reversed(_models.values()))
 
 
def _tl_not_in_catalog(exc: BaseException) -> bool:
    msg = str(exc).lower()
    return "not found" in msg or "valid official model names" in msg
 
 
def _tl_conversion_failed(exc: BaseException) -> bool:
    """True when TransformerLens cannot build/wrap this HF architecture (soft-fallback)."""
    if _tl_not_in_catalog(exc):
        return True
    msg = str(exc).lower()
    # e.g. GPTNeoXForCausalLM has no attribute 'embed_out' on some TL + transformers combos
    if "embed_out" in msg:
        return True
    if "has no attribute" in msg and (
        "neox" in msg or "gptneo" in msg or "gpt_neox" in msg or "pythia" in msg
    ):
        return True
    if isinstance(exc, AttributeError) and ("embed" in msg or "unembed" in msg):
        return True
    return False
 
 
def _wrap_hf_causal_lm(
    hf_name: str,
    cfg: dict[str, Any],
    hf_model: Any,
    *,
    dtype: Any,
    device: str,
) -> Any:
    """Wrap a loaded HuggingFace causal LM as HookedTransformer or HfLlmShim."""
    if cfg.get("hf_only"):
        from aquin.compute.hf_llm_shim import HfLlmShim
        from transformers import AutoTokenizer
 
        trust = bool(cfg.get("trust_remote_code", False))
        tokenizer = AutoTokenizer.from_pretrained(hf_name, trust_remote_code=trust)
        n_heads = int(cfg.get("n_heads", 32))
        return HfLlmShim(
            hf_model,
            tokenizer,
            hf_name=hf_name,
            n_layers=int(cfg["n_layers"]),
            d_model=int(cfg["d_model"]),
            n_heads=n_heads,
        )
 
    from transformer_lens import HookedTransformer
 
    trust = bool(cfg.get("trust_remote_code", False))
    try:
        tl = HookedTransformer.from_pretrained(
            hf_name,
            hf_model=hf_model,
            dtype=dtype,
            device=device,
            trust_remote_code=trust,
        )
    except (
        TypeError,
        ValueError,
        KeyError,
        RuntimeError,
        NotImplementedError,
        AttributeError,
    ) as exc:
        if not _tl_conversion_failed(exc):
            raise
        from aquin.compute.hf_llm_shim import HfLlmShim
        from transformers import AutoTokenizer
 
        print(
            f"[model] TransformerLens cannot wrap {hf_name} ({type(exc).__name__}: {exc}) "
            "— using HuggingFace shim…",
            flush=True,
        )
        tokenizer = AutoTokenizer.from_pretrained(hf_name, trust_remote_code=trust)
        n_heads = int(cfg.get("n_heads", 32))
        return HfLlmShim(
            hf_model,
            tokenizer,
            hf_name=hf_name,
            n_layers=int(cfg["n_layers"]),
            d_model=int(cfg["d_model"]),
            n_heads=n_heads,
        )
    tl.eval()
    return tl
 
 
def _build_hooked_transformer(cfg: dict[str, Any], *, dtype: Any, device: str) -> Any:
    """Load HookedTransformer, or HfLlmShim for custom HF-only architectures."""
    from aquin.compute.hf_auth import require_hf_hub_auth
 
    hf_name = cfg["hf_name"]
    require_hf_hub_auth(hf_repo=hf_name)
 
    if cfg.get("hf_only"):
        from aquin.compute.hf_llm_shim import HfLlmShim
 
        print(f"[model] {hf_name} — loading via HuggingFace (no TransformerLens wrapper)…", flush=True)
        return HfLlmShim.from_pretrained(hf_name, cfg, dtype=dtype, device=device)
 
    from transformers import AutoModelForCausalLM
 
    trust = bool(cfg.get("trust_remote_code", False))
    hf_first = bool(cfg.get("hf_first", False))
 
    if not hf_first:
        from transformer_lens import HookedTransformer
 
        try:
            return HookedTransformer.from_pretrained(
                hf_name,
                dtype=dtype,
                trust_remote_code=trust,
            )
        except (
            ValueError,
            KeyError,
            RuntimeError,
            NotImplementedError,
            AttributeError,
            TypeError,
        ) as exc:
            if not _tl_conversion_failed(exc):
                raise
            print(
                f"[model] {hf_name} cannot use TransformerLens directly "
                f"({type(exc).__name__}: {exc}) — loading via HuggingFace…",
                flush=True,
            )
 
    hf_model = AutoModelForCausalLM.from_pretrained(
        hf_name,
        torch_dtype=dtype,
        device_map=device,
        trust_remote_code=trust,
    )
    hf_model.eval()
    return _wrap_hf_causal_lm(hf_name, cfg, hf_model, dtype=dtype, device=device)
 
 
def load_model(model_id: str) -> Any:
    """
    Load a HookedTransformer model by short slug or HF name.
    Uses an LRU cache — at most MAX_LOADED_MODELS kept in VRAM.
    Raises ComputeNotAvailableError when no accelerator (unless AQUIN_ALLOW_CPU=1).
 
    Serialized: never interleave two builds (MPS unified RAM doubles fast).
    """
    with _model_io_lock:
        return _load_model_unlocked(model_id)
 
 
def _load_model_unlocked(model_id: str) -> Any:
    import torch
    from aquin.compute.device import (
        default_dtype_for_device,
        empty_device_cache,
        is_oom_error,
        require_load_device,
    )
 
    short = resolve_model_id(model_id)
 
    if short in _models:
        _models.move_to_end(short)
        _save_active_model(short)
        print(f"[model] cache hit: {short}", flush=True)
        return _models[short]
 
    # Avoid VRAM duplication: if a persistent daemon holds a model in another
    # process, free it before we allocate here. The daemon reloads lazily later.
    if os.environ.get("AQUIN_DAEMON") != "1":
        try:
            from aquin.compute.model_runtime import release_foreign_daemon
 
            release_foreign_daemon()
        except Exception:
            pass
 
    cfg = get_config(short)
    try:
        device = require_load_device(short, cfg)
    except RuntimeError as exc:
        raise ComputeNotAvailableError(str(exc)) from exc
    dtype = default_dtype_for_device(device)
 
    if not cfg.get("hf_only"):
        try:
            from transformer_lens import HookedTransformer  # noqa: F401
        except ImportError:
            raise ComputeNotAvailableError(
                "transformer_lens is not installed. "
                "Reinstall Aquin: pip install -U aquin"
            )
 
    from aquin.compute.hf_auth import gated_repo_hint, require_hf_hub_auth
 
    try:
        require_hf_hub_auth(hf_repo=cfg["hf_name"])
    except RuntimeError as exc:
        raise ComputeNotAvailableError(str(exc)) from exc
 
    import time
 
    print(f"[model] loading {cfg['hf_name']} ({dtype} on {device})...", flush=True)
    start = time.time()
    stop_ticker = threading.Event()
 
    def _tick() -> None:
        while not stop_ticker.wait(5):
            print(f"[model] still loading... {int(time.time() - start)}s", flush=True)
 
    threading.Thread(target=_tick, daemon=True).start()
    try:
        m = _build_hooked_transformer(cfg, dtype=dtype, device=device)
    except Exception as exc:
        msg = str(exc).lower()
        if "gated" in msg or "401" in msg or "unauthorized" in msg:
            raise ComputeNotAvailableError(gated_repo_hint(cfg["hf_name"])) from exc
        if is_oom_error(exc):
            hint = ""
            if short == "sarvam-30b" or int(cfg.get("d_model", 0)) >= 4096:
                hint = (
                    " Large MoE models like Sarvam 30B store all expert weights "
                    "(~60GB+ in bf16); use a GPU with enough VRAM."
                )
            raise ComputeNotAvailableError(
                f"Accelerator ran out of memory loading '{short}' on {device}.{hint}"
            ) from exc
        raise
    finally:
        stop_ticker.set()
 
    # Drop before install if the UI cancelled / unload raced us.
    try:
        from aquin.compute.model_runtime import is_load_cancelled
 
        if is_load_cancelled():
            del m
            empty_device_cache(device)
            print(f"[model] discarded {short} (load cancelled)", flush=True)
            raise RuntimeError(f"Model load cancelled ({short})")
    except ImportError:
        pass
 
    m.eval()
    from aquin.compute.hf_llm_shim import HfLlmShim
 
    if not isinstance(m, HfLlmShim):
        m.to(device)
    print(f"[model] loaded in {int(time.time() - start)}s", flush=True)
 
    _models[short] = m
    _save_active_model(short)
 
    # Evict previous model — only one LLM in VRAM at a time
    while len(_models) > MAX_LOADED_MODELS:
        evicted_id, evicted_model = _models.popitem(last=False)
        del evicted_model
        empty_device_cache(device)
        print(f"[model] unloaded {evicted_id} from VRAM", flush=True)
 
    print(f"[model] {short} ready. (only model in VRAM)", flush=True)
    return _models[short]
 
 
def load_sae(model: Any, layer: int, model_id: str, sae_dir: Path | None = None) -> Any:
    """
    Load a SparseAutoencoder for the given layer from ~/.aquin/sae/ or user bindings.
    Use `aquin load sae` (Aquin catalog) or `aquin load sae --path` for local files.
    """
    from aquin.compute.device import resolve_compute_device
 
    short = resolve_model_id(model_id)
    resolved_layer = int(layer)
 
    cache_key = (short, resolved_layer)
    cached = _sae_cache.get(cache_key)
    if cached is not None:
        return cached
 
    device = resolve_compute_device()
    sae = load_sae_from_disk(short, resolved_layer, device=device)
    print(f"[sae] loaded layer {resolved_layer} for {short}", flush=True)
    _sae_cache[cache_key] = sae
    return sae