extra_materials/scripts/Train_Transcoder.sh

#!/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