train_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.
 
import torch
import torch.nn.functional as F
from torch.optim import Adam
from pathlib import Path
from datasets import load_dataset
from transformer_lens import HookedTransformer
from aquin.compute.device import (
    accelerator_available,
    empty_device_cache,
    resolve_compute_device,
)
from aquin.compute.sae import SparseAutoencoder
import random
 
MODEL_NAME      = "meta-llama/Llama-3.2-1B-Instruct"
N_LAYERS        = 16
N_FEATURES      = 32768
D_MODEL         = 2048
L1_COEFF        = 10.0
LR              = 1e-4
BATCH_SIZE      = 4096
N_TOKENS        = 2_000_000
SEQ_LEN         = 64
COLLECT_BATCH   = 32        # sequences per forward pass during collection
CHUNK_SIZE      = 100_000   # activations per chunk file
SAVE_EVERY      = 2000
SAE_DIR       = Path(__file__).parent / "sae"
ACTS_DIR      = Path(__file__).parent / "sae" / "acts"
DEVICE = resolve_compute_device()
 
 
def chunk_paths(layer: int):
    return sorted(ACTS_DIR.glob(f"chunk_layer{layer}_*.pt"))
 
 
def norm_cache_path(layer: int) -> Path:
    return ACTS_DIR / f"norm_layer{layer}.pt"
 
 
def sae_save_path(layer: int) -> Path:
    return SAE_DIR / f"sae_layer{layer}.pt"
 
 
def collect_activations(model, layer: int):
    existing = chunk_paths(layer)
    existing_tokens = len(existing) * CHUNK_SIZE
    if existing_tokens >= N_TOKENS:
        print(f"[layer {layer}] found {len(existing)} chunks ({existing_tokens:,} tokens), skipping collection", flush=True)
        return
 
    print(f"[layer {layer}] collecting {N_TOKENS:,} activations...", flush=True)
    ACTS_DIR.mkdir(parents=True, exist_ok=True)
 
    dataset = load_dataset("Skylion007/openwebtext", split="train", streaming=True)
    tokenizer = model.tokenizer
    tokenizer.padding_side = "right"
 
    hook_name = f"blocks.{layer}.hook_resid_post"
    chunk, total, chunk_idx, buf = [], existing_tokens, len(existing), []
 
    model.eval()
    with torch.no_grad():
        for sample in dataset:
            if total >= N_TOKENS:
                break
            ids = tokenizer.encode(sample["text"][:1000], 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,
                )
                # shape: (COLLECT_BATCH, SEQ_LEN, D_MODEL) -> (COLLECT_BATCH * SEQ_LEN, D_MODEL)
                acts = cache[hook_name].reshape(-1, D_MODEL).cpu()
                chunk.append(acts)
                total += acts.shape[0]
 
                if sum(a.shape[0] for a in chunk) >= CHUNK_SIZE:
                    p = ACTS_DIR / f"chunk_layer{layer}_{chunk_idx}.pt"
                    torch.save(torch.cat(chunk, dim=0)[:CHUNK_SIZE], p)
                    chunk, chunk_idx = [], chunk_idx + 1
                    print(f"[layer {layer}] {total:,} / {N_TOKENS:,} tokens", flush=True)
 
    if chunk:
        p = ACTS_DIR / f"chunk_layer{layer}_{chunk_idx}.pt"
        torch.save(torch.cat(chunk, dim=0), p)
 
    print(f"[layer {layer}] collection complete", flush=True)
 
 
def compute_norm(layer: int):
    np = norm_cache_path(layer)
    if np.exists():
        d = torch.load(np)
        return d["mean"], d["std"]
 
    print(f"[layer {layer}] computing mean/std over chunks...", flush=True)
    running_mean = torch.zeros(D_MODEL)
    n_total = 0
    for p in chunk_paths(layer):
        chunk = torch.load(p, map_location="cpu")
        running_mean += chunk.sum(0)
        n_total += chunk.shape[0]
    mean = running_mean / n_total
 
    running_var = torch.zeros(D_MODEL)
    for p in chunk_paths(layer):
        chunk = torch.load(p, map_location="cpu")
        running_var += ((chunk - mean) ** 2).sum(0)
    std = (running_var / n_total).sqrt().clamp(min=1e-6)
 
    torch.save({"mean": mean, "std": std}, np)
    print(f"[layer {layer}] norm computed", flush=True)
    return mean, std
 
 
def train_layer(layer: int):
    save_path = sae_save_path(layer)
    if save_path.exists():
        print(f"[layer {layer}] checkpoint exists, skipping", flush=True)
        return
 
    mean, std = compute_norm(layer)
    SAE_DIR.mkdir(parents=True, exist_ok=True)
 
    sae = SparseAutoencoder(d_model=D_MODEL, n_features=N_FEATURES).to(DEVICE)
    opt = Adam(sae.parameters(), lr=LR)
    step = 0
 
    # init b_pre from first chunk
    first_chunk = torch.load(chunk_paths(layer)[0], map_location="cpu")
    with torch.no_grad():
        sae.b_pre.data = ((first_chunk[:10000] - mean) / std).to(DEVICE).mean(0)
    del first_chunk
 
    chunks = chunk_paths(layer)
    print(f"[layer {layer}] training on {len(chunks)} chunks (~{len(chunks) * CHUNK_SIZE:,} tokens)", flush=True)
 
    for epoch in range(1, 11):
        random.shuffle(chunks)
        for chunk_path in chunks:
            acts = torch.load(chunk_path, map_location="cpu")
            acts = (acts - mean) / std
            n = acts.shape[0]
            acts = acts[torch.randperm(n)]
 
            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 % 200 == 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"[layer {layer}] step={step} epoch={epoch} recon={recon:.4f} L0={l0:.1f} dead={dead}/{N_FEATURES}", flush=True)
 
                if step % SAVE_EVERY == 0:
                    sae.save(save_path)
 
            del acts
 
        sae.save(save_path)
        print(f"[layer {layer}] epoch {epoch} done", flush=True)
 
    print(f"[layer {layer}] training complete", flush=True)
 
 
if __name__ == "__main__":
    if not accelerator_available():
        raise RuntimeError("No GPU accelerator detected (CUDA/ROCm or Metal MPS).")
 
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--layers", type=str, default="all",
                        help="Comma-separated layer indices to train, or 'all' for all layers")
    args = parser.parse_args()
 
    layers = list(range(N_LAYERS)) if args.layers == "all" else [int(x) for x in args.layers.split(",")]
 
    print(f"[main] loading {MODEL_NAME}...", flush=True)
    model = HookedTransformer.from_pretrained(MODEL_NAME)
    model.eval()
    model.to(DEVICE)
 
    for layer in layers:
        collect_activations(model, layer)
 
    del model
    empty_device_cache()
    print("[main] collection done, starting training...", flush=True)
 
    for layer in layers:
        train_layer(layer)
 
    print("[main] all layers done.", flush=True)