Update FastPLMs files

README.md

@@ -237,16 +237,13 @@ sequence-group IDs. `-1` marks padding. Omit `sequence_id` to use
 
 ESM++ supports hidden-state SAEs from the official
 [Biohub ESMC SAE collection](https://huggingface.co/collections/biohub/esmc-saes-for-hidden-states-all-layers).
-Select an SAE for this ESMC scale. Load only required layers. Then attach them
-to the model:
+This artifact implements the SAE contract, so no Biohub runtime code is needed.
+Select an SAE for this ESMC scale, then load only the layers you need:
 
 ```python
 import torch
-from transformers import AutoModel
 
-sae = AutoModel.from_pretrained("biohub/ESMC-600M-sae-layer27-k64-codebook65536", device=model.device)
-sae.initialize_layers([27])
-model.add_sae_models([sae.layers["27"]])
+model.load_sae_models("biohub/ESMC-600M-sae-layer27-k64-codebook65536", [27])
 
 with torch.inference_mode():
     output = model(**batch, normalize_sae=True)
@@ -255,6 +252,12 @@ features = output.sae_outputs["layer27"]
 print(features.shape, features.layout)  # (valid_token_count, codebook_dim), sparse COO
 ```
 
+`load_sae_models` reads the shared `config.json` and one
+`layer_{index}.safetensors` shard per requested layer, from a Hub repository
+or a local directory, and attaches the layers on the model device in the model
+dtype. `add_sae_models` still accepts official Biohub `ESMCSAEModel.layers`
+entries.
+
 SAEs run after you attach them. Use `compute_sae=False` to skip SAE work.
 Outputs are detached sparse tensors with keys such as `layer{N}`. They omit
 padding. The model uses `sequence_id`, then `attention_mask`, for padding.

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import importlib
 import importlib.metadata
 import math
+import os
+from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass
 from functools import partial
@@ -25,6 +27,8 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
 )
 
+from .modeling_esm_plusplus_sae import ESMplusplusSAELayer, load_esmc_sae_layers
+
 
 try:
     from fastplms.attention import (
@@ -1115,14 +1119,48 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
         )
         return self._esmc_precision_status
 
+    def load_sae_models(
+        self,
+        repository: str | os.PathLike[str],
+        layers: Sequence[int],
+        *,
+        revision: str | None = None,
+        cache_dir: str | os.PathLike[str] | None = None,
+        token: str | bool | None = None,
+        local_files_only: bool = False,
+        dtype: torch.dtype | None = None,
+    ) -> dict[int, ESMplusplusSAELayer]:
+        """Load hidden-state SAE layers from a Hub repository or local directory, then attach them.
+
+        The layers land on this model's device and, unless ``dtype`` says otherwise, in this
+        model's parameter dtype, so they consume its hidden states without a dtype mismatch.
+        """
+
+        sae_layers = load_esmc_sae_layers(
+            repository,
+            layers,
+            revision=revision,
+            cache_dir=cache_dir,
+            token=token,
+            local_files_only=local_files_only,
+            device=self.device,
+            dtype=self.dtype if dtype is None else dtype,
+        )
+        self.add_sae_models(list(sae_layers.values()))
+        return sae_layers
+
     def add_sae_models(self, sae_models: list[nn.Module]) -> None:
-        """Attach official Biohub hidden-state SAE layers to this ESM++ model."""
+        """Attach hidden-state SAE layers to this ESM++ model.
+
+        Accepts layers from ``load_sae_models`` and official Biohub ``ESMCSAEModel.layers``
+        entries, which share one attachment contract.
+        """
 
         for sae_model in sae_models:
             if not isinstance(sae_model, nn.Module):
                 raise TypeError(
-                    "Each SAE must be an nn.Module obtained from an official Biohub "
-                    "ESMCSAEModel.layers entry."
+                    "Each SAE must be an nn.Module exposing the hidden-state SAE contract, such "
+                    "as a load_sae_models layer or an official Biohub ESMCSAEModel.layers entry."
                 )
             layer = getattr(sae_model, "layer", None)
             if isinstance(layer, bool) or not isinstance(layer, int):
@@ -1162,11 +1200,14 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
     ) -> tuple[tuple[int, ...], torch.Tensor | None]:
         if not compute_sae or not self._sae_models:
             return (), None
-        if input_ids is None:
+        if input_ids is None and attention_mask is None and sequence_id is None:
+            # Embedding inputs carry no token identities, so the caller must supply the mask.
+            # Rejecting masked tokens is then the caller's precondition rather than ours.
             raise ValueError(
-                "SAE computation requires input_ids so masked-token inputs can be rejected."
+                "SAE computation from inputs_embeds requires an explicit attention_mask or "
+                "sequence_id, because the token mask cannot be recovered from embeddings."
             )
