sae
1.0.0
ห้องสมุดนี้รถไฟ K -Sparse Autoencoders (SAES) ในการเปิดใช้งานสตรีมที่เหลือของโมเดลภาษา HuggingFace โดยประมาณตามสูตรที่มีรายละเอียดในการปรับขนาดและประเมินค่าอัตโนมัติแบบเบาบาง (Gao et al. 2024)
นี่คือไลบรารีที่เรียบง่ายที่มีตัวเลือกการกำหนดค่าน้อย ซึ่งแตกต่างจากไลบรารี SAE อื่น ๆ ส่วนใหญ่ (เช่น Saelens) มันไม่ได้แคชเปิดใช้งานบนดิสก์ แต่คำนวณพวกเขาทันที สิ่งนี้ช่วยให้เราสามารถปรับขนาดเป็นรุ่นและชุดข้อมูลขนาดใหญ่ที่มีค่าใช้จ่ายเป็นศูนย์ แต่มีข้อเสียที่ลองใช้พารามิเตอร์ไฮเปอร์พารามิเตอร์ที่แตกต่างกันสำหรับรุ่นเดียวกันและชุดข้อมูลจะช้ากว่าถ้าเราแคชเปิดใช้งาน (เนื่องจากการเปิดใช้งานจะถูกคำนวณใหม่) เราอาจเพิ่มแคชเป็นตัวเลือกในอนาคต
ต่อไปนี้ Gao et al. เราใช้ฟังก์ชั่นการเปิดใช้งาน TOPK ซึ่งบังคับใช้ระดับ sparsity ที่ต้องการโดยตรงในการเปิดใช้งาน สิ่งนี้ตรงกันข้ามกับห้องสมุดอื่น ๆ ที่ใช้การลงโทษ L1 ในฟังก์ชั่นการสูญเสีย เราเชื่อว่า Topk เป็นการปรับปรุง Pareto ในแนวทาง L1 และด้วยเหตุนี้จึงไม่ได้วางแผนที่จะสนับสนุน
ในการโหลด SAE ที่ได้รับการฝึกฝนจาก HuggingFace Hub คุณสามารถใช้วิธี Sae.load_from_hub ได้ดังนี้:
from sae import Sae
sae = Sae . load_from_hub ( "EleutherAI/sae-llama-3-8b-32x" , hookpoint = "layers.10" ) สิ่งนี้จะโหลด SAE สำหรับเลเยอร์กระแสที่เหลือ 10 ของ Llama 3 8B ซึ่งได้รับการฝึกฝนด้วยปัจจัยการขยายตัว 32 คุณยังสามารถโหลด SAEs สำหรับเลเยอร์ทั้งหมดในครั้งเดียวโดยใช้ Sae.load_many :
saes = Sae . load_many ( "EleutherAI/sae-llama-3-8b-32x" )
saes [ "layers.10" ] พจนานุกรมที่ส่งคืนโดย load_many รับประกันว่าจะถูกจัดเรียงตามธรรมชาติตามชื่อของจุดตะขอ สำหรับกรณีทั่วไปที่จุดตะขอมีชื่อว่า embed_tokens , layers.0 , ... , layers.n ซึ่งหมายความว่า SAES จะถูกจัดเรียงตามหมายเลขเลเยอร์ จากนั้นเราสามารถรวบรวมการเปิดใช้งาน SAE สำหรับ Model Forward Pass ดังนี้:
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 ในการฝึกอบรม SAES จากบรรทัดคำสั่งคุณสามารถใช้คำสั่งต่อไปนี้:
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 จะได้รับการฝึกอบรมเกี่ยวกับการเปิดใช้งานสตรีมที่เหลือของโมเดล อย่างไรก็ตามคุณยังสามารถฝึกอบรม SAEs ในการเปิดใช้งาน submodule อื่น ๆ โดยการระบุรูปแบบ hookpoint ที่กำหนดเอง รูปแบบเหล่านี้เป็นเหมือนชื่อโมดูล pytorch มาตรฐาน (เช่น h.0.ln_1 ) แต่ยังอนุญาตให้ใช้ไวยากรณ์การจับคู่รูปแบบ UNIX รวมถึงไวด์การ์ดและชุดอักขระ ตัวอย่างเช่นในการฝึกอบรม SAEs ในผลลัพธ์ของโมดูลความสนใจทุกอย่างและการเปิดใช้งานภายในของทุก MLP ใน GPT-2 คุณสามารถใช้รหัสต่อไปนี้:
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 " ขณะนี้เราไม่สนับสนุนการควบคุมด้วยตนเองอย่างละเอียดเกี่ยวกับอัตราการเรียนรู้จำนวนแฝงหรือพารามิเตอร์ไฮเปอร์พารามิเตอร์อื่น ๆ บนพื้นฐานของ hookpoint-by-hookpoint โดยค่าเริ่มต้นตัวเลือก expansion_ratio ใช้เพื่อเลือกจำนวนแฝงที่เหมาะสมสำหรับแต่ละ hookpoint ตามความกว้างของเอาต์พุตของ Hookpoint นั้น อัตราการเรียนรู้เริ่มต้นสำหรับแต่ละ hookpoint จะถูกตั้งค่าโดยใช้กฎหมายสเกลสแควร์สแควร์สแควร์ตามจำนวนแฝง หากคุณตั้งค่าจำนวนแฝงหรืออัตราการเรียนรู้ด้วยตนเองจะถูกนำไปใช้กับจุดเชื่อมต่อทั้งหมด
เราสนับสนุนการฝึกอบรมแบบกระจายผ่านคำสั่ง torchrun ของ Pytorch โดยค่าเริ่มต้นเราใช้วิธีการแบบขนานข้อมูลแบบกระจายซึ่งหมายความว่าน้ำหนักของแต่ละ 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 นี่เป็นเรื่องง่าย แต่หน่วยความจำไม่มีประสิทธิภาพมาก หากคุณต้องการฝึกอบรม SAES สำหรับหลายเลเยอร์ของแบบจำลองเราขอแนะนำให้ใช้ธง --distribute_modules ซึ่งจัดสรร SAE สำหรับเลเยอร์ที่แตกต่างกันไปยัง GPU ที่แตกต่างกัน ปัจจุบันเราต้องการให้จำนวน GPUs แบ่งจำนวนเลเยอร์ที่คุณฝึกซ้อมอย่างสม่ำเสมอ
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 คำสั่งด้านบนฝึกซ้อม SAE สำหรับทุก ชั้น ของ Llama 3 8B โดยใช้ GPU ที่มีอยู่ทั้งหมด มันสะสมการไล่ระดับสีมากกว่า 8 minibatches และแยกแต่ละ minibatch ออกเป็น 2 microbatches ก่อนที่จะป้อนมันลงในตัวเข้ารหัส SAE จึงช่วยประหยัดหน่วยความจำได้มาก นอกจากนี้ยังโหลดโมเดลในความแม่นยำ 8 บิตโดยใช้ bitsandbytes คำสั่งนี้ต้องใช้หน่วยความจำไม่เกิน 48GB ต่อ GPU บนโหนด 8 GPU
มีคุณสมบัติหลายอย่างที่เราต้องการเพิ่มในอนาคตอันใกล้:
หากคุณต้องการช่วยเหลือสิ่งเหล่านี้โปรดอย่าลังเลที่จะเปิดการประชาสัมพันธ์! คุณสามารถทำงานร่วมกับเราได้ในช่องทางที่กระจัดกระจาย Autoencoders ของ Eleutherai Discord
