Synthyra/ESMFold2
0423
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 model."""15 16import importlib17import math18import re19from dataclasses import dataclass20from typing import Optional, cast21 22import torch23import torch.nn as nn24from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss25from torch.nn import functional as F26 27from transformers.modeling_outputs import (28 MaskedLMOutput,29 ModelOutput,30 SequenceClassifierOutput,31 TokenClassifierOutput,32)33from transformers.modeling_utils import PreTrainedModel34from transformers.utils import (35 auto_docstring,36 can_return_tuple,37 is_flash_attn_2_available,38 logging,39)40from .configuration_esmc import ESMCConfig41from .modeling_esmc_sae import _ESMCSAELayer42 43logger = logging.get_logger(__name__)44 45_CONFIG_FOR_DOC = "ESMCConfig"46 47# Optional accelerated kernels. Pure-PyTorch fallbacks below if absent.48if is_flash_attn_2_available():49 flash_attn_module = importlib.import_module("flash_attn")50 flash_bert_padding = importlib.import_module("flash_attn.bert_padding")51 flash_attn_varlen_qkvpacked_func = (52 flash_attn_module.flash_attn_varlen_qkvpacked_func53 )54 pad_input = flash_bert_padding.pad_input55 unpad_input = flash_bert_padding.unpad_input56 57 _flash_attn_available = True58else:59 pad_input = unpad_input = flash_attn_varlen_qkvpacked_func = None60 _flash_attn_available = False61 62try:63 flash_rotary = importlib.import_module("flash_attn.ops.triton.rotary")64 apply_triton_rotary = flash_rotary.apply_rotary65 66 _flash_attn_rotary_available = torch.cuda.is_available()67except ImportError:68 apply_triton_rotary = None # type: ignore[assignment]69 _flash_attn_rotary_available = False70 71# Transformer Engine: fused LayerNorm+Linear / LayerNorm+MLP kernels with72# fp32 reduction inside the LayerNorm. Recommended on GPU for accurate bf1673# inference; without it the pure-PyTorch fallback drifts ~O(10) in fp32 and74# ~O(100) in bf16 on the unnormalized residual stream (perplexity stays75# within rounding noise).76try:77 te = importlib.import_module("transformer_engine.pytorch")78 79 _te_available = True80except ImportError:81 te = None # type: ignore[assignment]82 _te_available = False83 84# xformers: preferred SDPA implementation on GPU. Provides a fused85# bf16 attention kernel with deterministic reduction order. Flash86# Attention 2 and PyTorch's ``F.scaled_dot_product_attention`` are87# progressively-less-preferred fallbacks.88try:89 xops = importlib.import_module("xformers.ops")90 91 _xformers_available = True92except ImportError:93 xops = None # type: ignore[assignment]94 _xformers_available = False95 96# Flash Attention 2: secondary SDPA fallback. Used when xformers is not97# installed; fp16 / bf16 only.98if _flash_attn_available:99 flash_attn_func = flash_attn_module.flash_attn_func100else:101 flash_attn_func = None # type: ignore[assignment]102 103if not _te_available:104 logger.warning(105 "ESMC: transformer_engine is not installed; falling back to "106 "pure-PyTorch LayerNorm+Linear / LayerNorm+MLP. Outputs will differ "107 "numerically — measured on the unnormalized residual stream (before "108 "the final LayerNorm), ~O(10) max-diff in fp32 and ~O(100) in bf16; "109 "after the final LayerNorm these shrink to a few ULP and perplexity "110 "stays within rounding noise. Install with "111 "`pip install transformer-engine[pytorch]` to enable fused fp32-"112 "reduction LayerNorm."113 )114 115if not _xformers_available and not _flash_attn_available:116 logger.warning(117 "ESMC: neither xformers nor flash-attn is installed; falling back "118 "to PyTorch ``F.scaled_dot_product_attention``. The attention "119 "reduction order in bf16 differs from a fused kernel by ~1 bf16 "120 "ULP per attention block; compounded across the 80-block stack "121 "this reaches ~O(100) max-diff on the unnormalized residual stream. "122 "Install xformers (preferred) with `pip install xformers` for a "123 "fused attention kernel."124 )125 126if torch.cuda.is_available() and not _flash_attn_rotary_available:127 logger.warning(128 "ESMC: flash-attn rotary kernel not installed; falling back to "129 "pure-PyTorch RoPE. For faster GPU inference run `pip install flash-attn`."130 )131 132 133# ---------------------------------------------------------------------------134# Output dataclasses135# ---------------------------------------------------------------------------136 137 138@dataclass139class ESMCOutput(ModelOutput):140 """141 Args:142 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`):143 Sequence of hidden states at the output of the last layer, after layer normalisation.144 hidden_states (`torch.FloatTensor`, *optional*):145 Stacked hidden states for all encoder layers.146 Shape ``(n_layers, batch_size, sequence_length, d_model)``.147 Returned when ``output_hidden_states=True``.148 sae_outputs (`dict[str, torch.Tensor]`, *optional*):149 SAE feature magnitudes keyed by SAE model name (sparse tensors).150 Only populated when SAE models have been registered via151 ``add_sae_models`` and ``compute_sae=True``.152 attentions (`tuple(torch.FloatTensor)`, *optional*):153 Per-layer attention weights of shape154 ``(batch_size, num_heads, sequence_length, sequence_length)``.155 Returned when ``output_attentions=True``. Not available on the156 ``flash_attention_2`` path.157 """158 159 last_hidden_state: torch.FloatTensor | None = None160 hidden_states: torch.FloatTensor | None = None161 sae_outputs: dict[str, torch.Tensor] | None = None162 attentions: tuple[torch.FloatTensor, ...] | None = None163 164 165@dataclass166class ESMCMaskedLMOutput(MaskedLMOutput):167 """168 Args:169 loss (`torch.FloatTensor` of shape `(1,)`, *optional*):170 Masked language modelling loss. Returned when ``labels`` are provided.171 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, vocab_size)`):172 Prediction scores of the language modelling head.173 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`):174 Final hidden states after layer normalisation.175 hidden_states (`torch.FloatTensor`, *optional*):176 Stacked hidden states. Shape ``(n_layers, batch_size, sequence_length, d_model)``.177 sae_outputs (`dict[str, torch.Tensor]`, *optional*):178 SAE feature magnitudes keyed by SAE model name (sparse tensors).179 attentions (`tuple(torch.FloatTensor)`, *optional*):180 Per-layer attention weights of shape181 ``(batch_size, num_heads, sequence_length, sequence_length)``.182 Returned when ``output_attentions=True``.183 """184 185 loss: torch.FloatTensor | None = None186 logits: torch.FloatTensor | None = None187 last_hidden_state: torch.FloatTensor | None = None188 hidden_states: torch.FloatTensor | None = None189 sae_outputs: dict[str, torch.Tensor] | None = None190 attentions: tuple[torch.FloatTensor, ...] | None = None191 192 193@dataclass194class ESMCTokenClassifierOutput(TokenClassifierOutput):195 """196 Args:197 loss (`torch.FloatTensor` of shape `(1,)`, *optional*):198 Token classification loss. Returned when ``labels`` are provided.199 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_labels)`):200 Classification scores (before SoftMax).201 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`):202 Final hidden states after layer normalisation.203 hidden_states (`torch.FloatTensor`, *optional*):204 Stacked hidden states. Shape ``(n_layers, batch_size, sequence_length, d_model)``.205 sae_outputs (`dict[str, torch.Tensor]`, *optional*):206 SAE feature magnitudes keyed by SAE model name (sparse tensors).207 attentions (`tuple(torch.FloatTensor)`, *optional*):208 Per-layer attention weights of shape209 ``(batch_size, num_heads, sequence_length, sequence_length)``.210 Returned when ``output_attentions=True``.211 """212 213 loss: torch.FloatTensor | None = None214 logits: torch.FloatTensor | None = None215 last_hidden_state: torch.FloatTensor | None = None216 hidden_states: torch.FloatTensor | None = None217 sae_outputs: dict[str, torch.Tensor] | None = None218 attentions: tuple[torch.FloatTensor, ...] | None = None219 220 221@dataclass222class ESMCSequenceClassifierOutput(SequenceClassifierOutput):223 """224 Args:225 loss (`torch.FloatTensor` of shape `(1,)`, *optional*):226 Sequence classification loss. Returned when ``labels`` are provided.227 logits (`torch.FloatTensor` of shape `(batch_size, num_labels)`):228 Classification scores (before SoftMax).229 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, d_model)`):230 Final hidden states after layer normalisation.231 hidden_states (`torch.FloatTensor`, *optional*):232 Stacked hidden states. Shape ``(n_layers, batch_size, sequence_length, d_model)``.233 sae_outputs (`dict[str, torch.Tensor]`, *optional*):234 SAE feature magnitudes keyed by SAE model name (sparse tensors).235 attentions (`tuple(torch.FloatTensor)`, *optional*):236 Per-layer attention weights of shape237 ``(batch_size, num_heads, sequence_length, sequence_length)``.238 Returned when ``output_attentions=True``.239 """240 241 loss: torch.FloatTensor | None = None242 logits: torch.FloatTensor | None = None243 last_hidden_state: torch.FloatTensor | None = None244 hidden_states: torch.FloatTensor | None = None245 sae_outputs: dict[str, torch.Tensor] | None = None246 attentions: tuple[torch.FloatTensor, ...] | None = None247 248 249# ---------------------------------------------------------------------------250# Rotary position embedding helpers251# ---------------------------------------------------------------------------252 253 254def _rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:255 if not interleaved:256 x1, x2 = x.chunk(2, dim=-1)257 return torch.cat((-x2, x1), dim=-1)258 x1, x2 = x[..., ::2], x[..., 1::2]259 return torch.stack((-x2, x1), dim=-1).flatten(-2, -1)260 261 262def _apply_rotary_emb_torch(263 x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False264) -> torch.Tensor:265 """Apply rotary position embeddings (pure PyTorch, no Triton dependency).266 267 Args:268 x: ``(batch, seqlen, n_heads, head_dim)``269 cos: ``(seqlen, rotary_dim / 2)``270 sin: ``(seqlen, rotary_dim / 2)``271 """272 ro_dim = cos.shape[-1] * 2273 seqlen = x.size(1)274 cos = cos[:seqlen].unsqueeze(1).repeat(1, 1, 2)275 sin = sin[:seqlen].unsqueeze(1).repeat(1, 1, 2)276 return torch.cat(277 [278 x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim], interleaved) * sin,279 x[..., ro_dim:],280 ],281 dim=-1,282 )283 284 285class RotaryEmbedding(nn.Module):286 """Rotary position embeddings (RoPE) as described in `RoFormer`_.287 288 .. _RoFormer: https://arxiv.org/abs/2104.09864289 290 Args:291 dim: Size of a single attention head.292 base: Frequency base for the sinusoidal positions.293 interleaved: If ``True`` rotate adjacent pairs (GPT-J style) instead of294 splitting the head dimension in half (GPT-NeoX style).295 scaling_factor: Linear scaling factor applied to position indices.296 pos_idx_in_fp32: Compute position indices in float32 to avoid bf16297 rounding errors at large sequence lengths.298 """299 300 def __init__(301 self,302 dim: int,303 base: float = 10000.0,304 interleaved: bool = False,305 scale_base: float | None = None,306 scaling_factor: float = 1.0,307 pos_idx_in_fp32: bool = True,308 device=None,309 ):310 super().__init__()311 self.dim = dim312 self.base = base313 self.interleaved = interleaved314 self.scale_base = scale_base315 self.scaling_factor = scaling_factor316 self.pos_idx_in_fp32 = pos_idx_in_fp32317 318 self._seq_len_cached = 0319 self._cos_cached: torch.Tensor | None = None320 self._sin_cached: torch.Tensor | None = None321 self._cos_k_cached: torch.Tensor | None = None322 self._sin_k_cached: torch.Tensor | None = None323 324 self.reset_parameters(device=device)325 326 def reset_parameters(self, device=None):327 inv_freq = self._compute_inv_freq(device)328 self.register_buffer("inv_freq", inv_freq, persistent=False)329 arange = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)330 scale = (331 (arange + 0.4 * self.dim) / (1.4 * self.dim)332 if self.scale_base is not None333 else None334 )335 self.register_buffer("scale", scale, persistent=False)336 337 def _compute_inv_freq(self, device=None) -> torch.Tensor:338 return 1.0 / (339 self.base340 ** (341 torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)342 / self.dim343 )344 )345 346 def _update_cos_sin_cache(self, seqlen: int, device=None, dtype=None):347 if self.inv_freq.is_meta:348 self.reset_parameters(device=device)349 if (350 seqlen > self._seq_len_cached351 or self._cos_cached is None352 or self._cos_cached.device != device353 or self._cos_cached.dtype != dtype354 or (self.training and self._cos_cached.is_inference())355 ):356 self._seq_len_cached = seqlen357 if self.pos_idx_in_fp32:358 t = (359 torch.arange(seqlen, device=device, dtype=torch.float32)360 / self.scaling_factor361 )362 inv_freq = (363 self.inv_freq.to(torch.float32)364 if self.inv_freq.dtype != torch.float32365 else self.inv_freq366 )367 else:368 t = (369 torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) # type: ignore[call-overload]370 / self.scaling_factor371 )372 inv_freq = self.inv_freq373 freqs = torch.outer(t, inv_freq) # type: ignore[arg-type]374 375 if self.scale is None:376 self._cos_cached = torch.cos(freqs).to(dtype)377 self._sin_cached = torch.sin(freqs).to(dtype)378 else:379 _scale: torch.Tensor = self.scale # type: ignore[assignment]380 power = (381 torch.arange(seqlen, dtype=_scale.dtype, device=_scale.device)382 - seqlen // 2383 ) / self.scale_base # type: ignore[operator]384 scale = _scale.to(device=power.device) ** power.unsqueeze(-1)385 self._cos_cached = (torch.cos(freqs) * scale).to(dtype)386 self._sin_cached = (torch.sin(freqs) * scale).to(dtype)387 self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)388 self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)389 390 def _apply(self, fn, recurse=True):391 if self.inv_freq.is_meta:392 self.reset_parameters(device="cpu")393 result = super()._apply(fn, recurse=recurse)394 # Recompute inv_freq on the new device: CPU vs CUDA ``pow`` differ by395 # ~1 fp32 ULP, which compounds across attention layers. Keep this396 # buffer fp32 even when the module is cast to bf16/fp16; otherwise the397 # rounded RoPE frequencies drift from the internal ESMC path.398 new_inv_freq = self._compute_inv_freq(device=self.inv_freq.device)399 self.register_buffer("inv_freq", new_inv_freq, persistent=False)400 self._seq_len_cached = 0401 self._cos_cached = None402 self._sin_cached = None403 self._cos_k_cached = None404 self._sin_k_cached = None405 return result406 407 def forward(408 self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0409 ) -> tuple[torch.Tensor, torch.Tensor]:410 """Apply RoPE to query and key tensors.411 412 Args:413 q: ``(batch, seqlen, n_heads, head_dim)``414 k: ``(batch, seqlen, n_heads, head_dim)``415 seqlen_offset: Offset used in incremental decoding.416 417 Returns:418 Tuple of rotated ``(q, k)`` tensors with the same shape as the inputs.419 """420 self._update_cos_sin_cache(421 q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype422 )423 assert self._cos_cached is not None and self._sin_cached is not None424 425 if self.scale is not None:426 raise NotImplementedError("XPos scaling is not supported in this path.")427 428 cos = self._cos_cached[seqlen_offset:]429 sin = self._sin_cached[seqlen_offset:]430 431 if _flash_attn_rotary_available and q.device.type == "cuda":432 q_rot = apply_triton_rotary(q, cos, sin, interleaved=self.interleaved) # type: ignore[misc]433 k_rot = apply_triton_rotary(k, cos, sin, interleaved=self.interleaved) # type: ignore[misc]434 else:435 q_rot = _apply_rotary_emb_torch(q, cos, sin, self.interleaved)436 k_rot = _apply_rotary_emb_torch(k, cos, sin, self.interleaved)437 return q_rot, k_rot438 439 440class _TritonRotaryEmbedding(RotaryEmbedding):441 """RoPE variant that delegates to the Flash-Attention Triton kernel.442 443 Only used inside :class:`_FlashMultiHeadAttention` when Flash Attention 2444 is available. The ``forward`` signature differs from :class:`RotaryEmbedding`445 because Flash Attention packs Q, K, V together.446 """447 448 def forward(449 self, qkv: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int450 ) -> torch.Tensor: # type: ignore[override]451 """Apply RoPE in-place to a packed ``(N, 3, n_heads, head_dim)`` tensor."""452 self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)453 assert self._cos_cached is not None and self._sin_cached is not None454 assert apply_triton_rotary is not None455 456 apply_triton_rotary(457 qkv[:, 0],458 self._cos_cached,459 self._sin_cached,460 cu_seqlens=cu_seqlens,461 max_seqlen=max_seqlen,462 inplace=True,463 )464 apply_triton_rotary(465 qkv[:, 1],466 self._cos_cached,467 self._sin_cached,468 cu_seqlens=cu_seqlens,469 max_seqlen=max_seqlen,470 inplace=True,471 )472 return qkv473 474 475# ---------------------------------------------------------------------------476# Feed-forward network helpers477# ---------------------------------------------------------------------------478 479 480def _swiglu_hidden_dim(expansion_ratio: float, d_model: int) -> int:481 """Round hidden dim to the nearest multiple of 256 after applying expansion_ratio."""482 return int(((expansion_ratio * d_model) + 255) // 256 * 256)483 484 485class _SwiGLU(nn.Module):486 """SwiGLU activation: ``silu(x1) * x2`` where ``x`` is split along the last dim."""487 488 def forward(self, x: torch.Tensor) -> torch.Tensor:489 x1, x2 = x.chunk(2, dim=-1)490 return F.silu(x1) * x2491 492 493class _PyTorchLayerNormLinear(nn.Module):494 """LayerNorm followed by a Linear projection, sharing the parameter495 names ``layer_norm_weight``, ``layer_norm_bias`` and ``weight`` so the496 state-dict layout matches the accelerated TE module loaded on GPU.497 """498 499 def __init__(self, d_in: int, d_out: int, eps: float = 1e-5) -> None:500 super().__init__()501 self.d_in = d_in502 self.eps = eps503 self.layer_norm_weight = nn.Parameter(torch.ones(d_in))504 self.layer_norm_bias = nn.Parameter(torch.zeros(d_in))505 self.weight = nn.Parameter(torch.empty(d_out, d_in))506 nn.init.normal_(self.weight, std=0.02)507 508 def forward(self, x: torch.Tensor) -> torch.Tensor:509 x = F.layer_norm(510 x, (self.d_in,), self.layer_norm_weight, self.layer_norm_bias, self.eps511 )512 return F.linear(x, self.weight)513 514 515class _PyTorchLayerNormMLP(nn.Module):516 """LayerNorm + SwiGLU MLP, sharing the parameter names517 ``layer_norm_weight``, ``layer_norm_bias``, ``fc1_weight``,518 ``fc2_weight`` so the state-dict layout matches the accelerated TE519 module loaded on GPU.520 """521 522 def __init__(523 self, hidden_size: int, ffn_hidden_size: int, eps: float = 1e-5524 ) -> None:525 super().__init__()526 self.hidden_size = hidden_size527 self.ffn_hidden_size = ffn_hidden_size528 self.eps = eps529 self.layer_norm_weight = nn.Parameter(torch.ones(hidden_size))530 self.layer_norm_bias = nn.Parameter(torch.zeros(hidden_size))531 self.fc1_weight = nn.Parameter(torch.empty(2 * ffn_hidden_size, hidden_size))532 self.fc2_weight = nn.Parameter(torch.empty(hidden_size, ffn_hidden_size))533 nn.init.normal_(self.fc1_weight, std=0.02)534 nn.init.normal_(self.fc2_weight, std=0.02)535 536 def forward(self, x: torch.Tensor) -> torch.Tensor:537 x = F.layer_norm(538 x,539 (self.hidden_size,),540 self.layer_norm_weight,541 self.layer_norm_bias,542 self.eps,543 )544 x = F.linear(x, self.fc1_weight)545 x1, x2 = x.chunk(2, dim=-1)546 x = F.silu(x1) * x2547 return F.linear(x, self.fc2_weight)548 549 550def _swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Module:551 """LayerNorm + SwiGLU MLP. Uses Transformer Engine's fused LN+MLP when552 available; otherwise returns the pure-PyTorch fallback with matching553 state-dict layout."""554 assert not bias, "ESMC was trained with bias=False; bias=True not supported"555 hidden = _swiglu_hidden_dim(expansion_ratio, d_model)556 if _te_available:557 return te.LayerNormMLP( # type: ignore[union-attr]558 hidden_size=d_model,559 ffn_hidden_size=hidden,560 bias=bias,561 activation="swiglu",562 init_method=None,563 output_layer_init_method=None,564 )565 return _PyTorchLayerNormMLP(hidden_size=d_model, ffn_hidden_size=hidden)566 567 568def _make_attn_layernorm_qkv(d_model: int, bias: bool) -> nn.Module:569 """LayerNorm + fused QKV projection. Uses Transformer Engine when570 available; pure-PyTorch fallback otherwise."""571 assert not bias, "ESMC was trained with bias=False; bias=True not supported"572 if _te_available:573 return te.LayerNormLinear( # type: ignore[union-attr]574 d_model, d_model * 3, bias=bias, init_method=None575 )576 return _PyTorchLayerNormLinear(d_model, d_model * 3)577 578 579def _make_attn_out_proj(d_model: int, bias: bool) -> nn.Module:580 """Attention output projection. Uses Transformer Engine when available;581 pure-PyTorch ``nn.Linear`` otherwise."""582 if _te_available:583 return te.Linear( # type: ignore[union-attr]584 d_model, d_model, bias=bias, init_method=None585 )586 return nn.Linear(d_model, d_model, bias=bias)587 588 589def _gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Sequential:590 hidden = int(expansion_ratio * d_model)591 return nn.Sequential(592 nn.LayerNorm(d_model),593 nn.Linear(d_model, hidden, bias=bias),594 nn.GELU(),595 nn.Linear(hidden, d_model, bias=bias),596 )597 598 599# ---------------------------------------------------------------------------600# Attention601# ---------------------------------------------------------------------------602 603 604def _scaled_dot_product_attention(605 q: torch.Tensor,606 k: torch.Tensor,607 v: torch.Tensor,608 *,609 n_heads: int,610 d_head: int,611 seq_id: torch.Tensor | None,612) -> torch.Tensor:613 """Scaled dot-product attention with optional chain-aware mask.614 615 Dispatches in order of preference:616 1. xformers ``memory_efficient_attention`` — preferred fused kernel,617 requires ``xformers``, no chain mask.618 2. Flash Attention 2 (``flash_attn.flash_attn_func``) — secondary619 fused kernel, requires ``flash-attn``, no chain mask, fp16 /620 bf16 only.621 3. PyTorch's ``F.scaled_dot_product_attention`` — last-resort path;622 also handles the chain-aware mask when ``seq_id`` is present623 and the fp32 path that Flash Attention 2 does not support.624 """625 if seq_id is None and _xformers_available:626 b, s, _ = q.shape627 q4 = q.view(b, s, n_heads, d_head)628 k4 = k.view(b, s, n_heads, d_head)629 v4 = v.view(b, s, n_heads, d_head)630 context = xops.memory_efficient_attention( # type: ignore[union-attr]631 q4, k4, v4, attn_bias=None, scale=d_head**-0.5632 )633 return context.reshape(b, s, n_heads * d_head)634 if (635 seq_id is None636 and _flash_attn_available637 and q.dtype in (torch.float16, torch.bfloat16)638 ):639 b, s, _ = q.shape640 q4 = q.view(b, s, n_heads, d_head)641 k4 = k.view(b, s, n_heads, d_head)642 v4 = v.view(b, s, n_heads, d_head)643 context = flash_attn_func( # type: ignore[misc]644 q4, k4, v4, dropout_p=0.0, softmax_scale=d_head**-0.5645 )646 return context.reshape(b, s, n_heads * d_head) # type: ignore[union-attr]647 b, s, _ = q.shape648 q = q.view(b, s, n_heads, -1).transpose(1, 2)649 k = k.view(b, s, n_heads, -1).transpose(1, 2)650 v = v.view(b, s, n_heads, -1).transpose(1, 2)651 if seq_id is not None:652 mask = (seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)).unsqueeze(1)653 context = F.scaled_dot_product_attention(q, k, v, mask)654 else:655 context = F.scaled_dot_product_attention(q, k, v)656 _, h, _, d_out = context.shape657 return context.transpose(1, 2).reshape(b, s, h * d_out)658 659 660class MultiHeadAttention(nn.Module):661 """Multi-head self-attention with QK LayerNorm and RoPE.662 663 Args:664 d_model: Model hidden dimension.665 n_heads: Number of attention heads.666 bias: Whether to use bias in linear layers.667 qk_layernorm: Whether to apply LayerNorm to queries and keys before668 computing attention scores.669 """670 671 def __init__(672 self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True673 ):674 super().__init__()675 self.d_model = d_model676 self.n_heads = n_heads677 self.d_head = d_model // n_heads678 679 assert not bias, "ESMC was trained with bias=False; bias=True not supported"680 self.layernorm_qkv = _make_attn_layernorm_qkv(d_model, bias)681 self.out_proj = _make_attn_out_proj(d_model, bias)682 683 if qk_layernorm:684 self.q_ln = nn.LayerNorm(d_model, bias=bias)685 self.k_ln = nn.LayerNorm(d_model, bias=bias)686 else:687 self.q_ln = nn.Identity()688 self.k_ln = nn.Identity()689 690 self.rotary = RotaryEmbedding(d_model // n_heads)691 692 def _apply_rotary(693 self, q: torch.Tensor, k: torch.Tensor694 ) -> tuple[torch.Tensor, torch.Tensor]:695 q = q.unflatten(-1, (self.n_heads, self.d_head))696 k = k.unflatten(-1, (self.n_heads, self.d_head))697 q, k = self.rotary(q, k)698 q = q.flatten(-2, -1)699 k = k.flatten(-2, -1)700 return q, k701 702 def forward(703 self,704 x: torch.Tensor,705 seq_id: torch.Tensor | None,706 output_attentions: bool = False,707 ) -> tuple[torch.Tensor, torch.Tensor | None]:708 """Return ``(context, attn_weights)``.709 710 ``attn_weights`` is ``None`` unless ``output_attentions=True`` — the711 fused SDPA backends (xformers, flash-attn 2, ``F.scaled_dot_product_attention``)712 don't expose attention probabilities, so capturing them forces a713 materialized ``softmax(Q @ K.T / sqrt(d)) @ V`` path with shape714 ``(B, H, L, L)``.715 """716 qkv = self.layernorm_qkv(x)717 q, k, v = torch.chunk(qkv, 3, dim=-1)718 q = self.q_ln(q).to(q.dtype)719 k = self.k_ln(k).to(q.dtype)720 q, k = self._apply_rotary(q, k)721 722 b, s, _ = q.shape723 724 if output_attentions:725 # Manual SDPA so attention probabilities are observable.726 q4 = q.view(b, s, self.n_heads, self.d_head).transpose(1, 2)727 k4 = k.view(b, s, self.n_heads, self.d_head).transpose(1, 2)728 v4 = v.view(b, s, self.n_heads, self.d_head).transpose(1, 2)729 scale = self.d_head**-0.5730 attn_scores = (q4 @ k4.transpose(-2, -1)) * scale731 if seq_id is not None:732 mask = (seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)).unsqueeze(1)733 attn_scores = attn_scores.masked_fill(~mask, float("-inf"))734 attn_weights = torch.softmax(attn_scores, dim=-1)735 context = (attn_weights @ v4).transpose(1, 2).reshape(b, s, -1)736 return self.out_proj(context), attn_weights737 738 context = _scaled_dot_product_attention(739 q, k, v, n_heads=self.n_heads, d_head=self.d_head, seq_id=seq_id740 )741 return self.out_proj(context), None742 743 744class _FlashMultiHeadAttention(MultiHeadAttention):745 """Flash-Attention 2 variant of :class:`MultiHeadAttention`."""746 747 def __init__(748 self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True749 ):750 super().__init__(751 d_model=d_model, n_heads=n_heads, bias=bias, qk_layernorm=qk_layernorm752 )753 self.rotary = _TritonRotaryEmbedding(d_model // n_heads)754 755 def forward(756 self,757 x: torch.Tensor,758 seq_id: torch.Tensor | None,759 output_attentions: bool = False,760 ) -> tuple[torch.Tensor, torch.Tensor | None]:761 if output_attentions:762 raise ValueError(763 "output_attentions=True is not supported with "764 "attn_implementation='flash_attention_2'. "765 "Re-load the model with attn_implementation='sdpa' (or 'eager')."766 )767 assert seq_id is not None and seq_id.dtype == torch.bool768 769 seqlens = seq_id.sum(dim=-1, dtype=torch.int32)770 cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))771 max_seqlen = int(seqlens.max().item())772 773 qkv = self.layernorm_qkv(x)774 q, k, v = torch.chunk(qkv, 3, dim=-1)775 q = self.q_ln(q).to(q.dtype)776 k = self.k_ln(k).to(q.dtype)777 778 # ``q``/``k``/``v`` are 2D ``(T, D)`` here: the parent ``ESMCModel.forward``779 # calls ``unpad_input`` before the transformer stack to produce the780 # varlen-flat layout that ``flash_attn_varlen_qkvpacked_func`` requires.781 T = q.shape[0]782 qkv_packed = torch.stack([q, k, v], dim=1).view(T, 3, self.n_heads, self.d_head)783 qkv_packed = self.rotary(qkv_packed, cu_seqlens, max_seqlen)784 785 context = flash_attn_varlen_qkvpacked_func( # type: ignore[misc]786 qkv_packed, cu_seqlens, max_seqlen, softmax_scale=self.d_head**-0.5787 )788 n_out, h_out, d_out = context.shape # type: ignore[union-attr]789 return (790 self.out_proj(context.reshape(n_out, h_out * d_out)), # type: ignore[union-attr]791 None,792 )793 794 795# ---------------------------------------------------------------------------796# Transformer blocks797# ---------------------------------------------------------------------------798 799 800class UnifiedTransformerBlock(nn.Module):801 """Single transformer block: pre-norm attention + pre-norm FFN with residual scaling.802 803 Args:804 d_model: Hidden dimension.805 n_heads: Number of attention heads.806 use_flash_attn: Use Flash Attention 2 kernel if available.807 bias: Whether linear layers include bias terms.808 expansion_ratio: Hidden-dim expansion ratio for the FFN.809 residue_scaling_factor: Scales residual connections to stabilise deep810 networks (``1 / sqrt(n_layers / 36)`` is the ESM3 scheme).811 qk_layernorm: Whether to apply QK LayerNorm in attention.812 ffn_type: Feed-forward activation: ``"swiglu"`` or ``"gelu"``.813 """814 815 def __init__(816 self,817 d_model: int,818 n_heads: int,819 use_flash_attn: bool = False,820 bias: bool = False,821 expansion_ratio: float = 4.0,822 residue_scaling_factor: float = 1.0,823 qk_layernorm: bool = True,824 ffn_type: str = "swiglu",825 ):826 super().__init__()827 828 attn_cls = _FlashMultiHeadAttention if use_flash_attn else MultiHeadAttention829 self.attn = attn_cls(d_model, n_heads, bias=bias, qk_layernorm=qk_layernorm)830 831 if ffn_type == "swiglu":832 self.ffn = _swiglu_ln_ffn(d_model, expansion_ratio, bias)833 elif ffn_type == "gelu":834 self.ffn = _gelu_ln_ffn(d_model, expansion_ratio, bias)835 else:836 raise ValueError(837 f"Unknown ffn_type: {ffn_type!r}. Choose 'swiglu' or 'gelu'."838 )839 840 self.scaling_factor = residue_scaling_factor841 842 def forward(843 self,844 x: torch.Tensor,845 sequence_id: torch.Tensor | None,846 output_attentions: bool = False,847 ) -> tuple[torch.Tensor, torch.Tensor | None]:848 """849 Args:850 x: ``(batch, seq_len, d_model)``851 sequence_id: ``(batch, seq_len)`` chain-ID tensor used to restrict852 attention to tokens within the same chain. SDPA blocks accept853 an integer tensor (``-1`` marks padding); the flash-attn block854 takes a ``bool`` padding mask — the caller selects which.855 ``None`` skips chain-aware masking entirely (fast path).856 output_attentions: When ``True``, returns the per-head attention857 weights for this block alongside the residual output.858 859 Returns:860 ``(output, attn_weights_or_None)``. Shape of ``output`` is861 ``(batch, seq_len, d_model)``; ``attn_weights`` shape is862 ``(batch, num_heads, seq_len, seq_len)`` or ``None``.863 """864 attn_out, attn_weights = self.attn(865 x, sequence_id, output_attentions=output_attentions866 )867 x = x + attn_out / self.scaling_factor868 x = x + self.ffn(x) / self.scaling_factor869 return x, attn_weights870 871 872class TransformerStack(nn.Module):873 """Stack of :class:`UnifiedTransformerBlock` layers with a final LayerNorm.874 875 Args:876 d_model: Hidden dimension.877 n_heads: Number of attention heads.878 n_layers: Number of transformer blocks.879 scale_residue: When ``True`` apply ESM3 residue scaling880 ``sqrt(n_layers / 36)`` to each block.881 bias: Bias flag forwarded to every sub-module.882 qk_layernorm: QK LayerNorm flag forwarded to every block.883 ffn_type: FFN activation type (``"swiglu"`` or ``"gelu"``).884 expansion_ratio: FFN expansion ratio.885 use_flash_attn: Use Flash Attention 2 kernel when available.886 """887 888 def __init__(889 self,890 d_model: int,891 n_heads: int,892 n_layers: int,893 scale_residue: bool = True,894 bias: bool = False,895 qk_layernorm: bool = True,896 ffn_type: str = "swiglu",897 expansion_ratio: float = 8 / 3,898 use_flash_attn: bool = False,899 ):900 super().__init__()901 self.blocks = nn.ModuleList(902 [903 UnifiedTransformerBlock(904 d_model,905 n_heads,906 use_flash_attn=use_flash_attn,907 residue_scaling_factor=math.sqrt(n_layers / 36)908 if scale_residue909 else 1.0,910 expansion_ratio=expansion_ratio,911 bias=bias,912 qk_layernorm=qk_layernorm,913 ffn_type=ffn_type,914 )915 for _ in range(n_layers)916 ]917 )918 self.norm = nn.LayerNorm(d_model, bias=False)919 920 def forward(921 self,922 x: torch.Tensor,923 sequence_id: torch.Tensor | None = None,924 layers_to_collect: list[int] | None = None,925 output_attentions: bool = False,926 ) -> tuple[927 torch.Tensor,928 torch.Tensor,929 tuple[torch.Tensor, ...],930 tuple[torch.Tensor, ...] | None,931 ]:932 """Run the full transformer stack.933 934 Args:935 x: ``(batch, seq_len, d_model)``936 sequence_id: Optional chain-id tensor forwarded to each block.937 layers_to_collect: Layer indices (0-based pre-block inputs plus938 ``n_layers`` for the post-norm output) whose hidden states939 should be returned.940 output_attentions: When ``True``, collects the per-block attention941 weights and returns them as the fourth tuple element.942 943 Returns:944 ``(post_norm, pre_norm, hidden_states, attentions)`` where945 ``hidden_states`` is a (possibly empty) tuple of tensors and946 ``attentions`` is a tuple of per-block ``(B, H, L, L)`` tensors947 or ``None`` when ``output_attentions`` is ``False``.948 """949 if layers_to_collect is None:950 layers_to_collect = []951 952 collected: list[torch.Tensor] = []953 all_attentions: list[torch.Tensor] = []954 for layer_idx, block in enumerate(self.blocks):955 if layer_idx in layers_to_collect:956 collected.append(x)957 x, attn_weights = block(x, sequence_id, output_attentions=output_attentions)958 if output_attentions and attn_weights is not None:959 all_attentions.append(attn_weights)960 961 norm_x = self.norm(x)962 if len(self.blocks) in layers_to_collect:963 collected.append(norm_x)964 965 attentions = tuple(all_attentions) if output_attentions else None966 return norm_x, x, tuple(collected), attentions967 968 969# ---------------------------------------------------------------------------970# Pre-trained model base class971# ---------------------------------------------------------------------------972 973 974@auto_docstring975class ESMCPreTrainedModel(PreTrainedModel):976 """Base class for ESMC models.977 978 Handles weight initialisation and declares module-level capabilities.979 """980 981 config_class = ESMCConfig982 base_model_prefix = "esmc"983 supports_gradient_checkpointing = False984 _supports_sdpa = True985 _supports_flash_attn = True986 _supports_attention_backend = True987 _no_split_modules = ["UnifiedTransformerBlock"]988 _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"]989 990 def _init_weights(self, module: nn.Module):991 std = self.config.initializer_range992 if isinstance(module, nn.Linear):993 module.weight.data.normal_(mean=0.0, std=std)994 if module.bias is not None:995 module.bias.data.zero_()996 elif isinstance(module, RotaryEmbedding):997 module.reset_parameters(device=self.device)998 999 1000# ---------------------------------------------------------------------------1001# Base encoder model1002# ---------------------------------------------------------------------------1003 1004 1005@auto_docstring1006class ESMCModel(ESMCPreTrainedModel):1007 """The bare ESMC encoder outputting raw hidden states.1008 1009 ESMC is a protein language model trained by EvolutionaryScale using a1010 masked-token objective over amino acid sequences. The architecture is a1011 standard Transformer encoder with RoPE positional embeddings, QK LayerNorm,1012 and SwiGLU feed-forward networks.1013 1014 Args:1015 config: An :class:`ESMCConfig` instance.1016 """1017 1018 def __init__(self, config: ESMCConfig):1019 super().__init__(config)1020 self._use_flash_attn = (1021 _flash_attn_available and config._attn_implementation == "flash_attention_2"1022 )1023 self.embed = nn.Embedding(config.vocab_size, config.d_model)1024 self.transformer = TransformerStack(1025 config.d_model,1026 config.n_heads,1027 config.n_layers,1028 use_flash_attn=self._use_flash_attn,1029 )1030 self._sae_models: nn.ModuleDict = nn.ModuleDict()1031 self.post_init()1032 1033 def get_input_embeddings(self) -> nn.Embedding:1034 return self.embed1035 1036 def set_input_embeddings(self, value: nn.Embedding):1037 self.embed = value1038 1039 def add_sae_models(self, sae_models: list[_ESMCSAELayer]) -> None:1040 """Register one or more SAEs obtained from an :class:`ESMCSAEModel`.1041 1042 Each is keyed by ``f"layer{N}"`` (the backbone-layer index ``N`` the1043 SAE is trained against, set by1044 :meth:`ESMCSAEModel.initialize_layers`). Attaching two SAEs for the1045 same backbone layer raises — only one SAE per layer can be active.1046 1047 Example::1048 1049 sae = ESMCSAEModel.from_pretrained(1050 "biohub/esmc-600m-2024-12-sae-k64-codebook16384"1051 )1052 sae.initialize_layers([27, 33])1053 model.add_sae_models([sae.layers["27"], sae.layers["33"]])1054 """1055 for layer in sae_models:1056 assert isinstance(layer, _ESMCSAELayer), (1057 f"Expected an SAE layer (model.layers['<idx>']), got "1058 f"{type(layer).__name__}."1059 )1060 key = f"layer{int(layer.layer)}"1061 if key in self._sae_models:1062 raise ValueError(1063 f"An SAE is already registered at {key!r}. Only one SAE "1064 "per backbone layer can be active — pick a different "1065 "layer on one of them, or attach in a fresh model."1066 )1067 self._sae_models[key] = layer1068 1069 _SAE_KEY_RE = re.compile(r"layer(\d+)")1070 1071 def _get_sae_layer_num_requested(self, model_name: str) -> int:1072 """Recover the backbone-layer index from a key written by1073 :meth:`add_sae_models` (``"layer{N}"`` → ``N``)."""1074 match = self._SAE_KEY_RE.fullmatch(model_name)1075 assert (1076 match is not None1077 ), f"Unexpected SAE key {model_name!r}; expected 'layer{{N}}'."1078 return int(match.group(1))1079 1080 def _validate_sae_inputs(self, input_ids: torch.Tensor) -> None:1081 assert torch.all(input_ids != self.config.mask_token_id), (1082 "SAE inputs must not contain mask tokens. "1083 "SAEs were trained on unmasked sequences."1084 )1085 1086 def _get_sae_outputs(1087 self,1088 hidden_states: torch.Tensor,1089 layers_to_collect: list[int],1090 token_mask: torch.Tensor,1091 normalize_sae: bool = False,1092 ) -> dict[str, torch.Tensor]:1093 """Run all registered SAEs and return their feature magnitudes.1094 1095 Args:1096 hidden_states: Stacked tensor of shape1097 ``(len(layers_to_collect), batch, seq_len, d_model)``.1098 layers_to_collect: The ESMC layer indices that were collected,1099 in the same order as the first dim of ``hidden_states``.1100 token_mask: Boolean mask ``(batch, seq_len)`` — ``True`` for1101 real (non-padding) tokens.1102 normalize_sae: When ``True``, scale features by ``idf / max``1103 using the per-feature stats trained alongside each SAE.1104 """1105 layer_to_idx = {layer: idx for idx, layer in enumerate(layers_to_collect)}1106 sae_outputs: dict[str, torch.Tensor] = {}1107 1108 for model_name, sae_module in self._sae_models.items():1109 # `nn.ModuleDict` only stores `nn.Module`s at the type level;1110 # ``add_sae_models`` enforces that each entry is an ``_ESMCSAELayer``.1111 assert isinstance(sae_module, _ESMCSAELayer)1112 layer: _ESMCSAELayer = sae_module1113 requested_layer = self._get_sae_layer_num_requested(model_name)1114 layer_idx = layer_to_idx[requested_layer]1115 layer_states = hidden_states[layer_idx].clone().to(self.device)1116 1117 sae_out = layer.get_sae_output(layer_states, token_mask)1118 features = sae_out.feature_magnitudes.detach()1119 1120 if normalize_sae:1121 # ``register_buffer`` is typed as ``Tensor | Module`` on1122 # ``nn.Module``; narrow here since these are Tensors.1123 idf = cast(torch.Tensor, layer.idf)1124 max_val = cast(torch.Tensor, layer.max)1125 features = (features / max_val) * idf1126 1127 sae_outputs[model_name] = features.to_sparse()1128 1129 return sae_outputs1130 1131 @can_return_tuple1132 @auto_docstring1133 def forward(1134 self,1135 input_ids: Optional[torch.Tensor] = None,1136 attention_mask: Optional[torch.Tensor] = None,1137 sequence_id: Optional[torch.Tensor] = None,1138 output_hidden_states: Optional[bool] = None,1139 output_attentions: Optional[bool] = None,1140 return_dict: Optional[bool] = None,1141 compute_sae: bool = True,1142 normalize_sae: bool = False,1143 ) -> tuple[torch.Tensor, ...] | ESMCOutput:1144 r"""1145 sequence_id (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1146 Integer chain-ID tensor for chain-aware attention masking. Tokens with the same1147 non-negative integer value can attend to each other; tokens with different values1148 cannot (cross-chain masking). Padding positions should be set to ``-1``.1149 When provided, ``attention_mask`` is ignored. The ``flash_attention_2`` backend1150 only supports single-chain inputs (all non-padding values must be ``0``); pass1151 multi-chain ``sequence_id`` with ``attn_implementation='sdpa'`` (or ``'eager'``).1152 output_attentions (`bool`, *optional*):1153 Whether to return the per-block attention weights of shape1154 ``(batch_size, num_heads, sequence_length, sequence_length)``.1155 Forces a manual-SDPA path inside :class:`MultiHeadAttention` so the1156 attention probabilities are observable; raises on the1157 ``flash_attention_2`` path.1158 compute_sae (`bool`, *optional*, defaults to ``True``):1159 Whether to run any SAE models registered via :meth:`add_sae_models`.1160 Has no effect when no SAEs are registered.1161 normalize_sae (`bool`, *optional*, defaults to ``False``):1162 When ``True``, scale SAE feature magnitudes by ``idf / max`` (only1163 applied when the SAE's normalization buffers contain non-trivial values).1164 1165 Examples:1166 1167 ```python1168 >>> from transformers import AutoTokenizer, ESMCModel1169 1170 >>> model = ESMCModel.from_pretrained("Biohub/ESMC-600M-2024-12")1171 >>> tokenizer = AutoTokenizer.from_pretrained("Biohub/ESMC-600M-2024-12")1172 >>> inputs = tokenizer(["MLKNVQVQLV"], return_tensors="pt")1173 >>> outputs = model(**inputs)1174 >>> outputs.last_hidden_state.shape1175 torch.Size([1, 12, 960])1176 ```1177 """1178 output_hidden_states = (1179 output_hidden_states1180 if output_hidden_states is not None1181 else self.config.output_hidden_states1182 )1183 output_attentions = (1184 output_attentions1185 if output_attentions is not None1186 else self.config.output_attentions1187 )1188 return_dict = (1189 return_dict if return_dict is not None else self.config.use_return_dict1190 )1191 1192 output_sae = compute_sae and len(self._sae_models) > 01193 1194 # Determine which intermediate layers to collect. When SAEs are1195 # registered we must collect at least the layers they target, even if1196 # the caller did not ask for all hidden states.1197 if output_hidden_states:1198 layers_to_collect: list[int] = list(range(self.config.n_layers + 1))1199 elif output_sae:1200 layers_to_collect = sorted(