-        if torch.any(input_ids == self.config.mask_token_id):
+        if input_ids is not None and torch.any(input_ids == self.config.mask_token_id):
             raise ValueError("SAE inputs must not contain mask tokens; SAEs were trained unmasked.")
         if sequence_id is not None:
             token_mask = sequence_id >= 0
@@ -1183,7 +1224,15 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
         token_mask: torch.Tensor | None,
         *,
         normalize_sae: bool,
+        differentiable_sae: bool = False,
     ) -> dict[str, torch.Tensor] | None:
+        """Encode collected hidden states with every attached SAE.
+
+        The default detaches and sparsifies, which is right for interpretation and matches the
+        official implementation. ``differentiable_sae`` instead keeps the result attached to the
+        graph and dense, so a gradient-based sequence designer can optimize an objective built on
+        SAE features. The arithmetic is identical either way; only the tape and the layout differ.
+        """
         if not self._sae_models:
             return None
         if hidden_states is None or token_mask is None:
@@ -1193,14 +1242,18 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
             layer = int(key.removeprefix("layer"))
             if layer not in hidden_states:
                 raise RuntimeError(f"ESM++ did not collect the requested SAE layer {layer}.")
-            sae_output = sae_model.get_sae_output(hidden_states[layer].clone(), token_mask)
+            layer_states = hidden_states[layer]
+            sae_output = sae_model.get_sae_output(
+                layer_states if differentiable_sae else layer_states.clone(), token_mask
+            )
             features = getattr(sae_output, "feature_magnitudes", None)
             if not isinstance(features, torch.Tensor):
                 raise TypeError("SAE get_sae_output must return tensor feature_magnitudes.")
-            features = features.detach()
+            if not differentiable_sae:
+                features = features.detach()
             if normalize_sae:
                 features = (features / sae_model.max) * sae_model.idf
