hf_llm_shim.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.
 
"""HuggingFace causal LM wrapper with a minimal HookedTransformer-like surface."""
 
from __future__ import annotations
 
import re
from dataclasses import dataclass
from typing import Any, Callable
 
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
 
_RESID_POST = re.compile(r"^blocks\.(\d+)\.hook_resid_post$")
_RESID_PRE = re.compile(r"^blocks\.(\d+)\.hook_resid_pre$")
 
 
@dataclass
class _ShimCfg:
    n_layers: int
    d_model: int
    n_heads: int
    model_name: str
 
 
def _names_filter_match(names_filter: Any, name: str) -> bool:
    if names_filter is None:
        return True
    if callable(names_filter):
        return bool(names_filter(name))
    if isinstance(names_filter, str):
        return name == names_filter
    if isinstance(names_filter, (set, list, tuple, frozenset)):
        return name in names_filter
    return False
 
 
def _hf_load_device_map(cfg: dict[str, Any], device: str) -> str | dict[str, Any]:
    """Use device_map='auto' for large HF-native models (e.g. Sarvam 30B MoE)."""
    explicit = cfg.get("device_map")
    if explicit:
        return explicit
    if cfg.get("hf_only") and int(cfg.get("d_model", 0)) >= 4096:
        return "auto"
    return device
 
 
def _model_has_device_map(hf_model: Any) -> bool:
    return bool(getattr(hf_model, "hf_device_map", None))
 
 
def _resolve_modules(hf_model: Any) -> tuple[Any, Any, Any, Any]:
    """Return (backbone, layer_list, embed_module, lm_head)."""
    lm_head = getattr(hf_model, "lm_head", None)
    if lm_head is None and hasattr(hf_model, "get_output_embeddings"):
        lm_head = hf_model.get_output_embeddings()
 
    backbone = getattr(hf_model, "model", hf_model)
    if hasattr(backbone, "get_decoder"):
        backbone = backbone.get_decoder()
 
    layers = (
        getattr(backbone, "layers", None)
        or getattr(backbone, "h", None)
        or getattr(backbone, "decoder", None)
    )
    if layers is not None and hasattr(layers, "layers"):
        layers = layers.layers
 
    embed = (
        getattr(backbone, "embed_tokens", None)
        or getattr(backbone, "word_embeddings", None)
        or getattr(backbone, "wte", None)
        or getattr(backbone, "transformer", None)
    )
    if embed is None and hasattr(hf_model, "get_input_embeddings"):
        embed = hf_model.get_input_embeddings()
    if embed is not None and hasattr(embed, "wte"):
        embed = embed.wte
 
    if layers is None or embed is None:
        raise RuntimeError(
            "Could not locate decoder layers / embeddings on HuggingFace model "
            f"(type={type(hf_model).__name__}, backbone={type(backbone).__name__})"
        )
    return backbone, layers, embed, lm_head
 
 
