Synthyra/FastESMFold
065
1from __future__ import annotations
2
3import torch
4import torch._inductor.config as inductor_config
5import torch._dynamo as dynamo
6
7# Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs)
8# Provides significant speedup with minimal precision loss
9torch.set_float32_matmul_precision('high')
10
11# Enable TF32 for matrix multiplications and cuDNN operations
12torch.backends.cuda.matmul.allow_tf32 = True
13torch.backends.cudnn.allow_tf32 = True
14
15# Enable cuDNN autotuner - finds fastest algorithms for your hardware
16# Best when input sizes are consistent; may slow down first iterations
17torch.backends.cudnn.benchmark = True
18
19# Deterministic operations off for speed (set True if reproducibility needed)
20torch.backends.cudnn.deterministic = False
21inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM"
22
23dynamo.config.capture_scalar_outputs = True
24torch._dynamo.config.recompile_limit = 16
25
26"""Shared attention infrastructure for all FastPLMs models.
27
28Contains: AttentionBackend enum, backend resolution, mask creation,
29flex attention helpers, flash kernel detection/dispatch, and pad/unpad utilities.
30"""
31from enum import Enum
32from typing import Dict, List, Optional, Tuple
33
34import torch
35import torch.nn as nn
36from torch.nn import functional as F
37from einops import rearrange
38
39try:
40 from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask
41except ImportError:
42 create_block_mask = None
43 flex_attention = None
44 BlockMask = None
45
46_compiled_flex_attention = None
47
48
49def _get_flex_attention_fn():
50 """Return flex_attention callable: compiled (fused kernel) by default, or eager when debug flag is set."""
51 global _compiled_flex_attention
52 if flex_attention is None:
53 return None
54 flex_mod = torch.nn.attention.flex_attention
55 if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False):
56 return flex_attention
57 if _compiled_flex_attention is None:
58 _compiled_flex_attention = torch.compile(
59 flex_attention,
60 dynamic=False,
61 )
62 return _compiled_flex_attention
63
64
65# HuggingFace `kernels` exposes slightly different APIs for Flash Attention 2
66# and 3. Detect the loaded variant once so every caller uses the same dispatch.
67def _infer_kernels_flash_variant(kernel) -> Optional[str]:
68 if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"):
69 return "flash_attn2"
70 if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"):
71 return "flash_attn3"
72 return None
73
74
75def _try_get_kernels_flash():
76 try:
77 from kernels import get_kernel
78 except ImportError:
79 return None, None
80
81 flash_kernel = None
82 flash_kernel_variant = None
83 try:
84 flash_kernel = get_kernel("kernels-community/flash-attn3")
85 flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
86 assert flash_kernel_variant is not None, "Loaded flash-attn3 kernel does not expose a supported API."
87 except Exception:
88 try:
89 flash_kernel = get_kernel("kernels-community/flash-attn2")
90 flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
91 assert flash_kernel_variant is not None, "Loaded flash-attn2 kernel does not expose a supported API."
92 except Exception:
93 flash_kernel = None
94 flash_kernel_variant = None
95 return flash_kernel, flash_kernel_variant
96
97
98_FLASH_KERNELS_LOADED = False
99FLASH_KERNEL = None
100FLASH_KERNEL_VARIANT = None
101
102
103def _ensure_flash_kernels_loaded():
104 global _FLASH_KERNELS_LOADED, FLASH_KERNEL, FLASH_KERNEL_VARIANT
105 if _FLASH_KERNELS_LOADED:
106 return
107 _FLASH_KERNELS_LOADED = True
108 FLASH_KERNEL, FLASH_KERNEL_VARIANT = _try_get_kernels_flash()
109
110
111def _kernels_flash_forward(
112 query_states: torch.Tensor,
113 key_states: torch.Tensor,
114 value_states: torch.Tensor,
115 causal: bool = False,
116 softmax_scale: Optional[float] = None,
117) -> torch.Tensor:
118 """Flash-attention forward, optionally overriding the softmax scale.
119
120 When `softmax_scale is None`, the flash kernel applies its default
121 `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already
122 pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold).
123 Failing to override when Q is pre-scaled applies the scale twice. On
124 DPLM-150M, that produced pooled-embedding cosine around -0.12 and argmax
125 agreement around 0.27 vs SDPA.
126 """
127 assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
128 if FLASH_KERNEL_VARIANT == "flash_attn2":
129 return FLASH_KERNEL.fwd(
130 q=query_states, k=key_states, v=value_states,
131 softmax_scale=softmax_scale, is_causal=causal,
132 )[0]
133 if FLASH_KERNEL_VARIANT == "flash_attn3":
134 try:
135 output = FLASH_KERNEL.flash_attn_func(
136 q=query_states, k=key_states, v=value_states,
137 softmax_scale=softmax_scale, causal=causal,
138 )
139 except TypeError:
140 output = FLASH_KERNEL.flash_attn_func(
141 query_states, key_states, value_states,
142 0.0, softmax_scale, causal,
143 )
144 if isinstance(output, tuple):
145 return output[0]
146 return output
147 raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}")
148
149
150def _kernels_flash_varlen_forward(
151 query_states: torch.Tensor,
152 key_states: torch.Tensor,
153 value_states: torch.Tensor,
154 cu_seqlens_q: torch.Tensor,
155 cu_seqlens_k: torch.Tensor,
156 max_seqlen_in_batch_q: int,
157 max_seqlen_in_batch_k: int,
158 causal: bool = False,
159 softmax_scale: Optional[float] = None,
160) -> torch.Tensor:
161 """Varlen flash-attention forward, optionally overriding the softmax scale.
162
163 See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be
164 passed when Q has been pre-scaled by the caller.
165 """
166 assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
167 if FLASH_KERNEL_VARIANT == "flash_attn2":
168 return FLASH_KERNEL.varlen_fwd(
169 q=query_states, k=key_states, v=value_states,
170 cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
171 max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k,
172 softmax_scale=softmax_scale, is_causal=causal,
173 )[0]
174 if FLASH_KERNEL_VARIANT == "flash_attn3":
175 try:
176 output = FLASH_KERNEL.flash_attn_varlen_func(
177 q=query_states, k=key_states, v=value_states,
178 cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
179 max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k,
180 softmax_scale=softmax_scale, causal=causal,
181 )
182 except TypeError:
183 output = FLASH_KERNEL.flash_attn_varlen_func(
184 query_states, key_states, value_states,
185 cu_seqlens_q, cu_seqlens_k,
186 max_seqlen_in_batch_q, max_seqlen_in_batch_k,
187 0.0, softmax_scale, causal,
188 )
189 if isinstance(output, tuple):
190 return output[0]
191 return output
192 raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}")
193
194
195# Varlen flash attention runs only on real tokens. These helpers remove padding
196# before the kernel call and restore the original padded batch shape afterward.
197class IndexFirstAxis(torch.autograd.Function):
198 @staticmethod
199 def forward(ctx, input, indices) -> torch.Tensor:
200 ctx.save_for_backward(indices)
201 assert input.ndim >= 2
202 ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
203 second_dim = other_shape.numel()
204 return torch.gather(
205 rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim)
206 ).reshape(-1, *other_shape)
207
208 @staticmethod
209 def backward(ctx, grad_output) -> Tuple[torch.Tensor, None]:
210 (indices,) = ctx.saved_tensors
211 assert grad_output.ndim >= 2
212 other_shape = grad_output.shape[1:]
213 grad_output = rearrange(grad_output, "b ... -> b (...)")
214 grad_input = torch.zeros(
215 [ctx.first_axis_dim, grad_output.shape[1]], device=grad_output.device, dtype=grad_output.dtype
216 )
217 grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output)
218 return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
219
220
221class IndexPutFirstAxis(torch.autograd.Function):
222 @staticmethod
223 def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor:
224 ctx.save_for_backward(indices)
225 assert indices.ndim == 1
226 assert values.ndim >= 2
227 output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)
228 output[indices] = values
229 return output
230
231 @staticmethod
232 def backward(ctx, grad_output) -> Tuple[torch.Tensor, None, None]:
233 (indices,) = ctx.saved_tensors
234 return grad_output[indices], None, None
235
236
237index_first_axis = IndexFirstAxis.apply
238index_put_first_axis = IndexPutFirstAxis.apply
239
240
241def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
242 output = index_put_first_axis(hidden_states, indices, batch * seqlen)
243 return rearrange(output, "(b s) ... -> b s ...", b=batch)
244
245
246def _unpad_input(
247 query_layer: torch.Tensor,
248 key_layer: torch.Tensor,
249 value_layer: torch.Tensor,
250 attention_mask_2d: torch.Tensor,
251) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]:
252 batch_size, seq_len, num_heads, head_dim = query_layer.shape
253 seqlens = attention_mask_2d.sum(dim=1).int()
254 cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))
255 max_seqlen = int(seqlens.max().item())
256 indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()
257 query_layer = index_first_axis(query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
258 key_layer = index_first_axis(key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
259 value_layer = index_first_axis(value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
260 return query_layer, key_layer, value_layer, indices, (cu_seqlens, cu_seqlens), (max_seqlen, max_seqlen)
261
262
263def kernels_flash_attention_func(
264 query_states: torch.Tensor,
265 key_states: torch.Tensor,
266 value_states: torch.Tensor,
267 attention_mask_2d: Optional[torch.Tensor] = None,
268 causal: bool = False,
269 softmax_scale: Optional[float] = None,
270) -> torch.Tensor:
271 """Public flash-attention entry point with optional padding handling.
272
273 `softmax_scale`:
274 None -> kernel applies its default `1 / sqrt(head_dim)`.
275 float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled
276 by the caller).
277
278 Caller contract: if a model family pre-scales Q by `1/sqrt(head_dim)`
279 before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass
280 `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale
281 again, yielding an effective `1/head_dim` scale that drifts across layers.
282 """
283 assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
284 if not causal and attention_mask_2d is not None:
285 batch_size, q_len = query_states.shape[:2]
286 (
287 query_states, key_states, value_states,
288 indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k),
289 ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d)
290 attn_output_unpad = _kernels_flash_varlen_forward(
291 query_states=query_states, key_states=key_states, value_states=value_states,
292 cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
293 max_seqlen_in_batch_q=max_seqlen_q, max_seqlen_in_batch_k=max_seqlen_k,
294 softmax_scale=softmax_scale,
295 )
296 return pad_input(attn_output_unpad, indices_q, batch_size, q_len)
297 else:
298 return _kernels_flash_forward(
299 query_states=query_states, key_states=key_states, value_states=value_states,
300 causal=causal, softmax_scale=softmax_scale,
301 )
302
303
304# User-facing backend strings resolve to this enum before attention dispatch.
305class AttentionBackend(Enum):
306 AUTO = "auto"
307 KERNELS_FLASH = "kernels_flash"
308 FLEX = "flex"
309 SDPA = "sdpa"
310
311
312VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend)
313
314
315_BACKEND_CONFIRMED = False
316
317
318def resolve_attention_backend(requested_backend: str) -> AttentionBackend:
319 global _BACKEND_CONFIRMED
320 assert requested_backend in VALID_ATTENTION_BACKENDS, (
321 f"Unsupported attention backend: {requested_backend}. Expected one of {VALID_ATTENTION_BACKENDS}."
322 )
323 if requested_backend in (AttentionBackend.AUTO.value, AttentionBackend.KERNELS_FLASH.value):
324 _ensure_flash_kernels_loaded()
325 if requested_backend == AttentionBackend.AUTO.value:
326 if FLASH_KERNEL is not None:
327 resolved = AttentionBackend.KERNELS_FLASH
328 elif flex_attention is not None:
329 resolved = AttentionBackend.FLEX
330 else:
331 resolved = AttentionBackend.SDPA
332 elif requested_backend == AttentionBackend.KERNELS_FLASH.value:
333 assert FLASH_KERNEL is not None, "Kernels Flash Attention is not available in this environment."
334 resolved = AttentionBackend.KERNELS_FLASH
335 elif requested_backend == AttentionBackend.FLEX.value:
336 assert flex_attention is not None, "Flex Attention is not available in this environment."
337 resolved = AttentionBackend.FLEX
338 elif requested_backend == AttentionBackend.SDPA.value:
339 resolved = AttentionBackend.SDPA
340 else:
341 raise AssertionError(f"Unsupported attention backend: {requested_backend}")
342 if not _BACKEND_CONFIRMED:
343 print(f"Attention backend: config='{requested_backend}' -> resolved='{resolved.value}'")
344 _BACKEND_CONFIRMED = True
345 return resolved
346
347
348@torch.compiler.disable
349def get_attention_mask(
350 effective_backend: AttentionBackend,
351 batch_size: int,
352 seq_len: int,
353 device: torch.device,
354 attention_mask: Optional[torch.Tensor] = None,
355) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[BlockMask]]:
356 """Build padding masks once for all encoder layers.
357
358 Returns (attention_mask_2d, attention_mask_4d, flex_block_mask).
359 """
360 if attention_mask is None:
361 return None, None, None
362
363 attention_mask_2d = attention_mask.bool()
364
365 if effective_backend == AttentionBackend.KERNELS_FLASH:
366 return attention_mask_2d, None, None
367
368 if effective_backend == AttentionBackend.FLEX:
369 assert create_block_mask is not None, "Flex attention backend requested but torch.create_block_mask is unavailable."
370 valid_lens = attention_mask_2d.sum(dim=-1)
371
372 def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
373 return (q_idx < valid_lens[batch_idx]) & (kv_idx < valid_lens[batch_idx])
374
375 flex_block_mask = create_block_mask(mask_mod, batch_size, 1, seq_len, seq_len, device=device)
376 return attention_mask_2d, None, flex_block_mask
377
378 # SDPA/manual masks only keys. Padding queries still attend to real keys, so
379 # their outputs stay finite instead of softmaxing over all -inf scores.
380 attention_mask_4d = attention_mask_2d[:, None, None, :]
381 return attention_mask_2d, attention_mask_4d, None
382
383
384def bool_to_additive_mask(
385 bool_mask: torch.Tensor,
386 dtype: torch.dtype,
387) -> torch.Tensor:
388 """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid).
389
390 Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))`
391 directly on a bool tensor returns a bool tensor because `-inf` casts to `True`.
392 That silently drops the mask. Always allocate a float tensor first, then fill it.
393 This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask.
394 """
395 assert bool_mask.dtype == torch.bool, (
396 f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}"
397 )
398 additive = torch.zeros_like(bool_mask, dtype=dtype)
399 additive.masked_fill_(bool_mask.logical_not(), float("-inf"))
400 return additive
401
402import typing as T
403from dataclasses import dataclass, fields
404
405import torch
406import torch.nn as nn
407import torch.nn.functional as F
408
409
410@dataclass
411class TTTConfig:
412 lr: float = 4e-4
413 steps: int = 30
414 ags: int = 16
415 batch_size: int = 2
416 mask_ratio: float = 0.15
417 crop_size: int = 1024
418 bert_leave_prob: float = 0.1
419 bert_replace_prob: float = 0.1
420 optimizer: str = "sgd"
421 momentum: float = 0.0
422 weight_decay: float = 0.0
423 seed: int | None = 0
424 lora_rank: int = 8
425 lora_alpha: float = 32.0
426 lora_target_replace_module: str | None = None
427 lora_target_modules: tuple[str, ...] | None = None
428 initial_state_reset: bool = True
429 automatic_best_state_reset: bool = False
430 eval_each_step: bool = False
431 gradient_clip: bool = False
432 gradient_clip_max_norm: float = 1.0
433
434 @classmethod
435 def from_kwargs(cls, **kwargs: T.Any) -> "TTTConfig":
436 valid_names = {field.name for field in fields(cls)}
437 unknown_names = set(kwargs) - valid_names
438 assert len(unknown_names) == 0, f"Unknown TTTConfig fields: {sorted(unknown_names)}"
439 return cls(**kwargs)
440
441 def merged(self, overrides: T.Mapping[str, T.Any] | "TTTConfig" | None) -> "TTTConfig":
442 if overrides is None:
443 return self
444 if isinstance(overrides, TTTConfig):
445 return overrides
446 values = {field.name: self.__dict__[field.name] for field in fields(self)}
447 for name, value in overrides.items():
448 assert name in values, f"Unknown TTTConfig field: {name}"
449 values[name] = value
450 return TTTConfig(**values)
451
452 def verify(self) -> None:
453 assert self.lr > 0.0, "TTT learning rate must be positive."
454 assert self.steps >= 1, "TTT steps must be >= 1."
455 assert self.ags >= 1, "TTT gradient accumulation steps must be >= 1."
456 assert self.batch_size >= 1, "TTT batch_size must be >= 1."
457 assert 0.0 < self.mask_ratio <= 1.0, "TTT mask_ratio must be in (0, 1]."
458 assert self.crop_size >= 1, "TTT crop_size must be >= 1."
459 assert self.lora_rank >= 1, "TTT v1 is LoRA-only, so lora_rank must be >= 1."
460 assert self.lora_alpha > 0.0, "TTT lora_alpha must be positive."
461 assert self.optimizer in {"adamw", "sgd"}, "TTT optimizer must be 'adamw' or 'sgd'."
462 assert 0.0 <= self.bert_leave_prob <= 1.0, "bert_leave_prob must be in [0, 1]."
463 assert 0.0 <= self.bert_replace_prob <= 1.0, "bert_replace_prob must be in [0, 1]."
464 assert self.bert_leave_prob + self.bert_replace_prob <= 1.0, (
465 "bert_leave_prob + bert_replace_prob must be <= 1."
466 )
467 if self.gradient_clip:
468 assert self.gradient_clip_max_norm > 0.0, "gradient_clip_max_norm must be positive."
469
470
471class LoraInjectedLinear(nn.Module):
472 def __init__(self, linear: nn.Module, rank: int, alpha: float) -> None:
473 super().__init__()
474 weight = linear._parameters["weight"]
475 assert weight.ndim == 2, "LoRA can only wrap 2D linear weights."
476 self.linear = linear
477 self.linear.requires_grad_(False)
478 self.rank = rank
479 self.scale = alpha
480 in_features = weight.shape[1]
481 out_features = weight.shape[0]
482 self.lora_down = nn.Linear(in_features, rank, bias=False, dtype=torch.float32)
483 self.lora_up = nn.Linear(rank, out_features, bias=False, dtype=torch.float32)
484 self.lora_down.to(device=weight.device)
485 self.lora_up.to(device=weight.device)
486 nn.init.normal_(self.lora_down.weight, std=1.0 / rank)
487 nn.init.zeros_(self.lora_up.weight)
488
489 @property
490 def weight(self) -> torch.Tensor:
491 return self.linear._parameters["weight"]
492
493 @property
494 def bias(self) -> torch.Tensor | None:
495 return self.linear._parameters["bias"]
496
497 def forward(self, x: torch.Tensor) -> torch.Tensor:
498 base = self.linear(x)
499 delta = self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale
500 return base + delta.to(dtype=base.dtype)
501
502
503class FastPLMTestTimeTrainingMixin:
504 def init_ttt(self, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None) -> None:
505 base_config = TTTConfig()
506 self._ttt_cfg = base_config.merged(ttt_config)
507 self._ttt_cfg.verify()
508 self._ttt_initialized = False
509 self._ttt_initial_state: list[dict[str, torch.Tensor]] | None = None
510
511 @property
512 def ttt_config(self) -> TTTConfig:
513 if "_ttt_cfg" not in self.__dict__:
514 self.init_ttt()
515 return self._ttt_cfg
516
517 def _ttt_get_trainable_modules(self) -> list[nn.Module]:
518 return [self]
519
520 def _ttt_get_frozen_modules(self) -> list[nn.Module]:
521 return []
522
523 def _ttt_tokenize(
524 self,
525 seq: str | list[str] | None = None,
526 input_ids: torch.Tensor | None = None,
527 **kwargs: T.Any,
528 ) -> torch.Tensor | dict[str, torch.Tensor]:
529 del kwargs
530 if input_ids is not None:
531 return input_ids
532 assert seq is not None, "Pass either seq or input_ids for TTT."
533 tokenized = self.tokenizer(seq, return_tensors="pt", padding=True)
534 return tokenized["input_ids"]
535
536 def _ttt_mask_token(self) -> int:
537 return int(self.tokenizer.mask_token_id)
538
539 def _ttt_padding_token(self) -> int:
540 return int(self.tokenizer.pad_token_id)
541
542 def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
543 tokenizer = self.tokenizer
544 special_ids = set(tokenizer.all_special_ids)
545 vocab_size = int(self.config.vocab_size)
546 ids = [idx for idx in range(vocab_size) if idx not in special_ids]
547 assert len(ids) > 0, "TTT replacement token set is empty."
548 return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)
549
550 def _ttt_predict_logits(
551 self,
552 batch: torch.Tensor | dict[str, torch.Tensor],
553 **kwargs: T.Any,
554 ) -> torch.Tensor:
555 del kwargs
556 if isinstance(batch, dict):
557 output = self(**batch)
558 return output.logits
559 attention_mask = batch.ne(self._ttt_padding_token())
560 output = self(input_ids=batch, attention_mask=attention_mask)
561 return output.logits
562
563 def _ttt_eval_step(
564 self,
565 step: int,
566 loss: float,
567 seq: str | list[str] | None = None,
568 input_ids: torch.Tensor | None = None,
569 **kwargs: T.Any,
570 ) -> tuple[dict[str, T.Any], float | None]:
571 del step, loss, seq, input_ids, kwargs
572 return {}, None
573
574 def _ttt_is_lora_target(
575 self,
576 name: str,
577 full_name: str,
578 module: nn.Module,
579 active: bool,
580 target_modules: tuple[str, ...] | None,
581 ) -> bool:
582 if not active:
583 return False
584 if isinstance(module, LoraInjectedLinear):
585 return False
586 if (
587 target_modules is not None
588 and name not in target_modules
589 and full_name not in target_modules
590 ):
591 return False
592 if isinstance(module, nn.Linear):
593 return True
594 if "weight" not in module._parameters:
595 return False
596 weight = module._parameters["weight"]
597 if weight is None or weight.ndim != 2:
598 return False
599 return "Linear" in module.__class__.__name__
600
601 def _ttt_inject_lora(self) -> int:
602 cfg = self.ttt_config
603 cfg.verify()
604 target_class = cfg.lora_target_replace_module
605 target_modules = cfg.lora_target_modules
606 wrapped = 0
607
608 def inject(module: nn.Module, prefix: str, active: bool) -> None:
609 nonlocal wrapped
610 for name, child in list(module.named_children()):
611 full_name = f"{prefix}.{name}" if prefix else name
612 child_active = active
613 if target_class is not None:
614 child_active = active or child.__class__.__name__ == target_class
615 if self._ttt_is_lora_target(name, full_name, child, child_active, target_modules):
616 setattr(
617 module,
618 name,
619 LoraInjectedLinear(child, rank=cfg.lora_rank, alpha=cfg.lora_alpha),
620 )
621 wrapped += 1
622 continue
623 inject(child, full_name, child_active)
624
625 for trainable_module in self._ttt_get_trainable_modules():
626 inject(trainable_module, "", target_class is None)
627 assert wrapped > 0, "TTT LoRA injection did not find any target modules."
628 return wrapped
629
630 def _ttt_lora_modules(self) -> list[LoraInjectedLinear]:
631 return [module for module in self.modules() if isinstance(module, LoraInjectedLinear)]
632
633 def _ttt_lora_parameters(self) -> list[nn.Parameter]:
634 params: list[nn.Parameter] = []
635 for module in self._ttt_lora_modules():
636 params.extend(module.lora_down.parameters())
637 params.extend(module.lora_up.parameters())
638 assert len(params) > 0, "TTT has no LoRA parameters."
639 return params
640
641 def _ttt_snapshot_lora_state(self) -> list[dict[str, torch.Tensor]]:
642 snapshot = []
643 for module in self._ttt_lora_modules():
644 snapshot.append(
645 {
646 "lora_down.weight": module.lora_down.weight.detach().clone(),
647 "lora_up.weight": module.lora_up.weight.detach().clone(),
648 }
649 )
650 assert len(snapshot) > 0, "TTT has no LoRA state to snapshot."
651 return snapshot
652
653 def _ttt_restore_lora_state(self, state: list[dict[str, torch.Tensor]]) -> None:
654 modules = self._ttt_lora_modules()
655 assert len(modules) == len(state), "TTT LoRA state/module count mismatch."
656 with torch.no_grad():
657 for module, module_state in zip(modules, state):
658 module.lora_down.weight.copy_(module_state["lora_down.weight"])
659 module.lora_up.weight.copy_(module_state["lora_up.weight"])
660
661 def _ttt_ensure_initialized(self) -> None:
662 if "_ttt_cfg" not in self.__dict__:
663 self.init_ttt()
664 if self._ttt_initialized:
665 return
666 self._ttt_inject_lora()
667 self._ttt_initial_state = self._ttt_snapshot_lora_state()
668 self._ttt_initialized = True
669
670 def ttt_reset(self) -> None:
671 self._ttt_ensure_initialized()
672 assert self._ttt_initial_state is not None, "TTT initial state is not available."
673 self._ttt_restore_lora_state(self._ttt_initial_state)
674
675 def _ttt_make_optimizer(self) -> torch.optim.Optimizer:
676 cfg = self.ttt_config
677 params = self._ttt_lora_parameters()
678 if cfg.optimizer == "sgd":
679 return torch.optim.SGD(
680 params,
681 lr=cfg.lr,
682 momentum=cfg.momentum,
683 weight_decay=cfg.weight_decay,
684 )
685 return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay)
686
687 def _ttt_to_device(
688 self,
689 batch: torch.Tensor | dict[str, torch.Tensor],
690 device: torch.device,
691 ) -> torch.Tensor | dict[str, torch.Tensor]:
692 if isinstance(batch, dict):
693 return {name: tensor.to(device) for name, tensor in batch.items()}
694 return batch.to(device)
695
696 def _ttt_input_ids_from_batch(
697 self,
698 batch: torch.Tensor | dict[str, torch.Tensor],
699 ) -> torch.Tensor:
700 if isinstance(batch, dict):
701 return batch["input_ids"]
702 return batch
703
704 def _ttt_set_input_ids(
705 self,
706 batch: torch.Tensor | dict[str, torch.Tensor],
707 input_ids: torch.Tensor,
708 ) -> torch.Tensor | dict[str, torch.Tensor]:
709 if isinstance(batch, dict):
710 updated = dict(batch)
711 updated["input_ids"] = input_ids
712 return updated
713 return input_ids
714
715 def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
716 pad_token = self._ttt_padding_token()
717 mask = input_ids.ne(pad_token)
718 special_ids = set(self.tokenizer.all_special_ids)
719 for special_id in special_ids:
720 mask = mask & input_ids.ne(int(special_id))
721 return mask
722
723 def _ttt_sample_crop(
724 self,
725 batch: torch.Tensor | dict[str, torch.Tensor],
726 generator: torch.Generator,
727 ) -> torch.Tensor | dict[str, torch.Tensor]:
728 input_ids = self._ttt_input_ids_from_batch(batch)
729 cfg = self.ttt_config
730 if input_ids.shape[1] <= cfg.crop_size:
731 return batch
732 high = input_ids.shape[1] - cfg.crop_size + 1
733 start = int(
734 torch.randint(
735 high,
736 (1,),
737 generator=generator,
738 device=input_ids.device,
739 ).item()
740 )
741 end = start + cfg.crop_size
742 if isinstance(batch, dict):
743 cropped = {}
744 for name, tensor in batch.items():
745 if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]:
746 cropped[name] = tensor[:, start:end]
747 else:
748 cropped[name] = tensor
749 return cropped
750 return input_ids[:, start:end]
751
752 def _ttt_sample_batch(
753 self,
754 tokenized: torch.Tensor | dict[str, torch.Tensor],
755 generator: torch.Generator,
756 ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]:
757 cfg = self.ttt_config
758 batch = self._ttt_sample_crop(tokenized, generator)
759 input_ids = self._ttt_input_ids_from_batch(batch)
760 rows = torch.randint(
761 input_ids.shape[0],
762 (cfg.batch_size,),
763 generator=generator,
764 device=input_ids.device,
765 )
766 if isinstance(batch, dict):
767 sampled: torch.Tensor | dict[str, torch.Tensor] = {}
768 for name, tensor in batch.items():
769 if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]:
770 sampled[name] = tensor.index_select(0, rows)
771 else:
772 sampled[name] = tensor
773 else:
774 sampled = input_ids.index_select(0, rows)
775
776 sampled_ids = self._ttt_input_ids_from_batch(sampled)
777 labels = sampled_ids.clone()
778 non_special = self._ttt_non_special_mask(sampled_ids)
779 label_mask = torch.zeros_like(non_special)
780 for row_idx in range(sampled_ids.shape[0]):
781 candidate_positions = torch.where(non_special[row_idx])[0]
782 if candidate_positions.numel() == 0:
783 continue
784 num_mask = max(1, int(round(candidate_positions.numel() * cfg.mask_ratio)))
785 order = torch.randperm(
786 candidate_positions.numel(),
787 generator=generator,
788 device=sampled_ids.device,
789 )
790 chosen = candidate_positions[order[:num_mask]]
791 label_mask[row_idx, chosen] = True
792 labels = labels.masked_fill(~label_mask, -100)
793
794 masked_ids = sampled_ids.clone()
795 chosen_positions = torch.where(label_mask)
796 if chosen_positions[0].numel() > 0:
797 random_values = torch.rand(
798 chosen_positions[0].shape,
799 generator=generator,
800 device=sampled_ids.device,
801 )
802 leave = random_values < cfg.bert_leave_prob
803 replace = (random_values >= cfg.bert_leave_prob) & (
804 random_values < cfg.bert_leave_prob + cfg.bert_replace_prob
805 )
806 mask = ~(leave | replace)
807 if mask.any():
808 masked_ids[
809 chosen_positions[0][mask],
810 chosen_positions[1][mask],
811 ] = self._ttt_mask_token()
812 if replace.any():
813 replacement_tokens = self._ttt_replacement_tokens(sampled_ids)
814 replacement_idx = torch.randint(
815 replacement_tokens.shape[0],
816 (int(replace.sum().item()),),
817 generator=generator,
818 device=sampled_ids.device,
819 )
820 masked_ids[
821 chosen_positions[0][replace],
822 chosen_positions[1][replace],
823 ] = replacement_tokens[replacement_idx]
824
825 return self._ttt_set_input_ids(sampled, masked_ids), labels
826
827 def ttt(
828 self,
829 seq: str | list[str] | None = None,
830 input_ids: torch.Tensor | None = None,
831 ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None,
832 **kwargs: T.Any,
833 ) -> dict[str, T.Any]:
834 if ttt_config is not None:
835 if "_ttt_initialized" in self.__dict__ and self._ttt_initialized:
836 next_cfg = self.ttt_config.merged(ttt_config)
837 assert next_cfg.lora_rank == self.ttt_config.lora_rank, (
838 "Changing lora_rank after TTT initialization is not supported."
839 )
840 assert next_cfg.lora_alpha == self.ttt_config.lora_alpha, (
841 "Changing lora_alpha after TTT initialization is not supported."
842 )
843 assert (
844 next_cfg.lora_target_replace_module
845 == self.ttt_config.lora_target_replace_module
846 ), "Changing LoRA target class after TTT initialization is not supported."
847 assert next_cfg.lora_target_modules == self.ttt_config.lora_target_modules, (
848 "Changing LoRA target modules after TTT initialization is not supported."
849 )
850 self._ttt_cfg = next_cfg
851 else:
852 self.init_ttt(ttt_config)
853
854 self._ttt_ensure_initialized()
855 cfg = self.ttt_config
856 if cfg.initial_state_reset:
857 self.ttt_reset()
858
859 device = next(self.parameters()).device
860 tokenized = self._ttt_tokenize(seq=seq, input_ids=input_ids, **kwargs)
861 tokenized = self._ttt_to_device(tokenized, device)
862 generator_device = device if device.type == "cuda" else torch.device("cpu")
863 generator = torch.Generator(device=generator_device)
864 if cfg.seed is not None:
865 generator.manual_seed(cfg.seed)
866
867 module_modes = {module: module.training for module in self.modules()}
868 requires_grad = {param: param.requires_grad for param in self.parameters()}
869 losses: list[float] = []
870 step_metrics: list[dict[str, T.Any]] = []
871 best_state: list[dict[str, torch.Tensor]] | None = None
872 best_metric: float | None = None
873 best_step = 0
874
875 try:
876 self.train()
877 for param in self.parameters():
878 param.requires_grad_(False)
879 for param in self._ttt_lora_parameters():
880 param.requires_grad_(True)
881
882 optimizer = self._ttt_make_optimizer()
883 optimizer.zero_grad(set_to_none=True)
884 total_micro_steps = cfg.steps * cfg.ags
885 for micro_step in range(total_micro_steps):
886 batch, labels = self._ttt_sample_batch(tokenized, generator)
887 logits = self._ttt_predict_logits(batch, **kwargs)
888 labels = labels.to(device=logits.device)
889 loss = F.cross_entropy(
890 logits.reshape(-1, logits.shape[-1]),
891 labels.reshape(-1),
892 ignore_index=-100,
893 )
894 (loss / cfg.ags).backward()
895 if (micro_step + 1) % cfg.ags != 0:
896 continue
897
898 if cfg.gradient_clip:
899 torch.nn.utils.clip_grad_norm_(
900 self._ttt_lora_parameters(),
901 cfg.gradient_clip_max_norm,
902 )
903 optimizer.step()
904 optimizer.zero_grad(set_to_none=True)
905 step = (micro_step + 1) // cfg.ags
906 loss_value = float(loss.detach().item())
907 losses.append(loss_value)
908 if cfg.eval_each_step:
909 metrics, metric = self._ttt_eval_step(
910 step=step,
911 loss=loss_value,
912 seq=seq,
913 input_ids=input_ids,
914 **kwargs,
915 )
916 if len(metrics) > 0:
917 step_metrics.append(metrics)
918 if metric is not None and (
919 best_metric is None or metric > best_metric
920 ):
921 best_metric = metric
922 best_step = step
923 best_state = self._ttt_snapshot_lora_state()
924
925 if cfg.automatic_best_state_reset and best_state is not None:
926 self._ttt_restore_lora_state(best_state)
927 finally:
928 for param, value in requires_grad.items():
929 param.requires_grad_(value)
930 for module, training in module_modes.items():
931 module.train(training)
932
933 return {
934 "losses": losses,
935 "step_metrics": step_metrics,
936 "best_step": best_step,
937 "best_metric": best_metric,
938 }
939
940"""FastESMFold: self-contained ESMFold with FastESM2 attention and opt-in TTT.
941
942Usage:
943 from transformers import AutoModel
944 model = AutoModel.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True).cuda()
945
946 # Basic folding, no TTT
947 result = model.fold_protein("MKTLLILAVVA...")
948 print(result["plddt"], result["pdb_string"][:100])
949
950 # Experimental folding with TTT
951 result = model.fold_protein("MKTLLILAVVA...", ttt=True)
952
953Dependencies: torch, transformers, einops
954No dependency on: esm (fair-esm), proteinttt, openfold
955"""
956import copy
957from dataclasses import dataclass, field
958from functools import wraps
959from typing import Any, Callable, Dict, List, Optional, Tuple, Union
960
961import torch
962import torch.nn as nn
963from torch.nn import functional as F
964
965from einops import rearrange
966from transformers import EsmTokenizer, PretrainedConfig, PreTrainedModel
967from transformers.modeling_outputs import ModelOutput
968from transformers.models.esm.configuration_esm import EsmConfig
969from transformers.models.esm.modeling_esm import (
970 EsmContactPredictionHead,
971 EsmEmbeddings,
972 EsmIntermediate,
973 EsmLMHead,
974 EsmOutput,
975 EsmSelfOutput,
976 RotaryEmbedding,
977)
978from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
979
980
981
982
983# =============================================================================
984# Output Dataclass
985# =============================================================================
986
987@dataclass
988class FastEsmEncoderOutput(ModelOutput):
989 last_hidden_state: Optional[torch.Tensor] = None
990 hidden_states: Optional[Tuple[torch.Tensor, ...]] = None
991 attentions: Optional[Tuple[torch.Tensor, ...]] = None
992
993
994# =============================================================================
995# FastESM2 Attention Layers (multi-backend: SDPA, Flash, Flex)
996# =============================================================================
997
998class EsmSelfAttention(nn.Module):
999 def __init__(self, config, position_embedding_type: Optional[str] = None):
1000 super().__init__()
1001 assert config.hidden_size % config.num_attention_heads == 0, (
1002 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
1003 f"heads ({config.num_attention_heads})"
1004 )
1005 self.num_attention_heads = config.num_attention_heads
1006 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
1007 self.all_head_size = self.num_attention_heads * self.attention_head_size
1008
1009 self.query = nn.Linear(config.hidden_size, self.all_head_size)
1010 self.key = nn.Linear(config.hidden_size, self.all_head_size)
1011 self.value = nn.Linear(config.hidden_size, self.all_head_size)
1012 self.scale = self.attention_head_size**-0.5
1013
1014 self.dropout_prob = config.attention_probs_dropout_prob
1015 self.config = config
1016 self.attn_backend = resolve_attention_backend(config.attn_backend)
1017 self.position_embedding_type = position_embedding_type or config.position_embedding_type
1018 self.rotary_embeddings = None
1019 if self.position_embedding_type == "rotary":
1020 self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size)
1021
1022 def forward(
1023 self,
1024 hidden_states: torch.Tensor,
1025 attention_mask_2d: Optional[torch.Tensor] = None,
1026 attention_mask_4d: Optional[torch.Tensor] = None,
1027 flex_block_mask: Optional[BlockMask] = None,
1028 output_attentions: bool = False,
1029 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1030 batch_size, seq_length = hidden_states.shape[:-1]
1031 hidden_shape = (batch_size, seq_length, -1, self.attention_head_size)
1032 query_BHLD = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
1033 key_BHLD = self.key(hidden_states).view(hidden_shape).transpose(1, 2)
1034 value_BHLD = self.value(hidden_states).view(hidden_shape).transpose(1, 2)
1035
1036 query_BHLD = query_BHLD * self.scale
1037
1038 if self.position_embedding_type == "rotary":
1039 query_BHLD, key_BHLD = self.rotary_embeddings(query_BHLD, key_BHLD)
1040
1041 attn_output, attn_weights = self._attn(
1042 query_BHLD, key_BHLD, value_BHLD,
1043 attention_mask_2d=attention_mask_2d,
1044 attention_mask_4d=attention_mask_4d,
1045 flex_block_mask=flex_block_mask,
1046 output_attentions=output_attentions,
1047 )
1048 return attn_output, attn_weights
1049
1050 def _attn(
1051 self,
1052 query_BHLD: torch.Tensor,
1053 key_BHLD: torch.Tensor,
1054 value_BHLD: torch.Tensor,
1055 attention_mask_2d: Optional[torch.Tensor] = None,
1056 attention_mask_4d: Optional[torch.Tensor] = None,
1057 flex_block_mask: Optional[BlockMask] = None,
1058 output_attentions: bool = False,
1059 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1060 if output_attentions:
1061 return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d)
1062
1063 if self.attn_backend == AttentionBackend.KERNELS_FLASH:
1064 return self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d)
1065 elif self.attn_backend == AttentionBackend.FLEX:
1066 return self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask)
1067 elif self.attn_backend == AttentionBackend.SDPA:
1068 return self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d)
1069 else:
1070 raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}")
1071
1072 def _manual_attn(
1073 self,
1074 query_BHLD: torch.Tensor,
1075 key_BHLD: torch.Tensor,
1076 value_BHLD: torch.Tensor,
1077 attention_mask_4d: Optional[torch.Tensor] = None,
1078 ) -> Tuple[torch.Tensor, torch.Tensor]:
1079 attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2))
1080 if attention_mask_4d is not None:
1081 attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf"))
1082 attn_weights = F.softmax(attn_weights, dim=-1)
1083 if self.dropout_prob > 0 and self.training:
1084 attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training)
1085 context_BHLD = torch.matmul(attn_weights, value_BHLD)
1086 attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)")
1087 return attn_output, attn_weights
1088
1089 def _kernels_flash_attn(
1090 self,
1091 query_BHLD: torch.Tensor,
1092 key_BHLD: torch.Tensor,
1093 value_BHLD: torch.Tensor,
1094 attention_mask_2d: Optional[torch.Tensor] = None,
1095 ) -> Tuple[torch.Tensor, None]:
1096 query_BLHD = query_BHLD.transpose(1, 2).contiguous()
1097 key_BLHD = key_BHLD.transpose(1, 2).contiguous()
1098 value_BLHD = value_BHLD.transpose(1, 2).contiguous()
1099 # Q is pre-scaled by self.scale in forward() -- pass softmax_scale=1.0
1100 # to prevent the kernel from applying its default 1/sqrt(head_dim).
1101 attn_output = kernels_flash_attention_func(
1102 query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD,
1103 attention_mask_2d=attention_mask_2d, causal=False,
1104 softmax_scale=1.0,
1105 )
1106 return rearrange(attn_output, "b s h d -> b s (h d)"), None
1107
1108 def _flex_attn(
1109 self,
1110 query_BHLD: torch.Tensor,
1111 key_BHLD: torch.Tensor,
1112 value_BHLD: torch.Tensor,
1113 flex_block_mask: Optional[BlockMask] = None,
1114 ) -> Tuple[torch.Tensor, None]:
1115 assert flex_attention is not None, "Flex attention is not available in this environment."
1116 fn = _get_flex_attention_fn()
1117 context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=1.0)
1118 return rearrange(context_BHLD, "b h s d -> b s (h d)"), None
1119
1120 def _sdpa_attn(
1121 self,
1122 query_BHLD: torch.Tensor,
1123 key_BHLD: torch.Tensor,
1124 value_BHLD: torch.Tensor,
1125 attention_mask_4d: Optional[torch.Tensor] = None,
1126 ) -> Tuple[torch.Tensor, None]:
1127 context_BHLD = F.scaled_dot_product_attention(
1128 query_BHLD, key_BHLD, value_BHLD,
1129 attn_mask=attention_mask_4d,
1130 dropout_p=self.dropout_prob if self.training else 0.0,
1131 scale=1.0,
1132 )
1133 return rearrange(context_BHLD, "b h s d -> b s (h d)"), None
1134
1135
1136class EsmAttention(nn.Module):
1137 def __init__(self, config):
1138 super().__init__()
1139 self.self = EsmSelfAttention(config)
1140 self.output = EsmSelfOutput(config)
1141 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
1142
1143 def forward(
1144 self,
1145 hidden_states: torch.Tensor,
1146 attention_mask_2d: Optional[torch.Tensor] = None,
1147 attention_mask_4d: Optional[torch.Tensor] = None,
1148 flex_block_mask: Optional[BlockMask] = None,
1149 output_attentions: bool = False,
1150 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1151 hidden_states_ln = self.LayerNorm(hidden_states)
1152 attn_output, attn_weights = self.self(
1153 hidden_states_ln,
1154 attention_mask_2d=attention_mask_2d,
1155 attention_mask_4d=attention_mask_4d,
1156 flex_block_mask=flex_block_mask,
1157 output_attentions=output_attentions,
1158 )
1159 attention_output = self.output(attn_output, hidden_states)
1160 return attention_output, attn_weights
1161
1162
1163class EsmLayer(nn.Module):
1164 def __init__(self, config):
1165 super().__init__()
1166 self.attention = EsmAttention(config)
1167 self.intermediate = EsmIntermediate(config)
1168 self.output = EsmOutput(config)
1169 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
1170
1171 def forward(
1172 self,
1173 hidden_states: torch.Tensor,
1174 attention_mask_2d: Optional[torch.Tensor] = None,
1175 attention_mask_4d: Optional[torch.Tensor] = None,
1176 flex_block_mask: Optional[BlockMask] = None,
1177 output_attentions: bool = False,
1178 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1179 attention_output, attn_weights = self.attention(
1180 hidden_states,
1181 attention_mask_2d=attention_mask_2d,
1182 attention_mask_4d=attention_mask_4d,
1183 flex_block_mask=flex_block_mask,
1184 output_attentions=output_attentions,
1185 )
1186 layer_output = self._feed_forward(attention_output)
1187 return layer_output, attn_weights
1188
1189 def _feed_forward(self, attention_output: torch.Tensor) -> torch.Tensor:
1190 attention_output_ln = self.LayerNorm(attention_output)
1191 intermediate_output = self.intermediate(attention_output_ln)
1192 return self.output(intermediate_output, attention_output)
1193
1194
1195class FastEsmEncoder(nn.Module):
1196 def __init__(self, config):
1197 super().__init__()
1198 self.config = config
1199 self.attention_backend = resolve_attention_backend(config.attn_backend)
1200 self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)])