extra_materials/scripts/Train_SAE.sh

#!/bin/bash

# Define the parameters
model_name="gpt2"
tok_name="gpt2"
batch_size=5
max_tokens=500000000
sae_name="normalized"
lr=0.0001
expansion_factor=32
k=32
auxk=256
auxk_coef=0.1
device_id=7
max_epochs=1
dead_tokens_threshold=10000000
log_every_n_steps=50
save_every_n_training_steps=9700
save_step=true
use_loss_var=true
num_workers=63

cd ..

# Loop over 12 layers
for layer in {0..11}
do
    hook_name="blocks.${layer}.hook_mlp_out"

    python Train_SAE.py \
        --model_name "$model_name" \
        --tok_name "$tok_name" \
        --layer "$layer" \
        --batch_size "$batch_size" \
        --hook_name "$hook_name" \
        --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" \
        --save_step "$save_step" \
        --use_loss_var "$use_loss_var" \
        --max_tokens "$max_tokens" \
        --auxk_coef "$auxk_coef" \
        --num_workers "$num_workers"
done