sae

مقدمة

تقوم هذه المكتبة بتدريبات k -sparse Autoencoders (SAES) على تنشيط الدفق المتبقي لنماذج لغة العناق ، مما يتبع تقريبًا الوصفة المفصلة في تحجيم وتقييم العوامل التلقائية المتفرقة (Gao et al. 2024).

هذه مكتبة هزيلة وبسيطة مع خيارات تكوين قليلة. على عكس معظم مكتبات SAE الأخرى (على سبيل المثال Saelens) ، فإنه لا يقوم بتنشيط ذاكرة التخزين المؤقت على القرص ، ولكنه يحسبها على نحو حافل. يتيح لنا ذلك توسيع نطاق نماذج ومجموعات بيانات كبيرة جدًا بدون صفر تخزين ، ولكن يتمتع الجانب السلبي بأن محاولة فرط البرارمامات المختلفة لنفس النموذج ومجموعة البيانات ستكون أبطأ مما لو كنا نقوم بالتنشيط المخبأة (حيث سيتم إعادة حساب التنشيط). قد نضيف التخزين المؤقت كخيار في المستقبل.

بعد Gao et al. ، نستخدم وظيفة تنشيط TopK التي تفرض مباشرة المستوى المطلوب من التباين في التنشيطات. هذا على عكس المكتبات الأخرى التي تستخدم عقوبة L1 في وظيفة الخسارة. نعتقد أن Topk هو تحسن باريتو على نهج 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 لتمريرة إلى الأمام النموذجية على النحو التالي:

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

نقاط ربط مخصصة

بشكل افتراضي ، يتم تدريب SAEs على تنشيط الدفق المتبقي للنموذج. ومع ذلك ، يمكنك أيضًا تدريب SAEs على تنشيط أي وحدة فرعية أخرى (S) من خلال تحديد أنماط خط هوكس مخصصة. تشبه هذه الأنماط أسماء الوحدات النمطية القياسية (على سبيل المثال 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 "

نحن لا ندعم حاليًا التحكم اليدوي الدقيق في معدل التعلم ، أو عدد اللاتينات ، أو غيرها من المقاييس الفائقة على أساس نقطة خطاف. افتراضيًا ، يتم استخدام خيار expansion_ratio لتحديد العدد المناسب من اللاتينية لكل نقطة خطاف بناءً على عرض إخراج هذا الخطاب. ثم يتم تعيين معدل التعلم الافتراضي لكل نقطة خطاف باستخدام قانون تحجيم الجذر التربيعي العكسي بناءً على عدد اللاتينات. إذا قمت بتعيين عدد Latents يدويًا أو معدل التعلم ، فسيتم تطبيقه على جميع نقاط الخطاف.

التدريب الموزع

نحن ندعم التدريب الموزع عبر قيادة torchrun من Pytorch. بشكل افتراضي ، نستخدم الطريقة الموازية للبيانات الموزعة ، مما يعني أن أوزان كل SAE يتم نسخها على كل وحدة معالجة الرسومات.

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 ، والتي تخصص SAEs لطبقات مختلفة ل GPU مختلف. حاليًا ، نطلب من عدد وحدات معالجة الرسومات تقسيم عدد الطبقات التي تدرب عليها بشكل متساوٍ.

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 ، باستخدام جميع وحدات معالجة الرسومات المتاحة. يتراكم التدرجات أكثر من 8 وحدات صغيرة ، وتقسيم كل صغيرة إلى 2 microbatchs قبل إطعامها في تشفير SAE ، وبالتالي توفير الكثير من الذاكرة. كما أنه يقوم بتحميل النموذج بدقة 8 بت باستخدام bitsandbytes . لا يتطلب هذا الأمر أكثر من 48 جيجابايت من الذاكرة لكل وحدة معالجة الرسومات على عقدة GPU 8.

تودو

هناك العديد من الميزات التي نود إضافتها في المستقبل القريب:

• سازز المسبق
• دعم تنشيط التخزين المؤقت
• تقييم SAEs مع التباعد KL عند تطعيم في النموذج

إذا كنت ترغب في المساعدة في أي من هذه ، فلا تتردد في فتح العلاقات العامة! يمكنك التعاون معنا في قناة Autoencoders المتفجرة في Displicke the Eleutherai.