sae 다운로드 sae 소스 코드 다운로드

sae

기타 소스코드

1.0.0

다운로드

소개

이 라이브러리는 k -sparse autoencoders (SAES)를 조정하여 언어 모델의 잔차 스트림 활성화에 대해 훈련하며, 대략 드문 스파 스 자동 인코더 스케일링 및 평가에 자세히 설명 된 레시피를 따릅니다 (Gao et al. 2024).

구성 옵션이 거의없는 간단한 라이브러리입니다. 대부분의 다른 SAE 라이브러리 (예 : Saelens)와 달리 디스크에서 활성화를 캐시하는 것이 아니라 비행기를 계산합니다. 이를 통해 스토리지 오버 헤드가 0 인 매우 큰 모델 및 데이터 세트로 확장 할 수 있지만 동일한 모델 및 데이터 세트에 대해 다른 하이퍼 파라미터를 시도하는 것이 활성화 된 경우보다 느리게 진행되는 단점이 있습니다 (활성화가 재 계산되기 때문에). 우리는 향후 옵션으로 캐싱을 추가 할 수 있습니다.

Gao et al.에 따라, 우리는 활성화에서 원하는 수준의 희소성을 직접적으로 강화하는 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를로드 할 것입니다. LLAMA 3 8B는 32의 확장 계수로 훈련되었습니다. 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

명령 줄에서 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 옵션은 해당 후크 포인트 출력의 너비를 기반으로 각 훅 포인트에 대해 적절한 잠복 수를 선택하는 데 사용됩니다. 그런 다음 각 훅 포인트의 기본 학습 속도는 잠복 수를 기반으로 역 제곱근 스케일링 법을 사용하여 설정됩니다. 잠복 수 또는 학습 속도를 수동으로 설정하면 모든 훅 포인트에 적용됩니다.

분산 교육

우리는 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 플래그를 사용하는 것이 좋습니다. 현재 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 개의 미니 디바치에 걸쳐 그라디언트를 축적하고, SAE 인코더에 공급하기 전에 각 미니 배트를 2 개의 미생물로 나누어 많은 메모리를 저장합니다. 또한 bitsandbytes 사용하여 모델을 8 비트 정밀도로로드합니다. 이 명령은 8GPU 노드에서 GPU 당 48GB 이하의 메모리가 필요하지 않습니다.