sae-lens 5.11.0__py3-none-any.whl → 6.0.0__py3-none-any.whl

sae_lens/sae_training_runner.py

@@ -1,251 +0,0 @@
-import json
-import signal
-import sys
-from collections.abc import Sequence
-from pathlib import Path
-from typing import Any, cast
-
-import torch
-import wandb
-from simple_parsing import ArgumentParser
-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.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):
-    pass
-
-
-def interrupt_callback(sig_num: Any, stack_frame: Any):  # noqa: ARG001
-    raise InterruptedException()
-
-
-class SAETrainingRunner:
-    """
-    Class to run the training of a Sparse Autoencoder (SAE) on a TransformerLens model.
-    """
-
-    cfg: LanguageModelSAERunnerConfig
-    model: HookedRootModule
-    sae: TrainingSAE
-    activations_store: ActivationsStore
-
-    def __init__(
-        self,
-        cfg: LanguageModelSAERunnerConfig,
-        override_dataset: HfDataset | None = None,
-        override_model: HookedRootModule | None = None,
-        override_sae: TrainingSAE | None = None,
-    ):
-        if override_dataset is not None:
-            logger.warning(
-                f"You just passed in a dataset which will override the one specified in your configuration: {cfg.dataset_path}. As a consequence this run will not be reproducible via configuration alone."
-            )
-        if override_model is not None:
-            logger.warning(
-                f"You just passed in a model which will override the one specified in your configuration: {cfg.model_name}. As a consequence this run will not be reproducible via configuration alone."
-            )
-
-        self.cfg = cfg
-
-        if override_model is None:
-            self.model = load_model(
-                self.cfg.model_class_name,
-                self.cfg.model_name,
-                device=self.cfg.device,
-                model_from_pretrained_kwargs=self.cfg.model_from_pretrained_kwargs,
-            )
-        else:
-            self.model = override_model
-
-        self.activations_store = ActivationsStore.from_config(
-            self.model,
-            self.cfg,
-            override_dataset=override_dataset,
-        )
-
-        if override_sae is None:
-            if self.cfg.from_pretrained_path is not None:
-                self.sae = TrainingSAE.load_from_pretrained(
-                    self.cfg.from_pretrained_path, self.cfg.device
-                )
-            else:
-                self.sae = TrainingSAE(
-                    TrainingSAEConfig.from_dict(
-                        self.cfg.get_training_sae_cfg_dict(),
-                    )
-                )
-                self._init_sae_group_b_decs()
-        else:
-            self.sae = override_sae
-
-    def run(self):
-        """
-        Run the training of the SAE.
-        """
-
-        if self.cfg.log_to_wandb:
-            wandb.init(
-                project=self.cfg.wandb_project,
-                entity=self.cfg.wandb_entity,
-                config=cast(Any, self.cfg),
-                name=self.cfg.run_name,
-                id=self.cfg.wandb_id,
-            )
-
-        trainer = SAETrainer(
-            model=self.model,
-            sae=self.sae,
-            activation_store=self.activations_store,
-            save_checkpoint_fn=self.save_checkpoint,
-            cfg=self.cfg,
-        )
-
-        self._compile_if_needed()
-        sae = self.run_trainer_with_interruption_handling(trainer)
-
-        if self.cfg.log_to_wandb:
-            wandb.finish()
-
-        return sae
-
-    def _compile_if_needed(self):
-        # Compile model and SAE
-        #  torch.compile can provide significant speedups (10-20% in testing)
-        # using max-autotune gives the best speedups but:
-        # (a) increases VRAM usage,
-        # (b) can't be used on both SAE and LM (some issue with cudagraphs), and
-        # (c) takes some time to compile
-        # optimal settings seem to be:
-        # use max-autotune on SAE and max-autotune-no-cudagraphs on LM
-        # (also pylance seems to really hate this)
-        if self.cfg.compile_llm:
-            self.model = torch.compile(
-                self.model,
-                mode=self.cfg.llm_compilation_mode,
-            )  # type: ignore
-
-        if self.cfg.compile_sae:
-            backend = "aot_eager" if self.cfg.device == "mps" else "inductor"
-
-            self.sae.training_forward_pass = torch.compile(  # type: ignore
-                self.sae.training_forward_pass,
-                mode=self.cfg.sae_compilation_mode,
-                backend=backend,
-            )  # type: ignore
-
-    def run_trainer_with_interruption_handling(self, trainer: SAETrainer):
-        try:
-            # signal handlers (if preempted)
-            signal.signal(signal.SIGINT, interrupt_callback)
-            signal.signal(signal.SIGTERM, interrupt_callback)
-
-            # train SAE
-            sae = trainer.fit()
-
-        except (KeyboardInterrupt, InterruptedException):
-            logger.warning("interrupted, saving progress")
-            checkpoint_name = str(trainer.n_training_tokens)
-            self.save_checkpoint(trainer, checkpoint_name=checkpoint_name)
-            logger.info("done saving")
-            raise
-
-        return sae
-
-    # TODO: move this into the SAE trainer or Training SAE class
-    def _init_sae_group_b_decs(
-        self,
-    ) -> None:
-        """
-        extract all activations at a certain layer and use for sae b_dec initialization
-        """
-
-        if self.cfg.b_dec_init_method == "geometric_median":
-            self.activations_store.set_norm_scaling_factor_if_needed()
-            layer_acts = self.activations_store.storage_buffer.detach()[:, 0, :]
-            # get geometric median of the activations if we're using those.
-            median = compute_geometric_median(
-                layer_acts,
-                maxiter=100,
-            ).median
-            self.sae.initialize_b_dec_with_precalculated(median)  # type: ignore
-        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, :]
-            self.sae.initialize_b_dec_with_mean(layer_acts)  # type: ignore
-
-    @staticmethod
-    def save_checkpoint(
-        trainer: SAETrainer,
-        checkpoint_name: str,
-        wandb_aliases: list[str] | None = None,
-    ) -> None:
-        base_path = Path(trainer.cfg.checkpoint_path) / checkpoint_name
-        base_path.mkdir(exist_ok=True, parents=True)
-
-        trainer.activations_store.save(
-            str(base_path / "activations_store_state.safetensors")
-        )
-
-        if trainer.sae.cfg.normalize_sae_decoder:
-            trainer.sae.set_decoder_norm_to_unit_norm()
-
-        weights_path, cfg_path, sparsity_path = trainer.sae.save_model(
-            str(base_path),
-            trainer.log_feature_sparsity,
-        )
-
-        # let's over write the cfg file with the trainer cfg, which is a super set of the original cfg.
-        # and should not cause issues but give us more info about SAEs we trained in SAE Lens.
-        config = trainer.cfg.to_dict()
-        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__),
-            )
-            sparsity_artifact.add_file(str(sparsity_path))
-            wandb.log_artifact(sparsity_artifact)
-
-
-def _parse_cfg_args(args: Sequence[str]) -> LanguageModelSAERunnerConfig:
-    if len(args) == 0:
-        args = ["--help"]
-    parser = ArgumentParser(exit_on_error=False)
-    parser.add_arguments(LanguageModelSAERunnerConfig, dest="cfg")
-    return parser.parse_args(args).cfg
-
-
-# moved into its own function to make it easier to test
-def _run_cli(args: Sequence[str]):
-    cfg = _parse_cfg_args(args)
-    SAETrainingRunner(cfg=cfg).run()
-
-
-if __name__ == "__main__":
-    _run_cli(args=sys.argv[1:])