class _IdentityNorm(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x
 
 
class _UnembedShim:
    def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None:
        self.W_U = weight
        if bias is not None:
            self.b_U = bias
        else:
            self.b_U = torch.zeros(weight.shape[0], device=weight.device, dtype=weight.dtype)
 
 
def project_residual_to_logits(model: Any, resid: torch.Tensor) -> torch.Tensor:
    """Map a residual-space direction to vocab logits (TL or HF lm_head layout)."""
    if hasattr(model, "unembed"):
        w = model.unembed.W_U
        b = getattr(model.unembed, "b_U", None)
    else:
        w = model.W_U
        b = None
 
    d_model = resid.shape[-1]
    if w.shape[0] == d_model:
        logits = resid @ w
    elif w.shape[1] == d_model:
        logits = resid @ w.transpose(-1, -2)
    else:
        raise RuntimeError(
            f"Unexpected lm_head shape {tuple(w.shape)} for d_model={d_model}"
        )
 
    if b is not None and b.numel():
        logits = logits + b.to(dtype=logits.dtype, device=logits.device)
    return logits
 
 
class HfLlmShim:
    """Enough of HookedTransformer for SAE capture, steering, and eval logits."""
 
    def __init__(
        self,
        hf_model: Any,
        tokenizer: Any,
        *,
        hf_name: str,
        n_layers: int,
        d_model: int,
        n_heads: int,
    ) -> None:
        self.hf_model = hf_model
        self.tokenizer = tokenizer
        self.cfg = _ShimCfg(
            n_layers=n_layers,
            d_model=d_model,
            n_heads=n_heads,
            model_name=hf_name,
        )
        self._backbone, self._layers, self._embed, self._lm_head = _resolve_modules(hf_model)
        self.device = next(hf_model.parameters()).device
        self._final_norm = (
            getattr(self._backbone, "norm", None)
            or getattr(self._backbone, "ln_f", None)
            or getattr(self._backbone, "final_layer_norm", None)
        )
        head_bias = getattr(self._lm_head, "bias", None) if self._lm_head is not None else None
        self._unembed = _UnembedShim(self.W_U, head_bias)
 
    @property
    def ln_final(self) -> nn.Module:
        return self._final_norm if self._final_norm is not None else _IdentityNorm()
 
    @property
    def unembed(self) -> _UnembedShim:
        return self._unembed
 
    @classmethod
    def from_pretrained(
        cls,
        hf_name: str,
        cfg: dict[str, Any],
        *,
        dtype: torch.dtype,
        device: str,
        hf_kwargs: dict[str, Any] | None = None,
    ) -> HfLlmShim:
        trust = bool(cfg.get("trust_remote_code", False))
        kw = dict(hf_kwargs or {})
        load_device_map = _hf_load_device_map(cfg, device)
        common = dict(
            torch_dtype=dtype,
            trust_remote_code=trust,
            low_cpu_mem_usage=True,
            **kw,
        )
        if device in ("mps", "cpu"):
            hf_model = AutoModelForCausalLM.from_pretrained(hf_name, **common)
            hf_model = hf_model.to(device)
        else:
            hf_model = AutoModelForCausalLM.from_pretrained(
                hf_name,
                device_map=load_device_map,
                **common,
            )
        hf_model.eval()
        tokenizer = AutoTokenizer.from_pretrained(hf_name, trust_remote_code=trust, **kw)
        n_heads = int(cfg.get("n_heads", cfg.get("num_attention_heads", 32)))
        return cls(
            hf_model,
            tokenizer,
            hf_name=hf_name,
            n_layers=int(cfg["n_layers"]),
            d_model=int(cfg["d_model"]),
            n_heads=n_heads,
        )
 
    @property
    def W_E(self) -> torch.Tensor:
        return self._embed.weight
 
    @property
    def W_U(self) -> torch.Tensor:
        if self._lm_head is None:
            raise RuntimeError("Model has no lm_head")
        return self._lm_head.weight
 
    @property
    def blocks(self) -> list[Any]:
        return list(self._layers)
 
    def eval(self) -> HfLlmShim:
        self.hf_model.eval()
        return self
 
    def parameters(self, recurse: bool = True):
        return self.hf_model.parameters(recurse=recurse)
 
    def named_parameters(self, prefix: str = "", recurse: bool = True):
        return self.hf_model.named_parameters(prefix=prefix, recurse=recurse)
 
    def to(self, device: str | torch.device) -> HfLlmShim:
        if _model_has_device_map(self.hf_model):
            return self
        self.hf_model.to(device)
        self.device = next(self.hf_model.parameters()).device
        return self
 
    def to_tokens(self, text: str | torch.Tensor, prepend_bos: bool = False) -> torch.Tensor:
        if not isinstance(text, str):
            return text.to(self.device)
        encoded = self.tokenizer(
            text,
            return_tensors="pt",
            add_special_tokens=prepend_bos,
        )
        return encoded.input_ids.to(self.device)
 
    def to_string(self, tokens: Any) -> str:
        if isinstance(tokens, torch.Tensor):
            tokens = tokens.tolist()
        if isinstance(tokens, list) and tokens and isinstance(tokens[0], list):
            tokens = tokens[0]
        return self.tokenizer.decode(tokens, skip_special_tokens=True)
 
    def __call__(self, tokens: torch.Tensor) -> torch.Tensor:
        return self._forward_logits(tokens)
 
    def _forward_logits(self, tokens: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            out = self.hf_model(input_ids=tokens)
        return out.logits
 
    def _hooks_from_filter(self, names_filter: Any) -> tuple[set[int], set[int], bool]:
        if names_filter is None:
            return set(range(self.cfg.n_layers)), set(), False
 
        resid_post: set[int] = set()
        resid_pre: set[int] = set()
        need_embed = _names_filter_match(names_filter, "hook_embed")
        for layer in range(self.cfg.n_layers):
            if _names_filter_match(names_filter, f"blocks.{layer}.hook_resid_post"):
                resid_post.add(layer)
            if _names_filter_match(names_filter, f"blocks.{layer}.hook_resid_pre"):
                resid_pre.add(layer)
        return resid_post, resid_pre, need_embed
 
    def run_with_cache(
        self,
        tokens: torch.Tensor,
        names_filter: Any = None,
        return_type: str | None = "logits",
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        resid_post, resid_pre, need_embed = self._hooks_from_filter(names_filter)
        cache: dict[str, torch.Tensor] = {}
 
        if not resid_pre and not need_embed:
            with torch.no_grad():
                out = self.hf_model(input_ids=tokens, output_hidden_states=True)
            hs = out.hidden_states or ()
            for layer in resid_post:
                if layer + 1 < len(hs):
                    cache[f"blocks.{layer}.hook_resid_post"] = hs[layer + 1]
            logits = out.logits
            return logits, cache
 
        handles: list[Any] = []
 
        if need_embed:
            def _embed_hook(_module: Any, _inp: Any, output: torch.Tensor) -> torch.Tensor:
                cache["hook_embed"] = output
                return output
 
            handles.append(self._embed.register_forward_hook(_embed_hook))
 
        for layer in resid_pre:
            def _pre_hook(_module: Any, args: tuple[Any, ...], li: int = layer) -> None:
                if args:
                    cache[f"blocks.{li}.hook_resid_pre"] = args[0]
 
            handles.append(self._layers[layer].register_forward_pre_hook(_pre_hook))
 
        for layer in resid_post:
            def _post_hook(_module: Any, _inp: Any, output: Any, li: int = layer) -> Any:
                hidden = output[0] if isinstance(output, tuple) else output
                cache[f"blocks.{li}.hook_resid_post"] = hidden
                return output
 
            handles.append(self._layers[layer].register_forward_hook(_post_hook))
 
        try:
            with torch.no_grad():
                logits = self._forward_logits(tokens)
        finally:
            for handle in handles:
                handle.remove()
 
        return logits, cache
 
    def run_with_hooks(
        self,
        tokens: torch.Tensor,
        fwd_hooks: list[tuple[str, Callable[..., torch.Tensor]]] | None = None,
        return_type: str | None = "logits",
    ) -> torch.Tensor:
        handles: list[Any] = []
 
        for hook_name, hook_fn in fwd_hooks or []:
            if hook_name == "hook_embed":
 
                def _embed_hook(_module: Any, _inp: Any, output: torch.Tensor, fn=hook_fn) -> torch.Tensor:
                    return fn(output, None)
 
                handles.append(self._embed.register_forward_hook(_embed_hook))
                continue
 
            m_post = _RESID_POST.match(hook_name)
            if m_post:
                layer = int(m_post.group(1))
 
                def _post_hook(_module: Any, _inp: Any, output: Any, li: int = layer, fn=hook_fn) -> Any:
                    hidden = output[0] if isinstance(output, tuple) else output
                    modified = fn(hidden, None)
                    if isinstance(output, tuple):
                        return (modified,) + output[1:]
                    return modified
 
                handles.append(self._layers[layer].register_forward_hook(_post_hook))
                continue
 
            m_pre = _RESID_PRE.match(hook_name)
            if m_pre:
                layer = int(m_pre.group(1))
 
                def _pre_hook(_module: Any, args: tuple[Any, ...], fn=hook_fn) -> tuple[Any, ...] | None:
                    if not args:
                        return None
                    modified = fn(args[0], None)
                    return (modified,) + args[1:]
 
                handles.append(self._layers[layer].register_forward_pre_hook(_pre_hook))
 
        try:
            with torch.no_grad():
                logits = self._forward_logits(tokens)
        finally:
            for handle in handles:
                handle.remove()
 
        return logits