sae-lens 6.27.1__py3-none-any.whl → 6.27.3__py3-none-any.whl

sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.27.1"
+__version__ = "6.27.3"
 
 import logging
 

sae_lens/cache_activations_runner.py

@@ -263,14 +263,21 @@ class CacheActivationsRunner:
 
         for i in tqdm(range(self.cfg.n_buffers), desc="Caching activations"):
             try:
-                buffer = self.activations_store.get_raw_buffer(
-                    self.cfg.n_batches_in_buffer, shuffle=False
-                )
-                shard = self._create_shard(buffer)
+                # Accumulate n_batches_in_buffer batches into one shard
+                buffers: list[tuple[torch.Tensor, torch.Tensor | None]] = []
+                for _ in range(self.cfg.n_batches_in_buffer):
+                    buffers.append(self.activations_store.get_raw_llm_batch())
+                # Concatenate all batches
+                acts = torch.cat([b[0] for b in buffers], dim=0)
+                token_ids: torch.Tensor | None = None
+                if buffers[0][1] is not None:
+                    # All batches have token_ids if the first one does
+                    token_ids = torch.cat([b[1] for b in buffers], dim=0)  # type: ignore[arg-type]
+                shard = self._create_shard((acts, token_ids))
                 shard.save_to_disk(
                     f"{tmp_cached_activation_path}/shard_{i:05d}", num_shards=1
                 )
