sae-lens 6.9.1__py3-none-any.whl → 6.11.0__py3-none-any.whl

sae_lens/__init__.py

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

sae_lens/config.py

@@ -169,8 +169,10 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
         eval_batch_size_prompts (int, optional): The batch size for evaluation, in prompts. Useful if evals cause OOM.
         logger (LoggingConfig): Configuration for logging (e.g. W&B).
         n_checkpoints (int): The number of checkpoints to save during training. 0 means no checkpoints.
-        checkpoint_path (str): The path to save checkpoints. A unique ID will be appended to this path.
-        verbose (bool): Whether to print verbose output.
+        checkpoint_path (str | None): The path to save checkpoints. A unique ID will be appended to this path. Set to None to disable checkpoint saving. (default is "checkpoints")
+        save_final_checkpoint (bool): Whether to include an additional final checkpoint when training is finished. (default is False).
+        output_path (str | None): The path to save outputs. Set to None to disable output saving. (default is "output")
+        verbose (bool): Whether to print verbose output. (default is True)
         model_kwargs (dict[str, Any]): Keyword arguments for `model.run_with_cache`
         model_from_pretrained_kwargs (dict[str, Any], optional): Additional keyword arguments to pass to the model's `from_pretrained` method.
         sae_lens_version (str): The version of the sae_lens library.
@@ -254,9 +256,13 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
 
     logger: LoggingConfig = field(default_factory=LoggingConfig)
 
-    # Misc
+    # Outputs/Checkpoints
     n_checkpoints: int = 0
-    checkpoint_path: str = "checkpoints"
+    checkpoint_path: str | None = "checkpoints"
+    save_final_checkpoint: bool = False
+    output_path: str | None = "output"
+
+    # Misc
     verbose: bool = True
     model_kwargs: dict[str, Any] = dict_field(default={})
     model_from_pretrained_kwargs: dict[str, Any] | None = dict_field(default=None)
