Synthyra/ESMFold2

1# Copyright 2026 Biohub. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""PyTorch ESMC SAE (Sparse Autoencoder) model.15 16* :class:`ESMCSAEModel` — the published HF container, one repo per17  ``(backbone, codebook_dim, k)`` group. Each backbone layer ships as a18  ``layer_{i}.safetensors`` shard; ``from_pretrained`` downloads the whole19  snapshot but loads no weights — callers materialize the layers they need20  via :meth:`initialize_layers`. Single-layer repos auto-load so bare21  ``forward(x)`` works.22* :class:`_ESMCSAELayer` — internal ``nn.Module`` that holds the weights for23  one ``(backbone, codebook_dim, k, layer)`` SAE. Not a published HF artifact;24  obtained only via ``model.layers["<idx>"]``.25"""26 27from __future__ import annotations28 29import os30from dataclasses import dataclass31from pathlib import Path32from typing import Optional33 34import torch35import torch.nn as nn36import torch.nn.functional as F37from safetensors.torch import load_file, save_file38 39from transformers.modeling_outputs import ModelOutput40from transformers.modeling_utils import PreTrainedModel41from transformers.utils import auto_docstring42from .configuration_esmc_sae import ESMCSAEConfig, ESMCSAEParams43 44 45@dataclass46@auto_docstring(47    custom_intro="""48    Output type of [`ESMCSAEModel`].49    """50)51class ESMCSAEOutput(ModelOutput):52    feature_magnitudes: torch.Tensor53    reconstruction_loss: Optional[torch.Tensor] = None54 55    def to_sparse(self) -> None:56        self.feature_magnitudes = self.feature_magnitudes.to_sparse()57 58 59class _ESMCSAELayer(nn.Module):60    """One backbone layer's SAE — internal building block of :class:`ESMCSAEModel`.61 62    Not exposed via ``AutoModel`` and not loadable on its own. Obtain one63    via ``model.layers["<layer_idx>"]`` after calling ``initialize_layers``.64    """65 66    def __init__(self, params: ESMCSAEParams):67        super().__init__()68        self.params = params69 70        self.W_enc = nn.Parameter(torch.empty(params.d_model, params.codebook_dim))71        self.W_dec = nn.Parameter(torch.empty(params.codebook_dim, params.d_model))72        self.b_dec = nn.Parameter(torch.zeros(params.d_model))73        # Per-feature normalization stats. Trained alongside the SAE for some74        # variants; for variants that don't ship them, leaving these as ones75        # makes ``_get_sae_outputs``'s ``features / max * idf`` a no-op.76        self.register_buffer("idf", torch.ones(params.codebook_dim))77        self.register_buffer("max", torch.ones(params.codebook_dim))78 79    @property80    def layer(self) -> int:81        """Backbone-layer index this SAE is trained against."""82        return self.params.layer83 84    def forward(self, x: torch.Tensor, **_kwargs: object) -> ESMCSAEOutput:85        del _kwargs86        x = self._zscore_normalize_representation(x)87 88        x_with_pre_encoder_bias = x - self.b_dec89        preactivations = F.relu(x_with_pre_encoder_bias @ self.W_enc)90 91        topk = torch.topk(preactivations, self.params.k, dim=-1)92        feature_magnitudes = torch.zeros_like(preactivations).scatter(93            -1, topk.indices, topk.values94        )95 96        reconstructed = feature_magnitudes @ self.W_dec + self.b_dec97 98        reconstruction_loss = (reconstructed - x).pow(2).mean(dim=-1)99 100        return ESMCSAEOutput(101            feature_magnitudes=feature_magnitudes,102            reconstruction_loss=reconstruction_loss,103        )104 105    def get_sae_output(106        self, layer_states: torch.Tensor, token_mask: torch.Tensor107    ) -> ESMCSAEOutput:108        _, _, v_len = layer_states.shape109        nonpad_states = layer_states[token_mask].view(-1, v_len)110        return self(nonpad_states)111 112    def _zscore_normalize_representation(self, x: torch.Tensor) -> torch.Tensor:113        x_mean = x.mean(dim=-1, keepdim=True)114        x = x - x_mean115        x_std = x.std(dim=-1, keepdim=True)116        return x / (x_std + 1e-5)117 118 119@auto_docstring120class ESMCSAEPreTrainedModel(PreTrainedModel):121    config_class = ESMCSAEConfig122    base_model_prefix = "esmc_sae"123 124 125@auto_docstring(126    custom_intro="""127    HF container holding one SAE per backbone layer, all sharing the same128    ``(d_model, codebook_dim, k)``.129 130    ``from_pretrained`` downloads the entire repo (every ``layer_{i}.safetensors``)131    into the local HF cache but does **not** load any weights into memory.132    Callers materialize the layers they actually need by calling133    :meth:`initialize_layers`. The full set is available on disk after the134    first call, so subsequent layer switches read from the local cache without135    re-downloading.136 137    Examples::138 139        model = ESMCSAEModel.from_pretrained(140            "biohub/esmc-6b-2024-12-sae-k64-codebook16384"141        )142        model.initialize_layers([60])                  # ~2.5 GB into memory143        out = model(layer_states, layer=60)            # forward through layer 60144        model.initialize_layers([45])                  # add layer 45 (cached locally)145        model.release_layer(60)                        # free layer 60146    """147)148class ESMCSAEModel(ESMCSAEPreTrainedModel):149    def __init__(self, config: ESMCSAEConfig):150        super().__init__(config)151        # Layers are populated lazily by ``initialize_layers``; the container152        # starts empty so ``from_pretrained`` doesn't materialize hundreds of153        # GB of unused parameters.154        self.layers = nn.ModuleDict()155        # Zero-element buffer that rides along with ``.to(device/dtype)``.156        # ``initialize_layers`` reads its current device/dtype so SAEs added157        # after ``model.to("cuda")`` land on CUDA without re-passing ``device=``.158        self.register_buffer("_device_marker", torch.empty(0), persistent=False)159        self._snapshot_dir: Optional[str] = None160        self.post_init()161 162    @classmethod163    def from_pretrained(  # type: ignore[override]164        cls, pretrained_model_name_or_path: str | os.PathLike, *model_args, **kwargs165    ) -> "ESMCSAEModel":166        """Download (or reuse cached) the full repo and return the model.167 168        By default no weights are read into memory and the caller must invoke169        :meth:`initialize_layers` before running :meth:`forward`. The single170        exception is when the repo ships exactly one layer: that layer is171        auto-loaded (honoring ``torch_dtype`` / ``device`` if passed) so the172        bare ``forward(x)`` call just works.173 174        Honored kwargs: ``revision``, ``cache_dir``, ``token``,175        ``allow_patterns``, ``local_files_only``, ``force_download`` (forwarded176        to ``snapshot_download``); ``torch_dtype`` and ``device`` (used by the177        single-layer auto-load path; otherwise pass them to178        :meth:`initialize_layers`). Behavioral kwargs that imply work we do179        not perform (``device_map``, ``low_cpu_mem_usage``,180        ``quantization_config``, ``attn_implementation``) raise so the user181        isn't silently misled. Other HF housekeeping kwargs (``config``,182        ``trust_remote_code``, ``adapter_kwargs``, …) are accepted and183        ignored — they only matter for the standard loader, which we bypass.184        """185        del model_args186        torch_dtype = kwargs.pop("torch_dtype", None)187        device = kwargs.pop("device", None)188        local_dir = _resolve_snapshot_dir(pretrained_model_name_or_path, kwargs)189        unsupported = {190            "device_map",191            "low_cpu_mem_usage",192            "quantization_config",193            "attn_implementation",194            "max_memory",195            "offload_folder",196            "offload_state_dict",197        } & kwargs.keys()198        if unsupported:199            raise TypeError(200                f"Unsupported kwargs to ESMCSAEModel.from_pretrained: "201                f"{sorted(unsupported)}. The standard HF loader is bypassed —"202                " call initialize_layers(..., device=, dtype=) instead."203            )204        config = ESMCSAEConfig.from_pretrained(local_dir)205        model = cls(config)206        model._snapshot_dir = str(local_dir)207        if device is not None:208            model.to(device)209        if torch_dtype is not None:210            model.to(torch_dtype)211        if len(config.available_layers) == 1:212            model.initialize_layers(list(config.available_layers))213        return model214 215    def initialize_layers(216        self,217        layers: list[int],218        *,219        device: torch.device | str | None = None,220        dtype: torch.dtype | None = None,221    ) -> None:222        """Load the requested layers from the local snapshot into memory.223 224        Layers already present in :attr:`self.layers` are skipped — calling225        ``initialize_layers([23])`` twice is idempotent. ``device`` / ``dtype``226        default to wherever the model itself lives (via the ``_device_marker``227        buffer that moves with ``.to(...)``), so the common pattern of228        ``model.to("cuda"); model.initialize_layers([7])`` Just Works.229        """230        assert self._snapshot_dir is not None, (231            "ESMCSAEModel has no snapshot directory — call "232            "from_pretrained first, or set _snapshot_dir manually."233        )234        if device is None:235            device = self._device_marker.device236        if dtype is None:237            dtype = self._device_marker.dtype238        snapshot_dir = Path(self._snapshot_dir)239        available = set(self.config.available_layers)240        for layer_idx in layers:241            key = str(layer_idx)242            if key in self.layers:243                continue244            if layer_idx not in available:245                raise KeyError(246                    f"Layer {layer_idx} is not in this repo. "247                    f"available_layers={sorted(available)}"248                )249            shard = snapshot_dir / f"layer_{layer_idx}.safetensors"250            if not shard.exists():251                raise FileNotFoundError(252                    f"Missing layer file {shard} — config lists layer "253                    f"{layer_idx} as available but the shard is not on disk."254                )255            params = ESMCSAEParams(256                d_model=self.config.d_model,257                codebook_dim=self.config.codebook_dim,258                k=self.config.k,259                layer=layer_idx,260            )261            # Build on the meta device so we don't allocate weights that262            # ``load_state_dict`` would immediately overwrite.263            with torch.device("meta"):264                layer = _ESMCSAELayer(params)265            layer.to_empty(device=device)266            layer.load_state_dict(load_file(str(shard)))267            layer.to(dtype=dtype)268            self.layers[key] = layer269 270    def release_layer(self, layer: int) -> None:271        """Drop the named layer from memory. No-op if not loaded."""272        key = str(layer)273        if key in self.layers:274            del self.layers[key]275 276    def loaded_layers(self) -> list[int]:277        """Sorted list of layer indices currently materialized in memory."""278        return sorted(int(k) for k in self.layers.keys())279 280    def forward(281        self, x: torch.Tensor, layer: int | None = None, **kwargs: object282    ) -> ESMCSAEOutput:283        if layer is None:284            if len(self.layers) == 1:285                # Unambiguous: exactly one layer loaded → use it.286                ((_only_key, only_layer),) = self.layers.items()287                return only_layer(x, **kwargs)288            if len(self.layers) == 0:289                raise RuntimeError(290                    "No layers loaded — call "291                    f"initialize_layers([...]) first. "292                    f"available_layers={self.config.available_layers}"293                )294            raise RuntimeError(295                "Multiple layers are loaded — please select one via "296                f"forward(x, layer=<idx>). Loaded layers: {self.loaded_layers()}"297            )298        key = str(layer)299        if key not in self.layers:300            raise KeyError(301                f"Layer {layer} is not loaded. Call "302                f"initialize_layers([{layer}]) first. Loaded layers: "303                f"{self.loaded_layers()}"304            )305        return self.layers[key](x, **kwargs)306 307    def save_pretrained(  # type: ignore[override]308        self, save_directory: str | os.PathLike, *args, **kwargs309    ) -> None:310        """Write ``config.json`` plus one ``layer_{i}.safetensors`` per loaded layer.311 312        Only layers currently in :attr:`self.layers` are written.313        ``available_layers`` in the saved config is synced to what's actually314        on disk so a ``release_layer`` + ``save_pretrained`` round-trip never315        advertises a layer whose shard is missing.316        """317        del args, kwargs318        save_directory = Path(save_directory)319        save_directory.mkdir(parents=True, exist_ok=True)320        # Sync available_layers to what we're about to write — never advertise321        # a layer that isn't on disk in this repo.322        self.config.available_layers = self.loaded_layers()323        self.config.save_pretrained(str(save_directory))324        for key, layer in self.layers.items():325            shard = save_directory / f"layer_{key}.safetensors"326            save_file(327                {328                    k: v.detach().cpu().contiguous()329                    for k, v in layer.state_dict().items()330                },331                str(shard),332            )333 334 335def _resolve_snapshot_dir(336    pretrained_model_name_or_path: str | os.PathLike, kwargs: dict337) -> str:338    """Local dir → return as-is; hub id → ``snapshot_download`` it.339 340    A directory only counts as "local" if it actually contains ``config.json``,341    so a stale subdir named like a hub id (``./biohub/esmc-...``)342    doesn't accidentally shadow the hub fetch.343 344    Pops the standard ``snapshot_download`` keyword args from ``kwargs`` so345    callers can forward them via ``from_pretrained``.346    """347    path = Path(pretrained_model_name_or_path)348    if path.is_dir() and (path / "config.json").exists():349        return str(path)350    from huggingface_hub import snapshot_download351 352    return snapshot_download(353        repo_id=str(pretrained_model_name_or_path),354        revision=kwargs.pop("revision", None),355        cache_dir=kwargs.pop("cache_dir", None),356        token=kwargs.pop("token", None),357        allow_patterns=kwargs.pop("allow_patterns", None),358        local_files_only=kwargs.pop("local_files_only", False),359        force_download=kwargs.pop("force_download", False),360    )361 362 363__all__ = ["ESMCSAEModel", "ESMCSAEOutput", "ESMCSAEPreTrainedModel"]364