Finish training Gemma transcoders

neverix · Aug 5, 2024 · f20fc24 · f20fc24
1 parent 2e3dc93
commit f20fc24
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/scripts/train_gemma_sae.py b/scripts/train_gemma_sae.py
@@ -18,7 +18,7 @@ def train(
     cache_size = 2**16,
     cache_batch_size = 256,
     cache_ratio=1.0,
-    batch_size = 2**16,
+    batch_size = 2048,
     max_seq_len = 128,
     sparsity_coefficients=[4e-6],
     # save_steps=2500,

diff --git a/train_gemmas.py b/train_gemmas.py
@@ -1,5 +1,5 @@
 import os
-layers = [1, 2, 3, 4, 5]
+layers = [0]
 for layer_idx in range(len(layers)):
     layer = layers[layer_idx]
     restore = None  # if layer_idx == 0 else f"weights/phi-l{layers[layer_idx-1]}-gated.safetensors"