feat: add Apple Metal (MPS) support

README.md

@@ -89,7 +89,7 @@ from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer
 # example GFP sequence
 sequences = ["MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"]
 
-model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-6B", device="cuda").eval()
+model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-6B", device="auto").eval()
 tokenizer = EsmcTokenizer()
 
 inputs = tokenizer(sequences, return_tensors="pt", padding=True)
@@ -170,7 +170,7 @@ from esm.models.esmc import EsmcForMaskedLM, EsmcSaeModel, EsmcTokenizer
 
 sequence = "MGSNKSKPKDASQRRRSLEPAENVHGAGGGAFPASQTPSKPASADGHRGPSAAFAPAAAEPKLFGGFNSSDTVTSPQRAGPLAGGVTTFVALYDYESRTETDLSFKKGERLQIVNNTEGDWWLAHSLSTGQTGYIPSNYVAPSDSIQAEEWYFGKITRRESERLLLNAENPRGTFLVRESETTKGAYCLSVSDFDNAKGLNVKHYKIRKLDSGGFYITSRTQFNSLQQLVAYYSKHADGLCHRLTTVCPTSKPQTQGLAKDAWEIPRESLRLEVKLGQGCFGEVWMGTWNGTTRVAIKTLKPGTMSPEAFLQEAQVMKKLRHEKLVQLYAVVSEEPIYIVTEYMSKGSLLDFLKGETGKYLRLPQLVDMAAQIASGMAYVERMNYVHRDLRAANILVGENLVCKVADFGLARLIEDNEYTARQGAKFPIKWTAPEAALYGRFTIKSDVWSFGILLTELTTKGRVPYPGMVNREVLDQVERGYRMPCPPECPESLHDLMCQCWRKEPEERPTFEYLQAFLEDYFTSTEPQYQPGENL"
 
