sae.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
 
from pathlib import Path
 
import torch
import torch.nn as nn
 
 
class SparseAutoencoder(nn.Module):
    def __init__(self, d_model: int = 2048, n_features: int = 16384):
        super().__init__()
        self.d_model = d_model
        self.n_features = n_features
        self.b_pre = nn.Parameter(torch.zeros(d_model))
        self.W_enc = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(d_model, n_features)))
        self.b_enc = nn.Parameter(torch.zeros(n_features))
        self.W_dec = nn.Parameter(torch.nn.init.kaiming_uniform_(torch.empty(n_features, d_model)))
        self.b_dec = nn.Parameter(torch.zeros(d_model))
        self._normalise_decoder()
 
    def _normalise_decoder(self) -> None:
        with torch.no_grad():
            norms = self.W_dec.norm(dim=-1, keepdim=True).clamp(min=1e-8)
            self.W_dec.data = self.W_dec.data / norms
 
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        return torch.relu((x - self.b_pre) @ self.W_enc + self.b_enc)
 
    def decode(self, f: torch.Tensor) -> torch.Tensor:
        return f @ self.W_dec + self.b_dec
 
    def forward(self, x: torch.Tensor):
        f = self.encode(x)
        return f, self.decode(f)
 
    @torch.no_grad()
    def get_top_features(self, x: torch.Tensor, k: int = 10) -> list[tuple[int, float]]:
        f = self.encode(x.unsqueeze(0)).squeeze(0)
        topk = f.topk(k)
        return [(int(i), float(v)) for i, v in zip(topk.indices, topk.values)]
 
    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"d_model": self.d_model, "n_features": self.n_features, "state_dict": self.state_dict()}, path)
 
    @classmethod
    def _find_tensor(cls, tensors: dict, *names: str) -> torch.Tensor | None:
        for name in names:
            if name in tensors:
                return tensors[name]
        leaf_to_full: dict[str, str] = {}
        for key, val in tensors.items():
            if isinstance(val, torch.Tensor):
                leaf_to_full[key.split(".")[-1]] = key
        for name in names:
            leaf = name.split(".")[-1]
            if leaf in leaf_to_full:
                return tensors[leaf_to_full[leaf]]
        return None
 
    @classmethod
    def _infer_dims(cls, tensors: dict) -> tuple[int, int]:
        enc = cls._find_tensor(tensors, "W_enc", "encoder.weight")
        dec = cls._find_tensor(tensors, "W_dec", "decoder.weight")
        if enc is not None and enc.ndim == 2:
            if enc.shape[0] > enc.shape[1]:
                return int(enc.shape[1]), int(enc.shape[0])
            return int(enc.shape[0]), int(enc.shape[1])
        if dec is not None and dec.ndim == 2:
            if dec.shape[0] > dec.shape[1]:
                return int(dec.shape[1]), int(dec.shape[0])
            return int(dec.shape[1]), int(dec.shape[0])
        raise KeyError(
            "Cannot infer SAE dimensions — missing W_enc/encoder.weight or W_dec/decoder.weight. "
            f"Keys: {sorted(tensors.keys())}"
        )
 
    @classmethod
    def _coerce_checkpoint(cls, ckpt: Any) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
        if isinstance(ckpt, cls):
            return {"d_model": ckpt.d_model, "n_features": ckpt.n_features}, ckpt.state_dict()
 
        if not isinstance(ckpt, dict):
            raise ValueError(f"Unrecognised SAE checkpoint type: {type(ckpt)}")
 
        meta = dict(ckpt)
        if "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict):
            tensors = {k: v for k, v in ckpt["state_dict"].items() if isinstance(v, torch.Tensor)}
        elif "model_state" in ckpt and isinstance(ckpt["model_state"], dict):
            tensors = {k: v for k, v in ckpt["model_state"].items() if isinstance(v, torch.Tensor)}
        elif "model" in ckpt and isinstance(ckpt["model"], dict):
            tensors = {k: v for k, v in ckpt["model"].items() if isinstance(v, torch.Tensor)}
        else:
            tensors = {k: v for k, v in ckpt.items() if isinstance(v, torch.Tensor)}
 
        for nested_key in ("cfg", "config", "args"):
            nested = ckpt.get(nested_key)
            if isinstance(nested, dict):
                for key in ("d_model", "d_in", "n_features", "dict_size", "d_sae", "num_features"):
                    if key in nested and key not in meta:
                        meta[key] = nested[key]
 
        return meta, tensors
 
    @classmethod
    def _meta_int(cls, meta: dict[str, Any], *keys: str) -> int | None:
        for key in keys:
            if key in meta and meta[key] is not None:
                return int(meta[key])
        return None
 
    @classmethod
    def _remap_tensors(cls, tensors: dict[str, torch.Tensor], sae: "SparseAutoencoder") -> dict[str, torch.Tensor]:
        key_map = {
            "encoder.weight": "W_enc",
            "encoder.bias": "b_enc",
            "decoder.weight": "W_dec",
            "decoder.bias": "b_dec",
            "encoder_pre_bias": "b_pre",
            "pre_bias": "b_pre",
            "W_enc": "W_enc",
            "b_enc": "b_enc",
            "W_dec": "W_dec",
            "b_dec": "b_dec",
            "b_pre": "b_pre",
        }
        out = dict(sae.state_dict())
        enc_loaded = False
        for src_key, dst_key in key_map.items():
            t = cls._find_tensor(tensors, src_key)
            if t is None:
                continue
            if dst_key == "W_enc":
                if enc_loaded:
                    continue
                t = t.float()
                if t.shape[0] > t.shape[1]:
                    t = t.T
                enc_loaded = True
            elif dst_key == "W_dec":
                t = t.float()
                if t.shape == (sae.d_model, sae.n_features):
                    t = t.T
            else:
                t = t.float()
            out[dst_key] = t
        return out
 
    @classmethod
    def load(cls, path: str | Path, device: str = "cuda") -> "SparseAutoencoder":
        path = Path(path)
        if path.suffix == ".safetensors":
            return cls._load_safetensors(path, device)
        from aquin.compute.torch_io import is_corrupt_pytorch_zip, load_checkpoint
 
        try:
            ckpt = load_checkpoint(path, map_location=device)
        except Exception as exc:
            if is_corrupt_pytorch_zip(exc) or (
                exc.__cause__ is not None and is_corrupt_pytorch_zip(exc.__cause__)
            ):
                raise ValueError(
                    f"SAE checkpoint at {path} is corrupted or incomplete. "
                    "Delete the file and run: aquin load sae <model>-l<layer>"
                ) from exc
            raise
        from aquin.compute.torch_io import looks_like_catalog_metadata_dict
 
        if looks_like_catalog_metadata_dict(ckpt):
            raise ValueError(
                f"SAE checkpoint at {path} is catalog metadata, not weights. "
                "Re-download with: aquin load sae <model>-l<n>"
            )
        if isinstance(ckpt, cls):
            ckpt.to(device)
            ckpt.eval()
            return ckpt
 
        if not isinstance(ckpt, dict):
            raise ValueError(
                f"Not an SAE dictionary checkpoint ({type(ckpt)}): {path}. "
                "Use ~/.aquin/sae/user/<model>/<name>/sae_layerN.pt — "
                "not capture files like sae/sae_layer_N.pt or layers/layer_N.pt."
            )
 
        meta, tensors = cls._coerce_checkpoint(ckpt)
        if not tensors or (
            cls._find_tensor(tensors, "W_enc", "encoder.weight") is None
            and cls._find_tensor(tensors, "W_dec", "decoder.weight") is None
        ):
            raise ValueError(
                f"SAE checkpoint at {path} is not a valid PyTorch weights file. "
                "Re-download with: aquin load sae <model>-l<n>  "
                "or place weights with: aquin load sae --path <file.pt> --layer <n>"
            )
        d_model = cls._meta_int(meta, "d_model", "d_in")
        n_features = cls._meta_int(meta, "n_features", "dict_size", "d_sae", "num_features")
        if d_model is None or n_features is None:
            inferred_d, inferred_n = cls._infer_dims(tensors)
            d_model = d_model if d_model is not None else inferred_d
            n_features = n_features if n_features is not None else inferred_n
 
        sae = cls(d_model=d_model, n_features=n_features)
        remapped = cls._remap_tensors(tensors, sae)
        if cls._find_tensor(tensors, "W_enc", "encoder.weight") is not None:
            sae.load_state_dict(remapped, strict=False)
        else:
            sae.load_state_dict({k: v for k, v in remapped.items() if k in sae.state_dict()}, strict=False)
        sae.to(device)
        sae.eval()
        return sae
 
    @classmethod
    def _load_safetensors(cls, path: Path, device: str = "cuda") -> "SparseAutoencoder":
        from safetensors.torch import load_file
        tensors = load_file(str(path), device=device)
 
        # SAELens-style key layout:
        # W_enc, b_enc, W_dec, b_dec, b_pre (optional: scaling_factor)
        # Map to our parameter names
        key_map = {
            "W_enc": ["W_enc", "encoder.weight"],
            "b_enc": ["b_enc", "encoder.bias"],
            "W_dec": ["W_dec", "decoder.weight"],
            "b_dec": ["b_dec", "decoder.bias"],
            "b_pre": ["b_pre", "pre_bias", "encoder_pre_bias"],
        }
 
        def _get(candidates: list[str]) -> torch.Tensor | None:
            for k in candidates:
                if k in tensors:
                    return tensors[k]
            return None
 
        W_enc = _get(key_map["W_enc"])
        b_enc = _get(key_map["b_enc"])
        W_dec = _get(key_map["W_dec"])
        b_dec = _get(key_map["b_dec"])
        b_pre = _get(key_map["b_pre"])
 
        if W_enc is None or W_dec is None:
            raise ValueError(f"Unrecognised safetensors layout. Keys: {list(tensors.keys())}")
 
        # Infer dimensions — SAELens encoder.weight is [n_features, d_model]
        if W_enc.ndim != 2:
            raise ValueError(f"Unexpected W_enc shape: {W_enc.shape}")
        if W_enc.shape[0] > W_enc.shape[1]:
            n_features, d_model = W_enc.shape
            W_enc = W_enc.T
        else:
            d_model, n_features = W_enc.shape
 
        sae = cls(d_model=d_model, n_features=n_features)
        sae.W_enc.data = W_enc.float()
        if W_dec.shape == (n_features, d_model):
            sae.W_dec.data = W_dec.float()
        elif W_dec.shape == (d_model, n_features):
            sae.W_dec.data = W_dec.T.float()
        else:
            sae.W_dec.data = W_dec.T.float()
        if b_enc is not None:
            sae.b_enc.data = b_enc.float()
        if b_dec is not None:
            sae.b_dec.data = b_dec.float()
        if b_pre is not None:
            sae.b_pre.data = b_pre.float()
        sae.to(device)
        sae.eval()
        return sae