@@ -394,6 +400,7 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
         return SAETrainerConfig(
             n_checkpoints=self.n_checkpoints,
             checkpoint_path=self.checkpoint_path,
+            save_final_checkpoint=self.save_final_checkpoint,
             total_training_samples=self.total_training_tokens,
             device=self.device,
             autocast=self.autocast,
@@ -618,7 +625,8 @@ class PretokenizeRunnerConfig:
 @dataclass
 class SAETrainerConfig:
     n_checkpoints: int
-    checkpoint_path: str
+    checkpoint_path: str | None
+    save_final_checkpoint: bool
     total_training_samples: int
     device: str
     autocast: bool

sae_lens/llm_sae_training_runner.py

@@ -8,13 +8,18 @@ from typing import Any, Generic
 
 import torch
 import wandb
+from safetensors.torch import save_file
 from simple_parsing import ArgumentParser
 from transformer_lens.hook_points import HookedRootModule
 from typing_extensions import deprecated
 
 from sae_lens import logger
 from sae_lens.config import HfDataset, LanguageModelSAERunnerConfig
-from sae_lens.constants import ACTIVATIONS_STORE_STATE_FILENAME, RUNNER_CFG_FILENAME
+from sae_lens.constants import (
+    ACTIVATIONS_STORE_STATE_FILENAME,
+    RUNNER_CFG_FILENAME,
+    SPARSITY_FILENAME,
+)
 from sae_lens.evals import EvalConfig, run_evals
 from sae_lens.load_model import load_model
 from sae_lens.saes.batchtopk_sae import BatchTopKTrainingSAEConfig
@@ -185,11 +190,47 @@ class LanguageModelSAETrainingRunner:
         self._compile_if_needed()
         sae = self.run_trainer_with_interruption_handling(trainer)
 
+        if self.cfg.output_path is not None:
+            self.save_final_sae(
+                sae=sae,
+                output_path=self.cfg.output_path,
+                log_feature_sparsity=trainer.log_feature_sparsity,
+            )
+
         if self.cfg.logger.log_to_wandb:
             wandb.finish()
 
         return sae
 
+    def save_final_sae(
+        self,
+        sae: TrainingSAE[Any],
+        output_path: str,
+        log_feature_sparsity: torch.Tensor | None = None,
+    ):
+        base_output_path = Path(output_path)
+        base_output_path.mkdir(exist_ok=True, parents=True)
+
+        weights_path, cfg_path = sae.save_inference_model(str(base_output_path))
+
+        sparsity_path = None
+        if log_feature_sparsity is not None:
+            sparsity_path = base_output_path / SPARSITY_FILENAME
+            save_file({"sparsity": log_feature_sparsity}, sparsity_path)
+
+        runner_config = self.cfg.to_dict()
+        with open(base_output_path / RUNNER_CFG_FILENAME, "w") as f:
+            json.dump(runner_config, f)
+
+        if self.cfg.logger.log_to_wandb:
+            self.cfg.logger.log(
+                self,
+                weights_path,
+                cfg_path,
+                sparsity_path=sparsity_path,
+                wandb_aliases=["final_model"],
+            )
+
     def _set_sae_metadata(self):
         self.sae.cfg.metadata.dataset_path = self.cfg.dataset_path
         self.sae.cfg.metadata.hook_name = self.cfg.hook_name
@@ -247,20 +288,24 @@ class LanguageModelSAETrainingRunner:
             sae = trainer.fit()
 
         except (KeyboardInterrupt, InterruptedException):
-            logger.warning("interrupted, saving progress")
-            checkpoint_path = Path(self.cfg.checkpoint_path) / str(
-                trainer.n_training_samples
-            )
-            self.save_checkpoint(checkpoint_path)
-            logger.info("done saving")
+            if self.cfg.checkpoint_path is not None:
+                logger.warning("interrupted, saving progress")
+                checkpoint_path = Path(self.cfg.checkpoint_path) / str(
+                    trainer.n_training_samples
+                )
+                self.save_checkpoint(checkpoint_path)
+                logger.info("done saving")
             raise
 
         return sae
 
     def save_checkpoint(
         self,
-        checkpoint_path: Path,
+        checkpoint_path: Path | None,
     ) -> None:
+        if checkpoint_path is None:
+            return
+
         self.activations_store.save(
             str(checkpoint_path / ACTIVATIONS_STORE_STATE_FILENAME)
         )

sae_lens/saes/jumprelu_sae.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import torch
@@ -187,13 +187,29 @@ class JumpReLUSAE(SAE[JumpReLUSAEConfig]):
 class JumpReLUTrainingSAEConfig(TrainingSAEConfig):
     """
     Configuration class for training a JumpReLUTrainingSAE.
+
+    - jumprelu_init_threshold: initial threshold for the JumpReLU activation
+    - jumprelu_bandwidth: bandwidth for the JumpReLU activation
+    - jumprelu_sparsity_loss_mode: mode for the sparsity loss, either "step" or "tanh". "step" is Google Deepmind's L0 loss, "tanh" is Anthropic's sparsity loss.
+    - l0_coefficient: coefficient for the l0 sparsity loss
+    - l0_warm_up_steps: number of warm-up steps for the l0 sparsity loss
+    - pre_act_loss_coefficient: coefficient for the pre-activation loss. Set to None to disable. Set to 3e-6 to match Anthropic's setup. Default is None.
+    - jumprelu_tanh_scale: scale for the tanh sparsity loss. Only relevant for "tanh" sparsity loss mode. Default is 4.0.
     """
 
     jumprelu_init_threshold: float = 0.01
     jumprelu_bandwidth: float = 0.05
+    # step is Google Deepmind, tanh is Anthropic
+    jumprelu_sparsity_loss_mode: Literal["step", "tanh"] = "step"
     l0_coefficient: float = 1.0
     l0_warm_up_steps: int = 0
 
+    # anthropic's auxiliary loss to avoid dead features
+    pre_act_loss_coefficient: float | None = None
+
+    # only relevant for tanh sparsity loss mode
+    jumprelu_tanh_scale: float = 4.0
+
     @override
     @classmethod
     def architecture(cls) -> str:
@@ -267,9 +283,35 @@ class JumpReLUTrainingSAE(TrainingSAE[JumpReLUTrainingSAEConfig]):
         sae_out: torch.Tensor,
     ) -> dict[str, torch.Tensor]:
         """Calculate architecture-specific auxiliary loss terms."""
