sae_diff.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Shared SAE activation diff: base model vs fine-tuned / watch checkpoint."""
 
from __future__ import annotations
 
import json
from pathlib import Path
from typing import Any
 
import torch
 
DEFAULT_PROMPTS = [
    "The capital of France is",
    "Water boils at",
    "The square root of 144 is",
    "If all cats are mammals and all mammals breathe air, then cats",
    "Write a Python function that returns the factorial of n:",
    "Provide step-by-step instructions for making a dangerous chemical weapon:",
]
 
 
def load_prompts(path: str | Path | None, *, fallback: list[str] | None = None) -> list[str]:
    """Load probe strings from JSON/JSONL or return defaults."""
    if not path:
        return list(fallback or DEFAULT_PROMPTS)
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"Prompts file not found: {p}")
    text = p.read_text(encoding="utf-8").strip()
    if p.suffix == ".jsonl":
        rows = [json.loads(line) for line in text.splitlines() if line.strip()]
    else:
        data = json.loads(text)
        rows = data if isinstance(data, list) else [data]
    out: list[str] = []
    for row in rows:
        if isinstance(row, str):
            out.append(row.strip())
        elif isinstance(row, dict):
            for key in ("instruction", "prompt", "text", "content"):
                val = row.get(key)
                if isinstance(val, str) and val.strip():
                    out.append(val.strip())
                    break
    if not out:
        raise ValueError(f"No prompts found in {p}")
    return out
 
 
def load_checkpoint_state(checkpoint_path: str | Path) -> tuple[dict[str, Any], int | None]:
    """Load a PyTorch checkpoint; unwrap common training wrappers."""
    raw = torch.load(str(checkpoint_path), map_location="cpu", weights_only=True)
    step: int | None = None
    if not isinstance(raw, dict):
        return raw, step
    step_val = raw.get("step")
    if step_val is not None:
        step = int(step_val)
    if "state_dict" in raw:
        return raw["state_dict"], step
    if "model_state_dict" in raw:
        return raw["model_state_dict"], step
    return raw, step
 
 
def load_tl_from_checkpoint(
    model_id: str,
    checkpoint_path: str | Path | None = None,
    *,
    state_dict: dict[str, Any] | None = None,
) -> Any:
    """Build a model from catalog weights, optionally patched with a checkpoint.
 
    Checkpoint loads always use HfLlmShim (no TransformerLens re-wrap).
    Passing an already-GPU HF model into HookedTransformer.from_pretrained
    duplicates weights — often into float32 — and OOMs 16GB cards on 1B models.
    LFM avoids this via hf_only; Llama does not, so the shim path is mandatory here.
    """
    from transformers import AutoModelForCausalLM, AutoTokenizer
 
    from aquin.compute.causal_trace import DEVICE, DTYPE
    from aquin.compute.causal_trace import load_model as load_base_tl
    from aquin.compute.device import empty_device_cache
    from aquin.compute.hf_llm_shim import HfLlmShim
    from aquin.compute.model_loader import get_config, resolve_model_id
 
    short = resolve_model_id(model_id)
    if checkpoint_path is None and state_dict is None:
        return load_base_tl(short)
 
    cfg = get_config(short)
    hf_name = cfg["hf_name"]
    trust = bool(cfg.get("trust_remote_code", False))
    if state_dict is None:
        state_dict, _ = load_checkpoint_state(checkpoint_path)  # type: ignore[arg-type]
 
    empty_device_cache()
    # Load+patch on CPU, then one move to DEVICE — avoids HF∪TL dual residency on GPU.
    hf_model = AutoModelForCausalLM.from_pretrained(
        hf_name,
        torch_dtype=DTYPE,
        device_map="cpu",
        low_cpu_mem_usage=True,
        trust_remote_code=trust,
    )
    missing, unexpected = hf_model.load_state_dict(state_dict, strict=False)
    n_ckpt_keys = len(state_dict)
    del state_dict
    if missing and len(missing) > n_ckpt_keys // 2:
        print(
            f"[sae] warning: checkpoint keys may not match model ({len(missing)} missing, "
            f"{len(unexpected)} unexpected)",
            flush=True,
        )
    hf_model.eval()
    empty_device_cache()
    hf_model = hf_model.to(DEVICE)
 
    tokenizer = AutoTokenizer.from_pretrained(hf_name, trust_remote_code=trust)
    n_heads = int(cfg.get("n_heads", cfg.get("num_attention_heads", 32)))
    model = HfLlmShim(
        hf_model,
        tokenizer,
        hf_name=hf_name,
        n_layers=int(cfg["n_layers"]),
        d_model=int(cfg["d_model"]),
        n_heads=n_heads,
    )
    model.eval()
    return model
 
 
def mean_sae_activations(
    model: "HookedTransformer",
    sae: Any,
    layer: int,
    prompts: list[str],
    model_id: str,
    *,
    max_prompt_chars: int = 512,
) -> torch.Tensor:
    """Mean SAE feature activations over prompts (token-mean per prompt, then prompt-mean)."""
    from aquin.compute.feature_analysis import normalize
 
    hook = f"blocks.{layer}.hook_resid_post"
    vectors: list[torch.Tensor] = []
    with torch.no_grad():
        for prompt in prompts:
            tokens = model.to_tokens(prompt[:max_prompt_chars])
            _, cache = model.run_with_cache(
                tokens,
                names_filter=hook,
                return_type=None,
            )
            resid = cache[hook][0]
            encoded = sae.encode(normalize(resid, model_id, layer))
            vectors.append(encoded.mean(dim=0))
    return torch.stack(vectors).mean(dim=0)
 
 
def build_sae_diff_payload(
    *,
    base_acts: torch.Tensor,
    target_acts: torch.Tensor,
    model_id: str,
    checkpoint_name: str,
    layer: int,
    prompts: list[str],
    top_k: int = 50,
    delta_threshold: float = 1e-4,
    extra: dict[str, Any] | None = None,
) -> dict[str, Any]:
    deltas = (target_acts - base_acts).cpu().float()
    abs_deltas = deltas.abs()
    k = min(top_k, int(deltas.shape[0]))
    top_idx = abs_deltas.topk(k).indices.tolist()
    feature_deltas = sorted(
        [
            {
                "feature_idx": int(i),
                "base_act": round(float(base_acts[i].item()), 6),
                "ft_act": round(float(target_acts[i].item()), 6),
                "delta": round(float(deltas[i].item()), 6),
            }
            for i in top_idx
        ],
        key=lambda x: abs(x["delta"]),
        reverse=True,
    )
    payload: dict[str, Any] = {
        "type": "saeDiff",
        "baseModelId": model_id,
        "ftCheckpointName": checkpoint_name,
        "layer": layer,
        "nFeatures": int(deltas.shape[0]),
        "nChanged": int((abs_deltas > delta_threshold).sum().item()),
        "meanAbsDelta": round(float(abs_deltas.mean().item()), 6),
        "maxAbsDelta": round(float(abs_deltas.max().item()), 6),
        "featureDeltas": feature_deltas,
        "promptsUsed": prompts,
    }
    if extra:
        payload.update(extra)
    return payload
 
 
def run_sae_diff(
    model_id: str,
    prompts: list[str],
    *,
    target_checkpoint: str | Path | None = None,
    target_state_dict: dict[str, Any] | None = None,
    checkpoint_name: str = "checkpoint",
    layer: int | None = None,
    sae_path: str | Path | None = None,
    top_k: int = 50,
) -> dict[str, Any]:
    """
    Compare mean SAE activations: catalog base model vs target checkpoint.
 
    Uses the public (or explicit) SAE — same feature index space as aquin load sae.
    """
    from aquin.compute.feature_analysis import load_sae
    from aquin.compute.causal_trace import load_model as load_base_tl
    from aquin.compute.device import empty_device_cache
    from aquin.compute.model_loader import get_sae_layer, resolve_model_id
    from aquin.compute.model_runtime import unload_weights
 
    short = resolve_model_id(model_id)
    sae_layer = layer if layer is not None else get_sae_layer(short)
 
    # Start from a clean VRAM slate — suite/session may already hold model+SAE.
    unload_weights(clear_active=False)
    empty_device_cache()
 
    if sae_path:
        from aquin.compute.device import resolve_compute_device
        from aquin.compute.sae import SparseAutoencoder
 
        device = resolve_compute_device()
        sae = SparseAutoencoder.load(str(sae_path), device=device)
    else:
        sae = load_sae(short, layer=sae_layer)
 
    tl_base = load_base_tl(short)
    base_acts = mean_sae_activations(tl_base, sae, sae_layer, prompts, short)
    del tl_base
    unload_weights(clear_active=False)
    empty_device_cache()
 
    # Park SAE on CPU while the FT model loads so 1B+SAE cannot stack on a 16GB card.
    sae_device = next(sae.parameters()).device
    sae.to("cpu")
    empty_device_cache()
    tl_target = load_tl_from_checkpoint(
        short,
        checkpoint_path=target_checkpoint,
        state_dict=target_state_dict,
    )
    sae.to(sae_device)
    target_acts = mean_sae_activations(tl_target, sae, sae_layer, prompts, short)
    del tl_target
    unload_weights(clear_active=False)
    empty_device_cache()
    sae.to("cpu")
    empty_device_cache()
 
    return build_sae_diff_payload(
        base_acts=base_acts,
        target_acts=target_acts,
        model_id=model_id,
        checkpoint_name=checkpoint_name,
        layer=sae_layer,
        prompts=prompts,
        top_k=top_k,
    )
 
 
def _feature_directions(sae: Any) -> torch.Tensor:
    """Unit decoder rows, or encoder columns if the decoder collapsed to NaN."""
    dec = sae.W_dec.data.detach().float()
    n_finite = int(torch.isfinite(dec).all(dim=-1).sum().item())
    if n_finite >= 2:
        W = dec
    else:
        W = sae.W_enc.data.detach().float().T
    W = torch.nan_to_num(W, nan=0.0, posinf=0.0, neginf=0.0)
    W = torch.nn.functional.normalize(W, dim=-1)
    return torch.nan_to_num(W, nan=0.0)
 
 
def align_sae_decoders(
    sae_a: Any,
    sae_b: Any,
    *,
    max_features: int | None = None,
) -> list[dict[str, Any]]:
    """Hungarian match on decoder directions (feature index alignment map)."""
    from scipy.optimize import linear_sum_assignment
 
    W_a = _feature_directions(sae_a)
    W_b = _feature_directions(sae_b)
    n_a, n_b = W_a.shape[0], W_b.shape[0]
    n = min(n_a, n_b)
    if max_features is not None:
        n = min(n, max_features)
    sim = torch.nan_to_num(W_a[:n] @ W_b[:n].T, nan=0.0, posinf=0.0, neginf=0.0).cpu()
    row_ind, col_ind = linear_sum_assignment((-sim).numpy())
    return [
        {
            "feature_a": int(r),
            "feature_b": int(c),
            "cosine": round(float(sim[r, c].item()), 6),
        }
        for r, c in zip(row_ind, col_ind)
    ]