sae

Introdução

Esta biblioteca treina os autoencoders K -Sparse (SAES) nas ativações residuais de fluxo de modelos de linguagem Huggingface, seguindo aproximadamente a receita detalhada na escala e na avaliação de autoencoders esparsos (Gao et al. 2024).

Esta é uma biblioteca enxuta e simples, com poucas opções de configuração. Ao contrário da maioria das outras bibliotecas SAE (por exemplo, Saelens), ele não cache ativações no disco, mas as calcula na fly. Isso nos permite dimensionar para modelos e conjuntos de dados muito grandes com sobrecarga de armazenamento zero, mas tem a desvantagem de que experimentar diferentes hiperparâmetros para o mesmo modelo e conjunto de dados será mais lento do que se cactemos as ativações (pois as ativações serão recomputadas). Podemos adicionar o cache como uma opção no futuro.

Seguindo Gao et al., Utilizamos uma função de ativação TOPK que reforça diretamente um nível desejado de esparsidade nas ativações. Isso contrasta com outras bibliotecas que usam uma penalidade de L1 na função de perda. Acreditamos que Topk é uma melhoria de Pareto em relação à abordagem L1 e, portanto, não planeja apoiá -la.

Carregando SAEs pré -tenhados

Para carregar um SAE pré -traido no hub do Hugging Sace, você pode usar o método Sae.load_from_hub da seguinte forma:

 from sae import Sae

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

Isso carregará o SAE para a camada de fluxo residual 10 do llama 3 8b, que foi treinado com um fator de expansão de 32. Você também pode carregar o SAES para todas as camadas ao mesmo tempo usando Sae.load_many :

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

O dicionário retornado por load_many é garantido como naturalmente classificado pelo nome do ponto do gancho. Para o caso comum em que os pontos do gancho são nomeados embed_tokens , layers.0 , ..., layers.n , isso significa que os SAEs serão classificados pelo número da camada. Podemos então reunir as ativações do SAE para um passe para a frente da seguinte forma:

 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 

Treinando SAES

Para treinar SAES na linha de comando, você pode usar o seguinte comando:

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

A CLI suporta todas as opções de configuração fornecidas pela classe TrainConfig . Você pode vê -los executando python -m sae --help .

O uso programático é simples. Aqui está um exemplo:

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

Pontos de gancho personalizados

Por padrão, os SAEs são treinados nas ativações residuais do fluxo do modelo. No entanto, você também pode treinar o SAES nas ativações de qualquer outro submódulo (s) especificando padrões de gancho de gancho personalizados. Esses padrões são como nomes de módulos Pytorch padrão (por exemplo, h.0.ln_1 ), mas também permitem a sintaxe da correspondência de padrões do Unix, incluindo curingas e conjuntos de caracteres. Por exemplo, para treinar o SAES na saída de todos os módulos de atenção e as ativações internas de cada MLP no GPT-2, você pode usar o seguinte código:

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

Para restringir as três primeiras camadas:

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

Atualmente, não apoiamos o controle manual de grão fino sobre a taxa de aprendizagem, número de latentes ou outros hiperparâmetros com base no ponto de gancho por um ponto de gancho. Por padrão, a opção expansion_ratio é usada para selecionar o número apropriado de latentes para cada ponto de gancho com base na largura da saída desse gancho. A taxa de aprendizado padrão para cada ponto de gancho é então definida usando uma lei de escala de raiz quadrada inversa com base no número de latentes. Se você definir manualmente o número de latentes ou a taxa de aprendizado, ele será aplicado a todos os pontos de gancho.

Treinamento distribuído

Apoiamos o treinamento distribuído pelo comando torchrun da Pytorch. Por padrão, usamos o método paralelo de dados distribuídos, o que significa que os pesos de cada SAE são replicados em todas as 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

Isso é simples, mas muito ineficiente. Se você deseja treinar o SAES para muitas camadas de um modelo, recomendamos o uso do sinalizador --distribute_modules , que aloca o SAES para camadas diferentes para diferentes GPUs. Atualmente, exigimos que o número de GPUs divida uniformemente o número de camadas para as quais você está treinando SAES.

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

O comando acima treina um SAE para cada camada uniforme do llama 3 8b, usando todas as GPUs disponíveis. Ele acumula gradientes com mais de 8 minibatches e divide cada um minibatch em 2 microbatches antes de alimentá -los no codificador SAE, economizando muita memória. Ele também carrega o modelo em precisão de 8 bits usando bitsandbytes . Este comando não requer mais de 48 GB de memória por GPU em um nó de 8 GPU.

PENDÊNCIA

Existem vários recursos que gostaríamos de adicionar no futuro próximo:

• SAEs pré -tendidos fino
• Suporte para ativações de cache
• Avalie os SAEs com divergência de KL quando enxertados no modelo

Se você quiser ajudar com qualquer um deles, sinta -se à vontade para abrir um PR! Você pode colaborar conosco no canal esparso-autoencoder da Eleutherai Discord.