-        l0 = torch.sum(Step.apply(hidden_pre, self.threshold, self.bandwidth), dim=-1)  # type: ignore
-        l0_loss = (step_input.coefficients["l0"] * l0).mean()
-        return {"l0_loss": l0_loss}
+
+        threshold = self.threshold
+        W_dec_norm = self.W_dec.norm(dim=1)
+        if self.cfg.jumprelu_sparsity_loss_mode == "step":
+            l0 = torch.sum(
+                Step.apply(hidden_pre, threshold, self.bandwidth),  # type: ignore
+                dim=-1,
+            )
+            l0_loss = (step_input.coefficients["l0"] * l0).mean()
+        elif self.cfg.jumprelu_sparsity_loss_mode == "tanh":
+            per_item_l0_loss = torch.tanh(
+                self.cfg.jumprelu_tanh_scale * feature_acts * W_dec_norm
+            ).sum(dim=-1)
+            l0_loss = (step_input.coefficients["l0"] * per_item_l0_loss).mean()
+        else:
+            raise ValueError(
+                f"Invalid sparsity loss mode: {self.cfg.jumprelu_sparsity_loss_mode}"
+            )
+        losses = {"l0_loss": l0_loss}
+
+        if self.cfg.pre_act_loss_coefficient is not None:
+            losses["pre_act_loss"] = calculate_pre_act_loss(
+                self.cfg.pre_act_loss_coefficient,
+                threshold,
+                hidden_pre,
+                step_input.dead_neuron_mask,
+                W_dec_norm,
+            )
+        return losses
 
     @override
     def get_coefficients(self) -> dict[str, float | TrainCoefficientConfig]:
@@ -310,3 +352,21 @@ class JumpReLUTrainingSAE(TrainingSAE[JumpReLUTrainingSAEConfig]):
             threshold = state_dict["threshold"]
             del state_dict["threshold"]
             state_dict["log_threshold"] = torch.log(threshold).detach().contiguous()
