sae_train.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Train a temp SAE on activations from a catalog model or fine-tuned checkpoint."""
 
from __future__ import annotations
 
import json
import random
from pathlib import Path
from typing import Any, Iterator
 
import torch
import torch.nn.functional as F
from torch.optim import Adam
 
from aquin.compute.activation_store import (
    compute_norm,
    count_vectors,
    find_chunk_paths,
    resolve_training_acts_dir,
    validate_acts_manifest,
    write_collect_manifest,
)
from aquin.compute.sae import SparseAutoencoder
from aquin.compute.sae_diff import load_tl_from_checkpoint
from aquin.compute.device import empty_device_cache, resolve_compute_device
 
QUICK_TOKENS = 100_000
FULL_TOKENS = 2_000_000
QUICK_MAX_STEPS = 3_000
FULL_MAX_STEPS = 50_000
QUICK_MAX_EPOCHS = 10
FULL_MAX_EPOCHS = 200
SEQ_LEN = 64
COLLECT_BATCH = 32
CHUNK_SIZE = 50_000
BATCH_SIZE = 4096
LR = 1e-4
L1_COEFF = 10.0
USER_SAE_ROOT = Path.home() / ".aquin" / "sae" / "user"
 
 
def _iter_corpus_text(corpus_path: str | Path | None) -> Iterator[str]:
    if corpus_path:
        p = Path(corpus_path)
        if p.suffix == ".jsonl":
            for line in p.read_text(encoding="utf-8").splitlines():
                if not line.strip():
                    continue
                row = json.loads(line)
                if isinstance(row, str):
                    yield row
                elif isinstance(row, dict):
                    for key in ("text", "content", "instruction", "prompt"):
                        val = row.get(key)
                        if isinstance(val, str) and val.strip():
                            yield val[:2000]
                            break
        else:
            data = json.loads(p.read_text(encoding="utf-8"))
            rows = data if isinstance(data, list) else [data]
            for row in rows:
                if isinstance(row, str):
                    yield row
                elif isinstance(row, dict):
                    for key in ("text", "content", "instruction", "prompt"):
                        val = row.get(key)
                        if isinstance(val, str) and val.strip():
                            yield val[:2000]
                            break
        return
 
    from datasets import load_dataset
 
    dataset = load_dataset("Skylion007/openwebtext", split="train", streaming=True)
    for sample in dataset:
        yield sample["text"][:1000]
 
 
def collect_activations(
    model_id: str,
    layer: int,
    acts_dir: Path,
    n_tokens: int,
    *,
    checkpoint_path: str | Path | None = None,
    corpus_path: str | Path | None = None,
) -> None:
    device = resolve_compute_device()
    acts_dir.mkdir(parents=True, exist_ok=True)
 
    existing = sorted(acts_dir.glob("chunk_*.pt"))
    existing_tokens = len(existing) * CHUNK_SIZE
    if existing_tokens >= n_tokens:
        print(f"[sae-train] found {len(existing)} chunks ({existing_tokens:,} tokens), skipping collection", flush=True)
        return
 
    print(f"[sae-train] collecting {n_tokens:,} activations from layer {layer}...", flush=True)
    model = load_tl_from_checkpoint(model_id, checkpoint_path)
    model.eval()
 
    hook_name = f"blocks.{layer}.hook_resid_post"
    tokenizer = model.tokenizer
    if hasattr(tokenizer, "padding_side"):
        tokenizer.padding_side = "right"
 
    chunk: list[torch.Tensor] = []
    total = existing_tokens
    chunk_idx = len(existing)
    buf: list[int] = []
 
    with torch.no_grad():
        for sample in _iter_corpus_text(corpus_path):
            if total >= n_tokens:
                break
            ids = tokenizer.encode(sample, add_special_tokens=False)
            buf.extend(ids)
 
            while len(buf) >= SEQ_LEN * COLLECT_BATCH and total < n_tokens:
                seqs = [buf[i * SEQ_LEN:(i + 1) * SEQ_LEN] for i in range(COLLECT_BATCH)]
                buf = buf[SEQ_LEN * COLLECT_BATCH:]
                tokens = torch.tensor(seqs, device=device)
                _, cache = model.run_with_cache(
                    tokens,
                    names_filter=hook_name,
                    return_type=None,
                )
                acts = cache[hook_name].reshape(-1, cache[hook_name].shape[-1]).cpu()
                chunk.append(acts)
                total += acts.shape[0]
 
                if sum(a.shape[0] for a in chunk) >= CHUNK_SIZE:
                    stacked = torch.cat(chunk, dim=0)[:CHUNK_SIZE]
                    torch.save(stacked, acts_dir / f"chunk_{chunk_idx}.pt")
                    chunk, chunk_idx = [], chunk_idx + 1
                    print(f"[sae-train] {total:,} / {n_tokens:,} tokens", flush=True)
 
    if chunk:
        torch.save(torch.cat(chunk, dim=0), acts_dir / f"chunk_{chunk_idx}.pt")
 
    del model
    empty_device_cache()
    print("[sae-train] collection complete", flush=True)
 
 
def train_sae(
    model_id: str,
    layer: int,
    output_path: Path,
    *,
    checkpoint_path: str | Path | None = None,
    corpus_path: str | Path | None = None,
    quick: bool = False,
    d_model: int | None = None,
    n_features: int | None = None,
    max_steps: int | None = None,
    max_epochs: int | None = None,
    activations_dir: str | Path | None = None,
    balance: bool = False,
    balance_group: str | None = None,
) -> Path:
    from aquin.compute.feature_analysis import load_sae
    from aquin.compute.model_loader import get_config, resolve_model_id
 
    short = resolve_model_id(model_id)
    cfg = get_config(short)
    d_model = d_model or cfg["d_model"]
 
    if n_features is None:
        try:
            ref = load_sae(short, layer=layer)
            n_features = ref.n_features
            del ref
        except Exception:
            n_features = 32768
 
    n_tokens = QUICK_TOKENS if quick else FULL_TOKENS
    steps_cap = max_steps or (QUICK_MAX_STEPS if quick else FULL_MAX_STEPS)
    epochs_cap = max_epochs or (QUICK_MAX_EPOCHS if quick else FULL_MAX_EPOCHS)
 
    if activations_dir:
        work_dir = resolve_training_acts_dir(
            activations_dir,
            layer,
            balance=balance,
            balance_group=balance_group,
        )
        validate_acts_manifest(work_dir, model_id=short, layer=layer)
        n_vectors = count_vectors(work_dir)
        print(
            f"[sae-train] training from saved activations ({n_vectors:,} vectors) at {work_dir}",
            flush=True,
        )
        if n_vectors < 10_000:
            print(
                "[sae-train] warning: few activation vectors — quality may be poor; "
                "use corpus collection or a larger capture for production SAEs",
                flush=True,
            )
    else:
        work_dir = output_path.parent / f"_acts_layer{layer}"
        collect_activations(
            short,
            layer,
            work_dir,
            n_tokens,
            checkpoint_path=checkpoint_path,
            corpus_path=corpus_path,
        )
        write_collect_manifest(
            work_dir,
            model_id=short,
            layer=layer,
            model_mode="llm",
            d_model=d_model,
            source="corpus",
            n_vectors=count_vectors(work_dir),
            checkpoint_path=str(checkpoint_path) if checkpoint_path else None,
            corpus_path=str(corpus_path) if corpus_path else None,
        )
 
    mean, std = compute_norm(work_dir, d_model)
    device = resolve_compute_device()
    sae = SparseAutoencoder(d_model=d_model, n_features=n_features).to(device)
    opt = Adam(sae.parameters(), lr=LR)
 
    chunks = find_chunk_paths(work_dir)
    if not chunks:
        raise FileNotFoundError(f"No activation chunks in {work_dir}")
    first = torch.load(chunks[0], map_location="cpu", weights_only=False)
    with torch.no_grad():
        sae.b_pre.data = ((first[:10_000] - mean) / std).to(device).mean(0)
    del first
 
    step = 0
    print(
        f"[sae-train] training on {len(chunks)} chunks "
        f"(quick={quick}, steps_cap={steps_cap}, epochs_cap={epochs_cap})",
        flush=True,
    )
    for epoch in range(1, epochs_cap + 1):
        random.shuffle(chunks)
        for chunk_path in chunks:
            acts = torch.load(chunk_path, map_location="cpu", weights_only=False)
            acts = (acts - mean) / std
            n = acts.shape[0]
            perm = torch.randperm(n)
            acts = acts[perm]
 
            for start in range(0, n - BATCH_SIZE, BATCH_SIZE):
                batch = acts[start:start + BATCH_SIZE].to(device)
                f, x_hat = sae(batch)
                loss = F.mse_loss(x_hat, batch) + L1_COEFF * f.abs().mean()
                opt.zero_grad()
                loss.backward()
                opt.step()
                sae._normalise_decoder()
                step += 1
 
                if step % 500 == 0:
                    dead = (f.max(0).values == 0).sum().item()
                    l0 = (f > 0).float().sum(-1).mean().item()
                    recon = F.mse_loss(x_hat, batch).item()
                    print(
                        f"[sae-train] epoch={epoch} step={step} recon={recon:.4f} "
                        f"L0={l0:.1f} dead={dead}/{n_features}",
                        flush=True,
                    )
 
                if step >= steps_cap:
                    break
            del acts
            if step >= steps_cap:
                break
        if step >= steps_cap:
            break
 
    print(f"[sae-train] finished {step} steps ({epoch} epochs)", flush=True)
 
    output_path.parent.mkdir(parents=True, exist_ok=True)
    sae.save(output_path)
    norm_path = output_path.parent / f"norm_layer{layer}.pt"
    torch.save({"mean": mean.cpu(), "std": std.cpu()}, norm_path)
    meta = {
        "model_id": short,
        "layer": layer,
        "checkpoint_path": str(checkpoint_path) if checkpoint_path else None,
        "quick": quick,
        "steps": step,
        "max_steps": steps_cap,
        "epochs": epoch,
        "n_features": n_features,
        "d_model": d_model,
    }
    output_path.with_suffix(".meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"[sae-train] saved {output_path}", flush=True)
    return output_path
 
 
def default_user_sae_path(model_id: str, name: str, layer: int) -> Path:
    from aquin.compute.model_loader import resolve_model_id
 
    short = resolve_model_id(model_id)
    safe = name.replace("/", "--").replace(" ", "_")
    return USER_SAE_ROOT / short / safe / f"sae_layer{layer}.pt"