sae-lens 6.9.0__py3-none-any.whl → 6.10.0__py3-none-any.whl

sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.9.0"
+__version__ = "6.10.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/loading/pretrained_sae_loaders.py

@@ -1,12 +1,14 @@
 import json
 import re
+import warnings
 from pathlib import Path
 from typing import Any, Protocol
 
 import numpy as np
+import requests
 import torch
 import yaml
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, hf_hub_url
 from huggingface_hub.utils import EntryNotFoundError
 from packaging.version import Version
 from safetensors import safe_open
@@ -1330,6 +1332,48 @@ def mwhanna_transcoder_huggingface_loader(
     return cfg_dict, state_dict, None
 
 
+def get_safetensors_tensor_shapes(url: str) -> dict[str, list[int]]:
+    """
+    Get tensor shapes from a safetensors file using HTTP range requests
+    without downloading the entire file.
+
+    Args:
+        url: Direct URL to the safetensors file
+
+    Returns:
+        Dictionary mapping tensor names to their shapes
+    """
+    # Check if server supports range requests
+    response = requests.head(url, timeout=10)
+    response.raise_for_status()
+
+    accept_ranges = response.headers.get("Accept-Ranges", "")
+    if "bytes" not in accept_ranges:
+        raise ValueError("Server does not support range requests")
+
+    # Fetch first 8 bytes to get metadata size
+    headers = {"Range": "bytes=0-7"}
+    response = requests.get(url, headers=headers, timeout=10)
+    if response.status_code != 206:
+        raise ValueError("Failed to fetch initial bytes for metadata size")
+
+    meta_size = int.from_bytes(response.content, byteorder="little")
+
+    # Fetch the metadata header
+    headers = {"Range": f"bytes=8-{8 + meta_size - 1}"}
+    response = requests.get(url, headers=headers, timeout=10)
+    if response.status_code != 206:
+        raise ValueError("Failed to fetch metadata header")
+
+    metadata_json = response.content.decode("utf-8").strip()
+    metadata = json.loads(metadata_json)
+
+    # Extract tensor shapes, excluding the __metadata__ key
+    return {
+        name: info["shape"] for name, info in metadata.items() if name != "__metadata__"
+    }
+
+
 def mntss_clt_layer_huggingface_loader(
     repo_id: str,
     folder_name: str,
@@ -1341,11 +1385,20 @@ def mntss_clt_layer_huggingface_loader(
     Load a MNTSS CLT layer as a single layer transcoder.
     The assumption is that the `folder_name` is the layer to load as an int
     """
-    base_config_path = hf_hub_download(
-        repo_id, "config.yaml", force_download=force_download
+
+    # warn that this sums the decoders together, so should only be used to find feature activations, not for reconstruction
+    warnings.warn(
+        "This loads the CLT layer as a single layer transcoder by summing all decoders together. This should only be used to find feature activations, not for reconstruction",
+        UserWarning,
+    )
+
+    cfg_dict = get_mntss_clt_layer_config_from_hf(
+        repo_id,
+        folder_name,
+        device,
+        force_download,
+        cfg_overrides,
     )
-    with open(base_config_path) as f:
-        cfg_info: dict[str, Any] = yaml.safe_load(f)
 
     # We need to actually load the weights, since the config is missing most information
     encoder_path = hf_hub_download(
@@ -1370,11 +1423,39 @@ def mntss_clt_layer_huggingface_loader(
             "W_dec": decoder_state_dict[f"W_dec_{folder_name}"].sum(dim=1),  # type: ignore
         }
 
-    cfg_dict = {
+    return cfg_dict, state_dict, None
+
+
+def get_mntss_clt_layer_config_from_hf(
+    repo_id: str,
+    folder_name: str,
+    device: str,
+    force_download: bool = False,  # noqa: ARG001
+    cfg_overrides: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Load a MNTSS CLT layer as a single layer transcoder.
+    The assumption is that the `folder_name` is the layer to load as an int
+    """
+    base_config_path = hf_hub_download(
+        repo_id, "config.yaml", force_download=force_download
+    )
+    with open(base_config_path) as f:
+        cfg_info: dict[str, Any] = yaml.safe_load(f)
+
+    # Get tensor shapes without downloading full files using HTTP range requests
+    encoder_url = hf_hub_url(repo_id, f"W_enc_{folder_name}.safetensors")
+    encoder_shapes = get_safetensors_tensor_shapes(encoder_url)
+
+    # Extract shapes for the required tensors
+    b_dec_shape = encoder_shapes[f"b_dec_{folder_name}"]
+    b_enc_shape = encoder_shapes[f"b_enc_{folder_name}"]
+
+    return {
         "architecture": "transcoder",
-        "d_in": state_dict["b_dec"].shape[0],
-        "d_out": state_dict["b_dec"].shape[0],
-        "d_sae": state_dict["b_enc"].shape[0],
+        "d_in": b_dec_shape[0],
+        "d_out": b_dec_shape[0],
+        "d_sae": b_enc_shape[0],
         "dtype": "float32",
         "device": device if device is not None else "cpu",
         "activation_fn": "relu",
@@ -1387,8 +1468,6 @@ def mntss_clt_layer_huggingface_loader(
         **(cfg_overrides or {}),
     }
 
-    return cfg_dict, state_dict, None
-
 
 NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sae_lens": sae_lens_huggingface_loader,
@@ -1416,4 +1495,5 @@ NAMED_PRETRAINED_SAE_CONFIG_GETTERS: dict[str, PretrainedSaeConfigHuggingfaceLoa
     "sparsify": get_sparsify_config_from_hf,
     "gemma_2_transcoder": get_gemma_2_transcoder_config_from_hf,
     "mwhanna_transcoder": get_mwhanna_transcoder_config_from_hf,
+    "mntss_clt_layer_transcoder": get_mntss_clt_layer_config_from_hf,
 }

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.0.dist-info → sae_lens-6.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.9.0
+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.0.dist-info → sae_lens-6.10.0.dist-info}/RECORD

@@ -1,15 +1,15 @@
-sae_lens/__init__.py,sha256=TQtVpMEI_qo11x4o5w2cfOF9pqHeXXtyMCypr3lSRCI,3588
+sae_lens/__init__.py,sha256=k8M2SyKNE3KpipPxODICdLG8KJNVvf1Zab4KNJuGWMQ,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=IVtgxWN0w96ZORnWPYW2ndYWey7e5GpzlpedWF3NJ8k,46818
+sae_lens/loading/pretrained_sae_loaders.py,sha256=CVzHntSUKR1X3_gAqn8K_Ajq8D85qBrmrgEgU93IV4A,49609
 sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
 sae_lens/pretokenize_runner.py,sha256=w0f6SfZLAxbp5eAAKnet8RqUB_DKofZ9RGsoJwFnYbA,7058
 sae_lens/pretrained_saes.yaml,sha256=d6FYfWTdVAPlOCM55C1ICS6lF9nWPPVNwjlXCa9p7NU,600468
@@ -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.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.9.0.dist-info/METADATA,sha256=xqfc_6msZd2ZGEwnLFVBDAK2HdFrvllKwa78bx6iYhs,5244
-sae_lens-6.9.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-6.9.0.dist-info/RECORD,,
+sae_lens/util.py,sha256=lW7fBn_b8quvRYlen9PUmB7km60YhKyjmuelB1f6KzQ,2253
+sae_lens-6.10.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.10.0.dist-info/METADATA,sha256=7Yq4_hrZVc2CBB4nMvgy_BGFjT5FrF3SfOo8LnJ18Rg,5245
+sae_lens-6.10.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.10.0.dist-info/RECORD,,