OpenAI's GPT2-Small SAEs reformatted for easy loading from SAE Lens.
1import torch
2from transformer_lens import HookedTransformer
3from sae_lens import SAE, ActivationsStore
4
5torch.set_grad_enabled(False)
6model = HookedTransformer.from_pretrained("gpt2-small")
7sae, cfg, sparsity = SAE.from_pretrained(
8 "gpt2-small-resid-post-v5-128k", # to see the list of available releases, go to: https://github.com/jbloomAus/SAELens/blob/main/sae_lens/pretrained_saes.yaml
9 "blocks.11.hook_resid_post" # change this to another specific SAE ID in the release if desired.
10)
11
12# For loading activations or tokens from the training dataset.
13activation_store = ActivationsStore.from_sae(
14 model=model,
15 sae=sae,
16 streaming=True,
17 # fairly conservative parameters here so can use same for larger
18 # models without running out of memory.
19 store_batch_size_prompts=8,
20 train_batch_size_tokens=4096,
21 n_batches_in_buffer=4,
22 device=device,
23)
24
