File size: 1,258 Bytes
a2ffd07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | #!/bin/bash
# Define the parameters
model_name="gpt2-small"
tok_name="gpt2"
batch_size=1
max_tokens=500000000
sae_name="topk_crosscoder"
lr=0.0001
expansion_factor=32
k=32
auxk=256
auxk_coef=0.03125
device_id=7
max_epochs=1
dead_tokens_threshold=10000000
log_every_n_steps=50
save_every_n_training_steps=100
use_loss_var=true
num_workers=63
cd ..
input_hook_names=()
output_hook_names=()
for layer in {0..11}
do
input_hook_names+=("blocks.${layer}.ln2.hook_normalized")
output_hook_names+=("blocks.${layer}.hook_mlp_out")
done
python Train_Crosscoder.py \
--model_name "$model_name" \
--tok_name "$tok_name" \
--batch_size "$batch_size" \
--input_hook_names "${input_hook_names[@]}" \
--output_hook_names "${output_hook_names[@]}" \
--sae_name "$sae_name" \
--lr "$lr" \
--expansion_factor "$expansion_factor" \
--k "$k" \
--auxk "$auxk" \
--device_id "$device_id" \
--max_epochs "$max_epochs" \
--dead_tokens_threshold "$dead_tokens_threshold" \
--log_every_n_steps "$log_every_n_steps" \
--save_every_n_training_steps "$save_every_n_training_steps" \
--use_loss_var "$use_loss_var" \
--max_tokens "$max_tokens" \
--auxk_coef "$auxk_coef" \
--num_workers "$num_workers" |