File size: 1,192 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 | #!/bin/bash
# Define the parameters
model_name="gpt2-small"
tok_name="gpt2"
batch_size=10
max_tokens=500000000
sae_name="topk_transcoder"
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
use_loss_var=true
num_workers=63
cd ..
# Loop over 12 layers
for layer in {0..11}
do
input_hook="blocks.${layer}.ln2.hook_normalized"
output_hook="blocks.${layer}.hook_mlp_out"
python Train_Transcoder.py \
--model_name "$model_name" \
--tok_name "$tok_name" \
--layer "$layer" \
--batch_size "$batch_size" \
--input_hook "$input_hook" \
--output_hook "$output_hook" \
--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" \
--use_loss_var "$use_loss_var" \
--max_tokens "$max_tokens" \
--auxk_coef "$auxk_coef" \
--num_workers "$num_workers"
done |