sae-lens 6.0.0rc2__py3-none-any.whl → 6.0.0rc4__py3-none-any.whl

sae_lens/training/sae_trainer.py

@@ -1,16 +1,17 @@
 import contextlib
 from dataclasses import dataclass
-from typing import Any, Generic, Protocol, cast
+from pathlib import Path
+from typing import Any, Callable, Generic, Protocol
 
 import torch
 import wandb
+from safetensors.torch import save_file
 from torch.optim import Adam
-from tqdm import tqdm
-from transformer_lens.hook_points import HookedRootModule
+from tqdm.auto import tqdm
 
 from sae_lens import __version__
-from sae_lens.config import LanguageModelSAERunnerConfig
-from sae_lens.evals import EvalConfig, run_evals
+from sae_lens.config import SAETrainerConfig
+from sae_lens.constants import ACTIVATION_SCALER_CFG_FILENAME, SPARSITY_FILENAME
 from sae_lens.saes.sae import (
     T_TRAINING_SAE,
     T_TRAINING_SAE_CONFIG,
@@ -19,8 +20,9 @@ from sae_lens.saes.sae import (
     TrainStepInput,
     TrainStepOutput,
 )
-from sae_lens.training.activations_store import ActivationsStore
+from sae_lens.training.activation_scaler import ActivationScaler
 from sae_lens.training.optim import CoefficientScheduler, get_lr_scheduler
+from sae_lens.training.types import DataProvider
 
 
 def _log_feature_sparsity(
@@ -46,33 +48,39 @@ class TrainSAEOutput:
 class SaveCheckpointFn(Protocol):
     def __call__(
         self,
-        trainer: "SAETrainer[Any, Any]",
-        checkpoint_name: str,
-        wandb_aliases: list[str] | None = None,
+        checkpoint_path: Path,
     ) -> None: ...
 
 
+Evaluator = Callable[[T_TRAINING_SAE, DataProvider, ActivationScaler], dict[str, Any]]
+
+
 class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
     """
     Core SAE class used for inference. For training, see TrainingSAE.
     """
 
+    data_provider: DataProvider
+    activation_scaler: ActivationScaler
+    evaluator: Evaluator[T_TRAINING_SAE] | None
+
     def __init__(
         self,
-        model: HookedRootModule,
+        cfg: SAETrainerConfig,
         sae: T_TRAINING_SAE,
-        activation_store: ActivationsStore,
-        save_checkpoint_fn: SaveCheckpointFn,
-        cfg: LanguageModelSAERunnerConfig[T_TRAINING_SAE_CONFIG],
+        data_provider: DataProvider,
+        evaluator: Evaluator[T_TRAINING_SAE] | None = None,
+        save_checkpoint_fn: SaveCheckpointFn | None = None,
     ) -> None:
-        self.model = model
         self.sae = sae
-        self.activations_store = activation_store
-        self.save_checkpoint = save_checkpoint_fn
+        self.data_provider = data_provider
+        self.evaluator = evaluator
+        self.activation_scaler = ActivationScaler()
+        self.save_checkpoint_fn = save_checkpoint_fn
         self.cfg = cfg
 
         self.n_training_steps: int = 0
-        self.n_training_tokens: int = 0
+        self.n_training_samples: int = 0
         self.started_fine_tuning: bool = False
 
         _update_sae_lens_training_version(self.sae)
@@ -82,20 +90,16 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             self.checkpoint_thresholds = list(
                 range(
                     0,
-                    cfg.total_training_tokens,
-                    cfg.total_training_tokens // self.cfg.n_checkpoints,
+                    cfg.total_training_samples,
+                    cfg.total_training_samples // self.cfg.n_checkpoints,
                 )
             )[1:]
 
-        self.act_freq_scores = torch.zeros(
-            cast(int, cfg.sae.d_sae),
-            device=cfg.device,
-        )
+        self.act_freq_scores = torch.zeros(sae.cfg.d_sae, device=cfg.device)
         self.n_forward_passes_since_fired = torch.zeros(
-            cast(int, cfg.sae.d_sae),
-            device=cfg.device,
+            sae.cfg.d_sae, device=cfg.device
         )
-        self.n_frac_active_tokens = 0
+        self.n_frac_active_samples = 0
         # we don't train the scaling factor (initially)
         # set requires grad to false for the scaling factor
         for name, param in self.sae.named_parameters():
@@ -131,7 +135,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             )
 
         # Setup autocast if using
-        self.scaler = torch.amp.GradScaler(
+        self.grad_scaler = torch.amp.GradScaler(
             device=self.cfg.device, enabled=self.cfg.autocast
         )
 
@@ -144,23 +148,9 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
         else:
             self.autocast_if_enabled = contextlib.nullcontext()
 
-        # Set up eval config
-
-        self.trainer_eval_config = EvalConfig(
-            batch_size_prompts=self.cfg.eval_batch_size_prompts,
-            n_eval_reconstruction_batches=self.cfg.n_eval_batches,
-            n_eval_sparsity_variance_batches=self.cfg.n_eval_batches,
-            compute_ce_loss=True,
-            compute_l2_norms=True,
-            compute_sparsity_metrics=True,
-            compute_variance_metrics=True,
-            compute_kl=False,
-            compute_featurewise_weight_based_metrics=False,
-        )
-
     @property
     def feature_sparsity(self) -> torch.Tensor:
-        return self.act_freq_scores / self.n_frac_active_tokens
+        return self.act_freq_scores / self.n_frac_active_samples
 
     @property
     def log_feature_sparsity(self) -> torch.Tensor:
@@ -171,19 +161,24 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
         return (self.n_forward_passes_since_fired > self.cfg.dead_feature_window).bool()
 
     def fit(self) -> T_TRAINING_SAE:
-        pbar = tqdm(total=self.cfg.total_training_tokens, desc="Training SAE")
-
-        self.activations_store.set_norm_scaling_factor_if_needed()
+        self.sae.to(self.cfg.device)
+        pbar = tqdm(total=self.cfg.total_training_samples, desc="Training SAE")
+
+        if self.sae.cfg.normalize_activations == "expected_average_only_in":
+            self.activation_scaler.estimate_scaling_factor(
+                d_in=self.sae.cfg.d_in,
+                data_provider=self.data_provider,
+                n_batches_for_norm_estimate=int(1e3),
+            )
 
         # Train loop
-        while self.n_training_tokens < self.cfg.total_training_tokens:
+        while self.n_training_samples < self.cfg.total_training_samples:
             # Do a training step.
-            layer_acts = self.activations_store.next_batch()[:, 0, :].to(
-                self.sae.device
-            )
-            self.n_training_tokens += self.cfg.train_batch_size_tokens
+            batch = next(self.data_provider).to(self.sae.device)
+            self.n_training_samples += batch.shape[0]
+            scaled_batch = self.activation_scaler(batch)
 
-            step_output = self._train_step(sae=self.sae, sae_in=layer_acts)
+            step_output = self._train_step(sae=self.sae, sae_in=scaled_batch)
 
             if self.cfg.logger.log_to_wandb:
                 self._log_train_step(step_output)
@@ -194,22 +189,56 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             self._update_pbar(step_output, pbar)
 
         # fold the estimated norm scaling factor into the sae weights
-        if self.activations_store.estimated_norm_scaling_factor is not None:
+        if self.activation_scaler.scaling_factor is not None:
             self.sae.fold_activation_norm_scaling_factor(
-                self.activations_store.estimated_norm_scaling_factor
+                self.activation_scaler.scaling_factor
             )
-            self.activations_store.estimated_norm_scaling_factor = None
+            self.activation_scaler.scaling_factor = None
 
-        # save final sae group to checkpoints folder
+        # save final inference sae group to checkpoints folder
         self.save_checkpoint(
-            trainer=self,
-            checkpoint_name=f"final_{self.n_training_tokens}",
+            checkpoint_name=f"final_{self.n_training_samples}",
             wandb_aliases=["final_model"],
+            save_inference_model=True,
         )
 
         pbar.close()
         return self.sae
 
+    def save_checkpoint(
+        self,
+        checkpoint_name: str,
+        wandb_aliases: list[str] | None = None,
+        save_inference_model: bool = False,
+    ) -> None:
+        checkpoint_path = Path(self.cfg.checkpoint_path) / checkpoint_name
+        checkpoint_path.mkdir(exist_ok=True, parents=True)
+
+        save_fn = (
+            self.sae.save_inference_model
+            if save_inference_model
+            else self.sae.save_model
+        )
+        weights_path, cfg_path = save_fn(str(checkpoint_path))
+
+        sparsity_path = checkpoint_path / SPARSITY_FILENAME
+        save_file({"sparsity": self.log_feature_sparsity}, sparsity_path)
+
+        activation_scaler_path = checkpoint_path / ACTIVATION_SCALER_CFG_FILENAME
+        self.activation_scaler.save(str(activation_scaler_path))
+
+        if self.cfg.logger.log_to_wandb:
+            self.cfg.logger.log(
+                self,
+                weights_path,
+                cfg_path,
+                sparsity_path=sparsity_path,
+                wandb_aliases=wandb_aliases,
+            )
+
+        if self.save_checkpoint_fn is not None:
+            self.save_checkpoint_fn(checkpoint_path=checkpoint_path)
+
     def _train_step(
         self,
         sae: T_TRAINING_SAE,
@@ -242,17 +271,19 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
                 self.act_freq_scores += (
                     (train_step_output.feature_acts.abs() > 0).float().sum(0)
                 )
-                self.n_frac_active_tokens += self.cfg.train_batch_size_tokens
+                self.n_frac_active_samples += self.cfg.train_batch_size_samples
 
-        # Scaler will rescale gradients if autocast is enabled
-        self.scaler.scale(
+        # Grad scaler will rescale gradients if autocast is enabled
+        self.grad_scaler.scale(
             train_step_output.loss
         ).backward()  # loss.backward() if not autocasting
-        self.scaler.unscale_(self.optimizer)  # needed to clip correctly
+        self.grad_scaler.unscale_(self.optimizer)  # needed to clip correctly
         # TODO: Work out if grad norm clipping should be in config / how to test it.
         torch.nn.utils.clip_grad_norm_(sae.parameters(), 1.0)
-        self.scaler.step(self.optimizer)  # just ctx.optimizer.step() if not autocasting
-        self.scaler.update()
+        self.grad_scaler.step(
+            self.optimizer
+        )  # just ctx.optimizer.step() if not autocasting
+        self.grad_scaler.update()
 
         self.optimizer.zero_grad()
         self.lr_scheduler.step()
@@ -267,7 +298,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             wandb.log(
                 self._build_train_step_log_dict(
                     output=step_output,
-                    n_training_tokens=self.n_training_tokens,
+                    n_training_samples=self.n_training_samples,
                 ),
                 step=self.n_training_steps,
             )
@@ -283,7 +314,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
     def _build_train_step_log_dict(
         self,
         output: TrainStepOutput,
-        n_training_tokens: int,
+        n_training_samples: int,
     ) -> dict[str, Any]:
         sae_in = output.sae_in
         sae_out = output.sae_out
@@ -311,7 +342,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             "sparsity/mean_passes_since_fired": self.n_forward_passes_since_fired.mean().item(),
             "sparsity/dead_features": self.dead_neurons.sum().item(),
             "details/current_learning_rate": current_learning_rate,
-            "details/n_training_tokens": n_training_tokens,
+            "details/n_training_samples": n_training_samples,
             **{
                 f"details/{name}_coefficient": scheduler.value
                 for name, scheduler in self.coefficient_schedulers.items()
@@ -331,30 +362,11 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             * self.cfg.logger.eval_every_n_wandb_logs
         ) == 0:
             self.sae.eval()
-            ignore_tokens = set()
-            if self.activations_store.exclude_special_tokens is not None:
-                ignore_tokens = set(
-                    self.activations_store.exclude_special_tokens.tolist()
-                )
-            eval_metrics, _ = run_evals(
-                sae=self.sae,
-                activation_store=self.activations_store,
-                model=self.model,
-                eval_config=self.trainer_eval_config,
-                ignore_tokens=ignore_tokens,
-                model_kwargs=self.cfg.model_kwargs,
-            )  # not calculating featurwise metrics here.
-
-            # Remove eval metrics that are already logged during training
-            eval_metrics.pop("metrics/explained_variance", None)
-            eval_metrics.pop("metrics/explained_variance_std", None)
-            eval_metrics.pop("metrics/l0", None)
-            eval_metrics.pop("metrics/l1", None)
-            eval_metrics.pop("metrics/mse", None)
-
-            # Remove metrics that are not useful for wandb logging
-            eval_metrics.pop("metrics/total_tokens_evaluated", None)
-
+            eval_metrics = (
+                self.evaluator(self.sae, self.data_provider, self.activation_scaler)
+                if self.evaluator is not None
+                else {}
+            )
             for key, value in self.sae.log_histograms().items():
                 eval_metrics[key] = wandb.Histogram(value)  # type: ignore
 
@@ -378,21 +390,18 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
     @torch.no_grad()
     def _reset_running_sparsity_stats(self) -> None:
         self.act_freq_scores = torch.zeros(
-            self.cfg.sae.d_sae,  # type: ignore
+            self.sae.cfg.d_sae,  # type: ignore
             device=self.cfg.device,
         )
-        self.n_frac_active_tokens = 0
+        self.n_frac_active_samples = 0
 
     @torch.no_grad()
     def _checkpoint_if_needed(self):
         if (
             self.checkpoint_thresholds
-            and self.n_training_tokens > self.checkpoint_thresholds[0]
+            and self.n_training_samples > self.checkpoint_thresholds[0]
         ):
-            self.save_checkpoint(
-                trainer=self,
-                checkpoint_name=str(self.n_training_tokens),
-            )
+            self.save_checkpoint(checkpoint_name=str(self.n_training_samples))
             self.checkpoint_thresholds.pop(0)
 
     @torch.no_grad()
@@ -408,7 +417,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
                 for loss_name, loss_value in step_output.losses.items()
             )
             pbar.set_description(f"{self.n_training_steps}| {loss_strs}")
-            pbar.update(update_interval * self.cfg.train_batch_size_tokens)
+            pbar.update(update_interval * self.cfg.train_batch_size_samples)
 
 
 def _unwrap_item(item: float | torch.Tensor) -> float:

sae_lens/training/types.py

@@ -0,0 +1,5 @@
+from typing import Iterator
+
+import torch
+
+DataProvider = Iterator[torch.Tensor]

sae_lens/training/upload_saes_to_huggingface.py

@@ -88,7 +88,7 @@ def _create_default_readme(repo_id: str, sae_ids: Iterable[str]) -> str:
         ```python
         from sae_lens import SAE
 
-        sae, cfg_dict, sparsity = SAE.from_pretrained("{repo_id}", "<sae_id>")
+        sae = SAE.from_pretrained("{repo_id}", "<sae_id>")
         ```
         """
     )

sae_lens/util.py

@@ -1,3 +1,4 @@
+import re
 from dataclasses import asdict, fields, is_dataclass
 from typing import Sequence, TypeVar
 
@@ -26,3 +27,21 @@ def filter_valid_dataclass_fields(
     if whitelist_fields is not None:
         valid_field_names = valid_field_names.union(whitelist_fields)
     return {key: val for key, val in source_dict.items() if key in valid_field_names}
+
+
+def extract_stop_at_layer_from_tlens_hook_name(hook_name: str) -> int | None:
+    """Extract the stop_at layer from a HookedTransformer hook name.
+
+    Returns None if the hook name is not a valid HookedTransformer hook name.
+    """
+    layer = extract_layer_from_tlens_hook_name(hook_name)
+    return None if layer is None else layer + 1
+
+
+def extract_layer_from_tlens_hook_name(hook_name: str) -> int | None:
+    """Extract the layer from a HookedTransformer hook name.
+
+    Returns None if the hook name is not a valid HookedTransformer hook name.
+    """
+    hook_match = re.search(r"\.(\d+)\.", hook_name)
+    return None if hook_match is None else int(hook_match.group(1))

{sae_lens-6.0.0rc2.dist-info → sae_lens-6.0.0rc4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.0.0rc2
+Version: 6.0.0rc4
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch

sae_lens-6.0.0rc4.dist-info/RECORD

@@ -0,0 +1,37 @@
+sae_lens/__init__.py,sha256=dGZU3Y6iwiuW5oQVTfNvUmfnHO3bHWWbpU-nvXvw9M8,2861
+sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/analysis/hooked_sae_transformer.py,sha256=Eyg1Y2hVIHNuiiLOCTgzstOuW6iA-7hPHqaGR8y_vMs,13809
+sae_lens/analysis/neuronpedia_integration.py,sha256=MrENqc81Mc2SMbxGjbwHzpkGUCAFKSf0i4EdaUF2Oj4,18707
+sae_lens/cache_activations_runner.py,sha256=L5hhuU2-zPQr2S3L64GMKKLeMQfqXxwDl8NbuOtrybI,12567
+sae_lens/config.py,sha256=9Lg4HkQvj1t9QZJdmC071lyJMc_iqNQknosT7zOYfwM,27278
+sae_lens/constants.py,sha256=RJlzWx7wLNMNmrdI63naF7-M3enb55vYRN4x1hXx6vI,593
+sae_lens/evals.py,sha256=PIMGQobE9o2bHksDAtQe5bnTMYyHoZKB_elFhDOjrmo,38991
+sae_lens/llm_sae_training_runner.py,sha256=58XbDylw2fPOD7C-ZfSAjeNqJLXB05uHGTuiYVVbXXY,13354
+sae_lens/load_model.py,sha256=dBB_9gO6kWyQ4sXHq7qB8T3YUlXm3PGwYcpR4UVW4QY,8633
+sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/loading/pretrained_sae_loaders.py,sha256=kbirwfCg4Ks9Cg3rt78bYxIHMhz5h015n0UTRJQLJY0,31291
+sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
+sae_lens/pretokenize_runner.py,sha256=0nHQq3s_d80VS8iVK4-e6y_orAYVO8c4RrLGtIDfK_E,6885
+sae_lens/pretrained_saes.yaml,sha256=C_z-7Lxz6ZIy2V-c-4Xw45eAQ926O9aGjocSNuki0xs,573557
+sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
+sae_lens/saes/__init__.py,sha256=v6mfeDzyGYtT6x5SszAQtkldTXwPE-V_iwOlrT_pDwQ,1008
+sae_lens/saes/gated_sae.py,sha256=0zd66bH04nsaGk3bxHk10hsZofa2GrFbMo15LOsuqgU,9233
+sae_lens/saes/jumprelu_sae.py,sha256=iwmPQJ4XpIxzgosty680u8Zj7x1uVZhM75kPOT3obi0,12060
+sae_lens/saes/sae.py,sha256=HAGkJAj_FIDzbSR1dsG8b2AyMq8UauUU_yx-LvdfjuE,37465
+sae_lens/saes/standard_sae.py,sha256=PfkGLsw_6La3PXHOQL0u7qQsaZsXCJqYCeCcRDj5n64,6274
+sae_lens/saes/topk_sae.py,sha256=kmry1FE1H06OvCfn84V-j2JfWGKcU5b2urwAq_Oq5j4,9893
+sae_lens/tokenization_and_batching.py,sha256=oUAscjy_LPOrOb8_Ty6eLAcZ0B3HB_wiWjWktgolhG0,4314
+sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/training/activation_scaler.py,sha256=seEE-2Qd2JMHxqgnsNWPt-DGtYGZxWPnOwCGuVNSOtI,1719
+sae_lens/training/activations_store.py,sha256=s3Qvztv2siuuXSuXEUDZYSKq1QQCsqsGXY767kv6grc,32609
+sae_lens/training/mixing_buffer.py,sha256=vDpYG5ZE70szDvBsRKcNHEES3h_WTKJ16qDYk5jPOVA,2015
+sae_lens/training/optim.py,sha256=KXdOym-Ly3f2aFbndRc0JEH0Wa7u1BE5ljxGN3YtouQ,6836
+sae_lens/training/sae_trainer.py,sha256=9K0VudwSTJp9OlCVzaU_ngZ0WlYNrN6-ozTCCAxR9_k,15421
+sae_lens/training/types.py,sha256=qSjmGzXf3MLalygG0psnVjmhX_mpLmL47MQtZfe7qxg,81
+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=mCwLAilGMVo8Scm7CIsCafU7GsfmBvCcjwmloI4Ly7Y,1718
+sae_lens-6.0.0rc4.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.0.0rc4.dist-info/METADATA,sha256=wOQMSV4yNlpgpGxuE4DI0-q4KzTRYOg1m9ZxpdCsNjk,5326
+sae_lens-6.0.0rc4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.0.0rc4.dist-info/RECORD,,