-                del buffer, shard
+                del buffers, acts, token_ids, shard
             except StopIteration:
                 logger.warning(
                     f"Warning: Ran out of samples while filling the buffer at batch {i} before reaching {self.cfg.n_buffers} batches."

sae_lens/training/activations_store.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 import os
 import warnings
-from collections.abc import Generator, Iterator, Sequence
+from collections.abc import Generator, Iterator
 from pathlib import Path
 from typing import Any, Literal, cast
 
@@ -254,7 +254,6 @@ class ActivationsStore:
         self.context_size = context_size
         self.d_in = d_in
         self.n_batches_in_buffer = n_batches_in_buffer
-        self.half_buffer_size = n_batches_in_buffer // 2
         self.total_training_tokens = total_training_tokens
         self.store_batch_size_prompts = store_batch_size_prompts
         self.train_batch_size_tokens = train_batch_size_tokens
@@ -538,18 +537,15 @@ class ActivationsStore:
 
         return stacked_activations
 
-    def _load_buffer_from_cached(
+    def _load_raw_llm_batch_from_cached(
         self,
-        total_size: int,
-        context_size: int,
-        d_in: int,
         raise_on_epoch_end: bool,
     ) -> tuple[
         torch.Tensor,
         torch.Tensor | None,
     ]:
         """
-        Loads `total_size` activations from `cached_activation_dataset`
+        Loads a batch of activations from `cached_activation_dataset`
 
         The dataset has columns for each hook_name,
         each containing activations of shape (context_size, d_in).
@@ -557,6 +553,10 @@ class ActivationsStore:
         raises StopIteration
         """
         assert self.cached_activation_dataset is not None
+        context_size = self.context_size
+        batch_size = self.store_batch_size_prompts
+        d_in = self.d_in
+
         # In future, could be a list of multiple hook names
         if self.hook_name not in self.cached_activation_dataset.column_names:
             raise ValueError(
@@ -564,138 +564,100 @@ class ActivationsStore:
                 f"got {self.cached_activation_dataset.column_names}."
             )
 
-        if self.current_row_idx > len(self.cached_activation_dataset) - total_size:
+        if self.current_row_idx > len(self.cached_activation_dataset) - batch_size:
             self.current_row_idx = 0
             if raise_on_epoch_end:
                 raise StopIteration
 
-        new_buffer = []
         ds_slice = self.cached_activation_dataset[
-            self.current_row_idx : self.current_row_idx + total_size
+            self.current_row_idx : self.current_row_idx + batch_size
         ]
         # Load activations for each hook.
         # Usually faster to first slice dataset then pick column
-        new_buffer = ds_slice[self.hook_name]
-        if new_buffer.shape != (total_size, context_size, d_in):
+        acts_buffer = ds_slice[self.hook_name]
+        if acts_buffer.shape != (batch_size, context_size, d_in):
             raise ValueError(
-                f"new_buffer has shape {new_buffer.shape}, "
-                f"but expected ({total_size}, {context_size}, {d_in})."
+                f"acts_buffer has shape {acts_buffer.shape}, "
+                f"but expected ({batch_size}, {context_size}, {d_in})."
             )
 
-        self.current_row_idx += total_size
-        acts_buffer = new_buffer.reshape(total_size * context_size, d_in)
+        self.current_row_idx += batch_size
+        acts_buffer = acts_buffer.reshape(batch_size * context_size, d_in)
 
         if "token_ids" not in self.cached_activation_dataset.column_names:
             return acts_buffer, None
 
         token_ids_buffer = ds_slice["token_ids"]
-        if token_ids_buffer.shape != (total_size, context_size):
+        if token_ids_buffer.shape != (batch_size, context_size):
             raise ValueError(
                 f"token_ids_buffer has shape {token_ids_buffer.shape}, "
-                f"but expected ({total_size}, {context_size})."
+                f"but expected ({batch_size}, {context_size})."
             )
-        token_ids_buffer = token_ids_buffer.reshape(total_size * context_size)
+        token_ids_buffer = token_ids_buffer.reshape(batch_size * context_size)
         return acts_buffer, token_ids_buffer
 
     @torch.no_grad()
-    def get_raw_buffer(
+    def get_raw_llm_batch(
         self,
-        n_batches_in_buffer: int,
         raise_on_epoch_end: bool = False,
-        shuffle: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """
-        Loads the next n_batches_in_buffer batches of activations into a tensor and returns it.
+        Loads the next batch of activations from the LLM and returns it.
 
-        The primary purpose here is maintaining a shuffling buffer.
+        If raise_on_epoch_end is True, when the dataset is exhausted it will
+        automatically refill the dataset and then raise a StopIteration so that
+        the caller has a chance to react.
 
-        If raise_on_epoch_end is True, when the dataset it exhausted it will automatically refill the dataset and then raise a StopIteration so that the caller has a chance to react.
+        Returns:
+            Tuple of (activations, token_ids) where activations has shape
+            (batch_size * context_size, d_in) and token_ids has shape
+            (batch_size * context_size,).
         """
-        context_size = self.context_size
-        batch_size = self.store_batch_size_prompts
         d_in = self.d_in
-        total_size = batch_size * n_batches_in_buffer
 
         if self.cached_activation_dataset is not None:
-            return self._load_buffer_from_cached(
-                total_size, context_size, d_in, raise_on_epoch_end
-            )
+            return self._load_raw_llm_batch_from_cached(raise_on_epoch_end)
 
-        refill_iterator = range(0, total_size, batch_size)
-        # Initialize empty tensor buffer of the maximum required size with an additional dimension for layers
-        new_buffer_activations = torch.zeros(
-            (total_size, self.training_context_size, d_in),
-            dtype=self.dtype,  # type: ignore
-            device=self.device,
-        )
-        new_buffer_token_ids = torch.zeros(
-            (total_size, self.training_context_size),
-            dtype=torch.long,
-            device=self.device,
+        # move batch toks to gpu for model
+        batch_tokens = self.get_batch_tokens(raise_at_epoch_end=raise_on_epoch_end).to(
+            _get_model_device(self.model)
         )
+        activations = self.get_activations(batch_tokens).to(self.device)
 
-        for refill_batch_idx_start in tqdm(
-            refill_iterator, leave=False, desc="Refilling buffer"
-        ):
-            # move batch toks to gpu for model
-            refill_batch_tokens = self.get_batch_tokens(
-                raise_at_epoch_end=raise_on_epoch_end
-            ).to(_get_model_device(self.model))
-            refill_activations = self.get_activations(refill_batch_tokens)
-            # move acts back to cpu
-            refill_activations.to(self.device)
-            new_buffer_activations[
-                refill_batch_idx_start : refill_batch_idx_start + batch_size, ...
-            ] = refill_activations
-
-            # handle seqpos_slice, this is done for activations in get_activations
-            refill_batch_tokens = refill_batch_tokens[:, slice(*self.seqpos_slice)]
-            new_buffer_token_ids[
-                refill_batch_idx_start : refill_batch_idx_start + batch_size, ...
-            ] = refill_batch_tokens
-
-        new_buffer_activations = new_buffer_activations.reshape(-1, d_in)
-        new_buffer_token_ids = new_buffer_token_ids.reshape(-1)
-        if shuffle:
-            new_buffer_activations, new_buffer_token_ids = permute_together(
-                [new_buffer_activations, new_buffer_token_ids]
-            )
+        # handle seqpos_slice, this is done for activations in get_activations
+        batch_tokens = batch_tokens[:, slice(*self.seqpos_slice)]
 
-        return (
-            new_buffer_activations,
-            new_buffer_token_ids,
-        )
+        # reshape from (batch, context, d_in) to (batch * context, d_in)
+        activations = activations.reshape(-1, d_in)
+        token_ids = batch_tokens.reshape(-1)
 
-    def get_filtered_buffer(
+        return activations, token_ids
+
+    def get_filtered_llm_batch(
         self,
-        n_batches_in_buffer: int,
         raise_on_epoch_end: bool = False,
-        shuffle: bool = True,
     ) -> torch.Tensor:
+        """
+        Get a batch of LLM activations with special tokens filtered out.
+        """
         return _filter_buffer_acts(
-            self.get_raw_buffer(
-                n_batches_in_buffer=n_batches_in_buffer,
-                raise_on_epoch_end=raise_on_epoch_end,
-                shuffle=shuffle,
-            ),
+            self.get_raw_llm_batch(raise_on_epoch_end=raise_on_epoch_end),
             self.exclude_special_tokens,
         )
 
     def _iterate_filtered_activations(self) -> Generator[torch.Tensor, None, None]:
         """
-        Iterate over the filtered tokens in the buffer.
+        Iterate over filtered LLM activation batches.
         """
         while True:
             try:
-                yield self.get_filtered_buffer(
-                    self.half_buffer_size, raise_on_epoch_end=True
-                )
+                yield self.get_filtered_llm_batch(raise_on_epoch_end=True)
             except StopIteration:
                 warnings.warn(
                     "All samples in the training dataset have been exhausted, beginning new epoch."
                 )
                 try:
-                    yield self.get_filtered_buffer(self.half_buffer_size)
+                    yield self.get_filtered_llm_batch()
                 except StopIteration:
                     raise ValueError(
                         "Unable to fill buffer after starting new epoch. Dataset may be too small."
@@ -827,9 +789,3 @@ def _filter_buffer_acts(
 
     mask = torch.isin(tokens, exclude_tokens)
     return activations[~mask]
-
-
-def permute_together(tensors: Sequence[torch.Tensor]) -> tuple[torch.Tensor, ...]:
-    """Permute tensors together."""
-    permutation = torch.randperm(tensors[0].shape[0])
-    return tuple(t[permutation] for t in tensors)

sae_lens/training/mixing_buffer.py

@@ -44,9 +44,10 @@ def mixing_buffer(
             if mix_fraction > 0:
                 storage_buffer = storage_buffer[torch.randperm(storage_buffer.shape[0])]
 
-            num_serving_batches = max(
-                1, int(storage_buffer.shape[0] * (1 - mix_fraction)) // batch_size
-            )
+            # Keep a fixed amount for mixing, serve the rest
+            keep_for_mixing = int(buffer_size * mix_fraction)
+            num_to_serve = storage_buffer.shape[0] - keep_for_mixing
+            num_serving_batches = max(1, num_to_serve // batch_size)
             serving_cutoff = num_serving_batches * batch_size
             serving_buffer = storage_buffer[:serving_cutoff]
             storage_buffer = storage_buffer[serving_cutoff:]

{sae_lens-6.27.1.dist-info → sae_lens-6.27.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.27.1
+Version: 6.27.3
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE

{sae_lens-6.27.1.dist-info → sae_lens-6.27.3.dist-info}/RECORD

@@ -1,8 +1,8 @@
-sae_lens/__init__.py,sha256=MaSnckNx7ls1jVEQW1qgt7snwGbE62p91F6TAjaSR1k,4725
+sae_lens/__init__.py,sha256=ETLfd3PmdJ2aAaKyeTTHptBE2HaWY0OfzOKNk7dyhKE,4725
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=dQRgGVwce8XwylL2AzJE7l9elhtMRFCs2hdUj-Qyy4g,14038
 sae_lens/analysis/neuronpedia_integration.py,sha256=Gx1W7hUBEuMoasNcnOnZ1wmqbXDd1pSZ1nqKEya1HQc,4962
-sae_lens/cache_activations_runner.py,sha256=Lvlz-k5-3XxVRtUdC4b1CiKyx5s0ckLa8GDGv9_kcxs,12566
+sae_lens/cache_activations_runner.py,sha256=TjqNWIc46Nw09jHWFjzQzgzG5wdu_87Ahe-iFjI5_0Q,13117
 sae_lens/config.py,sha256=sseYcRMsAyopj8FICup1RGTXjFxzAithZ2OH7OpQV3Y,30839
 sae_lens/constants.py,sha256=CM-h9AjZNAl2aP7hVpKk7YsFHpu-_Lfhhmq2d5qPEVc,887
 sae_lens/evals.py,sha256=P0NUsJeGzYxFBiVKhbPzd72IFKY4gH40HHlEZ3jEAmg,39598
@@ -28,15 +28,15 @@ sae_lens/saes/transcoder.py,sha256=CTpJs8ASOK06npih7gZHygZuxqTR7HICWlOYfTiKjI4,1
 sae_lens/tokenization_and_batching.py,sha256=uoHtAs9z3XqG0Fh-iQVYVlrbyB_E3kFFhrKU30BosCo,5438
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/training/activation_scaler.py,sha256=FzNfgBplLWmyiSlZ6TUvE-nur3lOiGTrlvC97ys8S24,1973
-sae_lens/training/activations_store.py,sha256=2BVajHRcozKQFf1tkeraUCdFuut3spdk0hhgtdpizzI,34031
-sae_lens/training/mixing_buffer.py,sha256=DK22yPwEop4suG0K-8XFw5ZGNl0JrgCEjypmKEUAaGY,2394
+sae_lens/training/activations_store.py,sha256=kp4-6R4rTJUSt-g-Ifg5B1h7iIe7jZj-XQSKDvDpQMI,32187
+sae_lens/training/mixing_buffer.py,sha256=1Z-S2CcQXMWGxRZJFnXeZFxbZcALkO_fP6VO37XdJQQ,2519
 sae_lens/training/optim.py,sha256=bJpqqcK4enkcPvQAJkeH4Ci1LUOlfjIMTv6-IlaAbRA,5588
 sae_lens/training/sae_trainer.py,sha256=zhkabyIKxI_tZTV3_kwz6zMrHZ95Ecr97krmwc-9ffs,17600
 sae_lens/training/types.py,sha256=1FpLx_Doda9vZpmfm-x1e8wGBYpyhe9Kpb_JuM5nIFM,90
 sae_lens/training/upload_saes_to_huggingface.py,sha256=r_WzI1zLtGZ5TzAxuG3xa_8T09j3zXJrWd_vzPsPGkQ,4469
 sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
 sae_lens/util.py,sha256=spkcmQUsjVYFn5H2032nQYr1CKGVnv3tAdfIpY59-Mg,3919
-sae_lens-6.27.1.dist-info/METADATA,sha256=7jGERGOq6VhzilrpOOfVjWjxKEkpiok18GZH9eSzVTE,5361
-sae_lens-6.27.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-sae_lens-6.27.1.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.27.1.dist-info/RECORD,,
+sae_lens-6.27.3.dist-info/METADATA,sha256=c59mjyoausFHs1bd8n_4J6dA-2uDRPgY9Wwas52zydw,5361
+sae_lens-6.27.3.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+sae_lens-6.27.3.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.27.3.dist-info/RECORD,,