sae_lens/training/geometric_median.py

@@ -1,101 +0,0 @@
-from types import SimpleNamespace
-
-import torch
-import tqdm
-
-
-def weighted_average(points: torch.Tensor, weights: torch.Tensor):
-    weights = weights / weights.sum()
-    return (points * weights.view(-1, 1)).sum(dim=0)
-
-
-@torch.no_grad()
-def geometric_median_objective(
-    median: torch.Tensor, points: torch.Tensor, weights: torch.Tensor
-) -> torch.Tensor:
-    norms = torch.linalg.norm(points - median.view(1, -1), dim=1)  # type: ignore
-
-    return (norms * weights).sum()
-
-
-def compute_geometric_median(
-    points: torch.Tensor,
-    weights: torch.Tensor | None = None,
-    eps: float = 1e-6,
-    maxiter: int = 100,
-    ftol: float = 1e-20,
-    do_log: bool = False,
-):
-    """
-    :param points: ``torch.Tensor`` of shape ``(n, d)``
-    :param weights: Optional ``torch.Tensor`` of shape :math:``(n,)``.
-    :param eps: Smallest allowed value of denominator, to avoid divide by zero.
-        Equivalently, this is a smoothing parameter. Default 1e-6.
-    :param maxiter: Maximum number of Weiszfeld iterations. Default 100
-    :param ftol: If objective value does not improve by at least this `ftol` fraction, terminate the algorithm. Default 1e-20.
-    :param do_log: If true will return a log of function values encountered through the course of the algorithm
-    :return: SimpleNamespace object with fields
-        - `median`: estimate of the geometric median, which is a ``torch.Tensor`` object of shape :math:``(d,)``
-        - `termination`: string explaining how the algorithm terminated.
-        - `logs`: function values encountered through the course of the algorithm in a list (None if do_log is false).
-    """
-    with torch.no_grad():
-        if weights is None:
-            weights = torch.ones((points.shape[0],), device=points.device)
-        # initialize median estimate at mean
-        new_weights = weights
-        median = weighted_average(points, weights)
-        objective_value = geometric_median_objective(median, points, weights)
-        logs = [objective_value] if do_log else None
-
-        # Weiszfeld iterations
-        early_termination = False
-        pbar = tqdm.tqdm(range(maxiter))
-        for _ in pbar:
-            prev_obj_value = objective_value
-
-            norms = torch.linalg.norm(points - median.view(1, -1), dim=1)  # type: ignore
-            new_weights = weights / torch.clamp(norms, min=eps)
-            median = weighted_average(points, new_weights)
-            objective_value = geometric_median_objective(median, points, weights)
-
-            if logs is not None:
-                logs.append(objective_value)
-            if abs(prev_obj_value - objective_value) <= ftol * objective_value:
-                early_termination = True
-                break
-
-            pbar.set_description(f"Objective value: {objective_value:.4f}")
-
-    median = weighted_average(points, new_weights)  # allow autodiff to track it
-    return SimpleNamespace(
-        median=median,
-        new_weights=new_weights,
-        termination=(
-            "function value converged within tolerance"
-            if early_termination
-            else "maximum iterations reached"
-        ),
-        logs=logs,
-    )
-
-
-if __name__ == "__main__":
-    import time
-
-    TOLERANCE = 1e-2
-
-    dim1 = 10000
-    dim2 = 768
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    sample = (
-        torch.randn((dim1, dim2), device=device) * 100
-    )  # seems to be the order of magnitude of the actual use case
-    weights = torch.randn((dim1,), device=device)
-
-    torch.tensor(weights, device=device)
-
-    tic = time.perf_counter()
-    new = compute_geometric_median(sample, weights=weights, maxiter=100)
-    print(f"new code takes {time.perf_counter()-tic} seconds!")  # noqa: T201