sae-lens 6.22.1__py3-none-any.whl → 6.25.1__py3-none-any.whl

sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.22.1"
+__version__ = "6.25.1"
 
 import logging
 
@@ -15,6 +15,8 @@ from sae_lens.saes import (
     GatedTrainingSAEConfig,
     JumpReLUSAE,
     JumpReLUSAEConfig,
+    JumpReLUSkipTranscoder,
+    JumpReLUSkipTranscoderConfig,
     JumpReLUTrainingSAE,
     JumpReLUTrainingSAEConfig,
     JumpReLUTranscoder,
@@ -105,6 +107,8 @@ __all__ = [
     "SkipTranscoderConfig",
     "JumpReLUTranscoder",
     "JumpReLUTranscoderConfig",
+    "JumpReLUSkipTranscoder",
+    "JumpReLUSkipTranscoderConfig",
     "MatryoshkaBatchTopKTrainingSAE",
     "MatryoshkaBatchTopKTrainingSAEConfig",
     "TemporalSAE",
@@ -131,4 +135,7 @@ register_sae_training_class(
 register_sae_class("transcoder", Transcoder, TranscoderConfig)
 register_sae_class("skip_transcoder", SkipTranscoder, SkipTranscoderConfig)
 register_sae_class("jumprelu_transcoder", JumpReLUTranscoder, JumpReLUTranscoderConfig)
+register_sae_class(
+    "jumprelu_skip_transcoder", JumpReLUSkipTranscoder, JumpReLUSkipTranscoderConfig
+)
 register_sae_class("temporal", TemporalSAE, TemporalSAEConfig)

sae_lens/cache_activations_runner.py

@@ -14,9 +14,9 @@ from transformer_lens.HookedTransformer import HookedRootModule
 
 from sae_lens import logger
 from sae_lens.config import CacheActivationsRunnerConfig
-from sae_lens.constants import DTYPE_MAP
 from sae_lens.load_model import load_model
 from sae_lens.training.activations_store import ActivationsStore
+from sae_lens.util import str_to_dtype
 
 
 def _mk_activations_store(
@@ -97,7 +97,7 @@ class CacheActivationsRunner:
         bytes_per_token = (
             self.cfg.d_in * self.cfg.dtype.itemsize
             if isinstance(self.cfg.dtype, torch.dtype)
-            else DTYPE_MAP[self.cfg.dtype].itemsize
+            else str_to_dtype(self.cfg.dtype).itemsize
         )
         total_training_tokens = self.cfg.n_seq_in_dataset * self.context_size
         total_disk_space_gb = total_training_tokens * bytes_per_token / 10**9

sae_lens/config.py

@@ -17,9 +17,9 @@ from datasets import (
 )
 
 from sae_lens import __version__, logger
-from sae_lens.constants import DTYPE_MAP
 from sae_lens.registry import get_sae_training_class
 from sae_lens.saes.sae import TrainingSAEConfig
+from sae_lens.util import str_to_dtype
 
 if TYPE_CHECKING:
     pass
@@ -563,7 +563,7 @@ class CacheActivationsRunnerConfig:
 
     @property
     def bytes_per_token(self) -> int:
-        return self.d_in * DTYPE_MAP[self.dtype].itemsize
+        return self.d_in * str_to_dtype(self.dtype).itemsize
 
     @property
     def n_tokens_in_buffer(self) -> int:

sae_lens/constants.py

@@ -11,6 +11,14 @@ DTYPE_MAP = {
     "torch.bfloat16": torch.bfloat16,
 }
 
+# Reverse mapping from torch.dtype to canonical string format
+DTYPE_TO_STR = {
+    torch.float32: "float32",
+    torch.float64: "float64",
+    torch.float16: "float16",
+    torch.bfloat16: "bfloat16",
+}
+
 
 SPARSITY_FILENAME = "sparsity.safetensors"
 SAE_WEIGHTS_FILENAME = "sae_weights.safetensors"

sae_lens/loading/pretrained_sae_loaders.py

@@ -9,14 +9,12 @@ import requests
 import torch
 import yaml
 from huggingface_hub import hf_hub_download, hf_hub_url
-from huggingface_hub.utils import EntryNotFoundError
+from huggingface_hub.utils import EntryNotFoundError, build_hf_headers
 from packaging.version import Version
 from safetensors import safe_open
-from safetensors.torch import load_file
 
 from sae_lens import logger
 from sae_lens.constants import (
-    DTYPE_MAP,
     SAE_CFG_FILENAME,
     SAE_WEIGHTS_FILENAME,
     SPARSIFY_WEIGHTS_FILENAME,
@@ -28,7 +26,7 @@ from sae_lens.loading.pretrained_saes_directory import (
     get_repo_id_and_folder_name,
 )
 from sae_lens.registry import get_sae_class
-from sae_lens.util import filter_valid_dataclass_fields
+from sae_lens.util import filter_valid_dataclass_fields, str_to_dtype
 
 LLM_METADATA_KEYS = {
     "model_name",
@@ -46,9 +44,26 @@ LLM_METADATA_KEYS = {
     "sae_lens_training_version",
     "hook_name_out",
     "hook_head_index_out",
+    "hf_hook_name",
+    "hf_hook_name_out",
 }
 
 
+def load_safetensors_weights(
+    path: str | Path, device: str = "cpu", dtype: torch.dtype | str | None = None
+) -> dict[str, torch.Tensor]:
+    """Load safetensors weights and optionally convert to a different dtype"""
+    loaded_weights = {}
+    dtype = str_to_dtype(dtype) if isinstance(dtype, str) else dtype
+    with safe_open(path, framework="pt", device=device) as f:
+        for k in f.keys():  # noqa: SIM118
+            weight = f.get_tensor(k)
+            if dtype is not None:
+                weight = weight.to(dtype=dtype)
+            loaded_weights[k] = weight
+    return loaded_weights
+
+
 # loaders take in a release, sae_id, device, and whether to force download, and returns a tuple of config, state_dict, and log sparsity
 class PretrainedSaeHuggingfaceLoader(Protocol):
     def __call__(
@@ -339,7 +354,7 @@ def read_sae_components_from_disk(
     Given a loaded dictionary and a path to a weight file, load the weights and return the state_dict.
     """
     if dtype is None:
-        dtype = DTYPE_MAP[cfg_dict["dtype"]]
+        dtype = str_to_dtype(cfg_dict["dtype"])
 
     state_dict = {}
     with safe_open(weight_path, framework="pt", device=device) as f:  # type: ignore
@@ -523,6 +538,199 @@ def gemma_2_sae_huggingface_loader(
     return cfg_dict, state_dict, log_sparsity
 
 
+def _infer_gemma_3_raw_cfg_dict(repo_id: str, folder_name: str) -> dict[str, Any]:
+    """
+    Infer the raw config dict for Gemma 3 SAEs from the repo_id and folder_name.
+    This is used when config.json doesn't exist in the repo.
+    """
+    # Extract layer number from folder name
+    layer_match = re.search(r"layer_(\d+)", folder_name)
+    if layer_match is None:
+        raise ValueError(
+            f"Could not extract layer number from folder_name: {folder_name}"
+        )
+    layer = int(layer_match.group(1))
+
+    # Convert repo_id to model_name: google/gemma-scope-2-{size}-{suffix} -> google/gemma-3-{size}-{suffix}
+    model_name = repo_id.replace("gemma-scope-2", "gemma-3")
+
+    # Determine hook type and HF hook points based on folder_name
+    if "transcoder" in folder_name or "clt" in folder_name:
+        hf_hook_point_in = f"model.layers.{layer}.pre_feedforward_layernorm.output"
+        hf_hook_point_out = f"model.layers.{layer}.post_feedforward_layernorm.output"
+    elif "resid_post" in folder_name:
+        hf_hook_point_in = f"model.layers.{layer}.output"
+        hf_hook_point_out = None
+    elif "attn_out" in folder_name:
+        hf_hook_point_in = f"model.layers.{layer}.self_attn.o_proj.input"
+        hf_hook_point_out = None
+    elif "mlp_out" in folder_name:
+        hf_hook_point_in = f"model.layers.{layer}.post_feedforward_layernorm.output"
+        hf_hook_point_out = None
+    else:
+        raise ValueError(f"Could not infer hook type from folder_name: {folder_name}")
+
+    cfg: dict[str, Any] = {
+        "architecture": "jump_relu",
+        "model_name": model_name,
+        "hf_hook_point_in": hf_hook_point_in,
+    }
+    if hf_hook_point_out is not None:
+        cfg["hf_hook_point_out"] = hf_hook_point_out
+
+    return cfg
+
+
+def get_gemma_3_config_from_hf(
+    repo_id: str,
+    folder_name: str,
+    device: str,
+    force_download: bool = False,
+    cfg_overrides: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    # Try to load config.json from the repo, fall back to inferring if it doesn't exist
+    try:
+        config_path = hf_hub_download(
+            repo_id, f"{folder_name}/config.json", force_download=force_download
+        )
+        with open(config_path) as config_file:
+            raw_cfg_dict = json.load(config_file)
+    except EntryNotFoundError:
+        raw_cfg_dict = _infer_gemma_3_raw_cfg_dict(repo_id, folder_name)
+
+    if raw_cfg_dict.get("architecture") != "jump_relu":
+        raise ValueError(
+            f"Unexpected architecture in Gemma 3 config: {raw_cfg_dict.get('architecture')}"
+        )
+
+    layer_match = re.search(r"layer_(\d+)", folder_name)
+    if layer_match is None:
+        raise ValueError(
+            f"Could not extract layer number from folder_name: {folder_name}"
+        )
+    layer = int(layer_match.group(1))
+    hook_name_out = None
+    d_out = None
+    if "resid_post" in folder_name:
+        hook_name = f"blocks.{layer}.hook_resid_post"
+    elif "attn_out" in folder_name:
+        hook_name = f"blocks.{layer}.hook_attn_out"
+    elif "mlp_out" in folder_name:
+        hook_name = f"blocks.{layer}.hook_mlp_out"
+    elif "transcoder" in folder_name or "clt" in folder_name:
+        hook_name = f"blocks.{layer}.ln2.hook_normalized"
+        hook_name_out = f"blocks.{layer}.hook_mlp_out"
+    else:
+        raise ValueError("Hook name not found in folder_name.")
+
+    # hackily deal with clt file names
+    params_file_part = "/params.safetensors"
+    if "clt" in folder_name:
+        params_file_part = ".safetensors"
+
+    shapes_dict = get_safetensors_tensor_shapes(
+        repo_id, f"{folder_name}{params_file_part}"
+    )
+    d_in, d_sae = shapes_dict["w_enc"]
+    # TODO: update this for real model info
+    model_name = raw_cfg_dict["model_name"]
+    if "google" not in model_name:
+        model_name = "google/" + model_name
+    model_name = model_name.replace("-v3", "-3")
+    if "270m" in model_name:
+        # for some reason the 270m model on huggingface doesn't have the -pt suffix
+        model_name = model_name.replace("-pt", "")
+
+    architecture = "jumprelu"
+    if "transcoder" in folder_name or "clt" in folder_name:
+        architecture = "jumprelu_skip_transcoder"
+        d_out = shapes_dict["w_dec"][-1]
+
+    cfg = {
+        "architecture": architecture,
+        "d_in": d_in,
+        "d_sae": d_sae,
+        "dtype": "float32",
+        "model_name": model_name,
+        "hook_name": hook_name,
+        "hook_head_index": None,
+        "finetuning_scaling_factor": False,
+        "sae_lens_training_version": None,
+        "prepend_bos": True,
+        "dataset_path": "monology/pile-uncopyrighted",
+        "context_size": 1024,
+        "apply_b_dec_to_input": False,
+        "normalize_activations": None,
+        "hf_hook_name": raw_cfg_dict.get("hf_hook_point_in"),
+    }
+    if hook_name_out is not None:
+        cfg["hook_name_out"] = hook_name_out
+        cfg["hf_hook_name_out"] = raw_cfg_dict.get("hf_hook_point_out")
+    if d_out is not None:
+        cfg["d_out"] = d_out
+    if device is not None:
+        cfg["device"] = device
+
+    if cfg_overrides is not None:
+        cfg.update(cfg_overrides)
+
+    return cfg
+
+
+def gemma_3_sae_huggingface_loader(
+    repo_id: str,
+    folder_name: str,
+    device: str = "cpu",
+    force_download: bool = False,
+    cfg_overrides: dict[str, Any] | None = None,
+) -> tuple[dict[str, Any], dict[str, torch.Tensor], torch.Tensor | None]:
+    """
+    Custom loader for Gemma 3 SAEs.
+    """
+    cfg_dict = get_gemma_3_config_from_hf(
+        repo_id,
+        folder_name,
+        device,
+        force_download,
+        cfg_overrides,
+    )
+
+    params_file = "params.safetensors"
+    if "clt" in folder_name:
+        params_file = folder_name.split("/")[-1] + ".safetensors"
+        folder_name = "/".join(folder_name.split("/")[:-1])
+
+    # Download the SAE weights
+    sae_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=params_file,
+        subfolder=folder_name,
+        force_download=force_download,
+    )
+
+    raw_state_dict = load_safetensors_weights(
+        sae_path, device=device, dtype=cfg_dict.get("dtype")
+    )
+
+    with torch.no_grad():
+        w_dec = raw_state_dict["w_dec"]
+        if "clt" in folder_name:
+            w_dec = w_dec.sum(dim=1).contiguous()
+
+    state_dict = {
+        "W_enc": raw_state_dict["w_enc"],
+        "W_dec": w_dec,
+        "b_enc": raw_state_dict["b_enc"],
+        "b_dec": raw_state_dict["b_dec"],
+        "threshold": raw_state_dict["threshold"],
+    }
+
+    if "affine_skip_connection" in raw_state_dict:
+        state_dict["W_skip"] = raw_state_dict["affine_skip_connection"]
+
+    return cfg_dict, state_dict, None
+
+
 def get_goodfire_config_from_hf(
     repo_id: str,
     folder_name: str,  # noqa: ARG001
@@ -589,11 +797,13 @@ def get_goodfire_huggingface_loader(
     )
     raw_state_dict = torch.load(sae_path, map_location=device)
 
+    target_dtype = str_to_dtype(cfg_dict.get("dtype", "float32"))
+
     state_dict = {
-        "W_enc": raw_state_dict["encoder_linear.weight"].T,
-        "W_dec": raw_state_dict["decoder_linear.weight"].T,
-        "b_enc": raw_state_dict["encoder_linear.bias"],
-        "b_dec": raw_state_dict["decoder_linear.bias"],
+        "W_enc": raw_state_dict["encoder_linear.weight"].T.to(dtype=target_dtype),
+        "W_dec": raw_state_dict["decoder_linear.weight"].T.to(dtype=target_dtype),
+        "b_enc": raw_state_dict["encoder_linear.bias"].to(dtype=target_dtype),
+        "b_dec": raw_state_dict["decoder_linear.bias"].to(dtype=target_dtype),
     }
 
     return cfg_dict, state_dict, None
@@ -696,26 +906,19 @@ def llama_scope_sae_huggingface_loader(
         force_download=force_download,
     )
 
-    # Load the weights using load_file instead of safe_open
-    state_dict_loaded = load_file(sae_path, device=device)
+    state_dict_loaded = load_safetensors_weights(
+        sae_path, device=device, dtype=cfg_dict.get("dtype")
+    )
 
     # Convert and organize the weights
     state_dict = {
-        "W_enc": state_dict_loaded["encoder.weight"]
-        .to(dtype=DTYPE_MAP[cfg_dict["dtype"]])
-        .T,
-        "W_dec": state_dict_loaded["decoder.weight"]
-        .to(dtype=DTYPE_MAP[cfg_dict["dtype"]])
-        .T,
-        "b_enc": state_dict_loaded["encoder.bias"].to(
-            dtype=DTYPE_MAP[cfg_dict["dtype"]]
-        ),
-        "b_dec": state_dict_loaded["decoder.bias"].to(
-            dtype=DTYPE_MAP[cfg_dict["dtype"]]
-        ),
+        "W_enc": state_dict_loaded["encoder.weight"].T,
+        "W_dec": state_dict_loaded["decoder.weight"].T,
+        "b_enc": state_dict_loaded["encoder.bias"],
+        "b_dec": state_dict_loaded["decoder.bias"],
         "threshold": torch.ones(
             cfg_dict["d_sae"],
-            dtype=DTYPE_MAP[cfg_dict["dtype"]],
+            dtype=str_to_dtype(cfg_dict["dtype"]),
             device=cfg_dict["device"],
         )
         * cfg_dict["jump_relu_threshold"],
@@ -753,10 +956,14 @@ def get_dictionary_learning_config_1_from_hf(
     activation_fn = "topk" if trainer["dict_class"] == "AutoEncoderTopK" else "relu"
     activation_fn_kwargs = {"k": trainer["k"]} if activation_fn == "topk" else {}
 
+    architecture = "standard"
+    if trainer["dict_class"] == "GatedAutoEncoder":
+        architecture = "gated"
+    elif trainer["dict_class"] == "MatryoshkaBatchTopKSAE":
+        architecture = "jumprelu"
+
     return {
-        "architecture": (
-            "gated" if trainer["dict_class"] == "GatedAutoEncoder" else "standard"
-        ),
+        "architecture": architecture,
         "d_in": trainer["activation_dim"],
         "d_sae": trainer["dict_size"],
         "dtype": "float32",
@@ -905,9 +1112,12 @@ def dictionary_learning_sae_huggingface_loader_1(
     )
     encoder = torch.load(encoder_path, map_location="cpu")
 
+    W_enc = encoder["W_enc"] if "W_enc" in encoder else encoder["encoder.weight"].T
+    W_dec = encoder["W_dec"] if "W_dec" in encoder else encoder["decoder.weight"].T
+
     state_dict = {
-        "W_enc": encoder["encoder.weight"].T,
-        "W_dec": encoder["decoder.weight"].T,
+        "W_enc": W_enc,
+        "W_dec": W_dec,
         "b_dec": encoder.get(
             "b_dec", encoder.get("bias", encoder.get("decoder_bias", None))
         ),
@@ -915,6 +1125,8 @@ def dictionary_learning_sae_huggingface_loader_1(
 
     if "encoder.bias" in encoder:
         state_dict["b_enc"] = encoder["encoder.bias"]
+    if "b_enc" in encoder:
+        state_dict["b_enc"] = encoder["b_enc"]
 
     if "mag_bias" in encoder:
         state_dict["b_mag"] = encoder["mag_bias"]
@@ -923,6 +1135,12 @@ def dictionary_learning_sae_huggingface_loader_1(
     if "r_mag" in encoder:
         state_dict["r_mag"] = encoder["r_mag"]
 
+    if "threshold" in encoder:
+        threshold = encoder["threshold"]
+        if threshold.ndim == 0:
+            threshold = torch.full((W_enc.size(1),), threshold)
+        state_dict["threshold"] = threshold
+
     return cfg_dict, state_dict, None
 
 
@@ -1011,26 +1229,17 @@ def llama_scope_r1_distill_sae_huggingface_loader(
         force_download=force_download,
     )
 
-    # Load the weights using load_file instead of safe_open
-    state_dict_loaded = load_file(sae_path, device=device)
+    state_dict_loaded = load_safetensors_weights(
+        sae_path, device=device, dtype=cfg_dict.get("dtype")
+    )
 
     # Convert and organize the weights
     state_dict = {
-        "W_enc": state_dict_loaded["encoder.weight"]
-        .to(dtype=DTYPE_MAP[cfg_dict["dtype"]])
-        .T,
-        "W_dec": state_dict_loaded["decoder.weight"]
-        .to(dtype=DTYPE_MAP[cfg_dict["dtype"]])
-        .T,
-        "b_enc": state_dict_loaded["encoder.bias"].to(
-            dtype=DTYPE_MAP[cfg_dict["dtype"]]
-        ),
-        "b_dec": state_dict_loaded["decoder.bias"].to(
-            dtype=DTYPE_MAP[cfg_dict["dtype"]]
-        ),
-        "threshold": state_dict_loaded["log_jumprelu_threshold"]
-        .to(dtype=DTYPE_MAP[cfg_dict["dtype"]])
-        .exp(),
+        "W_enc": state_dict_loaded["encoder.weight"].T,
+        "W_dec": state_dict_loaded["decoder.weight"].T,
+        "b_enc": state_dict_loaded["encoder.bias"],
+        "b_dec": state_dict_loaded["decoder.bias"],
+        "threshold": state_dict_loaded["log_jumprelu_threshold"].exp(),
     }
 
     # No sparsity tensor for Llama Scope SAEs
@@ -1150,34 +1359,34 @@ def sparsify_disk_loader(
     cfg_dict = get_sparsify_config_from_disk(path, device, cfg_overrides)
 
     weight_path = Path(path) / SPARSIFY_WEIGHTS_FILENAME
-    state_dict_loaded = load_file(weight_path, device=device)
-
-    dtype = DTYPE_MAP[cfg_dict["dtype"]]
+    state_dict_loaded = load_safetensors_weights(
+        weight_path, device=device, dtype=cfg_dict.get("dtype")
+    )
 
     W_enc = (
         state_dict_loaded["W_enc"]
         if "W_enc" in state_dict_loaded
         else state_dict_loaded["encoder.weight"].T
-    ).to(dtype)
+    )
 
     if "W_dec" in state_dict_loaded:
-        W_dec = state_dict_loaded["W_dec"].T.to(dtype)
+        W_dec = state_dict_loaded["W_dec"].T
     else:
-        W_dec = state_dict_loaded["decoder.weight"].T.to(dtype)
+        W_dec = state_dict_loaded["decoder.weight"].T
 
     if "b_enc" in state_dict_loaded:
-        b_enc = state_dict_loaded["b_enc"].to(dtype)
+        b_enc = state_dict_loaded["b_enc"]
     elif "encoder.bias" in state_dict_loaded:
-        b_enc = state_dict_loaded["encoder.bias"].to(dtype)
+        b_enc = state_dict_loaded["encoder.bias"]
     else:
-        b_enc = torch.zeros(cfg_dict["d_sae"], dtype=dtype, device=device)
+        b_enc = torch.zeros(cfg_dict["d_sae"], dtype=W_dec.dtype, device=device)
 
     if "b_dec" in state_dict_loaded:
-        b_dec = state_dict_loaded["b_dec"].to(dtype)
+        b_dec = state_dict_loaded["b_dec"]
     elif "decoder.bias" in state_dict_loaded:
-        b_dec = state_dict_loaded["decoder.bias"].to(dtype)
+        b_dec = state_dict_loaded["decoder.bias"]
     else:
-        b_dec = torch.zeros(cfg_dict["d_in"], dtype=dtype, device=device)
+        b_dec = torch.zeros(cfg_dict["d_in"], dtype=W_dec.dtype, device=device)
 
     state_dict = {"W_enc": W_enc, "b_enc": b_enc, "W_dec": W_dec, "b_dec": b_dec}
     return cfg_dict, state_dict
@@ -1408,44 +1617,44 @@ def mwhanna_transcoder_huggingface_loader(
     )
 
     # Load weights from safetensors
-    state_dict = load_file(file_path, device=device)
+    state_dict = load_safetensors_weights(
+        file_path, device=device, dtype=cfg_dict.get("dtype")
+    )
     state_dict["W_enc"] = state_dict["W_enc"].T
 
     return cfg_dict, state_dict, None
 
 
-def get_safetensors_tensor_shapes(url: str) -> dict[str, list[int]]:
+def get_safetensors_tensor_shapes(repo_id: str, filename: str) -> dict[str, list[int]]:
     """
-    Get tensor shapes from a safetensors file using HTTP range requests
+    Get tensor shapes from a safetensors file on HuggingFace Hub
     without downloading the entire file.
 
+    Uses HTTP range requests to fetch only the metadata header.
+
     Args:
-        url: Direct URL to the safetensors file
+        repo_id: HuggingFace repo ID (e.g., "gg-gs/gemma-scope-2-1b-pt")
+        filename: Path to the safetensors file within the repo
 
     Returns:
         Dictionary mapping tensor names to their shapes
     """
-    # Check if server supports range requests
-    response = requests.head(url, timeout=10)
-    response.raise_for_status()
+    url = hf_hub_url(repo_id, filename)
 
-    accept_ranges = response.headers.get("Accept-Ranges", "")
-    if "bytes" not in accept_ranges:
-        raise ValueError("Server does not support range requests")
+    # Get HuggingFace headers (includes auth token if available)
+    hf_headers = build_hf_headers()
 
     # Fetch first 8 bytes to get metadata size
-    headers = {"Range": "bytes=0-7"}
+    headers = {**hf_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")
+    response.raise_for_status()
 
     meta_size = int.from_bytes(response.content, byteorder="little")
 
     # Fetch the metadata header
-    headers = {"Range": f"bytes=8-{8 + meta_size - 1}"}
+    headers = {**hf_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")
+    response.raise_for_status()
 
     metadata_json = response.content.decode("utf-8").strip()
     metadata = json.loads(metadata_json)
@@ -1494,8 +1703,12 @@ def mntss_clt_layer_huggingface_loader(
         force_download=force_download,
     )
 
-    encoder_state_dict = load_file(encoder_path, device=device)
-    decoder_state_dict = load_file(decoder_path, device=device)
+    encoder_state_dict = load_safetensors_weights(
+        encoder_path, device=device, dtype=cfg_dict.get("dtype")
+    )
+    decoder_state_dict = load_safetensors_weights(
+        decoder_path, device=device, dtype=cfg_dict.get("dtype")
+    )
 
     with torch.no_grad():
         state_dict = {
@@ -1525,9 +1738,10 @@ def get_mntss_clt_layer_config_from_hf(
     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)
+    # Get tensor shapes without downloading full files
+    encoder_shapes = get_safetensors_tensor_shapes(
+        repo_id, f"W_enc_{folder_name}.safetensors"
+    )
 
     # Extract shapes for the required tensors
     b_dec_shape = encoder_shapes[f"b_dec_{folder_name}"]
@@ -1637,7 +1851,9 @@ def temporal_sae_huggingface_loader(
     )
 
     # Load checkpoint from safetensors
-    state_dict_raw = load_file(ckpt_path, device=device)
+    state_dict_raw = load_safetensors_weights(
+        ckpt_path, device=device, dtype=cfg_dict.get("dtype")
+    )
 
     # Convert to SAELens naming convention
     # TemporalSAE uses: D (decoder), E (encoder), b (bias), attn_layers.*
@@ -1663,6 +1879,7 @@ NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sae_lens": sae_lens_huggingface_loader,
     "connor_rob_hook_z": connor_rob_hook_z_huggingface_loader,
     "gemma_2": gemma_2_sae_huggingface_loader,
+    "gemma_3": gemma_3_sae_huggingface_loader,
     "llama_scope": llama_scope_sae_huggingface_loader,
     "llama_scope_r1_distill": llama_scope_r1_distill_sae_huggingface_loader,
     "dictionary_learning_1": dictionary_learning_sae_huggingface_loader_1,
@@ -1680,6 +1897,7 @@ NAMED_PRETRAINED_SAE_CONFIG_GETTERS: dict[str, PretrainedSaeConfigHuggingfaceLoa
     "sae_lens": get_sae_lens_config_from_hf,
     "connor_rob_hook_z": get_connor_rob_hook_z_config_from_hf,
     "gemma_2": get_gemma_2_config_from_hf,
+    "gemma_3": get_gemma_3_config_from_hf,
     "llama_scope": get_llama_scope_config_from_hf,
     "llama_scope_r1_distill": get_llama_scope_r1_distill_config_from_hf,
     "dictionary_learning_1": get_dictionary_learning_config_1_from_hf,

sae_lens/pretokenize_runner.py

@@ -186,13 +186,13 @@ class PretokenizeRunner:
         """
         Load the dataset, tokenize it, and save it to disk and/or upload to Huggingface.
         """
-        dataset = load_dataset(
+        dataset = load_dataset(  # type: ignore
             self.cfg.dataset_path,
             name=self.cfg.dataset_name,
             data_dir=self.cfg.data_dir,
             data_files=self.cfg.data_files,
-            split=self.cfg.split,
-            streaming=self.cfg.streaming,
+            split=self.cfg.split,  # type: ignore
+            streaming=self.cfg.streaming,  # type: ignore
         )
         if isinstance(dataset, DatasetDict):
             raise ValueError(