sae-lens 5.9.0__py3-none-any.whl → 6.0.0rc1__py3-none-any.whl

sae_lens/__init__.py

@@ -1,10 +1,15 @@
 # ruff: noqa: E402
-__version__ = "5.9.0"
+__version__ = "6.0.0-rc.1"
 
 import logging
 
 logger = logging.getLogger(__name__)
 
+from sae_lens.saes.gated_sae import GatedSAE, GatedTrainingSAE
+from sae_lens.saes.jumprelu_sae import JumpReLUSAE, JumpReLUTrainingSAE
+from sae_lens.saes.standard_sae import StandardSAE, StandardTrainingSAE
+from sae_lens.saes.topk_sae import TopKSAE, TopKTrainingSAE
+
 from .analysis.hooked_sae_transformer import HookedSAETransformer
 from .cache_activations_runner import CacheActivationsRunner
 from .config import (
@@ -13,17 +18,26 @@ from .config import (
     PretokenizeRunnerConfig,
 )
 from .evals import run_evals
-from .pretokenize_runner import PretokenizeRunner, pretokenize_runner
-from .sae import SAE, SAEConfig
-from .sae_training_runner import SAETrainingRunner
-from .toolkit.pretrained_sae_loaders import (
+from .loading.pretrained_sae_loaders import (
     PretrainedSaeDiskLoader,
     PretrainedSaeHuggingfaceLoader,
 )
+from .pretokenize_runner import PretokenizeRunner, pretokenize_runner
+from .regsitry import register_sae_class, register_sae_training_class
+from .sae_training_runner import SAETrainingRunner
+from .saes.sae import SAE, SAEConfig, TrainingSAE, TrainingSAEConfig
 from .training.activations_store import ActivationsStore
-from .training.training_sae import TrainingSAE, TrainingSAEConfig
 from .training.upload_saes_to_huggingface import upload_saes_to_huggingface
 
+register_sae_class("standard", StandardSAE)
+register_sae_training_class("standard", StandardTrainingSAE)
+register_sae_class("gated", GatedSAE)
+register_sae_training_class("gated", GatedTrainingSAE)
+register_sae_class("topk", TopKSAE)
+register_sae_training_class("topk", TopKTrainingSAE)
+register_sae_class("jumprelu", JumpReLUSAE)
+register_sae_training_class("jumprelu", JumpReLUTrainingSAE)
+
 __all__ = [
     "SAE",
     "SAEConfig",
@@ -42,4 +56,6 @@ __all__ = [
     "upload_saes_to_huggingface",
     "PretrainedSaeHuggingfaceLoader",
     "PretrainedSaeDiskLoader",
+    "register_sae_class",
+    "register_sae_training_class",
 ]

sae_lens/analysis/hooked_sae_transformer.py

@@ -8,7 +8,7 @@ from transformer_lens.ActivationCache import ActivationCache
 from transformer_lens.hook_points import HookPoint  # Hooking utilities
 from transformer_lens.HookedTransformer import HookedTransformer
 
-from sae_lens.sae import SAE
+from sae_lens.saes.sae import SAE
 
 SingleLoss = Float[torch.Tensor, ""]  # Type alias for a single element tensor
 LossPerToken = Float[torch.Tensor, "batch pos-1"]
@@ -275,7 +275,7 @@ class HookedSAETransformer(HookedTransformer):
         .. code-block:: python
 
             from transformer_lens import HookedSAETransformer
-            from sae_lens.sae import SAE
+            from sae_lens.saes.sae import SAE
 
             model = HookedSAETransformer.from_pretrained('gpt2-small')
             sae_cfg = SAEConfig(...)

sae_lens/config.py

@@ -1,7 +1,8 @@
 import json
 import math
 import os
-from dataclasses import dataclass, field
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
 from typing import Any, Literal, cast
 
 import simple_parsing
@@ -53,6 +54,52 @@ def dict_field(default: dict[str, Any] | None, **kwargs: Any) -> Any:  # type: i
     return simple_parsing.helpers.dict_field(default, type=json_dict, **kwargs)
 
 
+@dataclass
+class LoggingConfig:
+    # WANDB
+    log_to_wandb: bool = True
+    log_activations_store_to_wandb: bool = False
+    log_optimizer_state_to_wandb: bool = False
+    wandb_project: str = "sae_lens_training"
+    wandb_id: str | None = None
+    run_name: str | None = None
+    wandb_entity: str | None = None
+    wandb_log_frequency: int = 10
+    eval_every_n_wandb_logs: int = 100  # logs every 100 steps.
+
+    def log(
+        self,
+        trainer: Any,  # avoid import cycle from importing SAETrainer
+        weights_path: Path | str,
+        cfg_path: Path | str,
+        sparsity_path: Path | str | None,
+        wandb_aliases: list[str] | None = None,
+    ) -> None:
+        # Avoid wandb saving errors such as:
+        #   ValueError: Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. Invalid name: sae_google/gemma-2b_etc
+        sae_name = trainer.sae.get_name().replace("/", "__")
+
+        # save model weights and cfg
+        model_artifact = wandb.Artifact(
+            sae_name,
+            type="model",
+            metadata=dict(trainer.cfg.__dict__),
+        )
+        model_artifact.add_file(str(weights_path))
+        model_artifact.add_file(str(cfg_path))
+        wandb.log_artifact(model_artifact, aliases=wandb_aliases)
+
+        # save log feature sparsity
+        sparsity_artifact = wandb.Artifact(
+            f"{sae_name}_log_feature_sparsity",
+            type="log_feature_sparsity",
+            metadata=dict(trainer.cfg.__dict__),
+        )
+        if sparsity_path is not None:
+            sparsity_artifact.add_file(str(sparsity_path))
+        wandb.log_artifact(sparsity_artifact)
+
+
 @dataclass
 class LanguageModelSAERunnerConfig:
     """
@@ -245,16 +292,7 @@ class LanguageModelSAERunnerConfig:
     n_eval_batches: int = 10
     eval_batch_size_prompts: int | None = None  # useful if evals cause OOM
 
-    # WANDB
-    log_to_wandb: bool = True
-    log_activations_store_to_wandb: bool = False
-    log_optimizer_state_to_wandb: bool = False
-    wandb_project: str = "mats_sae_training_language_model"
-    wandb_id: str | None = None
-    run_name: str | None = None
-    wandb_entity: str | None = None
-    wandb_log_frequency: int = 10
-    eval_every_n_wandb_logs: int = 100  # logs every 1000 steps.
+    logger: LoggingConfig = field(default_factory=LoggingConfig)
 
     # Misc
     resume: bool = False
@@ -310,8 +348,8 @@ class LanguageModelSAERunnerConfig:
             self.train_batch_size_tokens * self.context_size * self.n_batches_in_buffer
         )
 
-        if self.run_name is None:
-            self.run_name = f"{self.d_sae}-L1-{self.l1_coefficient}-LR-{self.lr}-Tokens-{self.training_tokens:3.3e}"
+        if self.logger.run_name is None:
+            self.logger.run_name = f"{self.d_sae}-L1-{self.l1_coefficient}-LR-{self.lr}-Tokens-{self.training_tokens:3.3e}"
 
         if self.model_from_pretrained_kwargs is None:
             if self.model_class_name == "HookedTransformer":
@@ -356,7 +394,7 @@ class LanguageModelSAERunnerConfig:
         if self.lr_end is None:
             self.lr_end = self.lr / 10
 
-        unique_id = self.wandb_id
+        unique_id = self.logger.wandb_id
         if unique_id is None:
             unique_id = cast(
                 Any, wandb
@@ -388,7 +426,9 @@ class LanguageModelSAERunnerConfig:
             ) // self.train_batch_size_tokens
             logger.info(f"Total training steps: {total_training_steps}")
 
-            total_wandb_updates = total_training_steps // self.wandb_log_frequency
+            total_wandb_updates = (
+                total_training_steps // self.logger.wandb_log_frequency
+            )
             logger.info(f"Total wandb updates: {total_wandb_updates}")
 
             # how many times will we sample dead neurons?
@@ -445,7 +485,7 @@ class LanguageModelSAERunnerConfig:
             "hook_name": self.hook_name,
             "hook_layer": self.hook_layer,
             "hook_head_index": self.hook_head_index,
-            "activation_fn_str": self.activation_fn,
+            "activation_fn": self.activation_fn,
             "apply_b_dec_to_input": self.apply_b_dec_to_input,
             "context_size": self.context_size,
             "prepend_bos": self.prepend_bos,
@@ -478,13 +518,16 @@ class LanguageModelSAERunnerConfig:
         }
 
     def to_dict(self) -> dict[str, Any]:
-        return {
-            **self.__dict__,
-            # some args may not be serializable by default
-            "dtype": str(self.dtype),
-            "device": str(self.device),
-            "act_store_device": str(self.act_store_device),
-        }
+        # Make a shallow copy of config’s dictionary
+        d = dict(self.__dict__)
+
+        d["logger"] = asdict(self.logger)
+
+        # Overwrite fields that might not be JSON-serializable
+        d["dtype"] = str(self.dtype)
+        d["device"] = str(self.device)
+        d["act_store_device"] = str(self.act_store_device)
+        return d
 
     def to_json(self, path: str) -> None:
         if not os.path.exists(os.path.dirname(path)):

sae_lens/evals.py

@@ -19,8 +19,8 @@ from tqdm import tqdm
 from transformer_lens import HookedTransformer
 from transformer_lens.hook_points import HookedRootModule
 
-from sae_lens.sae import SAE
-from sae_lens.toolkit.pretrained_saes_directory import get_pretrained_saes_directory
+from sae_lens.loading.pretrained_saes_directory import get_pretrained_saes_directory
+from sae_lens.saes.sae import SAE
 from sae_lens.training.activations_store import ActivationsStore
 
 
@@ -279,7 +279,6 @@ def get_featurewise_weight_based_metrics(sae: SAE) -> dict[str, Any]:
     unit_norm_decoder = (sae.W_dec.T / sae.W_dec.T.norm(dim=0, keepdim=True)).cpu()
 
     encoder_norms = sae.W_enc.norm(dim=-2).cpu().tolist()
-    encoder_bias = sae.b_enc.cpu().tolist()
     encoder_decoder_cosine_sim = (
         torch.nn.functional.cosine_similarity(
             unit_norm_decoder.T,
@@ -289,11 +288,13 @@ def get_featurewise_weight_based_metrics(sae: SAE) -> dict[str, Any]:
         .tolist()
     )
 
-    return {
-        "encoder_bias": encoder_bias,
+    metrics = {
         "encoder_norm": encoder_norms,
         "encoder_decoder_cosine_sim": encoder_decoder_cosine_sim,
     }
+    if hasattr(sae, "b_enc") and sae.b_enc is not None:
+        metrics["encoder_bias"] = sae.b_enc.cpu().tolist()  # type: ignore
+    return metrics
 
 
 def get_downstream_reconstruction_metrics(

sae_lens/{toolkit → loading}/pretrained_sae_loaders.py

@@ -17,7 +17,7 @@ from sae_lens.config import (
     SAE_WEIGHTS_FILENAME,
     SPARSITY_FILENAME,
 )
-from sae_lens.toolkit.pretrained_saes_directory import (
+from sae_lens.loading.pretrained_saes_directory import (
     get_config_overrides,
     get_pretrained_saes_directory,
     get_repo_id_and_folder_name,
@@ -174,30 +174,38 @@ def get_sae_lens_config_from_disk(
 
 
 def handle_config_defaulting(cfg_dict: dict[str, Any]) -> dict[str, Any]:
+    rename_keys_map = {
+        "hook_point": "hook_name",
+        "hook_point_layer": "hook_layer",
+        "hook_point_head_index": "hook_head_index",
+        "activation_fn_str": "activation_fn",
+    }
+    new_cfg = {rename_keys_map.get(k, k): v for k, v in cfg_dict.items()}
+
     # Set default values for backwards compatibility
-    cfg_dict.setdefault("prepend_bos", True)
-    cfg_dict.setdefault("dataset_trust_remote_code", True)
-    cfg_dict.setdefault("apply_b_dec_to_input", True)
-    cfg_dict.setdefault("finetuning_scaling_factor", False)
-    cfg_dict.setdefault("sae_lens_training_version", None)
-    cfg_dict.setdefault("activation_fn_str", cfg_dict.get("activation_fn", "relu"))
-    cfg_dict.setdefault("architecture", "standard")
-    cfg_dict.setdefault("neuronpedia_id", None)
-
-    if "normalize_activations" in cfg_dict and isinstance(
-        cfg_dict["normalize_activations"], bool
+    new_cfg.setdefault("prepend_bos", True)
+    new_cfg.setdefault("dataset_trust_remote_code", True)
+    new_cfg.setdefault("apply_b_dec_to_input", True)
+    new_cfg.setdefault("finetuning_scaling_factor", False)
+    new_cfg.setdefault("sae_lens_training_version", None)
+    new_cfg.setdefault("activation_fn", new_cfg.get("activation_fn", "relu"))
+    new_cfg.setdefault("architecture", "standard")
+    new_cfg.setdefault("neuronpedia_id", None)
+
+    if "normalize_activations" in new_cfg and isinstance(
+        new_cfg["normalize_activations"], bool
     ):
         # backwards compatibility
-        cfg_dict["normalize_activations"] = (
+        new_cfg["normalize_activations"] = (
             "none"
-            if not cfg_dict["normalize_activations"]
+            if not new_cfg["normalize_activations"]
             else "expected_average_only_in"
         )
 
-    cfg_dict.setdefault("normalize_activations", "none")
-    cfg_dict.setdefault("device", "cpu")
+    new_cfg.setdefault("normalize_activations", "none")
+    new_cfg.setdefault("device", "cpu")
 
-    return cfg_dict
+    return new_cfg
 
 
 def get_connor_rob_hook_z_config_from_hf(
@@ -223,7 +231,7 @@ def get_connor_rob_hook_z_config_from_hf(
         "hook_name": old_cfg_dict["act_name"],
         "hook_layer": old_cfg_dict["layer"],
         "hook_head_index": None,
-        "activation_fn_str": "relu",
+        "activation_fn": "relu",
         "apply_b_dec_to_input": True,
         "finetuning_scaling_factor": False,
         "sae_lens_training_version": None,
@@ -372,7 +380,7 @@ def get_gemma_2_config_from_hf(
         "hook_name": hook_name,
         "hook_layer": layer,
         "hook_head_index": None,
-        "activation_fn_str": "relu",
+        "activation_fn": "relu",
         "finetuning_scaling_factor": False,
         "sae_lens_training_version": None,
         "prepend_bos": True,
@@ -485,7 +493,7 @@ def get_llama_scope_config_from_hf(
         "hook_name": old_cfg_dict["hook_point_in"],
         "hook_layer": int(old_cfg_dict["hook_point_in"].split(".")[1]),
         "hook_head_index": None,
-        "activation_fn_str": "relu",
+        "activation_fn": "relu",
         "finetuning_scaling_factor": False,
         "sae_lens_training_version": None,
         "prepend_bos": True,
@@ -597,8 +605,8 @@ def get_dictionary_learning_config_1_from_hf(
 
     hook_point_name = f"blocks.{trainer['layer']}.hook_resid_post"
 
-    activation_fn_str = "topk" if trainer["dict_class"] == "AutoEncoderTopK" else "relu"
-    activation_fn_kwargs = {"k": trainer["k"]} if activation_fn_str == "topk" else {}
+    activation_fn = "topk" if trainer["dict_class"] == "AutoEncoderTopK" else "relu"
+    activation_fn_kwargs = {"k": trainer["k"]} if activation_fn == "topk" else {}
 
     return {
         "architecture": (
@@ -612,7 +620,7 @@ def get_dictionary_learning_config_1_from_hf(
         "hook_name": hook_point_name,
         "hook_layer": trainer["layer"],
         "hook_head_index": None,
-        "activation_fn_str": activation_fn_str,
+        "activation_fn": activation_fn,
         "activation_fn_kwargs": activation_fn_kwargs,
         "apply_b_dec_to_input": True,
         "finetuning_scaling_factor": False,
@@ -655,7 +663,7 @@ def get_deepseek_r1_config_from_hf(
         "dataset_path": "lmsys/lmsys-chat-1m",
         "dataset_trust_remote_code": True,
         "sae_lens_training_version": None,
-        "activation_fn_str": "relu",
+        "activation_fn": "relu",
         "normalize_activations": "none",
         "device": device,
         "apply_b_dec_to_input": False,
@@ -810,7 +818,7 @@ def get_llama_scope_r1_distill_config_from_hf(
         "hook_name": huggingface_cfg_dict["hook_point_in"],
         "hook_layer": int(huggingface_cfg_dict["hook_point_in"].split(".")[1]),
         "hook_head_index": None,
-        "activation_fn_str": "relu",
+        "activation_fn": "relu",
         "finetuning_scaling_factor": False,
         "sae_lens_training_version": None,
         "prepend_bos": True,

sae_lens/regsitry.py

@@ -0,0 +1,34 @@
+from typing import TYPE_CHECKING
+
+# avoid circular imports
+if TYPE_CHECKING:
+    from sae_lens.saes.sae import SAE, TrainingSAE
+
+SAE_CLASS_REGISTRY: dict[str, "type[SAE]"] = {}
+SAE_TRAINING_CLASS_REGISTRY: dict[str, "type[TrainingSAE]"] = {}
+
+
+def register_sae_class(architecture: str, sae_class: "type[SAE]") -> None:
+    if architecture in SAE_CLASS_REGISTRY:
+        raise ValueError(
+            f"SAE class for architecture {architecture} already registered."
+        )
+    SAE_CLASS_REGISTRY[architecture] = sae_class
+
+
+def register_sae_training_class(
+    architecture: str, sae_training_class: "type[TrainingSAE]"
+) -> None:
+    if architecture in SAE_TRAINING_CLASS_REGISTRY:
+        raise ValueError(
+            f"SAE training class for architecture {architecture} already registered."
+        )
+    SAE_TRAINING_CLASS_REGISTRY[architecture] = sae_training_class
+
+
+def get_sae_class(architecture: str) -> "type[SAE]":
+    return SAE_CLASS_REGISTRY[architecture]
+
+
+def get_sae_training_class(architecture: str) -> "type[TrainingSAE]":
+    return SAE_TRAINING_CLASS_REGISTRY[architecture]

sae_lens/sae_training_runner.py

@@ -13,10 +13,10 @@ from transformer_lens.hook_points import HookedRootModule
 from sae_lens import logger
 from sae_lens.config import HfDataset, LanguageModelSAERunnerConfig
 from sae_lens.load_model import load_model
+from sae_lens.saes.sae import TrainingSAE, TrainingSAEConfig
 from sae_lens.training.activations_store import ActivationsStore
 from sae_lens.training.geometric_median import compute_geometric_median
 from sae_lens.training.sae_trainer import SAETrainer
-from sae_lens.training.training_sae import TrainingSAE, TrainingSAEConfig
 
 
 class InterruptedException(Exception):
@@ -73,14 +73,14 @@ class SAETrainingRunner:
 
         if override_sae is None:
             if self.cfg.from_pretrained_path is not None:
-                self.sae = TrainingSAE.load_from_pretrained(
+                self.sae = TrainingSAE.load_from_disk(
                     self.cfg.from_pretrained_path, self.cfg.device
                 )
             else:
-                self.sae = TrainingSAE(
+                self.sae = TrainingSAE.from_dict(
                     TrainingSAEConfig.from_dict(
                         self.cfg.get_training_sae_cfg_dict(),
-                    )
+                    ).to_dict()
                 )
                 self._init_sae_group_b_decs()
         else:
@@ -91,13 +91,13 @@ class SAETrainingRunner:
         Run the training of the SAE.
         """
 
-        if self.cfg.log_to_wandb:
+        if self.cfg.logger.log_to_wandb:
             wandb.init(
-                project=self.cfg.wandb_project,
-                entity=self.cfg.wandb_entity,
+                project=self.cfg.logger.wandb_project,
+                entity=self.cfg.logger.wandb_entity,
                 config=cast(Any, self.cfg),
-                name=self.cfg.run_name,
-                id=self.cfg.wandb_id,
+                name=self.cfg.logger.run_name,
+                id=self.cfg.logger.wandb_id,
             )
 
         trainer = SAETrainer(
@@ -111,7 +111,7 @@ class SAETrainingRunner:
         self._compile_if_needed()
         sae = self.run_trainer_with_interruption_handling(trainer)
 
-        if self.cfg.log_to_wandb:
+        if self.cfg.logger.log_to_wandb:
             wandb.finish()
 
         return sae
@@ -175,7 +175,7 @@ class SAETrainingRunner:
                 layer_acts,
                 maxiter=100,
             ).median
-            self.sae.initialize_b_dec_with_precalculated(median)  # type: ignore
+            self.sae.initialize_b_dec_with_precalculated(median)
         elif self.cfg.b_dec_init_method == "mean":
             self.activations_store.set_norm_scaling_factor_if_needed()
             layer_acts = self.activations_store.storage_buffer.detach().cpu()[:, 0, :]
@@ -208,29 +208,14 @@ class SAETrainingRunner:
         with open(cfg_path, "w") as f:
             json.dump(config, f)
 
-        if trainer.cfg.log_to_wandb:
-            # Avoid wandb saving errors such as:
-            #   ValueError: Artifact name may only contain alphanumeric characters, dashes, underscores, and dots. Invalid name: sae_google/gemma-2b_etc
-            sae_name = trainer.sae.get_name().replace("/", "__")
-
-            # save model weights and cfg
-            model_artifact = wandb.Artifact(
-                sae_name,
-                type="model",
-                metadata=dict(trainer.cfg.__dict__),
-            )
-            model_artifact.add_file(str(weights_path))
-            model_artifact.add_file(str(cfg_path))
-            wandb.log_artifact(model_artifact, aliases=wandb_aliases)
-
-            # save log feature sparsity
-            sparsity_artifact = wandb.Artifact(
-                f"{sae_name}_log_feature_sparsity",
-                type="log_feature_sparsity",
-                metadata=dict(trainer.cfg.__dict__),
+        if trainer.cfg.logger.log_to_wandb:
+            trainer.cfg.logger.log(
+                trainer,
+                weights_path,
+                cfg_path,
+                sparsity_path=sparsity_path,
+                wandb_aliases=wandb_aliases,
             )
-            sparsity_artifact.add_file(str(sparsity_path))
-            wandb.log_artifact(sparsity_artifact)
 
 
 def _parse_cfg_args(args: Sequence[str]) -> LanguageModelSAERunnerConfig: