このライブラリは、スパース自動エンコーダーのスケーリングと評価で詳述されているレシピにほぼ追跡して、ハギングフェイス言語モデルの残留ストリームアクティベーションでk -sparse自動エンコーダー(SAE)を訓練します(Gao etal。2024)。
これは、構成オプションがほとんどない無駄のないシンプルなライブラリです。他のほとんどのSAEライブラリ(Saelensなど)とは異なり、ディスク上のアクティベーションをキャッシュするのではなく、フライでそれらを計算します。これにより、ストレージオーバーヘッドがゼロの非常に大きなモデルとデータセットにスケーリングすることができますが、同じモデルに対して異なるハイパーパラメーターを試すと、アクティブ化がキャッシュされた場合よりも遅くなるという欠点があります(アクティベーションが再計算されるため)。将来、キャッシュをオプションとして追加する場合があります。
Gao et al。に続いて、Activationsに望ましいレベルのスパース性を直接施行するTOPKアクティベーション関数を使用します。これは、損失関数にL1ペナルティを使用する他のライブラリとは対照的です。 TopkはL1アプローチよりもパレートの改善であると考えているため、サポートする予定はありません。
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 Sae.load_many
使用して、すべての層のSAEを一度にロードすることもできます。
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を訓練するには、次のコマンドを使用できます。
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 "
最初の3つのレイヤーに制限するには:
python -m sae gpt2 togethercomputer/RedPajama-Data-1T-Sample --hookpoints " h.[012].attn " " h.[012].mlp.act "
現在、フックポイントごとに学習率、潜水種の数、またはその他のハイパーパラメーターに対する細粒の手動制御はサポートしていません。デフォルトでは、 expansion_ratio
オプションを使用して、そのフックポイントの出力の幅に基づいて、各フックポイントに適切な数の潜在性を選択します。次に、各フックポイントのデフォルトの学習レートは、潜伏数に基づいて逆平方根スケーリング則を使用して設定されます。潜在性の数または学習率を手動で設定すると、すべてのフックポイントに適用されます。
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をトレーニングしたい場合は、異なるレイヤーにSAEを異なるGPUに割り当てる--distribute_modules
フラグを使用することをお勧めします。現在、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つのミニバッチを超える勾配を蓄積し、各ミニバッチを2つのマイクロバッチに分割してからSAEエンコーダーに供給し、多くのメモリを保存します。また、 bitsandbytes
を使用してモデルを8ビット精度でロードします。このコマンドには、8 GPUノードでGPUあたり48GB以下のメモリが必要です。
近い将来に追加したい機能がいくつかあります。
これらのいずれかを手伝いたい場合は、PRを自由に開いてください! Eleutherai Discordのスパースオートエンコーダーチャネルで私たちと協力できます。