sae-lens 6.9.1__tar.gz → 6.10.0__tar.gz

{sae_lens-6.9.1 → sae_lens-6.10.0}/PKG-INFO

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

{sae_lens-6.9.1 → sae_lens-6.10.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.9.1"
+version = "6.10.0"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"

{sae_lens-6.9.1 → sae_lens-6.10.0}/sae_lens/__init__.py

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

{sae_lens-6.9.1 → sae_lens-6.10.0}/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-6.9.1 → sae_lens-6.10.0}/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-6.9.1 → sae_lens-6.10.0}/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-6.9.1 → sae_lens-6.10.0}/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)