sae
1.0.0
Эта библиотека обучает k -sparse AutoEncoders (SAE) на активации остаточного потока моделей языка объятия, примерно следовать рецепту, подробно описанному в масштабировании и оценке разреженных автоэнкомодеров (Gao et al. 2024).
Это худой, простая библиотека с небольшим количеством параметров конфигурации. В отличие от большинства других библиотек SAE (например, Saelens), он не кэширует активации на диске, а скорее вычисляет их на лету. Это позволяет нам масштабировать до очень больших моделей и наборов данных с нулевыми накладными расходом хранения, но имеет недостаток, что попытка различных гиперпараметров для одной и той же модели и набора данных будет медленнее, чем если бы мы кэшировали активации (поскольку активации будут повторно выкомментировать). Мы можем добавить кэширование в качестве опции в будущем.
После Gao et al., Мы используем функцию активации Topk, которая непосредственно обеспечивает желаемый уровень редкости в активациях. Это в отличие от других библиотек, которые используют штраф L1 в функции потери. Мы считаем, что Topk - это улучшение Pareto по сравнению с подходом L1, и, следовательно, не планируем его поддержать.
Чтобы загрузить предварительно проведенный SAE из центра Huggingface, вы можете использовать метод 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. Вы также можете загрузить SAE для всех слоев одновременно с использованием 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 для прохода модели вперед следующим образом:
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, включая подстановочные знаки и наборы символов. Например, для обучения SAE на вывод каждого модуля внимания и внутренней активации каждого 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 используется для выбора соответствующего количества задержек для каждой подъемной точки на основе ширины выхода этой крючковой точки. Скорость обучения по умолчанию для каждой крючковой точки устанавливается с использованием закона об обратном квадратном корневом масштабировании на основе количества задержек. Если вы вручную установите количество задержек или скорость обучения, оно будет применено ко всем везде.
Мы поддерживаем распределенное обучение с помощью команды Pytorch torchrun . По умолчанию мы используем параллельный метод распределенных данных, что означает, что веса каждого 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 Это просто, но очень неэффективно память. Если вы хотите тренировать SAE для многих слоев модели, мы рекомендуем использовать флаг --distribute_modules , который выделяет SAE для разных слоев на разные графические процессоры. В настоящее время мы требуем, чтобы количество графических процессоров равномерно разделяло количество слоев, для которых вы тренируете 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 Приведенная выше команда обучает SAE для каждого четного слоя Llama 3 8b, используя все доступные графические процессоры. Он накапливает градиенты более 8 minibatches и разбивает каждый Minibatch на 2 микробатча, прежде чем подавать их в энкодер SAE, тем самым сохраняя большую память. Он также загружает модель в 8-битной точности с использованием bitsandbytes . Эта команда требует не более 48 ГБ памяти на графический процессор на узле 8 графических процессоров.
Есть несколько функций, которые мы хотели бы добавить в ближайшем будущем:
Если вы хотите помочь с любым из них, пожалуйста, не стесняйтесь открывать PR! Вы можете сотрудничать с нами на канале Sparse-AutoEcoders в Eleutherai Discord.
