sae下載 - sae源代碼下載

sae

其他源碼

1.0.0

下載

介紹

該庫訓練K -Sparse自動編碼器（SAE）在擁抱表語言模型的殘留流激活上，大致遵循縮放和評估稀疏自動編碼器中詳細介紹的食譜（Gao等，2024）。

這是一個精簡的庫，具有很少的配置選項。與大多數其他SAE庫（例如Saelens）不同，它不會在磁盤上緩存激活，而是在當時計算它們。這使我們可以擴展到非常大的模型和數據集，但存儲空間為零，但是與我們緩存激活相比，嘗試其他模型和數據集的不同超參數和數據集將慢（因為將重新計算激活）。將來，我們可能會添加緩存作為選項。

在Gao等人之後，我們使用TOPK激活函數，該功能直接在激活中實現了所需的稀疏度。這與其他在損失函數中使用L1罰款的庫形成對比。我們認為Topk是對L1方法的帕累托改進，因此不打算支持它。

加載預算的SAE

要從HuggingFace Hub加載預估計的SAE，您可以使用Sae.load_from_hub方法如下：

 from sae import Sae

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

這將為Llama 3 8B的殘留流層10加載SAE，該層的擴展係數為32。您還可以使用Sae.load_many ：

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

load_many返回的字典可以自然地按掛鉤點的名稱進行排序。對於常見的情況，掛鉤點命名為embed_tokens ， layers.0 ，...， layers.n ，這意味著SAE將按層編號進行排序。然後，我們可以收集SAE激活以進行模型前傳遞，如下所示：

 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

訓練SAE

要從命令行訓練SAE，您可以使用以下命令：

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

CLI支持TrainConfig類提供的所有配置選項。您可以通過運行python -m sae --help來查看它們。

程序化使用很簡單。這是一個示例：

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

自定義掛鉤點

默認情況下，SAE是在模型的殘留流激活上訓練的。但是，您也可以通過指定自定義掛鉤點模式來訓練SAE在任何其他子模塊的激活上的激活。這些模式就像標準的pytorch模塊名稱（例如h.0.ln_1 ），但也允許Unix模式匹配語法，包括通配符和字符集。例如，要在每個注意模塊的輸出和GPT-2中每個MLP的內部激活上訓練SAE，您可以使用以下代碼：

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

限製到前三層：

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

目前，我們不支持對學習率，潛伏數量或其他超參數的精細顆粒手動控制。默認情況下， expansion_ratio選項用於根據該Hookpoint輸出的寬度為每個掛鉤點選擇適當數量的潛在數量。然後，使用基於潛伏數量的逆平方根縮放定律設置每個掛鉤點的默認學習率。如果您手動設置潛在的數量或學習率，則將其應用於所有掛鉤點。

分佈式培訓

我們支持Pytorch的torchrun命令分佈式培訓。默認情況下，我們使用分佈式數據並行方法，這意味著每個SAE的權重都在每個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

這很簡單，但記憶力非常低。如果您想為模型的許多層訓練SAE，我們建議使用--distribute_modules標誌，該標誌將SAE的不同層的SAE分配給不同的GPU。目前，我們要求GPU的數量均勻地劃分您訓練SAE的層數。

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

上面的命令使用所有可用的GPU訓練每一層Llama 3 8B的SAE。它累積了8個小匹配的梯度，並將每個MiniBatch分成2個微匹配，然後將其饋入SAE編碼器，從而節省大量內存。它還使用bitsandbytes將模型加載到8位精度中。該命令在8 GPU節點上的每個GPU的內存不超過48GB。