diff --git a/README.md b/README.md index ec5f8658..235da2c8 100644 --- a/README.md +++ b/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=[ diff --git a/_assets/ESM3_README.md b/_assets/ESM3_README.md index 2e6421ee..481db27f 100644 --- a/_assets/ESM3_README.md +++ b/_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___________________________________________________________" diff --git a/cookbook/local/raw_forwards.py b/cookbook/local/raw_forwards.py index 35b7b094..8a6b2649 100644 --- a/cookbook/local/raw_forwards.py +++ b/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) diff --git a/esm/layers/attention.py b/esm/layers/attention.py index 718dea5f..be4e11f6 100644 --- a/esm/layers/attention.py +++ b/esm/layers/attention.py @@ -5,7 +5,9 @@ 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 @@ -13,6 +15,33 @@ 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. " diff --git a/esm/models/esm3.py b/esm/models/esm3.py index 05141f33..b7e796e6 100644 --- a/esm/models/esm3.py +++ b/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( diff --git a/esm/models/esmc/compatibility.py b/esm/models/esmc/compatibility.py index fea1e8b5..6823ae2b 100644 --- a/esm/models/esmc/compatibility.py +++ b/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) diff --git a/esm/models/esmc/model.py b/esm/models/esmc/model.py index dea2ee95..fdd093cf 100644 --- a/esm/models/esmc/model.py +++ b/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) diff --git a/esm/models/vqvae.py b/esm/models/vqvae.py index 9018251e..17ca1e09 100644 --- a/esm/models/vqvae.py +++ b/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 diff --git a/esm/pretrained.py b/esm/pretrained.py index 09389f06..9608ccd3 100644 --- a/esm/pretrained.py +++ b/esm/pretrained.py @@ -23,11 +23,13 @@ ESMC_300M, ESMC_600M, ) +from esm.utils.device import resolve_device ModelBuilder = Callable[[torch.device | str], nn.Module] -def ESM3_structure_encoder_v0(device: torch.device | str = "cpu"): +def ESM3_structure_encoder_v0(device: torch.device | str = "auto"): + device = resolve_device(device) with init_empty_weights(): model = StructureTokenEncoder( d_model=1024, n_heads=1, v_heads=128, n_layers=2, d_out=128, n_codes=4096 @@ -41,7 +43,8 @@ def ESM3_structure_encoder_v0(device: torch.device | str = "cpu"): return model -def ESM3_structure_decoder_v0(device: torch.device | str = "cpu"): +def ESM3_structure_decoder_v0(device: torch.device | str = "auto"): + device = resolve_device(device) with init_empty_weights(): model = StructureTokenDecoder(d_model=1280, n_heads=20, n_layers=30).eval() state_dict = torch.load( @@ -53,7 +56,8 @@ def ESM3_structure_decoder_v0(device: torch.device | str = "cpu"): return model -def ESM3_function_decoder_v0(device: torch.device | str = "cpu"): +def ESM3_function_decoder_v0(device: torch.device | str = "auto"): + device = resolve_device(device) with init_empty_weights(): model = FunctionTokenDecoder().eval() state_dict = torch.load( @@ -65,7 +69,8 @@ def ESM3_function_decoder_v0(device: torch.device | str = "cpu"): return model -def ESM3_sm_open_v0(device: torch.device | str = "cpu"): +def ESM3_sm_open_v0(device: torch.device | str = "auto"): + device = resolve_device(device) with init_empty_weights(): model = ESM3( d_model=1536, @@ -149,12 +154,11 @@ def ESMC_6B_202412( def load_local_model( - model_name: str, - device: torch.device = torch.device("cpu"), - use_flash_attn: bool = True, + model_name: str, device: torch.device | str = "auto", use_flash_attn: bool = True ) -> nn.Module: if model_name not in LOCAL_MODEL_REGISTRY: raise ValueError(f"Model {model_name} not found in local model registry.") + device = resolve_device(device) builder = LOCAL_MODEL_REGISTRY[model_name] kwargs = {} if "use_flash_attn" in inspect.signature(builder).parameters: diff --git a/esm/utils/device.py b/esm/utils/device.py new file mode 100644 index 00000000..56e513cf --- /dev/null +++ b/esm/utils/device.py @@ -0,0 +1,181 @@ +import os +import warnings + +import torch + +DeviceLike = torch.device | str | None + + +def is_mps_available() -> bool: + """Check if Apple Metal Performance Shaders (MPS) backend is available.""" + return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + + +def is_cuda_available() -> bool: + """Check if NVIDIA CUDA backend is available.""" + return torch.cuda.is_available() + + +def resolve_device(requested: str | torch.device | None = None) -> torch.device: + """Resolve a requested device string or torch.device to an available device. + + Supports 'auto', 'cuda', 'mps', and 'cpu', with graceful fallback + cuda -> mps -> cpu and clear warnings when falling back from an + explicitly requested device. + + Args: + requested: Device specification ('auto', 'cuda', 'mps', 'cpu', None, + or a torch.device instance). Defaults to 'auto'. + + Returns: + torch.device: The resolved device. + """ + if requested is None: + requested = "auto" + + if isinstance(requested, torch.device): + device_type = requested.type + device_index = requested.index + else: + req_str = str(requested).strip() + if ":" in req_str: + device_type, idx_str = req_str.split(":", 1) + try: + device_index = int(idx_str) + except ValueError: + device_index = None + else: + device_type = req_str + device_index = None + + device_type = device_type.lower() + + if device_type == "auto": + if is_cuda_available(): + return torch.device("cuda") + elif is_mps_available(): + return torch.device("mps") + else: + return torch.device("cpu") + + elif device_type == "cuda": + if is_cuda_available(): + if device_index is not None: + return torch.device(f"cuda:{device_index}") + return torch.device("cuda") + elif is_mps_available(): + warnings.warn( + f"CUDA device '{requested}' requested but CUDA is not available. " + "Falling back to Apple Metal (MPS).", + UserWarning, + stacklevel=2, + ) + return torch.device("mps") + else: + warnings.warn( + f"CUDA device '{requested}' requested but CUDA is not available. " + "Falling back to CPU.", + UserWarning, + stacklevel=2, + ) + return torch.device("cpu") + + elif device_type == "mps": + if is_mps_available(): + return torch.device("mps") + elif is_cuda_available(): + warnings.warn( + "Apple Metal (MPS) device requested but MPS is not available. " + "Falling back to CUDA.", + UserWarning, + stacklevel=2, + ) + return torch.device("cuda") + else: + warnings.warn( + "Apple Metal (MPS) device requested but MPS is not available. " + "Falling back to CPU.", + UserWarning, + stacklevel=2, + ) + return torch.device("cpu") + + elif device_type == "cpu": + return torch.device("cpu") + + else: + return torch.device(requested) + + +def get_default_model_dtype(device: torch.device | str) -> torch.dtype: + """Get the recommended default model dtype for a given device. + + On CUDA: Defaults to bfloat16 for efficiency. + On MPS: Defaults to float32 due to historical kernel coverage limitations. + Can be overridden via the ESM3_MPS_DTYPE environment variable + ('bfloat16'/'bf16' or 'float16'/'fp16'). + On CPU: Defaults to float32. + + Args: + device: Device to determine dtype for. + + Returns: + torch.dtype: The recommended torch dtype. + """ + if isinstance(device, str): + dev_type = device.lower().strip() + if dev_type == "auto": + dev = resolve_device("auto") + dev_type = dev.type + elif isinstance(device, torch.device): + dev_type = device.type + else: + dev_type = "cpu" + + if dev_type.startswith("cuda"): + return torch.bfloat16 + elif dev_type.startswith("mps"): + env_dtype = os.environ.get("ESM3_MPS_DTYPE", "").strip().lower() + if env_dtype in ("bfloat16", "bf16"): + return torch.bfloat16 + elif env_dtype in ("float16", "fp16", "half"): + return torch.float16 + return torch.float32 + else: + return torch.float32 + + +def empty_cache(device: torch.device | str | None = None) -> None: + """Clear cached memory for CUDA and/or MPS devices safely.""" + if is_cuda_available(): + try: + torch.cuda.empty_cache() + except RuntimeError: + pass + if ( + is_mps_available() + and hasattr(torch, "mps") + and hasattr(torch.mps, "empty_cache") + ): + try: + torch.mps.empty_cache() + except RuntimeError: + pass + + +def synchronize(device: torch.device | str | None = None) -> None: + """Synchronize device execution for CUDA and/or MPS devices safely.""" + if is_cuda_available(): + try: + torch.cuda.synchronize() + except RuntimeError: + pass + if ( + is_mps_available() + and hasattr(torch, "mps") + and hasattr(torch.mps, "synchronize") + ): + try: + torch.mps.synchronize() + except RuntimeError: + pass diff --git a/esm/widgets/utils/clients.py b/esm/widgets/utils/clients.py index 6c9846e9..e3daedcf 100644 --- a/esm/widgets/utils/clients.py +++ b/esm/widgets/utils/clients.py @@ -2,7 +2,6 @@ import huggingface_hub import huggingface_hub.errors -import torch from esm.models.esm3 import ESM3 from esm.sdk import ESM3ForgeInferenceClient @@ -14,7 +13,7 @@ def get_local_client() -> ESM3InferenceClient: huggingface_hub.whoami() except huggingface_hub.errors.LocalTokenNotFoundError: raise ValueError("Hugging Face token not found.") - return ESM3.from_pretrained(device=torch.device("cuda")) + return ESM3.from_pretrained(device="auto") def get_forge_client(model_name: str) -> ESM3InferenceClient: