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。