-model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-6B", device="cuda").eval()
+model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-6B", device="auto").eval()
 tokenizer = EsmcTokenizer()
 sae = EsmcSaeModel.from_pretrained(
     "biohub/ESMC-6B-sae-k64-codebook16384",
@@ -237,7 +237,7 @@ HHAI_SEQ = (
     "YKVHPSTSQAYKQFGNSVVINVLQYIAYNIGSSLNFKPY"
 )
 
-model = EsmFold2Model.from_pretrained("biohub/ESMFold2", device="cuda").eval()
+model = EsmFold2Model.from_pretrained("biohub/ESMFold2", device="auto").eval()
 
 spi = StructurePredictionInput(
     sequences=[

_assets/ESM3_README.md

@@ -68,7 +68,7 @@ from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
 login()
 
 # This will download the model weights and instantiate the model on your machine.
-model: ESM3InferenceClient = ESM3.from_pretrained("esm3-sm-open-v1").to("cuda") # or "cpu"
+model: ESM3InferenceClient = ESM3.from_pretrained("esm3-sm-open-v1").to("auto")
 
 # Generate a completion for a partial Carbonic Anhydrase (2vvb)
 prompt = "___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________"

cookbook/local/raw_forwards.py

@@ -14,24 +14,24 @@
     InterProQuantizedTokenizer as EsmFunctionTokenizer,
 )
 from esm.tokenization.sequence_tokenizer import EsmSequenceTokenizer
+from esm.utils.device import empty_cache, resolve_device
 from esm.utils.structure.protein_chain import ProteinChain
 from esm.utils.types import FunctionAnnotation
 
+device = resolve_device("auto")
+
 
-# The open weights are bf16, so calling `forward` directly needs autocast; the SDK
-# generate/logits paths set it up for you.
 @torch.no_grad()
-@torch.autocast("cuda", dtype=torch.bfloat16)
 def inverse_folding_example():
     tokenizer = EsmSequenceTokenizer()
-    encoder = ESM3_structure_encoder_v0("cuda")
-    model = ESM3_sm_open_v0("cuda")
+    encoder = ESM3_structure_encoder_v0(device)
+    model = ESM3_sm_open_v0(device)
 
     chain = ProteinChain.from_rcsb("1utn", "A")
     coords, plddt, residue_index = chain.to_structure_encoder_inputs()
-    coords = coords.cuda()
-    plddt = plddt.cuda()
-    residue_index = residue_index.cuda()
+    coords = coords.to(device)
+    plddt = plddt.to(device)
+    residue_index = residue_index.to(device)
     _, structure_tokens = encoder.encode(coords, residue_index=residue_index)
 
     # Add BOS/EOS padding
@@ -50,11 +50,10 @@ def inverse_folding_example():
 
 
 @torch.no_grad()
-@torch.autocast("cuda", dtype=torch.bfloat16)
 def conditioned_prediction_example():
     tokenizers = get_esm3_model_tokenizers()
 
-    model = ESM3_sm_open_v0("cuda")
+    model = ESM3_sm_open_v0(device)
 
     # PDB 1UTN
     sequence = "MKTFIFLALLGAAVAFPVDDDDKIVGGYTCGANTVPYQVSLNSGYHFCGGSLINSQWVVSAAHCYKSGIQVRLGEDNINVVEGNEQFISASKSIVHPSYNSNTLNNDIMLIKLKSAASLNSRVASISLPTSCASAGTQCLISGWGNTKSSGTSYPDVLKCLKAPILSDSSCKSAYPGQITSNMFCAGYLEGGKDSCQGDSGGPVVCSGKLQGIVSWGSGCAQKNKPGVYTKVCNYVSWIKQTIASN"
@@ -80,8 +79,8 @@ def conditioned_prediction_example():
     function_tokens = tokenizers.function.tokenize(function_annotations, len(sequence))
     function_tokens = tokenizers.function.encode(function_tokens)
 
-    function_tokens = function_tokens.cuda().unsqueeze(0)
-    sequence_tokens = sequence_tokens.cuda().unsqueeze(0)
+    function_tokens = function_tokens.to(device).unsqueeze(0)
+    sequence_tokens = sequence_tokens.to(device).unsqueeze(0)
 
     output = model.forward(
         sequence_tokens=sequence_tokens, function_tokens=function_tokens
@@ -92,8 +91,8 @@ def conditioned_prediction_example():
 @torch.no_grad()
 def decode(sequence, output, sequence_tokens):
     # To save on VRAM, we load these in separate functions
-    decoder = ESM3_structure_decoder_v0("cuda")
-    function_decoder = ESM3_function_decoder_v0("cuda")
+    decoder = ESM3_structure_decoder_v0(device)
+    function_decoder = ESM3_function_decoder_v0(device)
     function_tokenizer = EsmFunctionTokenizer()
 
     # Generally not recommended to just argmax the logits, decode iteratively!
@@ -150,6 +149,6 @@ def decode(sequence, output, sequence_tokens):
     inverse_folding_example()
 
     sequence, output, sequence_tokens = conditioned_prediction_example()
-    torch.cuda.empty_cache()
+    empty_cache()
     # And then decode from tokenized representation to outputs:
     decode(sequence, output, sequence_tokens)

esm/layers/attention.py

@@ -5,14 +5,43 @@
 import torch.nn.functional as F
 from torch import nn
 
+from esm.layers import rotary
 from esm.layers.rotary import RotaryEmbedding, TritonRotaryEmbedding
+from esm.utils.device import is_cuda_available
 
 try:
     from flash_attn import flash_attn_varlen_qkvpacked_func
 except (ImportError, RuntimeError):
     flash_attn_varlen_qkvpacked_func = None  # ty:ignore[invalid-assignment]
 
 
+def _eager_scaled_dot_product_attention(
+    query_BHLD: torch.Tensor,
+    key_BHLD: torch.Tensor,
+    value_BHLD: torch.Tensor,
+    mask_BHLL: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    scale = query_BHLD.shape[-1] ** -0.5
+    attn_scores = torch.einsum("bhld,bhsd->bhls", query_BHLD, key_BHLD) * scale
+    if mask_BHLL is not None:
+        attn_scores = attn_scores.masked_fill(~mask_BHLL, float("-inf"))
+    attn_weights = torch.softmax(attn_scores, dim=-1)
+    context_BHLD = torch.einsum("bhls,bhsd->bhld", attn_weights, value_BHLD)
+    return context_BHLD, attn_weights
+
+
+def _triton_rotary_available() -> bool:
+    return is_cuda_available() and rotary.apply_triton_rotary is not None
+
+
+def _flash_attention_available(x: torch.Tensor, rotary_module: nn.Module) -> bool:
+    return (
+        flash_attn_varlen_qkvpacked_func is not None
+        and x.is_cuda
+        and isinstance(rotary_module, TritonRotaryEmbedding)
+    )
+
+
 class MultiHeadAttention(nn.Module):
     def __init__(
         self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True
@@ -72,13 +101,10 @@ def forward(
             mask_BHLL = None
 
         if output_attentions:
-            scale = self.d_head**-0.5
-            attn_scores = torch.einsum("bhld,bhsd->bhls", query_BHLD, key_BHLD) * scale
-            if mask_BHLL is not None:
-                attn_scores = attn_scores.masked_fill(~mask_BHLL, float("-inf"))
-            attn_weights = torch.softmax(attn_scores, dim=-1)
-            context_BHLD = torch.einsum("bhls,bhsd->bhld", attn_weights, value_BHLD)
-        else:
+            context_BHLD, attn_weights = _eager_scaled_dot_product_attention(
+                query_BHLD, key_BHLD, value_BHLD, mask_BHLL
+            )
+        elif query_BHLD.is_cuda:
             attn_weights = None
             if mask_BHLL is None:
                 # Shortcut, if we don't use attention biases then torch
@@ -90,6 +116,11 @@ def forward(
                 context_BHLD = F.scaled_dot_product_attention(
                     query_BHLD, key_BHLD, value_BHLD, mask_BHLL
                 )
+        else:
+            attn_weights = None
+            context_BHLD, _ = _eager_scaled_dot_product_attention(
+                query_BHLD, key_BHLD, value_BHLD, mask_BHLL
+            )
 
         context_BLD = einops.rearrange(context_BHLD, "b h s d -> b s (h d)")
 
@@ -104,12 +135,17 @@ def __init__(
             d_model=d_model, n_heads=n_heads, bias=bias, qk_layernorm=qk_layernorm
         )
 
-        # Flash attention rotary.
-        self.rotary = TritonRotaryEmbedding(d_model // n_heads)
+        if _triton_rotary_available():
+            self.rotary = TritonRotaryEmbedding(d_model // n_heads)
+        else:
+            self.rotary = RotaryEmbedding(d_model // n_heads)
 
     def forward(
         self, x, seq_id, output_attentions: bool = False
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not _flash_attention_available(x, self.rotary):
+            return super().forward(x, seq_id, output_attentions=output_attentions)
+
         if output_attentions:
             raise ValueError(
                 "FlashMultiHeadAttention does not support output_attentions=True. "

esm/models/esm3.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import os
 from functools import partial
 from typing import Callable
 
@@ -31,6 +32,7 @@
 from esm.utils.constants import esm3 as C
 from esm.utils.constants.models import ESM3_OPEN_SMALL, normalize_model_name
 from esm.utils.decoding import decode_protein_tensor
+from esm.utils.device import DeviceLike, get_default_model_dtype, resolve_device
 from esm.utils.generation import (
     _batch_forward,
     _sample_per_prompt,
@@ -106,9 +108,13 @@ def forward(
         rbf_16_fn = partial(rbf, v_min=0.0, v_max=1.0, n_bins=16)
         # the `masked_fill(padding_mask.unsqueeze(2), 0)` for the two below is unnecessary
         # as pad tokens never even interact with the "real" tokens (due to sequence_id)
-        plddt_embed = self.plddt_projection(rbf_16_fn(average_plddt))
+        plddt_embed = self.plddt_projection(
+            rbf_16_fn(average_plddt).to(self.plddt_projection.weight.dtype)
+        )
         structure_per_res_plddt = self.structure_per_res_plddt_projection(
-            rbf_16_fn(per_res_plddt)
+            rbf_16_fn(per_res_plddt).to(
+                self.structure_per_res_plddt_projection.weight.dtype
+            )
         )
 
         # Structure + "structural features" embeds
@@ -226,21 +232,61 @@ def __init__(
 
     @classmethod
     def from_pretrained(
-        cls, model_name: str = ESM3_OPEN_SMALL, device: torch.device | None = None
+        cls, model_name: str = ESM3_OPEN_SMALL, device: DeviceLike = None
     ) -> ESM3:
         from esm.pretrained import load_local_model
 
         model_name = normalize_model_name(model_name)
         if not model_name:
             raise ValueError(f"Model name {model_name} is not a valid ESM3 model name.")
-        if device is None:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        device = resolve_device(device if device is not None else "auto")
         model = load_local_model(model_name, device=device)
-        if device.type != "cpu":
-            model = model.to(torch.bfloat16)
+        model_dtype = get_default_model_dtype(device)
+        model = model.to(model_dtype)
         assert isinstance(model, ESM3)
         return model
 
+    def to(self, *args, **kwargs):
+        device = None
+        dtype = None
+        new_args = list(args)
+
+        if new_args:
+            first_arg = new_args[0]
+            if isinstance(first_arg, (str, torch.device)):
+                resolved = resolve_device(first_arg)
+                new_args[0] = resolved
+                device = resolved
+            elif isinstance(first_arg, torch.dtype):
+                dtype = first_arg
+
+        if len(new_args) > 1 and isinstance(new_args[1], torch.dtype):
+            dtype = new_args[1]
+
+        if "device" in kwargs and kwargs["device"] is not None:
+            resolved = resolve_device(kwargs["device"])
+            kwargs["device"] = resolved
+            device = resolved
+
+        if "dtype" in kwargs:
+            dtype = kwargs["dtype"]
+
+        res = super().to(*new_args, **kwargs)
+
+        if device is not None:
+            if self._structure_encoder is not None:
+                self._structure_encoder.to(device)
+            if self._structure_decoder is not None:
+                self._structure_decoder.to(device)
+            if self._function_decoder is not None:
+                self._function_decoder.to(device)
+
+            if dtype is None:
+                target_dtype = get_default_model_dtype(device)
+                res = res.to(target_dtype)
+
+        return res
+
     @property
     def device(self):
         return next(self.parameters()).device
@@ -533,10 +579,26 @@ def logits(
             # 1.0 if all coordinates at specific indices have valid non-nan values.
             per_res_plddt = input.coordinates.isfinite().all(dim=-1).any(dim=-1).float()
 
+        autocast_enabled = False
+        autocast_dtype = torch.bfloat16
+        if device.type == "cuda":
+            autocast_enabled = True
+            autocast_dtype = torch.bfloat16
+        elif device.type == "mps":
+            mps_dtype_str = os.environ.get("ESM3_MPS_DTYPE", "").strip().lower()
+            if mps_dtype_str in ("bfloat16", "bf16"):
+                autocast_enabled = True
+                autocast_dtype = torch.bfloat16
+            elif mps_dtype_str in ("float16", "fp16", "half"):
+                autocast_enabled = True
+                autocast_dtype = torch.float16
+
         with (
             torch.no_grad(),  # Assume no gradients for now...
-            torch.autocast(enabled=True, device_type=device.type, dtype=torch.bfloat16)
-            if device.type == "cuda"
+            torch.autocast(
+                enabled=autocast_enabled, device_type=device.type, dtype=autocast_dtype
+            )
+            if autocast_enabled
             else contextlib.nullcontext(),
         ):
             output = self(

esm/models/esmc/compatibility.py

@@ -39,6 +39,7 @@
     LogitsOutput,
 )
 from esm.utils.constants.models import ESMC_6B, ESMC_300M, ESMC_600M
+from esm.utils.device import DeviceLike, get_default_model_dtype, resolve_device
 from esm.utils.sampling import _BatchedESMProteinTensor
 
 _DEPRECATION_MESSAGE = (
@@ -156,18 +157,20 @@ def __init__(
     def from_pretrained(
         cls,
         model_name: str = ESMC_600M,
-        device: torch.device | str | None = None,
+        device: DeviceLike = None,
         use_flash_attn: bool = True,
     ) -> "ESMC":
-        if device is None:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        device = torch.device(device)
+        device = resolve_device(device if device is not None else "auto")
+        model_dtype = get_default_model_dtype(device)
         model = EsmcForMaskedLM.from_pretrained(
             _legacy_name_to_repo(model_name),
             device=device,
-            # The old loader cast to bf16 whenever it left the CPU.
-            dtype=torch.bfloat16 if device.type != "cpu" else None,
-            attn_implementation="flash_attention_2" if use_flash_attn else "sdpa",
+            dtype=model_dtype if model_dtype != torch.float32 else None,
+            attn_implementation=(
+                "flash_attention_2"
+                if (use_flash_attn and device.type == "cuda")
+                else "sdpa"
+            ),
         )
         return cls(model=model)
 

esm/models/esmc/model.py

@@ -28,6 +28,7 @@
 from esm.models.esmc.layers import EsmcRotaryEmbedding, EsmcTransformerStack
 from esm.models.esmc.sae import EsmcSaeLayer
 from esm.models.hub import HubPreTrainedModel, resolve_model_dir
+from esm.utils.device import resolve_device
 
 _SAFETENSORS_INDEX = "model.safetensors.index.json"
 _SAFETENSORS_SINGLE = "model.safetensors"
@@ -252,8 +253,7 @@ def from_pretrained(
             config = replace(config, **config_overrides)
         if attn_implementation is not None:
             config.attn_implementation = attn_implementation
-        if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = resolve_device(device if device is not None else "auto")
         # Sets the default device the encoder reads to choose its kernels.
         with torch.device(device):
             return cls._load_pretrained(local_dir, config, device=device, dtype=dtype)

esm/models/vqvae.py

@@ -280,7 +280,10 @@ def find_knn_edges(
                 (coords.shape[0], coords.shape[1]), device=coords.device
             ).long()
 
-        with torch.no_grad(), torch.amp.autocast("cuda", enabled=False):
+        with (
+            torch.no_grad(),
+            torch.autocast(device_type=coords.device.type, enabled=False),
+        ):
             ca = coords[..., 1, :]
             edges, edge_mask = knn_graph(
                 ca, coord_mask, padding_mask, sequence_id, no_knn=knn