from __future__ import annotations import torch import torch._inductor.config as inductor_config import torch._dynamo as dynamo # Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs) # Provides significant speedup with minimal precision loss torch.set_float32_matmul_precision('high') # Enable TF32 for matrix multiplications and cuDNN operations torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True # Enable cuDNN autotuner - finds fastest algorithms for your hardware # Best when input sizes are consistent; may slow down first iterations torch.backends.cudnn.benchmark = True # Deterministic operations off for speed (set True if reproducibility needed) torch.backends.cudnn.deterministic = False inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM" dynamo.config.capture_scalar_outputs = True torch._dynamo.config.recompile_limit = 16 """Shared attention infrastructure for all FastPLMs models. Contains: AttentionBackend enum, backend resolution, mask creation, flex attention helpers, flash kernel detection/dispatch, and pad/unpad utilities. """ from enum import Enum from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn from torch.nn import functional as F from einops import rearrange try: from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask except ImportError: create_block_mask = None flex_attention = None BlockMask = None _compiled_flex_attention = None def _get_flex_attention_fn(): """Return flex_attention callable: compiled (fused kernel) by default, or eager when debug flag is set.""" global _compiled_flex_attention if flex_attention is None: return None flex_mod = torch.nn.attention.flex_attention if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False): return flex_attention if _compiled_flex_attention is None: _compiled_flex_attention = torch.compile( flex_attention, dynamic=False, ) return _compiled_flex_attention # HuggingFace `kernels` exposes slightly different APIs for Flash Attention 2 # and 3. Detect the loaded variant once so every caller uses the same dispatch. def _infer_kernels_flash_variant(kernel) -> Optional[str]: if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"): return "flash_attn2" if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"): return "flash_attn3" return None def _try_get_kernels_flash(): try: from kernels import get_kernel except ImportError: return None, None flash_kernel = None flash_kernel_variant = None try: flash_kernel = get_kernel("kernels-community/flash-attn3") flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel) assert flash_kernel_variant is not None, "Loaded flash-attn3 kernel does not expose a supported API." except Exception: try: flash_kernel = get_kernel("kernels-community/flash-attn2") flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel) assert flash_kernel_variant is not None, "Loaded flash-attn2 kernel does not expose a supported API." except Exception: flash_kernel = None flash_kernel_variant = None return flash_kernel, flash_kernel_variant _FLASH_KERNELS_LOADED = False FLASH_KERNEL = None FLASH_KERNEL_VARIANT = None def _ensure_flash_kernels_loaded(): global _FLASH_KERNELS_LOADED, FLASH_KERNEL, FLASH_KERNEL_VARIANT if _FLASH_KERNELS_LOADED: return _FLASH_KERNELS_LOADED = True FLASH_KERNEL, FLASH_KERNEL_VARIANT = _try_get_kernels_flash() def _kernels_flash_forward( query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, causal: bool = False, softmax_scale: Optional[float] = None, ) -> torch.Tensor: """Flash-attention forward, optionally overriding the softmax scale. When `softmax_scale is None`, the flash kernel applies its default `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold). Failing to override when Q is pre-scaled applies the scale twice. On DPLM-150M, that produced pooled-embedding cosine around -0.12 and argmax agreement around 0.27 vs SDPA. """ assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." if FLASH_KERNEL_VARIANT == "flash_attn2": return FLASH_KERNEL.fwd( q=query_states, k=key_states, v=value_states, softmax_scale=softmax_scale, is_causal=causal, )[0] if FLASH_KERNEL_VARIANT == "flash_attn3": try: output = FLASH_KERNEL.flash_attn_func( q=query_states, k=key_states, v=value_states, softmax_scale=softmax_scale, causal=causal, ) except TypeError: output = FLASH_KERNEL.flash_attn_func( query_states, key_states, value_states, 0.0, softmax_scale, causal, ) if isinstance(output, tuple): return output[0] return output raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}") def _kernels_flash_varlen_forward( query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, max_seqlen_in_batch_q: int, max_seqlen_in_batch_k: int, causal: bool = False, softmax_scale: Optional[float] = None, ) -> torch.Tensor: """Varlen flash-attention forward, optionally overriding the softmax scale. See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be passed when Q has been pre-scaled by the caller. """ assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." if FLASH_KERNEL_VARIANT == "flash_attn2": return FLASH_KERNEL.varlen_fwd( q=query_states, k=key_states, v=value_states, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k, softmax_scale=softmax_scale, is_causal=causal, )[0] if FLASH_KERNEL_VARIANT == "flash_attn3": try: output = FLASH_KERNEL.flash_attn_varlen_func( q=query_states, k=key_states, v=value_states, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k, softmax_scale=softmax_scale, causal=causal, ) except TypeError: output = FLASH_KERNEL.flash_attn_varlen_func( query_states, key_states, value_states, cu_seqlens_q, cu_seqlens_k, max_seqlen_in_batch_q, max_seqlen_in_batch_k, 0.0, softmax_scale, causal, ) if isinstance(output, tuple): return output[0] return output raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}") # Varlen flash attention runs only on real tokens. These helpers remove padding # before the kernel call and restore the original padded batch shape afterward. class IndexFirstAxis(torch.autograd.Function): @staticmethod def forward(ctx, input, indices) -> torch.Tensor: ctx.save_for_backward(indices) assert input.ndim >= 2 ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] second_dim = other_shape.numel() return torch.gather( rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim) ).reshape(-1, *other_shape) @staticmethod def backward(ctx, grad_output) -> Tuple[torch.Tensor, None]: (indices,) = ctx.saved_tensors assert grad_output.ndim >= 2 other_shape = grad_output.shape[1:] grad_output = rearrange(grad_output, "b ... -> b (...)") grad_input = torch.zeros( [ctx.first_axis_dim, grad_output.shape[1]], device=grad_output.device, dtype=grad_output.dtype ) grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output) return grad_input.reshape(ctx.first_axis_dim, *other_shape), None class IndexPutFirstAxis(torch.autograd.Function): @staticmethod def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor: ctx.save_for_backward(indices) assert indices.ndim == 1 assert values.ndim >= 2 output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype) output[indices] = values return output @staticmethod def backward(ctx, grad_output) -> Tuple[torch.Tensor, None, None]: (indices,) = ctx.saved_tensors return grad_output[indices], None, None index_first_axis = IndexFirstAxis.apply index_put_first_axis = IndexPutFirstAxis.apply def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor: output = index_put_first_axis(hidden_states, indices, batch * seqlen) return rearrange(output, "(b s) ... -> b s ...", b=batch) def _unpad_input( query_layer: torch.Tensor, key_layer: torch.Tensor, value_layer: torch.Tensor, attention_mask_2d: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]: batch_size, seq_len, num_heads, head_dim = query_layer.shape seqlens = attention_mask_2d.sum(dim=1).int() cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0)) max_seqlen = int(seqlens.max().item()) indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten() query_layer = index_first_axis(query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices) key_layer = index_first_axis(key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices) value_layer = index_first_axis(value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices) return query_layer, key_layer, value_layer, indices, (cu_seqlens, cu_seqlens), (max_seqlen, max_seqlen) def kernels_flash_attention_func( query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, attention_mask_2d: Optional[torch.Tensor] = None, causal: bool = False, softmax_scale: Optional[float] = None, ) -> torch.Tensor: """Public flash-attention entry point with optional padding handling. `softmax_scale`: None -> kernel applies its default `1 / sqrt(head_dim)`. float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled by the caller). Caller contract: if a model family pre-scales Q by `1/sqrt(head_dim)` before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale again, yielding an effective `1/head_dim` scale that drifts across layers. """ assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment." if not causal and attention_mask_2d is not None: batch_size, q_len = query_states.shape[:2] ( query_states, key_states, value_states, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k), ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d) attn_output_unpad = _kernels_flash_varlen_forward( query_states=query_states, key_states=key_states, value_states=value_states, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_in_batch_q=max_seqlen_q, max_seqlen_in_batch_k=max_seqlen_k, softmax_scale=softmax_scale, ) return pad_input(attn_output_unpad, indices_q, batch_size, q_len) else: return _kernels_flash_forward( query_states=query_states, key_states=key_states, value_states=value_states, causal=causal, softmax_scale=softmax_scale, ) # User-facing backend strings resolve to this enum before attention dispatch. class AttentionBackend(Enum): AUTO = "auto" KERNELS_FLASH = "kernels_flash" FLEX = "flex" SDPA = "sdpa" VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend) _BACKEND_CONFIRMED = False def resolve_attention_backend(requested_backend: str) -> AttentionBackend: global _BACKEND_CONFIRMED assert requested_backend in VALID_ATTENTION_BACKENDS, ( f"Unsupported attention backend: {requested_backend}. Expected one of {VALID_ATTENTION_BACKENDS}." ) if requested_backend in (AttentionBackend.AUTO.value, AttentionBackend.KERNELS_FLASH.value): _ensure_flash_kernels_loaded() if requested_backend == AttentionBackend.AUTO.value: if FLASH_KERNEL is not None: resolved = AttentionBackend.KERNELS_FLASH elif flex_attention is not None: resolved = AttentionBackend.FLEX else: resolved = AttentionBackend.SDPA elif requested_backend == AttentionBackend.KERNELS_FLASH.value: assert FLASH_KERNEL is not None, "Kernels Flash Attention is not available in this environment." resolved = AttentionBackend.KERNELS_FLASH elif requested_backend == AttentionBackend.FLEX.value: assert flex_attention is not None, "Flex Attention is not available in this environment." resolved = AttentionBackend.FLEX elif requested_backend == AttentionBackend.SDPA.value: resolved = AttentionBackend.SDPA else: raise AssertionError(f"Unsupported attention backend: {requested_backend}") if not _BACKEND_CONFIRMED: print(f"Attention backend: config='{requested_backend}' -> resolved='{resolved.value}'") _BACKEND_CONFIRMED = True return resolved @torch.compiler.disable def get_attention_mask( effective_backend: AttentionBackend, batch_size: int, seq_len: int, device: torch.device, attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[BlockMask]]: """Build padding masks once for all encoder layers. Returns (attention_mask_2d, attention_mask_4d, flex_block_mask). """ if attention_mask is None: return None, None, None attention_mask_2d = attention_mask.bool() if effective_backend == AttentionBackend.KERNELS_FLASH: return attention_mask_2d, None, None if effective_backend == AttentionBackend.FLEX: assert create_block_mask is not None, "Flex attention backend requested but torch.create_block_mask is unavailable." valid_lens = attention_mask_2d.sum(dim=-1) def mask_mod(batch_idx, head_idx, q_idx, kv_idx): return (q_idx < valid_lens[batch_idx]) & (kv_idx < valid_lens[batch_idx]) flex_block_mask = create_block_mask(mask_mod, batch_size, 1, seq_len, seq_len, device=device) return attention_mask_2d, None, flex_block_mask # SDPA/manual masks only keys. Padding queries still attend to real keys, so # their outputs stay finite instead of softmaxing over all -inf scores. attention_mask_4d = attention_mask_2d[:, None, None, :] return attention_mask_2d, attention_mask_4d, None def bool_to_additive_mask( bool_mask: torch.Tensor, dtype: torch.dtype, ) -> torch.Tensor: """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid). Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))` directly on a bool tensor returns a bool tensor because `-inf` casts to `True`. That silently drops the mask. Always allocate a float tensor first, then fill it. This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask. """ assert bool_mask.dtype == torch.bool, ( f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}" ) additive = torch.zeros_like(bool_mask, dtype=dtype) additive.masked_fill_(bool_mask.logical_not(), float("-inf")) return additive import typing as T from dataclasses import dataclass, fields import torch import torch.nn as nn import torch.nn.functional as F @dataclass class TTTConfig: lr: float = 4e-4 steps: int = 30 ags: int = 16 batch_size: int = 2 mask_ratio: float = 0.15 crop_size: int = 1024 bert_leave_prob: float = 0.1 bert_replace_prob: float = 0.1 optimizer: str = "sgd" momentum: float = 0.0 weight_decay: float = 0.0 seed: int | None = 0 lora_rank: int = 8 lora_alpha: float = 32.0 lora_target_replace_module: str | None = None lora_target_modules: tuple[str, ...] | None = None initial_state_reset: bool = True automatic_best_state_reset: bool = False eval_each_step: bool = False gradient_clip: bool = False gradient_clip_max_norm: float = 1.0 @classmethod def from_kwargs(cls, **kwargs: T.Any) -> "TTTConfig": valid_names = {field.name for field in fields(cls)} unknown_names = set(kwargs) - valid_names assert len(unknown_names) == 0, f"Unknown TTTConfig fields: {sorted(unknown_names)}" return cls(**kwargs) def merged(self, overrides: T.Mapping[str, T.Any] | "TTTConfig" | None) -> "TTTConfig": if overrides is None: return self if isinstance(overrides, TTTConfig): return overrides values = {field.name: self.__dict__[field.name] for field in fields(self)} for name, value in overrides.items(): assert name in values, f"Unknown TTTConfig field: {name}" values[name] = value return TTTConfig(**values) def verify(self) -> None: assert self.lr > 0.0, "TTT learning rate must be positive." assert self.steps >= 1, "TTT steps must be >= 1." assert self.ags >= 1, "TTT gradient accumulation steps must be >= 1." assert self.batch_size >= 1, "TTT batch_size must be >= 1." assert 0.0 < self.mask_ratio <= 1.0, "TTT mask_ratio must be in (0, 1]." assert self.crop_size >= 1, "TTT crop_size must be >= 1." assert self.lora_rank >= 1, "TTT v1 is LoRA-only, so lora_rank must be >= 1." assert self.lora_alpha > 0.0, "TTT lora_alpha must be positive." assert self.optimizer in {"adamw", "sgd"}, "TTT optimizer must be 'adamw' or 'sgd'." assert 0.0 <= self.bert_leave_prob <= 1.0, "bert_leave_prob must be in [0, 1]." assert 0.0 <= self.bert_replace_prob <= 1.0, "bert_replace_prob must be in [0, 1]." assert self.bert_leave_prob + self.bert_replace_prob <= 1.0, ( "bert_leave_prob + bert_replace_prob must be <= 1." ) if self.gradient_clip: assert self.gradient_clip_max_norm > 0.0, "gradient_clip_max_norm must be positive." class LoraInjectedLinear(nn.Module): def __init__(self, linear: nn.Module, rank: int, alpha: float) -> None: super().__init__() weight = linear._parameters["weight"] assert weight.ndim == 2, "LoRA can only wrap 2D linear weights." self.linear = linear self.linear.requires_grad_(False) self.rank = rank self.scale = alpha in_features = weight.shape[1] out_features = weight.shape[0] self.lora_down = nn.Linear(in_features, rank, bias=False, dtype=torch.float32) self.lora_up = nn.Linear(rank, out_features, bias=False, dtype=torch.float32) self.lora_down.to(device=weight.device) self.lora_up.to(device=weight.device) nn.init.normal_(self.lora_down.weight, std=1.0 / rank) nn.init.zeros_(self.lora_up.weight) @property def weight(self) -> torch.Tensor: return self.linear._parameters["weight"] @property def bias(self) -> torch.Tensor | None: return self.linear._parameters["bias"] def forward(self, x: torch.Tensor) -> torch.Tensor: base = self.linear(x) delta = self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale return base + delta.to(dtype=base.dtype) class FastPLMTestTimeTrainingMixin: def init_ttt(self, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None) -> None: base_config = TTTConfig() self._ttt_cfg = base_config.merged(ttt_config) self._ttt_cfg.verify() self._ttt_initialized = False self._ttt_initial_state: list[dict[str, torch.Tensor]] | None = None @property def ttt_config(self) -> TTTConfig: if "_ttt_cfg" not in self.__dict__: self.init_ttt() return self._ttt_cfg def _ttt_get_trainable_modules(self) -> list[nn.Module]: return [self] def _ttt_get_frozen_modules(self) -> list[nn.Module]: return [] def _ttt_tokenize( self, seq: str | list[str] | None = None, input_ids: torch.Tensor | None = None, **kwargs: T.Any, ) -> torch.Tensor | dict[str, torch.Tensor]: del kwargs if input_ids is not None: return input_ids assert seq is not None, "Pass either seq or input_ids for TTT." tokenized = self.tokenizer(seq, return_tensors="pt", padding=True) return tokenized["input_ids"] def _ttt_mask_token(self) -> int: return int(self.tokenizer.mask_token_id) def _ttt_padding_token(self) -> int: return int(self.tokenizer.pad_token_id) def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: tokenizer = self.tokenizer special_ids = set(tokenizer.all_special_ids) vocab_size = int(self.config.vocab_size) ids = [idx for idx in range(vocab_size) if idx not in special_ids] assert len(ids) > 0, "TTT replacement token set is empty." return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) def _ttt_predict_logits( self, batch: torch.Tensor | dict[str, torch.Tensor], **kwargs: T.Any, ) -> torch.Tensor: del kwargs if isinstance(batch, dict): output = self(**batch) return output.logits attention_mask = batch.ne(self._ttt_padding_token()) output = self(input_ids=batch, attention_mask=attention_mask) return output.logits def _ttt_eval_step( self, step: int, loss: float, seq: str | list[str] | None = None, input_ids: torch.Tensor | None = None, **kwargs: T.Any, ) -> tuple[dict[str, T.Any], float | None]: del step, loss, seq, input_ids, kwargs return {}, None def _ttt_is_lora_target( self, name: str, full_name: str, module: nn.Module, active: bool, target_modules: tuple[str, ...] | None, ) -> bool: if not active: return False if isinstance(module, LoraInjectedLinear): return False if ( target_modules is not None and name not in target_modules and full_name not in target_modules ): return False if isinstance(module, nn.Linear): return True if "weight" not in module._parameters: return False weight = module._parameters["weight"] if weight is None or weight.ndim != 2: return False return "Linear" in module.__class__.__name__ def _ttt_inject_lora(self) -> int: cfg = self.ttt_config cfg.verify() target_class = cfg.lora_target_replace_module target_modules = cfg.lora_target_modules wrapped = 0 def inject(module: nn.Module, prefix: str, active: bool) -> None: nonlocal wrapped for name, child in list(module.named_children()): full_name = f"{prefix}.{name}" if prefix else name child_active = active if target_class is not None: child_active = active or child.__class__.__name__ == target_class if self._ttt_is_lora_target(name, full_name, child, child_active, target_modules): setattr( module, name, LoraInjectedLinear(child, rank=cfg.lora_rank, alpha=cfg.lora_alpha), ) wrapped += 1 continue inject(child, full_name, child_active) for trainable_module in self._ttt_get_trainable_modules(): inject(trainable_module, "", target_class is None) assert wrapped > 0, "TTT LoRA injection did not find any target modules." return wrapped def _ttt_lora_modules(self) -> list[LoraInjectedLinear]: return [module for module in self.modules() if isinstance(module, LoraInjectedLinear)] def _ttt_lora_parameters(self) -> list[nn.Parameter]: params: list[nn.Parameter] = [] for module in self._ttt_lora_modules(): params.extend(module.lora_down.parameters()) params.extend(module.lora_up.parameters()) assert len(params) > 0, "TTT has no LoRA parameters." return params def _ttt_snapshot_lora_state(self) -> list[dict[str, torch.Tensor]]: snapshot = [] for module in self._ttt_lora_modules(): snapshot.append( { "lora_down.weight": module.lora_down.weight.detach().clone(), "lora_up.weight": module.lora_up.weight.detach().clone(), } ) assert len(snapshot) > 0, "TTT has no LoRA state to snapshot." return snapshot def _ttt_restore_lora_state(self, state: list[dict[str, torch.Tensor]]) -> None: modules = self._ttt_lora_modules() assert len(modules) == len(state), "TTT LoRA state/module count mismatch." with torch.no_grad(): for module, module_state in zip(modules, state): module.lora_down.weight.copy_(module_state["lora_down.weight"]) module.lora_up.weight.copy_(module_state["lora_up.weight"]) def _ttt_ensure_initialized(self) -> None: if "_ttt_cfg" not in self.__dict__: self.init_ttt() if self._ttt_initialized: return self._ttt_inject_lora() self._ttt_initial_state = self._ttt_snapshot_lora_state() self._ttt_initialized = True def ttt_reset(self) -> None: self._ttt_ensure_initialized() assert self._ttt_initial_state is not None, "TTT initial state is not available." self._ttt_restore_lora_state(self._ttt_initial_state) def _ttt_make_optimizer(self) -> torch.optim.Optimizer: cfg = self.ttt_config params = self._ttt_lora_parameters() if cfg.optimizer == "sgd": return torch.optim.SGD( params, lr=cfg.lr, momentum=cfg.momentum, weight_decay=cfg.weight_decay, ) return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay) def _ttt_to_device( self, batch: torch.Tensor | dict[str, torch.Tensor], device: torch.device, ) -> torch.Tensor | dict[str, torch.Tensor]: if isinstance(batch, dict): return {name: tensor.to(device) for name, tensor in batch.items()} return batch.to(device) def _ttt_input_ids_from_batch( self, batch: torch.Tensor | dict[str, torch.Tensor], ) -> torch.Tensor: if isinstance(batch, dict): return batch["input_ids"] return batch def _ttt_set_input_ids( self, batch: torch.Tensor | dict[str, torch.Tensor], input_ids: torch.Tensor, ) -> torch.Tensor | dict[str, torch.Tensor]: if isinstance(batch, dict): updated = dict(batch) updated["input_ids"] = input_ids return updated return input_ids def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: pad_token = self._ttt_padding_token() mask = input_ids.ne(pad_token) special_ids = set(self.tokenizer.all_special_ids) for special_id in special_ids: mask = mask & input_ids.ne(int(special_id)) return mask def _ttt_sample_crop( self, batch: torch.Tensor | dict[str, torch.Tensor], generator: torch.Generator, ) -> torch.Tensor | dict[str, torch.Tensor]: input_ids = self._ttt_input_ids_from_batch(batch) cfg = self.ttt_config if input_ids.shape[1] <= cfg.crop_size: return batch high = input_ids.shape[1] - cfg.crop_size + 1 start = int( torch.randint( high, (1,), generator=generator, device=input_ids.device, ).item() ) end = start + cfg.crop_size if isinstance(batch, dict): cropped = {} for name, tensor in batch.items(): if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]: cropped[name] = tensor[:, start:end] else: cropped[name] = tensor return cropped return input_ids[:, start:end] def _ttt_sample_batch( self, tokenized: torch.Tensor | dict[str, torch.Tensor], generator: torch.Generator, ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]: cfg = self.ttt_config batch = self._ttt_sample_crop(tokenized, generator) input_ids = self._ttt_input_ids_from_batch(batch) rows = torch.randint( input_ids.shape[0], (cfg.batch_size,), generator=generator, device=input_ids.device, ) if isinstance(batch, dict): sampled: torch.Tensor | dict[str, torch.Tensor] = {} for name, tensor in batch.items(): if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]: sampled[name] = tensor.index_select(0, rows) else: sampled[name] = tensor else: sampled = input_ids.index_select(0, rows) sampled_ids = self._ttt_input_ids_from_batch(sampled) labels = sampled_ids.clone() non_special = self._ttt_non_special_mask(sampled_ids) label_mask = torch.zeros_like(non_special) for row_idx in range(sampled_ids.shape[0]): candidate_positions = torch.where(non_special[row_idx])[0] if candidate_positions.numel() == 0: continue num_mask = max(1, int(round(candidate_positions.numel() * cfg.mask_ratio))) order = torch.randperm( candidate_positions.numel(), generator=generator, device=sampled_ids.device, ) chosen = candidate_positions[order[:num_mask]] label_mask[row_idx, chosen] = True labels = labels.masked_fill(~label_mask, -100) masked_ids = sampled_ids.clone() chosen_positions = torch.where(label_mask) if chosen_positions[0].numel() > 0: random_values = torch.rand( chosen_positions[0].shape, generator=generator, device=sampled_ids.device, ) leave = random_values < cfg.bert_leave_prob replace = (random_values >= cfg.bert_leave_prob) & ( random_values < cfg.bert_leave_prob + cfg.bert_replace_prob ) mask = ~(leave | replace) if mask.any(): masked_ids[ chosen_positions[0][mask], chosen_positions[1][mask], ] = self._ttt_mask_token() if replace.any(): replacement_tokens = self._ttt_replacement_tokens(sampled_ids) replacement_idx = torch.randint( replacement_tokens.shape[0], (int(replace.sum().item()),), generator=generator, device=sampled_ids.device, ) masked_ids[ chosen_positions[0][replace], chosen_positions[1][replace], ] = replacement_tokens[replacement_idx] return self._ttt_set_input_ids(sampled, masked_ids), labels def ttt( self, seq: str | list[str] | None = None, input_ids: torch.Tensor | None = None, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None, **kwargs: T.Any, ) -> dict[str, T.Any]: if ttt_config is not None: if "_ttt_initialized" in self.__dict__ and self._ttt_initialized: next_cfg = self.ttt_config.merged(ttt_config) assert next_cfg.lora_rank == self.ttt_config.lora_rank, ( "Changing lora_rank after TTT initialization is not supported." ) assert next_cfg.lora_alpha == self.ttt_config.lora_alpha, ( "Changing lora_alpha after TTT initialization is not supported." ) assert ( next_cfg.lora_target_replace_module == self.ttt_config.lora_target_replace_module ), "Changing LoRA target class after TTT initialization is not supported." assert next_cfg.lora_target_modules == self.ttt_config.lora_target_modules, ( "Changing LoRA target modules after TTT initialization is not supported." ) self._ttt_cfg = next_cfg else: self.init_ttt(ttt_config) self._ttt_ensure_initialized() cfg = self.ttt_config if cfg.initial_state_reset: self.ttt_reset() device = next(self.parameters()).device tokenized = self._ttt_tokenize(seq=seq, input_ids=input_ids, **kwargs) tokenized = self._ttt_to_device(tokenized, device) generator_device = device if device.type == "cuda" else torch.device("cpu") generator = torch.Generator(device=generator_device) if cfg.seed is not None: generator.manual_seed(cfg.seed) module_modes = {module: module.training for module in self.modules()} requires_grad = {param: param.requires_grad for param in self.parameters()} losses: list[float] = [] step_metrics: list[dict[str, T.Any]] = [] best_state: list[dict[str, torch.Tensor]] | None = None best_metric: float | None = None best_step = 0 try: self.train() for param in self.parameters(): param.requires_grad_(False) for param in self._ttt_lora_parameters(): param.requires_grad_(True) optimizer = self._ttt_make_optimizer() optimizer.zero_grad(set_to_none=True) total_micro_steps = cfg.steps * cfg.ags for micro_step in range(total_micro_steps): batch, labels = self._ttt_sample_batch(tokenized, generator) logits = self._ttt_predict_logits(batch, **kwargs) labels = labels.to(device=logits.device) loss = F.cross_entropy( logits.reshape(-1, logits.shape[-1]), labels.reshape(-1), ignore_index=-100, ) (loss / cfg.ags).backward() if (micro_step + 1) % cfg.ags != 0: continue if cfg.gradient_clip: torch.nn.utils.clip_grad_norm_( self._ttt_lora_parameters(), cfg.gradient_clip_max_norm, ) optimizer.step() optimizer.zero_grad(set_to_none=True) step = (micro_step + 1) // cfg.ags loss_value = float(loss.detach().item()) losses.append(loss_value) if cfg.eval_each_step: metrics, metric = self._ttt_eval_step( step=step, loss=loss_value, seq=seq, input_ids=input_ids, **kwargs, ) if len(metrics) > 0: step_metrics.append(metrics) if metric is not None and ( best_metric is None or metric > best_metric ): best_metric = metric best_step = step best_state = self._ttt_snapshot_lora_state() if cfg.automatic_best_state_reset and best_state is not None: self._ttt_restore_lora_state(best_state) finally: for param, value in requires_grad.items(): param.requires_grad_(value) for module, training in module_modes.items(): module.train(training) return { "losses": losses, "step_metrics": step_metrics, "best_step": best_step, "best_metric": best_metric, } """FastESMFold: self-contained ESMFold with FastESM2 attention and opt-in TTT. Usage: from transformers import AutoModel model = AutoModel.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True).cuda() # Basic folding, no TTT result = model.fold_protein("MKTLLILAVVA...") print(result["plddt"], result["pdb_string"][:100]) # Experimental folding with TTT result = model.fold_protein("MKTLLILAVVA...", ttt=True) Dependencies: torch, transformers, einops No dependency on: esm (fair-esm), proteinttt, openfold """ import copy from dataclasses import dataclass, field from functools import wraps from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn from torch.nn import functional as F from einops import rearrange from transformers import EsmTokenizer, PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import ModelOutput from transformers.models.esm.configuration_esm import EsmConfig from transformers.models.esm.modeling_esm import ( EsmContactPredictionHead, EsmEmbeddings, EsmIntermediate, EsmLMHead, EsmOutput, EsmSelfOutput, RotaryEmbedding, ) from transformers.models.esm.modeling_esmfold import EsmForProteinFolding # ============================================================================= # Output Dataclass # ============================================================================= @dataclass class FastEsmEncoderOutput(ModelOutput): last_hidden_state: Optional[torch.Tensor] = None hidden_states: Optional[Tuple[torch.Tensor, ...]] = None attentions: Optional[Tuple[torch.Tensor, ...]] = None # ============================================================================= # FastESM2 Attention Layers (multi-backend: SDPA, Flash, Flex) # ============================================================================= class EsmSelfAttention(nn.Module): def __init__(self, config, position_embedding_type: Optional[str] = None): super().__init__() assert config.hidden_size % config.num_attention_heads == 0, ( f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " f"heads ({config.num_attention_heads})" ) self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.all_head_size = self.num_attention_heads * self.attention_head_size self.query = nn.Linear(config.hidden_size, self.all_head_size) self.key = nn.Linear(config.hidden_size, self.all_head_size) self.value = nn.Linear(config.hidden_size, self.all_head_size) self.scale = self.attention_head_size**-0.5 self.dropout_prob = config.attention_probs_dropout_prob self.config = config self.attn_backend = resolve_attention_backend(config.attn_backend) self.position_embedding_type = position_embedding_type or config.position_embedding_type self.rotary_embeddings = None if self.position_embedding_type == "rotary": self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size) def forward( self, hidden_states: torch.Tensor, attention_mask_2d: Optional[torch.Tensor] = None, attention_mask_4d: Optional[torch.Tensor] = None, flex_block_mask: Optional[BlockMask] = None, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: batch_size, seq_length = hidden_states.shape[:-1] hidden_shape = (batch_size, seq_length, -1, self.attention_head_size) query_BHLD = self.query(hidden_states).view(hidden_shape).transpose(1, 2) key_BHLD = self.key(hidden_states).view(hidden_shape).transpose(1, 2) value_BHLD = self.value(hidden_states).view(hidden_shape).transpose(1, 2) query_BHLD = query_BHLD * self.scale if self.position_embedding_type == "rotary": query_BHLD, key_BHLD = self.rotary_embeddings(query_BHLD, key_BHLD) attn_output, attn_weights = self._attn( query_BHLD, key_BHLD, value_BHLD, attention_mask_2d=attention_mask_2d, attention_mask_4d=attention_mask_4d, flex_block_mask=flex_block_mask, output_attentions=output_attentions, ) return attn_output, attn_weights def _attn( self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor, value_BHLD: torch.Tensor, attention_mask_2d: Optional[torch.Tensor] = None, attention_mask_4d: Optional[torch.Tensor] = None, flex_block_mask: Optional[BlockMask] = None, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if output_attentions: return self._manual_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) if self.attn_backend == AttentionBackend.KERNELS_FLASH: return self._kernels_flash_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_2d) elif self.attn_backend == AttentionBackend.FLEX: return self._flex_attn(query_BHLD, key_BHLD, value_BHLD, flex_block_mask) elif self.attn_backend == AttentionBackend.SDPA: return self._sdpa_attn(query_BHLD, key_BHLD, value_BHLD, attention_mask_4d) else: raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") def _manual_attn( self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor, value_BHLD: torch.Tensor, attention_mask_4d: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: attn_weights = torch.matmul(query_BHLD, key_BHLD.transpose(-1, -2)) if attention_mask_4d is not None: attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) attn_weights = F.softmax(attn_weights, dim=-1) if self.dropout_prob > 0 and self.training: attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training) context_BHLD = torch.matmul(attn_weights, value_BHLD) attn_output = rearrange(context_BHLD, "b h s d -> b s (h d)") return attn_output, attn_weights def _kernels_flash_attn( self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor, value_BHLD: torch.Tensor, attention_mask_2d: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, None]: query_BLHD = query_BHLD.transpose(1, 2).contiguous() key_BLHD = key_BHLD.transpose(1, 2).contiguous() value_BLHD = value_BHLD.transpose(1, 2).contiguous() # Q is pre-scaled by self.scale in forward() -- pass softmax_scale=1.0 # to prevent the kernel from applying its default 1/sqrt(head_dim). attn_output = kernels_flash_attention_func( query_states=query_BLHD, key_states=key_BLHD, value_states=value_BLHD, attention_mask_2d=attention_mask_2d, causal=False, softmax_scale=1.0, ) return rearrange(attn_output, "b s h d -> b s (h d)"), None def _flex_attn( self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor, value_BHLD: torch.Tensor, flex_block_mask: Optional[BlockMask] = None, ) -> Tuple[torch.Tensor, None]: assert flex_attention is not None, "Flex attention is not available in this environment." fn = _get_flex_attention_fn() context_BHLD = fn(query_BHLD, key_BHLD, value_BHLD, block_mask=flex_block_mask, scale=1.0) return rearrange(context_BHLD, "b h s d -> b s (h d)"), None def _sdpa_attn( self, query_BHLD: torch.Tensor, key_BHLD: torch.Tensor, value_BHLD: torch.Tensor, attention_mask_4d: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, None]: context_BHLD = F.scaled_dot_product_attention( query_BHLD, key_BHLD, value_BHLD, attn_mask=attention_mask_4d, dropout_p=self.dropout_prob if self.training else 0.0, scale=1.0, ) return rearrange(context_BHLD, "b h s d -> b s (h d)"), None class EsmAttention(nn.Module): def __init__(self, config): super().__init__() self.self = EsmSelfAttention(config) self.output = EsmSelfOutput(config) self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask_2d: Optional[torch.Tensor] = None, attention_mask_4d: Optional[torch.Tensor] = None, flex_block_mask: Optional[BlockMask] = None, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: hidden_states_ln = self.LayerNorm(hidden_states) attn_output, attn_weights = self.self( hidden_states_ln, attention_mask_2d=attention_mask_2d, attention_mask_4d=attention_mask_4d, flex_block_mask=flex_block_mask, output_attentions=output_attentions, ) attention_output = self.output(attn_output, hidden_states) return attention_output, attn_weights class EsmLayer(nn.Module): def __init__(self, config): super().__init__() self.attention = EsmAttention(config) self.intermediate = EsmIntermediate(config) self.output = EsmOutput(config) self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask_2d: Optional[torch.Tensor] = None, attention_mask_4d: Optional[torch.Tensor] = None, flex_block_mask: Optional[BlockMask] = None, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: attention_output, attn_weights = self.attention( hidden_states, attention_mask_2d=attention_mask_2d, attention_mask_4d=attention_mask_4d, flex_block_mask=flex_block_mask, output_attentions=output_attentions, ) layer_output = self._feed_forward(attention_output) return layer_output, attn_weights def _feed_forward(self, attention_output: torch.Tensor) -> torch.Tensor: attention_output_ln = self.LayerNorm(attention_output) intermediate_output = self.intermediate(attention_output_ln) return self.output(intermediate_output, attention_output) class FastEsmEncoder(nn.Module): def __init__(self, config): super().__init__() self.config = config self.attention_backend = resolve_attention_backend(config.attn_backend) self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)]) self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_hidden_states: bool = False, output_attentions: bool = False, ) -> FastEsmEncoderOutput: all_hidden_states = () if output_hidden_states else None all_attentions = () if output_attentions else None attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask( effective_backend=self.attention_backend, batch_size=hidden_states.shape[0], seq_len=hidden_states.shape[1], device=hidden_states.device, attention_mask=attention_mask, ) for layer_module in self.layer: if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) hidden_states, attn_weights = layer_module( hidden_states, attention_mask_2d=attention_mask_2d, attention_mask_4d=attention_mask_4d, flex_block_mask=flex_block_mask, output_attentions=output_attentions, ) if all_attentions is not None: all_attentions = all_attentions + (attn_weights,) if self.emb_layer_norm_after: hidden_states = self.emb_layer_norm_after(hidden_states) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) return FastEsmEncoderOutput( last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions, ) # ============================================================================= # FastESM Backbone (replaces EsmModel inside ESMFold) # ============================================================================= class FastEsmBackbone(nn.Module): """FastESM2 backbone with multi-backend attention. Drop-in replacement for transformers.EsmModel inside EsmForProteinFolding. State dict keys match HuggingFace EsmModel exactly, so pretrained weights load without any key remapping. """ def __init__(self, config): super().__init__() self.config = config self.embeddings = EsmEmbeddings(config) self.encoder = FastEsmEncoder(config) self.contact_head = EsmContactPredictionHead( in_features=config.num_hidden_layers * config.num_attention_heads, bias=True ) def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> FastEsmEncoderOutput: output_attentions = output_attentions if output_attentions is not None else False output_hidden_states = output_hidden_states if output_hidden_states is not None else False token_embedding_output = self.embeddings( input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, ) encoder_outputs = self.encoder( token_embedding_output, attention_mask=attention_mask, output_hidden_states=output_hidden_states, output_attentions=output_attentions, ) return FastEsmEncoderOutput( last_hidden_state=encoder_outputs.last_hidden_state, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, ) # ============================================================================= # TTT (Test-Time Training) Configuration and Utilities # ============================================================================= _ESM_STANDARD_AA = list("ACDEFGHIKLMNPQRSTVWY") class LoraInjectedLinear(nn.Module): """LoRA-augmented linear layer matching lora_diffusion's behavior. Replaces an existing nn.Linear with base(x) + lora_up(lora_down(x)) * scale. Initialization follows cloneofsimo/lora: down=Normal(0, 1/r), up=zeros. """ def __init__(self, original_linear: nn.Linear, r: int = 4, scale: float = 1.0): super().__init__() self.linear = original_linear in_features = original_linear.in_features out_features = original_linear.out_features assert r <= min(in_features, out_features), f"LoRA rank {r} exceeds dimensions ({in_features}, {out_features})" self.lora_down = nn.Linear(in_features, r, bias=False) self.lora_up = nn.Linear(r, out_features, bias=False) self.scale = scale nn.init.normal_(self.lora_down.weight, std=1.0 / r) nn.init.zeros_(self.lora_up.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear(x) + self.lora_up(self.lora_down(x)) * self.scale def inject_trainable_lora( model: nn.Module, target_class_name: str, r: int, scale: float, ) -> List[nn.Parameter]: """Replace nn.Linear layers inside modules matching target_class_name with LoRA. Matches lora_diffusion's inject_trainable_lora behavior: finds all modules whose class name matches target_class_name, then replaces their nn.Linear children with LoraInjectedLinear. Returns the list of trainable LoRA parameters. """ lora_params: List[nn.Parameter] = [] for _parent_name, parent_module in model.named_modules(): if parent_module.__class__.__name__ != target_class_name: continue for child_name, child_module in list(parent_module.named_children()): if not isinstance(child_module, nn.Linear): continue lora_linear = LoraInjectedLinear(child_module, r=r, scale=scale) lora_linear = lora_linear.to( device=child_module.weight.device, dtype=child_module.weight.dtype, ) setattr(parent_module, child_name, lora_linear) lora_params.extend(lora_linear.lora_down.parameters()) lora_params.extend(lora_linear.lora_up.parameters()) return lora_params @dataclass class TTTConfig: lr: float = 4e-4 ags: int = 4 steps: int = 10 batch_size: int = 4 mask_ratio: float = 0.15 crop_size: int = 1024 bert_leave_prob: float = 0.1 bert_replace_prob: float = 0.1 optimizer: str = "sgd" momentum: float = 0.0 weight_decay: float = 0.0 seed: Optional[int] = 0 initial_state_reset: bool = True freeze_embeddings: bool = True lora_rank: int = 8 lora_alpha: float = 32.0 lora_target_class: str = "EsmSelfAttention" def verify(self) -> None: assert self.lr > 0.0, "TTT learning rate must be positive." assert self.ags > 0, "TTT ags must be positive." assert self.steps >= 0, "TTT steps must be non-negative." assert self.batch_size > 0, "TTT batch_size must be positive." assert 0.0 < self.mask_ratio <= 1.0, "TTT mask_ratio must be in (0, 1]." assert self.crop_size > 0, "TTT crop_size must be positive." assert 0.0 <= self.bert_leave_prob <= 1.0 assert 0.0 <= self.bert_replace_prob <= 1.0 assert self.bert_leave_prob + self.bert_replace_prob <= 1.0 assert self.optimizer in {"sgd", "adamw"} assert self.lora_rank >= 0 assert self.lora_alpha > 0.0 def preserve_model_state(func: Callable[..., Any]) -> Callable[..., Any]: @wraps(func) def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: was_training = self.training original_device = next(self.parameters()).device original_requires_grad = { name: parameter.requires_grad for name, parameter in self.named_parameters() } try: return func(self, *args, **kwargs) finally: self.train(was_training) self.to(original_device) for name, parameter in self.named_parameters(): if name in original_requires_grad: parameter.requires_grad = original_requires_grad[name] else: parameter.requires_grad = False return wrapper # ============================================================================= # FastEsmFoldConfig # ============================================================================= class FastEsmFoldConfig(EsmConfig): model_type = "fast_esmfold" def __init__(self, attn_backend: str = "sdpa", ttt_config: Optional[Dict[str, Any]] = None, **kwargs): super().__init__(**kwargs) self.attn_backend = attn_backend self.ttt_config = ttt_config or { "lr": 4e-4, "steps": 10, "lora_rank": 8, "lora_alpha": 32.0, } # ============================================================================= # FastEsmForProteinFolding # ============================================================================= class FastEsmForProteinFolding(EsmForProteinFolding): """ESMFold with FastESM2 attention backends and opt-in experimental TTT. Inherits all folding logic (trunk, structure module, output_to_pdb, infer) from transformers.EsmForProteinFolding. Replaces the ESM2 backbone with FastESM2 for optimized attention and adds opt-in TTT for difficult targets. Key API: result = model.fold_protein("MKTL...", ttt=True) # result = {"plddt": float, "ptm": float, "pdb_string": str} """ config_class = FastEsmFoldConfig def __init__(self, config: FastEsmFoldConfig): super().__init__(config) # Replace standard ESM2 backbone with FastESM2 (multi-backend attention) # unless use_standard_backbone is set (for TTT debugging/compatibility) if not config.ttt_config.get("use_standard_backbone", False): self.esm = FastEsmBackbone(config) self.esm.requires_grad_(False) if config.esmfold_config.fp16_esm: self.esm.half() # MLM head for TTT (pretrained EsmLMHead: Dense -> GELU -> LN -> Linear) self.mlm_head = EsmLMHead(config) # TTT state (lazy initialization) ttt_kwargs = {k: v for k, v in config.ttt_config.items() if k != "use_standard_backbone"} self._ttt_cfg = TTTConfig(**ttt_kwargs) self._ttt_cfg.verify() self._ttt_initialized = False self._ttt_initial_state = None self._ttt_generator = torch.Generator() if self._ttt_cfg.seed is not None: self._ttt_generator.manual_seed(self._ttt_cfg.seed) self._non_special_tokens_cache = None self._ttt_tokenizer = None def _get_ttt_tokenizer(self) -> EsmTokenizer: if self._ttt_tokenizer is None: self._ttt_tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D") return self._ttt_tokenizer def _ensure_ttt_ready(self) -> None: """Lazy TTT initialization. Injects LoRA adapters and saves initial state. Must be called after weights are loaded (not in __init__).""" if self._ttt_initialized: return self._ttt_initialized = True tokenizer = self._get_ttt_tokenizer() vocab = tokenizer.get_vocab() self._non_special_tokens_cache = [vocab[c] for c in _ESM_STANDARD_AA if c in vocab] if self._ttt_cfg.lora_rank > 0: self.mlm_head.eval() for p in self.mlm_head.parameters(): p.requires_grad = False # Seed global state before LoRA init for reproducible weight initialization if self._ttt_cfg.seed is not None: torch.manual_seed(self._ttt_cfg.seed) self._inject_lora() else: # Legacy path: jointly-trained random linear projection head H = self.config.hidden_size V = self.config.vocab_size device = next(self.esm.parameters()).device self._ttt_lm_proj = nn.Linear(H, V, bias=True).to(device) if self._ttt_cfg.initial_state_reset: self._ttt_initial_state = self._ttt_get_state() @property def _uses_lora(self) -> bool: return self._ttt_cfg.lora_rank > 0 def _inject_lora(self) -> None: """Inject LoRA adapters into ESM2 attention layers (matching lora_diffusion behavior).""" self._lora_params = inject_trainable_lora( self.esm, target_class_name=self._ttt_cfg.lora_target_class, r=self._ttt_cfg.lora_rank, scale=self._ttt_cfg.lora_alpha, ) assert len(self._lora_params) > 0, ( f"No LoRA params injected. Check target_class_name='{self._ttt_cfg.lora_target_class}' " f"matches attention modules in the backbone." ) # ---- TTT State Management ---- def _get_lora_modules(self) -> List[LoraInjectedLinear]: """Find all LoraInjectedLinear modules in the backbone.""" return [m for m in self.esm.modules() if isinstance(m, LoraInjectedLinear)] def _ttt_get_state(self) -> Dict[str, Any]: if self._uses_lora: lora_state = [] for m in self._get_lora_modules(): lora_state.append({ "down": m.lora_down.weight.data.clone(), "up": m.lora_up.weight.data.clone(), }) return {"_lora_state": lora_state} return { "esm": copy.deepcopy(self.esm), "_ttt_lm_proj": copy.deepcopy(self._ttt_lm_proj), } def _ttt_set_state(self, state: Dict[str, Any]) -> None: if "_lora_state" in state: modules = self._get_lora_modules() assert len(modules) == len(state["_lora_state"]) for m, saved in zip(modules, state["_lora_state"]): m.lora_down.weight.data.copy_(saved["down"]) m.lora_up.weight.data.copy_(saved["up"]) return if "esm" in state: self.esm = copy.deepcopy(state["esm"]) if "_ttt_lm_proj" in state: self._ttt_lm_proj = copy.deepcopy(state["_ttt_lm_proj"]) def ttt_reset(self) -> None: """Reset model to pre-TTT state (restore initial LoRA or backbone weights).""" assert self._ttt_initial_state is not None, "TTT reset requires initial_state_reset=True." self._ttt_set_state(self._ttt_initial_state) # ---- TTT Core ---- def _ttt_tokenize(self, seq: str) -> torch.Tensor: tokenizer = self._get_ttt_tokenizer() out = tokenizer( seq, return_tensors="pt", add_special_tokens=self._uses_lora, padding=False, truncation=False, ) return out["input_ids"] def _ttt_mask_token(self) -> int: return self._get_ttt_tokenizer().mask_token_id def _ttt_get_non_special_tokens(self) -> List[int]: if self._non_special_tokens_cache is not None: return self._non_special_tokens_cache tokenizer = self._get_ttt_tokenizer() vocab = tokenizer.get_vocab() self._non_special_tokens_cache = [vocab[c] for c in _ESM_STANDARD_AA if c in vocab] return self._non_special_tokens_cache def _ttt_predict_logits(self, batch: torch.Tensor) -> torch.Tensor: """Run ESM2 backbone + LM head to get MLM logits.""" # Temporarily unfreeze backbone for gradient flow during TTT output = self.esm(input_ids=batch) hidden = output.last_hidden_state if self._uses_lora: return self.mlm_head(hidden) return self._ttt_lm_proj(hidden) def _ttt_sample_batch( self, x: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: _, seq_len = x.shape batch_size = self._ttt_cfg.batch_size crop_size = min(self._ttt_cfg.crop_size, seq_len) x_expanded = x.expand(batch_size, -1) if seq_len == crop_size: start_indices = torch.zeros(batch_size, dtype=torch.long) else: start_indices = torch.randint( 0, seq_len - crop_size + 1, (batch_size,), generator=self._ttt_generator, ).to(torch.long) batch_cropped = torch.stack([ x_expanded[index, start : start + crop_size] for index, start in enumerate(start_indices) ]) non_special_tokens = set(self._ttt_get_non_special_tokens()) mask = torch.zeros((batch_size, crop_size), dtype=torch.bool) mask_token_id = self._ttt_mask_token() for row_index in range(batch_size): non_special_positions = [ col for col in range(crop_size) if batch_cropped[row_index, col].item() in non_special_tokens ] assert len(non_special_positions) > 0, "Sequence must contain at least one non-special token." num_to_mask = max(1, int(round(len(non_special_positions) * self._ttt_cfg.mask_ratio))) sampled_indices = torch.randperm( len(non_special_positions), generator=self._ttt_generator, )[:num_to_mask] positions_to_mask = torch.tensor(non_special_positions, dtype=torch.long)[sampled_indices] mask[row_index, positions_to_mask] = True batch_masked = batch_cropped.clone() for row_index in range(batch_size): masked_positions = torch.nonzero(mask[row_index], as_tuple=True)[0] for masked_position in masked_positions: probability = float(torch.rand(1, generator=self._ttt_generator).item()) if probability < 1.0 - self._ttt_cfg.bert_leave_prob - self._ttt_cfg.bert_replace_prob: batch_masked[row_index, masked_position] = mask_token_id continue if probability < 1.0 - self._ttt_cfg.bert_leave_prob: replacement_candidates = self._ttt_get_non_special_tokens() replacement_index = int(torch.randint( 0, len(replacement_candidates), (1,), generator=self._ttt_generator, ).item()) batch_masked[row_index, masked_position] = replacement_candidates[replacement_index] return batch_masked, batch_cropped, mask, start_indices def _ttt_cross_entropy_loss( self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor, ) -> torch.Tensor: assert logits.ndim == 3, "Logits must be [batch, seq, vocab]." _, _, vocab_size = logits.shape logits_flat = logits.reshape(-1, vocab_size) targets_flat = targets.reshape(-1) mask_flat = mask.reshape(-1) assert int(mask_flat.sum().item()) > 0, "TTT mask must select at least one token." loss = F.cross_entropy( logits_flat[mask_flat], targets_flat[mask_flat], reduction="none", ) masked_tokens_per_seq = mask.sum(dim=1).tolist() per_sequence_losses = torch.split(loss, masked_tokens_per_seq) return torch.stack([sl.mean() for sl in per_sequence_losses]).mean() def _ttt_get_optimizer(self, parameters) -> torch.optim.Optimizer: if self._ttt_cfg.optimizer == "sgd": return torch.optim.SGD( parameters, lr=self._ttt_cfg.lr, momentum=self._ttt_cfg.momentum, weight_decay=self._ttt_cfg.weight_decay, ) return torch.optim.AdamW( parameters, lr=self._ttt_cfg.lr, weight_decay=self._ttt_cfg.weight_decay, ) def _lora_ttt(self, seq: str) -> Dict[str, List[float]]: """LoRA TTT: only LoRA adapter weights are trained, mlm_head is frozen.""" x = self._ttt_tokenize(seq) device = next(self.parameters()).device non_blocking = device.type == "cuda" losses = [] if self._ttt_cfg.steps == 0: return {"losses": losses} for parameter in self.parameters(): parameter.requires_grad = False for p in self._lora_params: p.requires_grad = True optimizer = self._ttt_get_optimizer(self._lora_params) optimizer.zero_grad(set_to_none=True) self.eval() for step in range(self._ttt_cfg.steps * self._ttt_cfg.ags): batch_masked, targets, mask, start_indices = self._ttt_sample_batch(x) batch_masked = batch_masked.to(device, non_blocking=non_blocking) targets = targets.to(device, non_blocking=non_blocking) mask = mask.to(device, non_blocking=non_blocking) self.train() logits = self._ttt_predict_logits(batch_masked) loss = self._ttt_cross_entropy_loss(logits, targets, mask) loss.backward() losses.append(float(loss.detach().cpu().item())) if (step + 1) % self._ttt_cfg.ags == 0: optimizer.step() optimizer.zero_grad(set_to_none=True) self.eval() return {"losses": losses} def _legacy_ttt(self, seq: str) -> Dict[str, List[float]]: """Legacy TTT: full fine-tuning of ESM2 backbone with random linear projection head.""" x = self._ttt_tokenize(seq) device = next(self.parameters()).device non_blocking = device.type == "cuda" losses = [] if self._ttt_cfg.steps == 0: return {"losses": losses} # Full fine-tune: all backbone params trainable for parameter in self.parameters(): parameter.requires_grad = False for parameter in self.esm.parameters(): parameter.requires_grad = True if self._ttt_cfg.freeze_embeddings: for parameter in self.esm.embeddings.parameters(): parameter.requires_grad = False for parameter in self._ttt_lm_proj.parameters(): parameter.requires_grad = True trainable_params = filter(lambda p: p.requires_grad, self.parameters()) optimizer = self._ttt_get_optimizer(trainable_params) optimizer.zero_grad(set_to_none=True) self.eval() for step in range(self._ttt_cfg.steps * self._ttt_cfg.ags): batch_masked, targets, mask, start_indices = self._ttt_sample_batch(x) batch_masked = batch_masked.to(device, non_blocking=non_blocking) targets = targets.to(device, non_blocking=non_blocking) mask = mask.to(device, non_blocking=non_blocking) self.train() logits = self._ttt_predict_logits(batch_masked) loss = self._ttt_cross_entropy_loss(logits, targets, mask) loss.backward() losses.append(float(loss.detach().cpu().item())) if (step + 1) % self._ttt_cfg.ags == 0: optimizer.step() optimizer.zero_grad(set_to_none=True) self.eval() return {"losses": losses} @preserve_model_state def ttt(self, seq: str) -> Dict[str, List[float]]: """Run test-time training on a single sequence using masked language modeling. Adapts the ESM2 backbone (via LoRA or full fine-tuning) to the input sequence before structure prediction. Call fold_protein(seq, ttt=True) for the full pipeline. Args: seq: Protein sequence (single-letter amino acid codes) Returns: Dict with "losses" key containing per-step MLM loss values """ self._ensure_ttt_ready() # TTT requires fp32 for stable gradient computation. ESMFold typically # runs the backbone in fp16, but small LoRA updates vanish in half precision. esm_dtype = next(self.esm.parameters()).dtype if esm_dtype != torch.float32: self.esm.float() self.mlm_head.float() if self._uses_lora: result = self._lora_ttt(seq) else: result = self._legacy_ttt(seq) # Restore original dtype (backbone back to fp16 for inference) if esm_dtype != torch.float32: self.esm.to(esm_dtype) self.mlm_head.to(esm_dtype) return result # ---- High-Level API ---- def _fold_single(self, sequence: str, return_pdb_string: bool = True) -> Dict[str, Any]: """Fold a sequence once and return pLDDT, ptm, and optionally PDB string.""" with torch.no_grad(): output = self.infer(sequence) plddt = output["plddt"] # plddt shape is (batch, L, 37) - per-atom across atom37 types. # Use CA atom (index 1) only, matching PDB B-factor output. if plddt.dim() == 3: mean_plddt = float(plddt[:, :, 1].mean().item()) elif plddt.dim() == 2: mean_plddt = float(plddt[:, 1].mean().item()) else: mean_plddt = float(plddt.mean().item()) result = { "plddt": mean_plddt, "ptm": float(output["ptm"].item()) if "ptm" in output else None, } if return_pdb_string: pdb_strings = self.output_to_pdb(output) result["pdb_string"] = pdb_strings[0] if isinstance(pdb_strings, list) else pdb_strings return result def fold_protein( self, sequence: str, return_pdb_string: bool = True, ttt: bool = False, ) -> Dict[str, Any]: """Fold a protein sequence. Test-time training is disabled by default. Pass ``ttt=True`` or call ``fold_protein_ttt`` to opt in to experimental TTT. Args: sequence: Protein sequence (single-letter amino acid codes) return_pdb_string: If True, include PDB string in output ttt: If True, run experimental LoRA TTT before returning the best fold Returns: Dict with keys: - plddt: float, mean pLDDT - ptm: float, predicted TM-score - pdb_string: str (if return_pdb_string=True), PDB from best step - step_plddts: list[float], baseline pLDDT when TTT is disabled - best_step: int, 0 when TTT is disabled """ if ttt: return self.fold_protein_ttt( sequence=sequence, return_pdb_string=return_pdb_string, ) result = self._fold_single(sequence, return_pdb_string=return_pdb_string) return { "plddt": result["plddt"], "ptm": result["ptm"], "pdb_string": result.get("pdb_string"), "step_plddts": [result["plddt"]], "best_step": 0, } def fold_protein_ttt( self, sequence: str, return_pdb_string: bool = True, ) -> Dict[str, Any]: """Fold a protein sequence with experimental test-time training. Runs TTT (masked language model adaptation via LoRA) for the configured number of steps, folding after each optimizer step to track pLDDT. Returns the structure with the highest pLDDT across all steps (including baseline). Args: sequence: Protein sequence (single-letter amino acid codes) return_pdb_string: If True, include PDB string in output Returns: Dict with keys: - plddt: float, best mean pLDDT across all TTT steps - ptm: float, predicted TM-score from best step - pdb_string: str (if return_pdb_string=True), PDB from best step - step_plddts: list[float], pLDDT at each step [baseline, s1, ..., s10] - best_step: int, which step produced best structure (0=baseline) """ self._ensure_ttt_ready() # Cast to fp32 for TTT stability esm_dtype = next(self.esm.parameters()).dtype if esm_dtype != torch.float32: self.esm.float() self.mlm_head.float() device = next(self.parameters()).device non_blocking = device.type == "cuda" # Step 0: baseline fold (no TTT adaptation) best = self._fold_single(sequence, return_pdb_string=return_pdb_string) step_plddts = [best["plddt"]] if self._ttt_cfg.steps > 0: # Tokenize for masked LM training x = self._ttt_tokenize(sequence) # Freeze all, unfreeze LoRA for p in self.parameters(): p.requires_grad = False if self._uses_lora: for p in self._lora_params: p.requires_grad = True optimizer = self._ttt_get_optimizer(self._lora_params) else: for p in self.esm.parameters(): p.requires_grad = True if self._ttt_cfg.freeze_embeddings: for p in self.esm.embeddings.parameters(): p.requires_grad = False for p in self._ttt_lm_proj.parameters(): p.requires_grad = True trainable = [p for p in self.parameters() if p.requires_grad] optimizer = self._ttt_get_optimizer(trainable) optimizer.zero_grad(set_to_none=True) self.eval() for step in range(self._ttt_cfg.steps * self._ttt_cfg.ags): batch_masked, targets, mask, _start = self._ttt_sample_batch(x) batch_masked = batch_masked.to(device, non_blocking=non_blocking) targets = targets.to(device, non_blocking=non_blocking) mask = mask.to(device, non_blocking=non_blocking) self.train() logits = self._ttt_predict_logits(batch_masked) loss = self._ttt_cross_entropy_loss(logits, targets, mask) loss.backward() if (step + 1) % self._ttt_cfg.ags == 0: optimizer.step() optimizer.zero_grad(set_to_none=True) # Fold after this optimizer step self.eval() current = self._fold_single(sequence, return_pdb_string=return_pdb_string) step_plddts.append(current["plddt"]) if current["plddt"] > best["plddt"]: best = current self.eval() # Restore requires_grad for p in self.parameters(): p.requires_grad = False # Reset LoRA weights for next sequence self.ttt_reset() # Restore dtype if esm_dtype != torch.float32: self.esm.to(esm_dtype) self.mlm_head.to(esm_dtype) return { "plddt": best["plddt"], "ptm": best["ptm"], "pdb_string": best.get("pdb_string"), "step_plddts": step_plddts, "best_step": step_plddts.index(max(step_plddts)), }