sae

Einführung

Diese Bibliothek zeichnet K -Sparse AutoCodher (SAEs) zu den Reststreamaktivierungen von Harmgingface -Sprachmodellen, die ungefähr dem Rezept folgen, das bei der Skalierung und Bewertung von spärlichen Autoencodern detailliert detailliert ist (Gao et al. 2024).

Dies ist eine schlanke, einfache Bibliothek mit wenigen Konfigurationsoptionen. Im Gegensatz zu den meisten anderen SAE-Bibliotheken (z. B. SAELENS) wird die Aktivierungen auf der Festplatte nicht geändert, sondern sie im Fliege berechnet. Dies ermöglicht es uns, auf sehr große Modelle und Datensätze mit Nullspeicher-Overhead zu skalieren, hat jedoch den Nachteil, dass es langsamer ist, verschiedene Hyperparameter für dasselbe Modell und Datensatz auszuprobieren, als wenn wir Aktivierungen zwischengespeichert werden (da Aktivierungen erneut erfunden werden). In Zukunft können wir das Caching als Option hinzufügen.

Nach Gao et al. Verwenden wir eine TOPK -Aktivierungsfunktion, die in den Aktivierungen direkt einen gewünschten Sparsity -Grad erzwingt. Dies steht im Gegensatz zu anderen Bibliotheken, die eine L1 -Strafe in der Verlustfunktion verwenden. Wir glauben, dass Topk eine Pareto -Verbesserung gegenüber dem L1 -Ansatz ist und daher nicht vorhat, ihn zu unterstützen.

Laden vorbereitet SAES

Um eine vorgezogene SAE aus dem Hubface -Hub zu laden, können Sie die Methode Sae.load_from_hub wie folgt verwenden:

 from sae import Sae

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

Dadurch wird die SAE für die Reststromschicht 10 von Lama 3 8b geladen, die mit einem Expansionsfaktor von 32 trainiert wurde. Sie können die SAES auch für alle Schichten gleichzeitig mit Sae.load_many laden:

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

Das von load_many zurückgegebene Wörterbuch wird garantiert mit dem Namen des Hakenpunkts sortiert. Für den gemeinsamen Fall, in dem die Hakenpunkte mit dem Namen embed_tokens , layers.0 , ..., layers.n bezeichnet werden, bedeutet dies, dass die SAES nach Ebenennummer sortiert werden. Anschließend können wir die SAE -Aktivierungen für einen Vorwärtspass wie folgt sammeln:

 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 

Trainingssaes

Um SAES aus der Befehlszeile zu trainieren, können Sie den folgenden Befehl verwenden:

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

Die CLI unterstützt alle Konfigurationsoptionen, die von der TrainConfig -Klasse bereitgestellt werden. Sie können sie sehen, indem Sie python -m sae --help ausführen.

Der programmatische Gebrauch ist einfach. Hier ist ein Beispiel:

 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 ()

Benutzerdefinierte Hakenpunkte

Standardmäßig werden die SAEs auf den Reststromaktivierungen des Modells geschult. Sie können jedoch auch SAEs an den Aktivierungen eines anderen Submoduls trainieren, indem Sie benutzerdefinierte Hakenpunktmuster angeben. Diese Muster sind wie Standardnamen des Pytorch -Moduls (z. B. h.0.ln_1 ), ermöglichen jedoch auch die UNIX -Muster -Matching -Syntax, einschließlich Wildcards und Charaktersätze. Zum Beispiel können Sie den folgenden Code verwenden:

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

Auf die ersten drei Schichten einschränken:

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

Derzeit unterstützen wir die feinkörnige manuelle Kontrolle über die Lernrate, die Anzahl der Latenten oder andere Hyperparameter auf der Basis von Hakenpunkten. Standardmäßig wird die Option expansion_ratio verwendet, um die entsprechende Anzahl von Lattern für jeden Hakenpunkt basierend auf der Breite des Ausgangs dieses Hookpoint auszuwählen. Die Standard -Lernrate für jeden Hookpoint wird dann unter Verwendung eines inversen Quadratwurzel -Skalierungsgesetzes festgelegt, der auf der Anzahl der Latients basiert. Wenn Sie die Anzahl der Leitungen oder die Lernrate manuell festlegen, wird sie auf alle Hakenpunkte angewendet.

Verteiltes Training

Wir unterstützen das verteilte Training über Pytorchs torchrun -Befehl. Standardmäßig verwenden wir die verteilte Daten parallele Methode, was bedeutet, dass die Gewichte jeder SAE auf jeder GPU repliziert werden.

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

Dies ist einfach, aber sehr ineffizient. Wenn Sie SAES für viele Schichten eines Modells trainieren möchten, empfehlen wir, das Flag --distribute_modules zu verwenden, das die SAES für verschiedene Schichten an verschiedene GPUs zuteilt. Derzeit verlangen wir, dass die Anzahl der GPUs gleichmäßig die Anzahl der Ebenen, für die Sie SAES trainieren, gleichmäßig teilt.

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

Der obige Befehl trainiert eine SAE für jede gleichmäßige Schicht Lama 3 8b unter Verwendung aller verfügbaren GPUs. Es akkumuliert Gradienten über 8 Minibatches und spaltet jede Minibatch in 2 Mikrobatches, bevor sie in den SAE -Encoder füttert, wodurch viel Speicher gespeichert wird. Es lädt das Modell auch in 8-Bit-Präzision unter Verwendung von bitsandbytes . Dieser Befehl erfordert nicht mehr als 48 GB Speicher pro GPU auf einem 8 -GPU -Knoten.

Todo

Es gibt mehrere Funktionen, die wir in naher Zukunft hinzufügen möchten:

• Fülle vorbereitet SAES
• Unterstützung für das zwischengespeicherte Aktivierungen
• Bewerten Sie SAES mit KL -Divergenz, wenn Sie in das Modell gepfropft sind

Wenn Sie mit einem dieser Helfer helfen möchten, können Sie bitte eine PR eröffnen! Sie können mit uns im spärlichen Kanal der Eleutherai-Zwietracht zusammenarbeiten.