1for name in ["blocks.0.hook_resid_post", "blocks.4.hook_resid_post", "blocks.8.hook_resid_post", "blocks.12.hook_resid_post", "blocks.16.hook_resid_post"]: #["blocks.0.hook_resid_post",
2 total_training_steps = 30_000 # probably we should do more
3 batch_size = 4096
4 total_training_tokens = total_training_steps * batch_size
5 num = int(name.split(".")[1])
6
7 lr_warm_up_steps = 0
8 lr_decay_steps = total_training_steps // 5 # 20% of training
9 l1_warm_up_steps = total_training_steps // 20 # 5% of training
10
11 cfg = LanguageModelSAERunnerConfig(
12 # Data Generating Function (Model + Training Distibuion)
13 model_name="Qwen/Qwen2.5-0.5B", # our model (more options here: https://neelnanda-io.github.io/TransformerLens/generated/model_properties_table.html)
14 hook_name=name, # A valid hook point (see more details here: https://neelnanda-io.github.io/TransformerLens/generated/demos/Main_Demo.html#Hook-Points)
15 hook_layer=num, # Only one layer in the model.
16 d_in=896, # the width of the mlp output.
17 dataset_path="NoamDiamant52/TinyStories_tokenized", # this is a tokenized language dataset on Huggingface for the Tiny Stories corpus.
18 is_dataset_tokenized=True,
19 streaming=True, # we could pre-download the token dataset if it was small.
20 # SAE Parameters
21 mse_loss_normalization=None, # We won't normalize the mse loss,
22 expansion_factor=16, # the width of the SAE. Larger will result in better stats but slower training.
23 b_dec_init_method="zeros", # The geometric median can be used to initialize the decoder weights.
24 apply_b_dec_to_input=False, # We won't apply the decoder weights to the input.
25 normalize_sae_decoder=False,
26 scale_sparsity_penalty_by_decoder_norm=True,
27 decoder_heuristic_init=True,
28 init_encoder_as_decoder_transpose=True,
29 normalize_activations="expected_average_only_in",
30 # Training Parameters
31 lr=5e-5, # lower the better, we'll go fairly high to speed up the tutorial.
32 adam_beta1=0.9, # adam params (default, but once upon a time we experimented with these.)
33 adam_beta2=0.999,
34 lr_scheduler_name="constant", # constant learning rate with warmup. Could be better schedules out there.
35 lr_warm_up_steps=lr_warm_up_steps, # this can help avoid too many dead features initially.
36 lr_decay_steps=lr_decay_steps, # this will help us avoid overfitting.
37 l1_coefficient=5, # will control how sparse the feature activations are
38 l1_warm_up_steps=l1_warm_up_steps, # this can help avoid too many dead features initially.
39 lp_norm=1.0, # the L1 penalty (and not a Lp for p < 1)
40 train_batch_size_tokens=batch_size,
41 context_size=512, # will control the lenght of the prompts we feed to the model. Larger is better but slower. so for the tutorial we'll use a short one.
42 # Activation Store Parameters
43 n_batches_in_buffer=64, # controls how many activations we store / shuffle.
44 training_tokens=total_training_tokens, # 100 million tokens is quite a few, but we want to see good stats. Get a coffee, come back.
45 store_batch_size_prompts=16,
46 # Resampling protocol
47 use_ghost_grads=False, # we don't use ghost grads anymore.
48 feature_sampling_window=1000, # this controls our reporting of feature sparsity stats
49 dead_feature_window=1000, # would effect resampling or ghost grads if we were using it.
50 dead_feature_threshold=1e-4, # would effect resampling or ghost grads if we were using it.
51 # WANDB
52 log_to_wandb=True, # always use wandb unless you are just testing code.
53 wandb_project="sae_lens_tutorial",
54 wandb_log_frequency=30,
55 eval_every_n_wandb_logs=20,
56 # Misc
57 device=device,
58 seed=42,
59 n_checkpoints=0,
60 checkpoint_path="checkpoints",
61 dtype="float32",
62 )
63 # look at the next cell to see some instruction for what to do while this is running.
64 sparse_autoencoder = SAETrainingRunner(cfg).run()
65 sparse_autoencoder.save_model(f"post_residual_layer_{num}")