-            outputs[key] = features.to_sparse()
+            outputs[key] = features if differentiable_sae else features.to_sparse()
         return outputs
 
     def _pad_fp8_inputs(
@@ -1422,6 +1475,7 @@ class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
         return_dict: bool | None = None,
         compute_sae: bool = True,
         normalize_sae: bool = False,
+        differentiable_sae: bool = False,
     ) -> TransformerOutput | tuple[torch.Tensor, ...]:
         """Run ESMC inference with the pinned Biohub mask precedence.
 
@@ -1475,6 +1529,7 @@ class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
                 transformer_output.sae_hidden_states,
                 sae_token_mask,
                 normalize_sae=normalize_sae,
+                differentiable_sae=differentiable_sae,
             )
             if sae_layers
             else None
@@ -1571,6 +1626,7 @@ class ESMplusplusForMaskedLM(
         compute_logits: bool = True,
         compute_sae: bool = True,
         normalize_sae: bool = False,
+        differentiable_sae: bool = False,
     ) -> ESMplusplusOutput | tuple[torch.Tensor, ...]:
         if input_ids is None and inputs_embeds is None:
             raise ValueError("You have to specify either input_ids or inputs_embeds")
@@ -1617,6 +1673,7 @@ class ESMplusplusForMaskedLM(
                 output.sae_hidden_states,
                 sae_token_mask,
                 normalize_sae=normalize_sae,
+                differentiable_sae=differentiable_sae,
             )
             if sae_layers
             else None

fastplms/models/esm_plusplus/modeling_esm_plusplus_sae.py

@@ -0,0 +1,266 @@
+"""Hidden-state sparse autoencoders for ESM++ (ESMC) models.
+
+FastPLMs implements the published Biohub hidden-state SAE contract directly, so attaching an SAE
+needs only PyTorch, Transformers, and the checkpoint itself. Biohub still owns the SAE weights:
+this module reads their published repository layout, one shared ``config.json`` plus one
+``layer_{index}.safetensors`` shard per backbone layer, and never redistributes those tensors.
+
+A layer produced here satisfies the same attachment contract as an official Biohub
+``ESMCSAEModel.layers`` entry, so ``PreTrainedESMplusplusModel.add_sae_models`` accepts either.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from torch import Tensor
+
+
+SAE_CONFIG_FILE = "config.json"
+SAE_REQUIRED_CONFIG_FIELDS = ("d_model", "codebook_dim", "k")
+# Trained per-feature statistics. Shards that never trained them omit them, and the ones default
+# makes the (features / max) * idf normalization an identity.
+SAE_OPTIONAL_STATE_NAMES = ("idf", "max")
+STANDARDIZATION_EPSILON = 1e-5
+
+
+@dataclass(frozen=True, slots=True)
+class ESMplusplusSAEParams:
+    """Shape contract of one hidden-state SAE.
+
+    ``layer`` indexes the ESM++ hidden state the SAE reads, where ``0`` is the embedding output and
+    ``num_hidden_layers`` is the final normalized state.
+    """
+
+    d_model: int
+    codebook_dim: int
+    k: int
+    layer: int
+
+
+@dataclass(slots=True)
+class ESMplusplusSAEOutput:
+    """Sparse features, and the optional reconstruction error, for one batch of residues."""
+
+    feature_magnitudes: Tensor
+    reconstruction_loss: Tensor | None = None
+
+
+def _standardize_residue_states(residue_states: Tensor) -> Tensor:
+    """Center and scale each residue vector the way the SAEs were trained."""
+
+    # residue_states: (n, d)
+    centered = residue_states - residue_states.mean(dim=-1, keepdim=True)  # (n, d)
+    return centered / (centered.std(dim=-1, keepdim=True) + STANDARDIZATION_EPSILON)  # (n, d)
+
+
+class ESMplusplusSAELayer(nn.Module):
+    """Top-k sparse autoencoder over one ESM++ hidden state.
+
+    Encoding standardizes each residue vector, projects it into a wide codebook, and keeps only the
+    ``k`` largest activations. The decoder is used only to measure reconstruction error, which is
+    why it is opt-in: interpretation and gradient-based design read ``feature_magnitudes`` alone.
+    """
+
+    idf: Tensor
+    max: Tensor
+
+    def __init__(self, params: ESMplusplusSAEParams) -> None:
+        super().__init__()
+        self.params = params
+        self.W_enc = nn.Parameter(torch.empty(params.d_model, params.codebook_dim))  # (d, c)
+        self.W_dec = nn.Parameter(torch.empty(params.codebook_dim, params.d_model))  # (c, d)
+        self.b_dec = nn.Parameter(torch.zeros(params.d_model))  # (d,)
+        self.register_buffer("idf", torch.ones(params.codebook_dim))  # (c,)
+        self.register_buffer("max", torch.ones(params.codebook_dim))  # (c,)
+
+    @property
+    def layer(self) -> int:
+        """ESM++ hidden-state index this SAE was trained against."""
+
+        return self.params.layer
+
+    def forward(
+        self,
+        residue_states: Tensor,
+        *,
+        compute_reconstruction_loss: bool = False,
+    ) -> ESMplusplusSAEOutput:
+        # residue_states: (n, d) for n residues; c is the codebook width, k the retained features
+        standardized = _standardize_residue_states(residue_states)  # (n, d)
+        preactivations = F.relu((standardized - self.b_dec) @ self.W_enc)  # (n, c)
+        retained = torch.topk(preactivations, self.params.k, dim=-1)  # values, indices: (n, k)
+        feature_magnitudes = torch.zeros_like(preactivations).scatter(
+            -1, retained.indices, retained.values
+        )  # (n, c)
+        if not compute_reconstruction_loss:
+            return ESMplusplusSAEOutput(feature_magnitudes=feature_magnitudes)
+
+        reconstructed = feature_magnitudes @ self.W_dec + self.b_dec  # (n, d)
+        return ESMplusplusSAEOutput(
+            feature_magnitudes=feature_magnitudes,
+            reconstruction_loss=(reconstructed - standardized).pow(2).mean(dim=-1),  # (n,)
+        )
+
+    def get_sae_output(self, layer_states: Tensor, token_mask: Tensor) -> ESMplusplusSAEOutput:
+        """Encode the unpadded residues of one ESM++ hidden state.
+
+        This name and signature are the attachment contract that
+        ``PreTrainedESMplusplusModel.add_sae_models`` validates and calls.
+        """
+
+        # layer_states: (b, l, d); token_mask: (b, l)
+        residue_states = layer_states[token_mask]  # (n, d) for n valid tokens
+        encoded: ESMplusplusSAEOutput = self(residue_states)
+        return encoded
+
+
+def _repository_file(
+    repository: str | os.PathLike[str],
+    filename: str,
+    *,
+    revision: str | None,
+    cache_dir: str | os.PathLike[str] | None,
+    token: str | bool | None,
+    local_files_only: bool,
+) -> Path:
+    """Resolve one repository file, from a local directory or the Hub cache."""
+
+    # A directory counts as local only when it holds the shared config, so a stale directory named
+    # like a Hub identifier cannot shadow the download.
+    local_directory = Path(repository)
+    if (local_directory / SAE_CONFIG_FILE).is_file():
+        path = local_directory / filename
+        if not path.is_file():
+            raise FileNotFoundError(f"SAE repository {local_directory} has no {filename}.")
+        return path
+
+    return Path(
+        hf_hub_download(
+            repo_id=str(repository),
+            filename=filename,
+            revision=revision,
+            cache_dir=None if cache_dir is None else str(cache_dir),
+            token=token,
+            local_files_only=local_files_only,
+        )
+    )
+
+
+def _load_sae_layer(
+    shard_path: Path,
+    *,
+    params: ESMplusplusSAEParams,
+    device: torch.device,
+    dtype: torch.dtype | None,
+) -> ESMplusplusSAELayer:
+    state = load_file(str(shard_path), device=str(device))
+    encoder = state.get("W_enc")
+    if encoder is None:
+        raise ValueError(f"{shard_path} is not an ESMC SAE shard; it has no 'W_enc' entry.")
+
+    # Build on the meta device so the shard tensors are the only materialized copy of a codebook
+    # that reaches roughly one gigabyte for the widest published SAEs.
+    with torch.device("meta"):
+        sae_layer = ESMplusplusSAELayer(params)
+    sae_layer.to(dtype=encoder.dtype if dtype is None else dtype)
+    sae_layer.to_empty(device=device)
+    # to_empty leaves the statistics buffers uninitialized, so restore the identity defaults that
+    # shards without trained statistics rely on.
+    sae_layer.idf.fill_(1.0)
+    sae_layer.max.fill_(1.0)
+
+    incompatible = sae_layer.load_state_dict(state, strict=False)
+    missing = tuple(
+        name for name in incompatible.missing_keys if name not in SAE_OPTIONAL_STATE_NAMES
+    )
+    unexpected = tuple(incompatible.unexpected_keys)
+    if missing or unexpected:
+        raise ValueError(
+            f"{shard_path} does not match the ESMC SAE state contract; "
+            f"missing {list(missing)}, unexpected {list(unexpected)}."
+        )
+    return sae_layer
+
+
+def load_esmc_sae_layers(
+    repository: str | os.PathLike[str],
+    layers: Sequence[int],
+    *,
+    revision: str | None = None,
+    cache_dir: str | os.PathLike[str] | None = None,
+    token: str | bool | None = None,
+    local_files_only: bool = False,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype | None = None,
+) -> dict[int, ESMplusplusSAELayer]:
+    """Load the requested hidden-state SAE layers from a Hub repository or a local directory.
+
+    Only the shared config and the requested shards are read, so a repository that publishes every
+    backbone layer costs one shard per requested layer. ``dtype`` defaults to the dtype stored in
+    the shard; pass the ESM++ model dtype when the SAE has to consume its hidden states directly.
+    """
+
+    requested = tuple(dict.fromkeys(int(layer) for layer in layers))
+    if not requested:
+        raise ValueError("Loading SAE layers requires at least one backbone layer index.")
+
+    config_path = _repository_file(
+        repository,
+        SAE_CONFIG_FILE,
+        revision=revision,
+        cache_dir=cache_dir,
+        token=token,
+        local_files_only=local_files_only,
+    )
+    config = json.loads(config_path.read_text(encoding="utf-8"))
+    missing_fields = [name for name in SAE_REQUIRED_CONFIG_FIELDS if name not in config]
+    if missing_fields:
+        raise ValueError(f"{config_path} is not an ESMC SAE config; it omits {missing_fields}.")
+    available = tuple(int(index) for index in config.get("available_layers", ()))
+
+    target_device = torch.device(device)
+    sae_layers: dict[int, ESMplusplusSAELayer] = {}
+    for layer in requested:
+        if available and layer not in available:
+            raise ValueError(
+                f"SAE repository {repository} does not publish layer {layer}; "
+                f"available layers are {list(available)}."
+            )
+        shard_path = _repository_file(
+            repository,
+            f"layer_{layer}.safetensors",
+            revision=revision,
+            cache_dir=cache_dir,
+            token=token,
+            local_files_only=local_files_only,
+        )
+        sae_layers[layer] = _load_sae_layer(
+            shard_path,
+            params=ESMplusplusSAEParams(
+                d_model=int(config["d_model"]),
+                codebook_dim=int(config["codebook_dim"]),
+                k=int(config["k"]),
+                layer=layer,
+            ),
+            device=target_device,
+            dtype=dtype,
+        )
+    return sae_layers
+
+
+__all__ = [
+    "ESMplusplusSAELayer",
+    "ESMplusplusSAEOutput",
+    "ESMplusplusSAEParams",
+    "load_esmc_sae_layers",
+]

modeling_fastplms.py

@@ -13,7 +13,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "65b4cb38129822c36a448e295f77702ee64ca00fcfb9265b8c60f1c3ef169097":
+if RUNTIME_HASH != "80b9bd43aa70400483b3f35121ed5039a743154f5b44b872578b5a12bf95bdd0":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []