sae

Perkenalan

Perpustakaan ini melatih K -sparse Autoencoders (SAE) pada aktivasi aliran residu dari model bahasa pelukan, secara kasar mengikuti resep yang dirinci dalam penskalaan dan mengevaluasi autoencoders yang jarang (Gao et al. 2024).

Ini adalah pustaka yang ramping dan sederhana dengan beberapa opsi konfigurasi. Tidak seperti kebanyakan perpustakaan SAE lainnya (misalnya Saelens), itu tidak cache aktivasi pada disk, melainkan menghitungnya saat terbang. Ini memungkinkan kami untuk skala ke model dan set data yang sangat besar dengan overhead penyimpanan nol, tetapi memiliki sisi kelemahan bahwa mencoba hyperparameter yang berbeda untuk model dan dataset yang sama akan lebih lambat daripada jika kami mengacak aktivasi (karena aktivasi akan dikomputasi ulang). Kami dapat menambahkan caching sebagai opsi di masa depan.

Mengikuti Gao et al., Kami menggunakan fungsi aktivasi TOPK yang secara langsung menegakkan tingkat sparsity yang diinginkan dalam aktivasi. Ini berbeda dengan perpustakaan lain yang menggunakan penalti L1 dalam fungsi kerugian. Kami percaya Topk adalah peningkatan Pareto atas pendekatan L1, dan karenanya tidak berencana untuk mendukungnya.

Memuat saus pretrained

Untuk memuat SAE pretrained dari hub Huggingface, Anda dapat menggunakan metode Sae.load_from_hub sebagai berikut:

 from sae import Sae

sae = Sae . load_from_hub ( "EleutherAI/sae-llama-3-8b-32x" , hookpoint = "layers.10" )

Ini akan memuat SAE untuk sisa aliran aliran 10 dari Llama 3 8b, yang dilatih dengan faktor ekspansi 32. Anda juga dapat memuat SAE untuk semua lapisan sekaligus menggunakan Sae.load_many :

 saes = Sae . load_many ( "EleutherAI/sae-llama-3-8b-32x" )
saes [ "layers.10" ]

Kamus yang dikembalikan oleh load_many dijamin akan diurutkan secara alami dengan nama titik kait. Untuk kasus umum di mana titik -titik kait dinamai embed_tokens , layers.0 , ..., layers.n , ini berarti bahwa SAES akan diurutkan berdasarkan nomor layer. Kami kemudian dapat mengumpulkan aktivasi SAE untuk Model Forward Pass sebagai berikut:

 from transformers import AutoModelForCausalLM , AutoTokenizer
import torch

tokenizer = AutoTokenizer . from_pretrained ( "meta-llama/Meta-Llama-3-8B" )
inputs = tokenizer ( "Hello, world!" , return_tensors = "pt" )

with torch . inference_mode ():
    model = AutoModelForCausalLM . from_pretrained ( "meta-llama/Meta-Llama-3-8B" )
    outputs = model ( ** inputs , output_hidden_states = True )

    latent_acts = []
    for sae , hidden_state in zip ( saes . values (), outputs . hidden_states ):
        latent_acts . append ( sae . encode ( hidden_state ))

# Do stuff with the latent activations 

Pelatihan SAES

Untuk melatih SAE dari baris perintah, Anda dapat menggunakan perintah berikut:

python -m sae EleutherAI/pythia-160m togethercomputer/RedPajama-Data-1T-Sample

CLI mendukung semua opsi konfigurasi yang disediakan oleh kelas TrainConfig . Anda dapat melihatnya dengan menjalankan python -m sae --help .

Penggunaan terprogram sederhana. Inilah contohnya:

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

from sae import SaeConfig , SaeTrainer , TrainConfig
from sae . data import chunk_and_tokenize

MODEL = "EleutherAI/pythia-160m"
dataset = load_dataset (
    "togethercomputer/RedPajama-Data-1T-Sample" ,
    split = "train" ,
    trust_remote_code = True ,
)
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 ( gpt . config . hidden_size ), batch_size = 16
)
trainer = SaeTrainer ( cfg , tokenized , gpt )

trainer . fit ()

Hookpoint khusus

Secara default, SAE dilatih pada aktivasi aliran residual model. Namun, Anda juga dapat melatih SAES pada aktivasi submodule lainnya dengan menentukan pola hookpoint khusus. Pola -pola ini seperti nama modul pytorch standar (misalnya h.0.ln_1 ) tetapi juga memungkinkan sintaks pencocokan pola UNIX, termasuk wildcard dan set karakter. Misalnya, untuk melatih SAES pada output dari setiap modul perhatian dan aktivasi dalam setiap MLP di GPT-2, Anda dapat menggunakan kode berikut:

python -m sae gpt2 togethercomputer/RedPajama-Data-1T-Sample --hookpoints " h.*.attn " " h.*.mlp.act "

Untuk membatasi tiga lapisan pertama:

python -m sae gpt2 togethercomputer/RedPajama-Data-1T-Sample --hookpoints " h.[012].attn " " h.[012].mlp.act "

Kami saat ini tidak mendukung kontrol manual berbutir halus atas tingkat pembelajaran, jumlah laten, atau hiperparameter lainnya berdasarkan hookpoint-by-hookpoint. Secara default, opsi expansion_ratio digunakan untuk memilih jumlah laten yang sesuai untuk setiap hookpoint berdasarkan lebar output hookpoint itu. Tingkat pembelajaran default untuk setiap hookpoint kemudian diatur menggunakan hukum penskalaan akar kuadrat terbalik berdasarkan jumlah laten. Jika Anda secara manual mengatur jumlah laten atau tingkat pembelajaran, itu akan diterapkan pada semua hookpoints.

Pelatihan Terdistribusi

Kami mendukung pelatihan terdistribusi melalui perintah Pytorch's torchrun . Secara default kami menggunakan metode paralel data terdistribusi, yang berarti bahwa bobot masing -masing SAE direplikasi pada setiap GPU.

torchrun --nproc_per_node gpu -m sae meta-llama/Meta-Llama-3-8B --batch_size 1 --layers 16 24 --k 192 --grad_acc_steps 8 --ctx_len 2048

Ini sederhana, tetapi sangat tidak efisien. Jika Anda ingin melatih SAES untuk banyak lapisan model, kami sarankan menggunakan flag --distribute_modules , yang mengalokasikan SAE untuk lapisan yang berbeda ke GPU yang berbeda. Saat ini, kami mensyaratkan bahwa jumlah GPU secara merata membagi jumlah lapisan yang Anda latih.

torchrun --nproc_per_node gpu -m sae meta-llama/Meta-Llama-3-8B --distribute_modules --batch_size 1 --layer_stride 2 --grad_acc_steps 8 --ctx_len 2048 --k 192 --load_in_8bit --micro_acc_steps 2

Perintah di atas melatih SAE untuk setiap lapisan Llama 3 8b, menggunakan semua GPU yang tersedia. Ini mengumpulkan gradien lebih dari 8 minibatch, dan membagi setiap minibatch menjadi 2 microbatch sebelum memasukkannya ke dalam encoder SAE, sehingga menghemat banyak memori. Ini juga memuat model dalam presisi 8-bit menggunakan bitsandbytes . Perintah ini membutuhkan tidak lebih dari 48GB memori per GPU pada simpul 8 GPU.

Todo

Ada beberapa fitur yang ingin kami tambahkan dalam waktu dekat:

• Finetuning Saes pretrained
• Dukungan untuk aktivasi caching
• Evaluasi SAE dengan divergensi KL saat dicangkokkan ke dalam model

Jika Anda ingin membantu dengan semua ini, jangan ragu untuk membuka PR! Anda dapat berkolaborasi dengan kami di saluran Sparse-Autoencoders dari Eleutherai Discord.