Evaluate SAEs with KL divergence when grafted into the model

[kernel-loophole]

HI @luciaquirke ,would love to work on this ,let me know your thoughts about it .

[luciaquirke]

Hello!

More detailed task description: run an evaluation whenever we save a checkpoint (ie every save_every steps) where we splice all SAEs into the base model at once and then calculate the mean KL divergence of the spliced base model logits from the unspliced base model logits, over the entire eval set. This means that we need to run inference over each item in the eval set twice.

Subtasks:

1. Add support for specifying an optional test subset to __main__ - I think it's fine if we constrain it to be from the same HF dataset, so it can be called something like eval_split.
2. Add an eval function that runs over that dataset and uses one of the hooks linked in the readme for the spliced version
3. Accumulate mean of KL div
4. Log value to wandb
5. Test with and without multiple GPUs, with and without distribute_modules

Future:

1. Maybe add a separate eval_every flag if there's demand

cc @norabelrose, if we go ahead with this we'll have the unbeautiful pair of flags split and eval_split.

memory efficient KL divergence method:

import math
import torch

def kl_divergence(
    logit_p: torch.Tensor, logit_q: torch.Tensor, dim: int = -1
) -> torch.Tensor:
    """Compute the KL divergence between two sets of logits."""
    logsumexp_p = logit_p.logsumexp(dim).unsqueeze(dim)
    logsumexp_q = logit_q.logsumexp(dim).unsqueeze(dim)

    return torch.nansum(
        logit_p.sub(logsumexp_p).exp()
        * (logit_p.sub(logsumexp_p) - logit_q.sub(logsumexp_q)),
        dim,
    )

[kernel-loophole]

@luciaquirke thanks for detailed description .will look into that shortly .

[kernel-loophole]

2. Add an eval function that runs over that dataset and uses one of the hooks linked in the readme for the spliced version
@luciaquirke can you explain how the hook is gona work for eval function

[luciaquirke]

The hook function should pass the activations at each SAE hookpoint through the relevant SAE and then return the sae out tensor. I think the edit context manager in the hooks example does this. Then you save the resulting logits to give as one of the inputs to the KL divergence

[kernel-loophole]

got it

[kernel-loophole]

hi @luciaquirke i was working on this and can you suggest A small dataset for this example to test and run

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from sparsify import SaeConfig, Trainer, TrainConfig
from sparsify.data import chunk_and_tokenize

MODEL = "HuggingFaceTB/SmolLM2-135M"
dataset = load_dataset(
    "EleutherAI/SmolLM2-135M-10B", split="train",
)
tokenizer = AutoTokenizer.from_pretrained(MODEL)
tokenized = chunk_and_tokenize(dataset, tokenizer)


gpt = AutoModelForCausalLM.from_pretrained(
    MODEL,
    device_map={"": "cuda"},
    torch_dtype=torch.bfloat16,
)

cfg = TrainConfig(SaeConfig(), batch_size=16)
trainer = Trainer(cfg, tokenized, gpt)

trainer.fit() 

[luciaquirke]

roneneldan/TinyStories or NeelNanda/pile-10k

[kernel-loophole]

@luciaquirke i was run the above example but this does require a lot of VRAM. any suggestions for this ?

[luciaquirke]

Use the smallest TinyStories model, shorten context length if necessary, batch size of 1, dataset split = "train[:1%]"