+
+
+def calculate_pre_act_loss(
+    pre_act_loss_coefficient: float,
+    threshold: torch.Tensor,
+    hidden_pre: torch.Tensor,
+    dead_neuron_mask: torch.Tensor | None,
+    W_dec_norm: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Calculate Anthropic's pre-activation loss, except we only calculate this for latents that are actually dead.
+    """
+    if dead_neuron_mask is None or not dead_neuron_mask.any():
+        return hidden_pre.new_tensor(0.0)
+    per_item_loss = (
+        (threshold - hidden_pre).relu() * dead_neuron_mask * W_dec_norm
+    ).sum(dim=-1)
+    return pre_act_loss_coefficient * per_item_loss.mean()

sae_lens/training/sae_trainer.py

@@ -22,6 +22,7 @@ from sae_lens.saes.sae import (
 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
+from sae_lens.util import path_or_tmp_dir
 
 
 def _log_feature_sparsity(
@@ -40,7 +41,7 @@ def _update_sae_lens_training_version(sae: TrainingSAE[Any]) -> None:
 class SaveCheckpointFn(Protocol):
     def __call__(
         self,
-        checkpoint_path: Path,
+        checkpoint_path: Path | None,
     ) -> None: ...
 
 
@@ -187,12 +188,8 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             )
             self.activation_scaler.scaling_factor = None
 
-        # save final inference sae group to checkpoints folder
-        self.save_checkpoint(
-            checkpoint_name=f"final_{self.n_training_samples}",
-            wandb_aliases=["final_model"],
-            save_inference_model=True,
-        )
+        if self.cfg.save_final_checkpoint:
+            self.save_checkpoint(checkpoint_name=f"final_{self.n_training_samples}")
 
         pbar.close()
         return self.sae
@@ -201,32 +198,31 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
         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))
+        checkpoint_path = None
+        if self.cfg.checkpoint_path is not None or self.cfg.logger.log_to_wandb:
+            with path_or_tmp_dir(self.cfg.checkpoint_path) as base_checkpoint_path:
+                checkpoint_path = base_checkpoint_path / checkpoint_name
+                checkpoint_path.mkdir(exist_ok=True, parents=True)
 
-        sparsity_path = checkpoint_path / SPARSITY_FILENAME
-        save_file({"sparsity": self.log_feature_sparsity}, sparsity_path)
+                weights_path, cfg_path = self.sae.save_model(str(checkpoint_path))
 
-        activation_scaler_path = checkpoint_path / ACTIVATION_SCALER_CFG_FILENAME
-        self.activation_scaler.save(str(activation_scaler_path))
+                sparsity_path = checkpoint_path / SPARSITY_FILENAME
+                save_file({"sparsity": self.log_feature_sparsity}, sparsity_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,
-            )
+                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)

sae_lens/util.py

@@ -1,5 +1,8 @@
 import re
+import tempfile
+from contextlib import contextmanager
 from dataclasses import asdict, fields, is_dataclass
+from pathlib import Path
 from typing import Sequence, TypeVar
 
 K = TypeVar("K")
@@ -45,3 +48,18 @@ def extract_layer_from_tlens_hook_name(hook_name: str) -> int | None:
     """
     hook_match = re.search(r"\.(\d+)\.", hook_name)
     return None if hook_match is None else int(hook_match.group(1))
+
+
+@contextmanager
+def path_or_tmp_dir(path: str | Path | None):
+    """Context manager that yields a concrete Path for path.
+
+    - If path is None, creates a TemporaryDirectory and yields its Path.
+      The directory is cleaned up on context exit.
+    - Otherwise, yields Path(path) without creating or cleaning.
+    """
+    if path is None:
+        with tempfile.TemporaryDirectory() as td:
+            yield Path(td)
+    else:
+        yield Path(path)

{sae_lens-6.9.1.dist-info → sae_lens-6.11.0.dist-info}/METADATA

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

{sae_lens-6.9.1.dist-info → sae_lens-6.11.0.dist-info}/RECORD

@@ -1,12 +1,12 @@
-sae_lens/__init__.py,sha256=YTdinwk5L_K9nK5pBJIKiJTT_bLudiY_uvF_osio3gs,3588
+sae_lens/__init__.py,sha256=87qzm9CROR5ywD55DSxpt_ut24G4O8MR2zGYDtHEA0Y,3589
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=vRu6JseH1lZaEeILD5bEkQEQ1wYHHDcxD-f2olKmE9Y,14275
 sae_lens/analysis/neuronpedia_integration.py,sha256=Gx1W7hUBEuMoasNcnOnZ1wmqbXDd1pSZ1nqKEya1HQc,4962
 sae_lens/cache_activations_runner.py,sha256=cNeAtp2JQ_vKbeddZVM-tcPLYyyfTWL8NDna5KQpkLI,12583
-sae_lens/config.py,sha256=IrjbsKBbaZoFXYrsPJ5xBwIqi9uZJIIFXjV_uoErJaE,28176
+sae_lens/config.py,sha256=IdRXSKPfYY3hwUovj-u83eep8z52gkJHII0mY0KseYY,28739
 sae_lens/constants.py,sha256=CSjmiZ-bhjQeVLyRvWxAjBokCgkfM8mnvd7-vxLIWTY,639
 sae_lens/evals.py,sha256=4hanbyG8qZLItWqft94F4ZjUoytPVB7fw5s0P4Oi0VE,39504
-sae_lens/llm_sae_training_runner.py,sha256=exxNX_OEhdiUrlgmBP9bjX9DOf0HUcNQGO4unKeDjKM,13713
+sae_lens/llm_sae_training_runner.py,sha256=sJTcDX1bUJJ_jZLUT88-8KUYIAPeUGoXktX68PsBqw0,15137
 sae_lens/load_model.py,sha256=C8AMykctj6H7tz_xRwB06-EXj6TfW64PtSJZR5Jxn1Y,8649
 sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/loading/pretrained_sae_loaders.py,sha256=CVzHntSUKR1X3_gAqn8K_Ajq8D85qBrmrgEgU93IV4A,49609
@@ -17,7 +17,7 @@ sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
 sae_lens/saes/__init__.py,sha256=jVwazK8Q6dW5J6_zFXPoNAuBvSxgziQ8eMOjGM3t-X8,1475
 sae_lens/saes/batchtopk_sae.py,sha256=CyaFG2hMyyDaEaXXrAMJC8wQDW1JoddTKF5mvxxBQKY,3395
 sae_lens/saes/gated_sae.py,sha256=qcmM9JwBA8aZR8z_IRHV1_gQX-q_63tKewWXRnhdXuo,8986
-sae_lens/saes/jumprelu_sae.py,sha256=3xkhBcCol2mEpIBLceymCpudocm2ypOjTeTXbpiXoA4,10794
+sae_lens/saes/jumprelu_sae.py,sha256=HHBF1sJ95lZvxwP5vwLSQFKdnJN2KKYK0WAEaLTrta0,13399
 sae_lens/saes/sae.py,sha256=gdUZuLaOHQrPjbDj-nZI813B6-_mNAnV9i9z4qTnpHk,38255
 sae_lens/saes/standard_sae.py,sha256=9UqYyYtQuThYxXKNaDjYcyowpOx2-7cShG-TeUP6JCQ,5940
 sae_lens/saes/topk_sae.py,sha256=CXMBI6CFvI5829bOhoQ350VXR9d8uFHUDlULTIWHXoU,8686
@@ -28,12 +28,12 @@ sae_lens/training/activation_scaler.py,sha256=seEE-2Qd2JMHxqgnsNWPt-DGtYGZxWPnOw
 sae_lens/training/activations_store.py,sha256=2EUY2abqpT5El3T95sypM_JRDgiKL3VeT73U9SQIFGY,32903
 sae_lens/training/mixing_buffer.py,sha256=vDpYG5ZE70szDvBsRKcNHEES3h_WTKJ16qDYk5jPOVA,2015
 sae_lens/training/optim.py,sha256=TiI9nbffzXNsI8WjcIsqa2uheW6suxqL_KDDmWXobWI,5312
-sae_lens/training/sae_trainer.py,sha256=6HPf5wtmY1wMUTkLFRg9DujNMMXJkVMPdAhB2svvlkk,15368
+sae_lens/training/sae_trainer.py,sha256=Jh5AyBGtfZjnprv9H3k0p_luWWnM7YFjlmHuO1W_J6U,15465
 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.9.1.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.9.1.dist-info/METADATA,sha256=ELrgBpTtDALXmLtQuJSS7-ui_pf-DJbAyL4uUr5N-UU,5244
-sae_lens-6.9.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-6.9.1.dist-info/RECORD,,
+sae_lens/util.py,sha256=lW7fBn_b8quvRYlen9PUmB7km60YhKyjmuelB1f6KzQ,2253
+sae_lens-6.11.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.11.0.dist-info/METADATA,sha256=yEZL6kzNBUKjzd0d73Hiup6SpLNtNEGUOgFu0MwgVpo,5245
+sae_lens-6.11.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.11.0.dist-info/RECORD,,