minor changes

jbloomAus · Apr 18, 2024 · bc766e4 · bc766e4
1 parent 822882c
commit bc766e4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/docs/training_saes.md b/docs/training_saes.md
@@ -49,7 +49,7 @@ cfg = LanguageModelSAERunnerConfig(
 
     # Activation Store Parameters
     n_batches_in_buffer = 128,
-    total_training_tokens = 1_000_000 * 300,
+    training_tokens = 1_000_000 * 300,
     store_batch_size = 32,
 
     # Dead Neurons and Sparsity
@@ -60,7 +60,7 @@ cfg = LanguageModelSAERunnerConfig(
 
     # WANDB
     log_to_wandb = True,
-    wandb_project= "mats_sae_training_gpt2",
+    wandb_project= "gpt2",
     wandb_entity = None,
     wandb_log_frequency=100,
 

diff --git a/sae_lens/training/config.py b/sae_lens/training/config.py
@@ -174,7 +174,9 @@ def __post_init__(self):
                 f"Lower bound: n_contexts_per_buffer (millions): {n_contexts_per_buffer / 10 **6}"
             )
 
-            total_training_steps = self.training_tokens // self.train_batch_size
+            total_training_steps = (
+                self.training_tokens + self.finetuning_tokens
+            ) // self.train_batch_size
             print(f"Total training steps: {total_training_steps}")
 
             total_wandb_updates = total_training_steps // self.wandb_log_frequency