From ba98fc1608a42d003cd5b808c6cb346a6185aa33 Mon Sep 17 00:00:00 2001 From: recurse-agent Date: Fri, 14 Aug 2026 20:22:59 +0000 Subject: [PATCH 1/3] Sync ESM models for OSS FoldBench evaluation Co-Authored-By: OpenAI Codex --- src/transformers/models/esmc/__init__.py | 6 +- .../models/esmc/configuration_esmc.py | 2 +- .../models/esmc/configuration_esmc_sae.py | 2 +- src/transformers/models/esmc/modeling_esmc.py | 77 ++++----- .../models/esmc/modeling_esmc_sae.py | 10 +- .../models/esmc/tokenization_esmc.py | 6 +- src/transformers/models/esmfold2/__init__.py | 6 +- .../models/esmfold2/configuration_esmfold2.py | 8 +- .../models/esmfold2/distributed/manager.py | 36 ++--- .../distributed/model/layers/layernorm.py | 8 +- .../distributed/model/layers/linear.py | 18 +-- .../distributed/model/layers/pairformer.py | 2 +- .../models/esmfold2/distributed/utils.py | 25 +-- .../kernels/fused_attention_pair_bias.py | 13 +- .../kernels/fused_dropout_residual.py | 26 +-- .../esmfold2/kernels/fused_dual_gemm.py | 26 +-- .../esmfold2/kernels/fused_ln_residual.py | 32 ++-- .../esmfold2/kernels/fused_lnlin_swiglu.py | 18 +-- .../esmfold2/kernels/trimul_einsum_triton.py | 6 +- .../esmfold2/kernels/trimul_with_residual.py | 24 +-- .../models/esmfold2/modeling_esmfold2.py | 80 +++++----- .../esmfold2/modeling_esmfold2_common.py | 149 +++++++++--------- .../modeling_esmfold2_experimental.py | 24 +-- 23 files changed, 303 insertions(+), 301 deletions(-) diff --git a/src/transformers/models/esmc/__init__.py b/src/transformers/models/esmc/__init__.py index 07d913a152..28dc2fbb53 100644 --- a/src/transformers/models/esmc/__init__.py +++ b/src/transformers/models/esmc/__init__.py @@ -13,8 +13,10 @@ # limitations under the License. from typing import TYPE_CHECKING -from ...utils import _LazyModule # type: ignore[import] -from ...utils.import_utils import define_import_structure # type: ignore[import] +from ...utils import _LazyModule # ty:ignore[unresolved-import] +from ...utils.import_utils import ( # ty:ignore[unresolved-import] + define_import_structure, +) if TYPE_CHECKING: from .configuration_esmc import * # noqa: F403 diff --git a/src/transformers/models/esmc/configuration_esmc.py b/src/transformers/models/esmc/configuration_esmc.py index 962a23e461..fa7c473540 100644 --- a/src/transformers/models/esmc/configuration_esmc.py +++ b/src/transformers/models/esmc/configuration_esmc.py @@ -13,7 +13,7 @@ # limitations under the License. """ESMC model configuration.""" -from ...configuration_utils import PretrainedConfig # type: ignore[import] +from ...configuration_utils import PretrainedConfig # ty:ignore[unresolved-import] class ESMCConfig(PretrainedConfig): diff --git a/src/transformers/models/esmc/configuration_esmc_sae.py b/src/transformers/models/esmc/configuration_esmc_sae.py index 42ecc1f32c..4fcc14b6a6 100644 --- a/src/transformers/models/esmc/configuration_esmc_sae.py +++ b/src/transformers/models/esmc/configuration_esmc_sae.py @@ -15,7 +15,7 @@ from dataclasses import dataclass -from ...configuration_utils import PretrainedConfig # type: ignore[import] +from ...configuration_utils import PretrainedConfig # ty:ignore[unresolved-import] @dataclass diff --git a/src/transformers/models/esmc/modeling_esmc.py b/src/transformers/models/esmc/modeling_esmc.py index f281871690..8f249f6493 100644 --- a/src/transformers/models/esmc/modeling_esmc.py +++ b/src/transformers/models/esmc/modeling_esmc.py @@ -23,14 +23,14 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from torch.nn import functional as F -from ...modeling_outputs import ( # type: ignore[import] +from ...modeling_outputs import ( # ty:ignore[unresolved-import] MaskedLMOutput, ModelOutput, SequenceClassifierOutput, TokenClassifierOutput, ) -from ...modeling_utils import PreTrainedModel # type: ignore[import] -from ...utils import ( # type: ignore[import] +from ...modeling_utils import PreTrainedModel # ty:ignore[unresolved-import] +from ...utils import ( # ty:ignore[unresolved-import] auto_docstring, can_return_tuple, is_flash_attn_2_available, @@ -58,7 +58,7 @@ _flash_attn_rotary_available = torch.cuda.is_available() except ImportError: - apply_triton_rotary = None # type: ignore[assignment] + apply_triton_rotary = None # ty:ignore[invalid-assignment] _flash_attn_rotary_available = False # Transformer Engine: fused LayerNorm+Linear / LayerNorm+MLP kernels with @@ -67,11 +67,11 @@ # ~O(100) in bf16 on the unnormalized residual stream (perplexity stays # within rounding noise). try: - import transformer_engine.pytorch as te # type: ignore[import-untyped] + import transformer_engine.pytorch as te _te_available = True except ImportError: - te = None # type: ignore[assignment] + te = None # ty:ignore[invalid-assignment] _te_available = False # xformers: preferred SDPA implementation on GPU. Provides a fused @@ -79,11 +79,11 @@ # Attention 2 and PyTorch's ``F.scaled_dot_product_attention`` are # progressively-less-preferred fallbacks. try: - import xformers.ops as xops # type: ignore[import-untyped] + import xformers.ops as xops _xformers_available = True except ImportError: - xops = None # type: ignore[assignment] + xops = None # ty:ignore[invalid-assignment] _xformers_available = False # Flash Attention 2: secondary SDPA fallback. Used when xformers is not @@ -91,7 +91,7 @@ if _flash_attn_available: from flash_attn import flash_attn_func else: - flash_attn_func = None # type: ignore[assignment] + flash_attn_func = None if not _te_available: logger.warning( @@ -359,21 +359,21 @@ def _update_cos_sin_cache(self, seqlen: int, device=None, dtype=None): ) else: t = ( - torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) # type: ignore[call-overload] + torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) # ty:ignore[no-matching-overload] / self.scaling_factor ) inv_freq = self.inv_freq - freqs = torch.outer(t, inv_freq) # type: ignore[arg-type] + freqs = torch.outer(t, inv_freq) # ty:ignore[invalid-argument-type] if self.scale is None: self._cos_cached = torch.cos(freqs).to(dtype) self._sin_cached = torch.sin(freqs).to(dtype) else: - _scale: torch.Tensor = self.scale # type: ignore[assignment] + _scale: torch.Tensor = self.scale # ty:ignore[invalid-assignment] power = ( torch.arange(seqlen, dtype=_scale.dtype, device=_scale.device) - seqlen // 2 - ) / self.scale_base # type: ignore[operator] + ) / self.scale_base # ty:ignore[unsupported-operator] scale = _scale.to(device=power.device) ** power.unsqueeze(-1) self._cos_cached = (torch.cos(freqs) * scale).to(dtype) self._sin_cached = (torch.sin(freqs) * scale).to(dtype) @@ -422,8 +422,8 @@ def forward( sin = self._sin_cached[seqlen_offset:] if _flash_attn_rotary_available and q.device.type == "cuda": - q_rot = apply_triton_rotary(q, cos, sin, interleaved=self.interleaved) # type: ignore[misc] - k_rot = apply_triton_rotary(k, cos, sin, interleaved=self.interleaved) # type: ignore[misc] + q_rot = apply_triton_rotary(q, cos, sin, interleaved=self.interleaved) + k_rot = apply_triton_rotary(k, cos, sin, interleaved=self.interleaved) else: q_rot = _apply_rotary_emb_torch(q, cos, sin, self.interleaved) k_rot = _apply_rotary_emb_torch(k, cos, sin, self.interleaved) @@ -440,7 +440,7 @@ class _TritonRotaryEmbedding(RotaryEmbedding): def forward( self, qkv: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int - ) -> torch.Tensor: # type: ignore[override] + ) -> torch.Tensor: """Apply RoPE in-place to a packed ``(N, 3, n_heads, head_dim)`` tensor.""" self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype) assert self._cos_cached is not None and self._sin_cached is not None @@ -547,7 +547,7 @@ def _swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Modul assert not bias, "ESMC was trained with bias=False; bias=True not supported" hidden = _swiglu_hidden_dim(expansion_ratio, d_model) if _te_available: - return te.LayerNormMLP( # type: ignore[union-attr] + return te.LayerNormMLP( hidden_size=d_model, ffn_hidden_size=hidden, bias=bias, @@ -563,9 +563,7 @@ def _make_attn_layernorm_qkv(d_model: int, bias: bool) -> nn.Module: available; pure-PyTorch fallback otherwise.""" assert not bias, "ESMC was trained with bias=False; bias=True not supported" if _te_available: - return te.LayerNormLinear( # type: ignore[union-attr] - d_model, d_model * 3, bias=bias, init_method=None - ) + return te.LayerNormLinear(d_model, d_model * 3, bias=bias, init_method=None) return _PyTorchLayerNormLinear(d_model, d_model * 3) @@ -573,9 +571,7 @@ def _make_attn_out_proj(d_model: int, bias: bool) -> nn.Module: """Attention output projection. Uses Transformer Engine when available; pure-PyTorch ``nn.Linear`` otherwise.""" if _te_available: - return te.Linear( # type: ignore[union-attr] - d_model, d_model, bias=bias, init_method=None - ) + return te.Linear(d_model, d_model, bias=bias, init_method=None) return nn.Linear(d_model, d_model, bias=bias) @@ -630,7 +626,7 @@ def _scaled_dot_product_attention( q4 = q.view(b, s, n_heads, d_head) k4 = k.view(b, s, n_heads, d_head) v4 = v.view(b, s, n_heads, d_head) - context = xops.memory_efficient_attention( # type: ignore[union-attr] + context = xops.memory_efficient_attention( q4, k4, v4, attn_bias=None, scale=d_head**-0.5 ) return context.reshape(b, s, n_heads * d_head) @@ -643,10 +639,8 @@ def _scaled_dot_product_attention( q4 = q.view(b, s, n_heads, d_head) k4 = k.view(b, s, n_heads, d_head) v4 = v.view(b, s, n_heads, d_head) - context = flash_attn_func( # type: ignore[misc] - q4, k4, v4, dropout_p=0.0, softmax_scale=d_head**-0.5 - ) - return context.reshape(b, s, n_heads * d_head) # type: ignore[union-attr] + context = flash_attn_func(q4, k4, v4, dropout_p=0.0, softmax_scale=d_head**-0.5) # ty:ignore[call-non-callable] + return context.reshape(b, s, n_heads * d_head) b, s, _ = q.shape q = q.view(b, s, n_heads, -1).transpose(1, 2) k = k.view(b, s, n_heads, -1).transpose(1, 2) @@ -785,14 +779,11 @@ def forward( qkv_packed = torch.stack([q, k, v], dim=1).view(T, 3, self.n_heads, self.d_head) qkv_packed = self.rotary(qkv_packed, cu_seqlens, max_seqlen) - context = flash_attn_varlen_qkvpacked_func( # type: ignore[misc] + context = flash_attn_varlen_qkvpacked_func( qkv_packed, cu_seqlens, max_seqlen, softmax_scale=self.d_head**-0.5 - ) - n_out, h_out, d_out = context.shape # type: ignore[union-attr] - return ( - self.out_proj(context.reshape(n_out, h_out * d_out)), # type: ignore[union-attr] - None, - ) + ) # ty:ignore[call-non-callable] + n_out, h_out, d_out = context.shape + return (self.out_proj(context.reshape(n_out, h_out * d_out)), None) # --------------------------------------------------------------------------- @@ -1079,9 +1070,9 @@ def _get_sae_layer_num_requested(self, model_name: str) -> int: """Recover the backbone-layer index from a key written by :meth:`add_sae_models` (``"layer{N}"`` → ``N``).""" match = self._SAE_KEY_RE.fullmatch(model_name) - assert ( - match is not None - ), f"Unexpected SAE key {model_name!r}; expected 'layer{{N}}'." + assert match is not None, ( + f"Unexpected SAE key {model_name!r}; expected 'layer{{N}}'." + ) return int(match.group(1)) def _validate_sae_inputs(self, input_ids: torch.Tensor) -> None: @@ -1215,7 +1206,7 @@ def forward( bool_mask = sequence_id >= 0 else: if attention_mask is None: - attention_mask = input_ids != self.config.pad_token_id + attention_mask = input_ids != self.config.pad_token_id # ty:ignore[invalid-assignment] assert attention_mask is not None bool_mask = attention_mask.bool() sequence_id = bool_mask.to(torch.long) - 1 @@ -1260,7 +1251,7 @@ def forward( # Stack once; reused for both SAE and hidden-state output. collected_tensor: torch.Tensor | None = ( - torch.stack(collected, dim=0) if collected else None # type: ignore[arg-type] + torch.stack(collected, dim=0) if collected else None ) sae_outputs: dict[str, torch.Tensor] | None = None @@ -1283,11 +1274,11 @@ def forward( attentions, ] if v is not None - ) + ) # ty:ignore[invalid-return-type] return ESMCOutput( last_hidden_state=last_hidden_state, - hidden_states=hidden_states_tensor, + hidden_states=hidden_states_tensor, # ty:ignore[invalid-argument-type] sae_outputs=sae_outputs, attentions=attentions, ) @@ -1332,7 +1323,7 @@ def __init__(self, config: ESMCConfig): self.post_init() def get_output_embeddings(self) -> nn.Linear: - return self.lm_head[-1] # type: ignore[return-value] + return self.lm_head[-1] def set_output_embeddings(self, new_embeddings: nn.Linear): self.lm_head[-1] = new_embeddings diff --git a/src/transformers/models/esmc/modeling_esmc_sae.py b/src/transformers/models/esmc/modeling_esmc_sae.py index b706e50e91..6e549b070c 100644 --- a/src/transformers/models/esmc/modeling_esmc_sae.py +++ b/src/transformers/models/esmc/modeling_esmc_sae.py @@ -36,9 +36,9 @@ import torch.nn.functional as F from safetensors.torch import load_file, save_file -from ...modeling_outputs import ModelOutput # type: ignore[import] -from ...modeling_utils import PreTrainedModel # type: ignore[import] -from ...utils import auto_docstring # type: ignore[import] +from ...modeling_outputs import ModelOutput # ty:ignore[unresolved-import] +from ...modeling_utils import PreTrainedModel # ty:ignore[unresolved-import] +from ...utils import auto_docstring # ty:ignore[unresolved-import] from .configuration_esmc_sae import ESMCSAEConfig, ESMCSAEParams @@ -160,7 +160,7 @@ def __init__(self, config: ESMCSAEConfig): self.post_init() @classmethod - def from_pretrained( # type: ignore[override] + def from_pretrained( cls, pretrained_model_name_or_path: str | os.PathLike, *model_args, **kwargs ) -> "ESMCSAEModel": """Download (or reuse cached) the full repo and return the model. @@ -304,7 +304,7 @@ def forward( ) return self.layers[key](x, **kwargs) - def save_pretrained( # type: ignore[override] + def save_pretrained( self, save_directory: str | os.PathLike, *args, **kwargs ) -> None: """Write ``config.json`` plus one ``layer_{i}.safetensors`` per loaded layer. diff --git a/src/transformers/models/esmc/tokenization_esmc.py b/src/transformers/models/esmc/tokenization_esmc.py index 2485779194..d74e105a53 100644 --- a/src/transformers/models/esmc/tokenization_esmc.py +++ b/src/transformers/models/esmc/tokenization_esmc.py @@ -17,8 +17,10 @@ from tokenizers.models import BPE from tokenizers.processors import TemplateProcessing -from ...tokenization_utils_fast import PreTrainedTokenizerFast # type: ignore[import] -from ...utils import logging # type: ignore[import] +from ...tokenization_utils_fast import ( # ty:ignore[unresolved-import] + PreTrainedTokenizerFast, +) +from ...utils import logging # ty:ignore[unresolved-import] logger = logging.get_logger(__name__) diff --git a/src/transformers/models/esmfold2/__init__.py b/src/transformers/models/esmfold2/__init__.py index 6b245e85ea..32eba6722d 100644 --- a/src/transformers/models/esmfold2/__init__.py +++ b/src/transformers/models/esmfold2/__init__.py @@ -13,8 +13,10 @@ # limitations under the License. from typing import TYPE_CHECKING -from ...utils import _LazyModule # type: ignore[import] -from ...utils.import_utils import define_import_structure # type: ignore[import] +from ...utils import _LazyModule # ty:ignore[unresolved-import] +from ...utils.import_utils import ( # ty:ignore[unresolved-import] + define_import_structure, +) if TYPE_CHECKING: from .configuration_esmfold2 import * # noqa: F403 diff --git a/src/transformers/models/esmfold2/configuration_esmfold2.py b/src/transformers/models/esmfold2/configuration_esmfold2.py index 0d8fb6c082..3ea4c86aaa 100644 --- a/src/transformers/models/esmfold2/configuration_esmfold2.py +++ b/src/transformers/models/esmfold2/configuration_esmfold2.py @@ -17,7 +17,7 @@ from dataclasses import asdict, dataclass, field -from ...configuration_utils import PretrainedConfig # type: ignore[import] +from ...configuration_utils import PretrainedConfig # ty:ignore[unresolved-import] # --------------------------------------------------------------------------- # Nested dataclass configs @@ -90,7 +90,7 @@ class InputsEmbedderConfig: def __post_init__(self): if isinstance(self.atom_encoder, dict): - self.atom_encoder = AtomAttentionConfig(**self.atom_encoder) + self.atom_encoder = AtomAttentionConfig(**self.atom_encoder) # ty:ignore[invalid-argument-type] @dataclass @@ -139,7 +139,7 @@ class DiffusionStructureHeadConfig: def __post_init__(self): if isinstance(self.diffusion_module, dict): - self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module) + self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module) # ty:ignore[invalid-argument-type] @dataclass @@ -157,7 +157,7 @@ class ConfidenceHeadConfig: def __post_init__(self): if isinstance(self.folding_trunk, dict): - self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk) + self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk) # ty:ignore[invalid-argument-type] # --------------------------------------------------------------------------- diff --git a/src/transformers/models/esmfold2/distributed/manager.py b/src/transformers/models/esmfold2/distributed/manager.py index f2b2c4b7e1..9fcdedd736 100644 --- a/src/transformers/models/esmfold2/distributed/manager.py +++ b/src/transformers/models/esmfold2/distributed/manager.py @@ -187,10 +187,10 @@ def _setup( DistributedManager._state["_initialized"] = True manager = DistributedManager() - manager._has_dist = torch.distributed.is_available() # type: ignore[assignment] - manager._rank = rank # type: ignore[assignment] - manager._world_size = world_size # type: ignore[assignment] - manager._node_rank = node_rank # type: ignore[assignment] + manager._has_dist = torch.distributed.is_available() # ty:ignore[unresolved-attribute] + manager._rank = rank # ty:ignore[unresolved-attribute] + manager._world_size = world_size # ty:ignore[unresolved-attribute] + manager._node_rank = node_rank # ty:ignore[unresolved-attribute] if device_type == "cuda": if ( @@ -199,14 +199,14 @@ def _setup( ): warn("world_size is not a multiple of torch.cuda.device_count()") if local_rank is None: - manager._local_rank = manager.rank % torch.cuda.device_count() # type: ignore[assignment] + manager._local_rank = manager.rank % torch.cuda.device_count() # ty:ignore[unresolved-attribute] else: - manager._local_rank = local_rank # type: ignore[assignment] - manager._device = torch.device(f"cuda:{manager.local_rank}") # type: ignore[assignment] + manager._local_rank = local_rank # ty:ignore[unresolved-attribute] + manager._device = torch.device(f"cuda:{manager.local_rank}") # ty:ignore[unresolved-attribute] else: if local_rank is not None: - manager._local_rank = local_rank # type: ignore[assignment] - manager._device = torch.device("cpu") # type: ignore[assignment] + manager._local_rank = local_rank # ty:ignore[unresolved-attribute] + manager._device = torch.device("cpu") # ty:ignore[unresolved-attribute] if not manager.has_dist: warn("DistributedManager initialized without torch.distributed package") @@ -217,7 +217,7 @@ def _setup( torch.cuda.device(manager.device) torch.cuda.empty_cache() - manager._backend = backend # type: ignore[assignment] + manager._backend = backend # ty:ignore[unresolved-attribute] if manager.device.type == "cuda" and backend == "nccl": try: @@ -248,7 +248,7 @@ def _setup( manager._group_ranks["world"] = torch.distributed.get_process_group_ranks( manager.group["world"] ) - manager._method_init = method_init # type: ignore[assignment] + manager._method_init = method_init # ty:ignore[unresolved-attribute] if grid_group_sizes is not None: DistributedManager.create_grid_group(grid_group_sizes) @@ -430,8 +430,8 @@ def _initialize_env(*args, **kwargs): group_rank = os.environ.get("GROUP_RANK", 0) node_rank = int(os.environ.get("NODE_RANK", group_rank)) try: - rank = int(rank) # type: ignore[arg-type] - world_size = int(world_size) # type: ignore[arg-type] + rank = int(rank) # ty:ignore[invalid-argument-type] + world_size = int(world_size) # ty:ignore[invalid-argument-type] if local_rank is not None: local_rank = int(local_rank) except TypeError: @@ -444,8 +444,8 @@ def _initialize_env(*args, **kwargs): node_rank=node_rank, world_size=world_size, local_rank=local_rank, - addr=os.environ.get("MASTER_ADDR"), # type: ignore[arg-type] - port=os.environ.get("MASTER_PORT"), # type: ignore[arg-type] + addr=os.environ.get("MASTER_ADDR"), # ty:ignore[invalid-argument-type] + port=os.environ.get("MASTER_PORT"), # ty:ignore[invalid-argument-type] method_init="ENV", **kwargs, ) @@ -468,8 +468,8 @@ def _initialize_slurm(*args, **kwargs): local_rank = os.environ.get("SLURM_LOCALID") addr = os.environ.get("SLURM_LAUNCH_NODE_IPADDR") try: - rank = int(rank) # type: ignore[arg-type] - world_size = int(world_size) # type: ignore[arg-type] + rank = int(rank) # ty:ignore[invalid-argument-type] + world_size = int(world_size) # ty:ignore[invalid-argument-type] if local_rank is not None: local_rank = int(local_rank) except TypeError: @@ -482,7 +482,7 @@ def _initialize_slurm(*args, **kwargs): node_rank=node_rank, world_size=world_size, local_rank=local_rank, - addr=addr, # type: ignore[arg-type] + addr=addr, # ty:ignore[invalid-argument-type] method_init="SLURM", **kwargs, ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py b/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py index 88fec07fa0..9db298391a 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py @@ -128,7 +128,7 @@ def backward(ctx, d_out: DTensor): replicate = [Replicate()] * ctx.device_mesh.ndim if ctx.needs_input_grad[2]: - dw_work.wait() # type: ignore[union-attr] + dw_work.wait() dw_dtensor = DTensor.from_local( dw, device_mesh=ctx.device_mesh, @@ -137,7 +137,7 @@ def backward(ctx, d_out: DTensor): stride=dw.stride(), ) if ctx.needs_input_grad[3]: - db_work.wait() # type: ignore[union-attr] + db_work.wait() db_dtensor = DTensor.from_local( db, device_mesh=ctx.device_mesh, @@ -189,13 +189,13 @@ def __init__(self, layer_local: nn.LayerNorm, device_mesh: DeviceMesh) -> None: else: self.bias = None - if "cp" in device_mesh.mesh_dim_names: # type: ignore[operator] + if "cp" in device_mesh.mesh_dim_names: # ty:ignore[unsupported-operator] self._reduce_group = device_mesh.get_group("cp") else: self._reduce_group = dist.group.WORLD def forward(self, x: DTensor) -> DTensor: - return _LayerNormParamsReplicatedImpl.apply( # type: ignore[return-value] + return _LayerNormParamsReplicatedImpl.apply( x, self.normalized_shape, self.weight, diff --git a/src/transformers/models/esmfold2/distributed/model/layers/linear.py b/src/transformers/models/esmfold2/distributed/model/layers/linear.py index ba195f0509..c58a0efd68 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/linear.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/linear.py @@ -96,7 +96,7 @@ def backward(ctx, d_out: DTensor): if ctx.needs_input_grad[1]: # Aggregate over all but the last two dims (batch + seq dims) dw = torch.einsum("...i,...j->ij", d_out_local, x_saved) - dw = dw.contiguous() # type: ignore[union-attr] + dw = dw.contiguous() op = dist.ReduceOp.AVG if ctx.avg_reduce else dist.ReduceOp.SUM dw_work = dist.all_reduce(dw, op=op, group=ctx.reduce_group, async_op=True) @@ -127,22 +127,22 @@ def backward(ctx, d_out: DTensor): if dw_work is not None: dw_work.wait() dw_dtensor = DTensor.from_local( - dw, # type: ignore[arg-type] + dw, # ty:ignore[invalid-argument-type] device_mesh=ctx.device_mesh, placements=replicate, - shape=dw.shape, # type: ignore[union-attr] - stride=dw.stride(), # type: ignore[union-attr] + shape=dw.shape, # ty:ignore[unresolved-attribute] + stride=dw.stride(), # ty:ignore[unresolved-attribute] ) db_dtensor: Optional[DTensor] = None if db_work is not None: db_work.wait() db_dtensor = DTensor.from_local( - db, # type: ignore[arg-type] + db, # ty:ignore[invalid-argument-type] device_mesh=ctx.device_mesh, placements=replicate, - shape=db.shape, # type: ignore[union-attr] - stride=db.stride(), # type: ignore[union-attr] + shape=db.shape, # ty:ignore[unresolved-attribute] + stride=db.stride(), # ty:ignore[unresolved-attribute] ) return dx_dtensor, dw_dtensor, db_dtensor, None, None @@ -188,12 +188,12 @@ def __init__( self.bias = None # Choose reduce group: use cp group if present, otherwise world - if "cp" in device_mesh.mesh_dim_names: # type: ignore[operator] + if "cp" in device_mesh.mesh_dim_names: # ty:ignore[unsupported-operator] self._reduce_group = device_mesh.get_group("cp") else: self._reduce_group = dist.group.WORLD def forward(self, x: DTensor) -> DTensor: - return _LinearParamsReplicatedImpl.apply( # type: ignore[return-value] + return _LinearParamsReplicatedImpl.apply( x, self.weight, self.bias, self._reduce_group, self.avg_reduce ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py b/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py index 038f1471b1..a9e3322691 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py @@ -170,7 +170,7 @@ def __init__( raise TypeError(f"trunk must be FoldingTrunk, got {type(trunk).__name__}") self.blocks = nn.ModuleList( - [PairUpdateBlockDistributed(block, dist_manager) for block in trunk.blocks] # type: ignore[arg-type] + [PairUpdateBlockDistributed(block, dist_manager) for block in trunk.blocks] # ty:ignore[invalid-argument-type] ) def forward( diff --git a/src/transformers/models/esmfold2/distributed/utils.py b/src/transformers/models/esmfold2/distributed/utils.py index a4ba6e7e60..bcd4d024cf 100644 --- a/src/transformers/models/esmfold2/distributed/utils.py +++ b/src/transformers/models/esmfold2/distributed/utils.py @@ -50,7 +50,7 @@ class LayoutMap: def __init__( self, strides: tuple[int, ...], shape: tuple[int, ...], offset: int = 0 ): - if not all(isinstance(s, (int, np.int64)) and s > 0 for s in strides): # type: ignore[arg-type] + if not all(isinstance(s, (int, np.int64)) and s > 0 for s in strides): raise ValueError(f"Strides must be positive integers: {strides}") if any(s < 0 for s in shape): raise ValueError(f"Shape must be non-negative: {shape}") @@ -181,8 +181,8 @@ def __getitem__(self, slices) -> "LayoutMap": new_offset = self.offset for axis, s in enumerate(slices): - if isinstance(s, (int, np.int64)): # type: ignore[arg-type] - new_offset += s * self.strides[axis] # type: ignore[operator] + if isinstance(s, (int, np.int64)): + new_offset += s * self.strides[axis] elif isinstance(s, slice): start, stop, step = s.indices(self.shape[axis]) if step <= 0: @@ -231,7 +231,7 @@ def get_group_rank_from_axial_shift( raise ValueError(f"Axis {axis} out of range for coord {coord}") coord_shifted = list(coord) coord_shifted[axis] = (coord_shifted[axis] + delta) % layout_group.shape[axis] - return layout_group(coord_shifted) # type: ignore[arg-type] + return layout_group(coord_shifted) # ty:ignore[invalid-argument-type] def update_exhaustive_strides( @@ -323,19 +323,22 @@ def tiled_softmax_attention_update( if is_initial_chunk: return o_chunk, lse_m_chunk, amax_chunk + assert o is not None and lse_m is not None + if has_amax: - d_lse_m = lse_m - lse_m_chunk # type: ignore[operator] - amax_next = torch.maximum(amax_chunk, amax) # type: ignore[arg-type] - delta_lse = amax_chunk - amax - d_lse_m # type: ignore[operator] + assert amax is not None + d_lse_m = lse_m - lse_m_chunk + amax_next = torch.maximum(amax_chunk, amax) + delta_lse = amax_chunk - amax - d_lse_m o_new = o - torch.sigmoid(delta_lse) * (o - o_chunk) lse_m_new = lse_m_chunk + torch.logsumexp( - torch.cat([(amax - amax_next) + d_lse_m, amax_chunk - amax_next], dim=-1), # type: ignore[operator] + torch.cat([(amax - amax_next) + d_lse_m, amax_chunk - amax_next], dim=-1), dim=-1, keepdim=True, ).to(dtype=lse_m_chunk.dtype) return o_new, lse_m_new, amax_next else: - d_lse_m = lse_m - lse_m_chunk # type: ignore[operator] + d_lse_m = lse_m - lse_m_chunk o_new = o - torch.sigmoid(-d_lse_m) * (o - o_chunk) lse_m_new = lse_m_chunk + torch.log1p(torch.exp(d_lse_m)).to( dtype=lse_m_chunk.dtype @@ -405,8 +408,8 @@ def __init__(self, serial_trunk: nn.Module, dist_manager) -> None: # symbol is imported lazily inside the function to avoid a circular # import with pairformer.py); pyright can't narrow through the # lazy-import isinstance check. - serial_trunk.set_kernel_backend(None) # type: ignore[operator] - serial_trunk.set_chunk_size(None) # type: ignore[operator] + serial_trunk.set_kernel_backend(None) + serial_trunk.set_chunk_size(None) self.dist_trunk = FoldingTrunkDistributed(serial_trunk, dist_manager) self.dist_manager = dist_manager diff --git a/src/transformers/models/esmfold2/kernels/fused_attention_pair_bias.py b/src/transformers/models/esmfold2/kernels/fused_attention_pair_bias.py index 98b722397b..df35e0ab06 100644 --- a/src/transformers/models/esmfold2/kernels/fused_attention_pair_bias.py +++ b/src/transformers/models/esmfold2/kernels/fused_attention_pair_bias.py @@ -345,10 +345,9 @@ def _launch_forward( """ assert z.dim() == 4, f"z must be (B,Q,K,DIM_Z); got {z.shape}" B, Q, K, DIM_Z = z.shape - assert w_proj_z.shape == ( - num_heads, - DIM_Z, - ), f"w_proj_z {w_proj_z.shape} ≠ ({num_heads}, {DIM_Z})" + assert w_proj_z.shape == (num_heads, DIM_Z), ( + f"w_proj_z {w_proj_z.shape} ≠ ({num_heads}, {DIM_Z})" + ) z = z.contiguous() w_proj_z = w_proj_z.contiguous() @@ -656,9 +655,9 @@ def fused_attention_pair_bias( d_model = H * D if precomputed_bias is not None: - assert ( - not torch.is_grad_enabled() - ), "precomputed_bias path is inference-only; autograd is not supported." + assert not torch.is_grad_enabled(), ( + "precomputed_bias path is inference-only; autograd is not supported." + ) bias = precomputed_bias else: if z is None or w_proj_z is None: diff --git a/src/transformers/models/esmfold2/kernels/fused_dropout_residual.py b/src/transformers/models/esmfold2/kernels/fused_dropout_residual.py index 2335d66fe8..02f7ed3c38 100644 --- a/src/transformers/models/esmfold2/kernels/fused_dropout_residual.py +++ b/src/transformers/models/esmfold2/kernels/fused_dropout_residual.py @@ -129,12 +129,12 @@ def _fused_dropout_residual_fwd( mask_2d, out, M, - D, - n_col, - n_row * n_col, - BLOCK_M=BLOCK_M, - BLOCK_D=BLOCK_D, - num_warps=4, # pyright: ignore[reportCallIssue] + D, # ty:ignore[invalid-argument-type] + n_col, # ty:ignore[invalid-argument-type] + n_row * n_col, # ty:ignore[invalid-argument-type] + BLOCK_M=BLOCK_M, # ty:ignore[invalid-argument-type] + BLOCK_D=BLOCK_D, # ty:ignore[invalid-argument-type] + num_warps=4, # ty:ignore[unknown-argument] ) return out @@ -152,12 +152,12 @@ def _fused_dropout_residual_bwd( mask_2d, ddelta, M, - D, - n_col, - n_row * n_col, - BLOCK_M=BLOCK_M, - BLOCK_D=BLOCK_D, - num_warps=4, # pyright: ignore[reportCallIssue] + D, # ty:ignore[invalid-argument-type] + n_col, # ty:ignore[invalid-argument-type] + n_row * n_col, # ty:ignore[invalid-argument-type] + BLOCK_M=BLOCK_M, # ty:ignore[invalid-argument-type] + BLOCK_D=BLOCK_D, # ty:ignore[invalid-argument-type] + num_warps=4, # ty:ignore[unknown-argument] ) return ddelta @@ -244,4 +244,4 @@ def forward(self, pair: torch.Tensor, delta: torch.Tensor) -> torch.Tensor: shape[1] = 1 # row-shared mask: [B, 1, N_col, D] ones = delta.new_ones(shape) mask = torch.nn.functional.dropout(ones, p=self.r, training=True) - return FusedDropoutResidualFn.apply(pair, delta, mask) # type: ignore[return-value] + return FusedDropoutResidualFn.apply(pair, delta, mask) diff --git a/src/transformers/models/esmfold2/kernels/fused_dual_gemm.py b/src/transformers/models/esmfold2/kernels/fused_dual_gemm.py index 66ae88ec42..e642275be9 100644 --- a/src/transformers/models/esmfold2/kernels/fused_dual_gemm.py +++ b/src/transformers/models/esmfold2/kernels/fused_dual_gemm.py @@ -54,7 +54,7 @@ def _gated_dual_gemm_kernel( GROUP_M: tl.constexpr, HAS_MASK: tl.constexpr, TRANSPOSE_OUT: tl.constexpr, # store (N, M) instead of (M, N) - NEEDS_INT64: tl.constexpr = True, # type: ignore[assignment] + NEEDS_INT64: tl.constexpr = True, # ty:ignore[invalid-parameter-default] ): """Per (TILE_M, TILE_N) output tile: gate_acc = Σ_K (x[:, k] @ w1[:, k]) over k @@ -170,7 +170,7 @@ def _gated_dual_gemm_backward_kernel( HAS_MASK: tl.constexpr, GRAD_OUT_TRANSPOSED: tl.constexpr, # 1: load grad from (N, M) layout GRAD_OUT_SPLIT: tl.constexpr, # 1: read from two (N/2, M) tensors (chunk-free path) - NEEDS_INT64: tl.constexpr = True, # type: ignore[assignment] + NEEDS_INT64: tl.constexpr = True, # ty:ignore[invalid-parameter-default] ): """Per (TILE_M, TILE_N) output tile: gate_acc = Σ_K x[:, k] @ w1[:, k] @@ -303,9 +303,9 @@ def _fused_gated_dual_gemm_bwd( M = x_2d.shape[0] N = w1.shape[0] - assert ( - w1.dtype == w2.dtype == torch.bfloat16 - ), f"weights must be bf16; got {w1.dtype}/{w2.dtype}" + assert w1.dtype == w2.dtype == torch.bfloat16, ( + f"weights must be bf16; got {w1.dtype}/{w2.dtype}" + ) assert x_2d.dtype == torch.bfloat16, "bwd only supports bf16 x" if grad_out_split is not None: @@ -425,7 +425,7 @@ def forward( return out @staticmethod - def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + def backward(ctx, grad_out: torch.Tensor): if ctx.has_mask: x, w1, w2, mask = ctx.saved_tensors else: @@ -478,7 +478,7 @@ def forward( return a, b_t @staticmethod - def backward(ctx, grad_a: torch.Tensor, grad_b_t: torch.Tensor): # type: ignore[override] + def backward(ctx, grad_a: torch.Tensor, grad_b_t: torch.Tensor): if ctx.has_mask: x, w1, w2, mask = ctx.saved_tensors else: @@ -487,7 +487,7 @@ def backward(ctx, grad_a: torch.Tensor, grad_b_t: torch.Tensor): # type: ignore # Don't call .contiguous() — _fused_gated_dual_gemm_bwd validates instead # (avoids a full-tensor copy on the common contig path). d_x, d_w1, d_w2, d_mask = _fused_gated_dual_gemm_bwd( - None, # type: ignore[arg-type] + None, # ty:ignore[invalid-argument-type] x, w1, w2, @@ -514,7 +514,7 @@ def fused_gated_dual_gemm_split( if torch.is_grad_enabled() and ( x.requires_grad or w1.requires_grad or w2.requires_grad ): - return FusedGatedDualGEMMSplit.apply(x, w1, w2, mask, trailing_shape) # type: ignore[return-value] + return FusedGatedDualGEMMSplit.apply(x, w1, w2, mask, trailing_shape) out = _fused_gated_dual_gemm_fwd(x, w1, w2, mask=mask, transpose_out=True) N = w1.shape[0] out_view = out.view((N,) + trailing_shape) @@ -548,9 +548,9 @@ def _fused_gated_dual_gemm_fwd( assert w1.shape == w2.shape, f"w1 {w1.shape} ≠ w2 {w2.shape}" assert w1.shape[1] == K - assert ( - w1.dtype == w2.dtype == torch.bfloat16 - ), f"weights must be bf16; got {w1.dtype}/{w2.dtype}" + assert w1.dtype == w2.dtype == torch.bfloat16, ( + f"weights must be bf16; got {w1.dtype}/{w2.dtype}" + ) out_dtype = torch.bfloat16 if transpose_out: @@ -606,5 +606,5 @@ def fused_gated_dual_gemm( "transpose_out=True is inference-only; train path must use the " "non-transposed output (post-stage-3 einsum already handles layout)" ) - return FusedGatedDualGEMM.apply(x, w1, w2, mask) # type: ignore[return-value] + return FusedGatedDualGEMM.apply(x, w1, w2, mask) return _fused_gated_dual_gemm_fwd(x, w1, w2, mask=mask, transpose_out=transpose_out) diff --git a/src/transformers/models/esmfold2/kernels/fused_ln_residual.py b/src/transformers/models/esmfold2/kernels/fused_ln_residual.py index ec349e64a8..3ab8d4add5 100644 --- a/src/transformers/models/esmfold2/kernels/fused_ln_residual.py +++ b/src/transformers/models/esmfold2/kernels/fused_ln_residual.py @@ -210,12 +210,12 @@ def _ln_fwd( mean, rstd, M, - D=D, - EPS=eps, - LAYOUT=layout_int, - TILE_M=_FWD_TILE_M, - num_warps=_FWD_NUM_WARPS, # type: ignore[call-arg] - num_stages=_FWD_NUM_STAGES, # type: ignore[call-arg] + D=D, # ty:ignore[invalid-argument-type] + EPS=eps, # ty:ignore[invalid-argument-type] + LAYOUT=layout_int, # ty:ignore[invalid-argument-type] + TILE_M=_FWD_TILE_M, # ty:ignore[invalid-argument-type] + num_warps=_FWD_NUM_WARPS, # ty:ignore[unknown-argument] + num_stages=_FWD_NUM_STAGES, # ty:ignore[unknown-argument] ) return out, mean, rstd @@ -258,12 +258,12 @@ def _ln_bwd( grad_b_partial, grad_residual if has_residual else _dummy, M, - D=D, - LAYOUT=layout_int, - HAS_RESIDUAL=has_residual, - TILE_M=_BWD_TILE_M, - num_warps=_BWD_NUM_WARPS, # type: ignore[call-arg] - num_stages=_BWD_NUM_STAGES, # type: ignore[call-arg] + D=D, # ty:ignore[invalid-argument-type] + LAYOUT=layout_int, # ty:ignore[invalid-argument-type] + HAS_RESIDUAL=has_residual, # ty:ignore[invalid-argument-type] + TILE_M=_BWD_TILE_M, # ty:ignore[invalid-argument-type] + num_warps=_BWD_NUM_WARPS, # ty:ignore[unknown-argument] + num_stages=_BWD_NUM_STAGES, # ty:ignore[unknown-argument] ) grad_w = grad_w_partial.sum(dim=0).to(w.dtype) @@ -287,7 +287,7 @@ def forward( return out_bnd.view(*out_shape) @staticmethod - def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + def backward(ctx, grad_out: torch.Tensor): x_view, w, mean, rstd = ctx.saved_tensors grad_y = grad_out.contiguous().view(ctx.M, ctx.D) grad_x, grad_w, grad_b = _ln_bwd( @@ -305,7 +305,7 @@ def fused_ln_transpose( layout: str = "bijd->bijd", ) -> torch.Tensor: """Plain LN replacement for ``layer_norm_transpose`` (no residual fusion).""" - return _LayerNormTransposeFn.apply(x, w, b, eps, layout) # type: ignore[return-value] + return _LayerNormTransposeFn.apply(x, w, b, eps, layout) # Stage-1 LN with residual-add folded into the bwd kernel. The Function returns @@ -347,7 +347,7 @@ def forward( return ln_out, residual_link.view_as(residual_link) @staticmethod - def backward(ctx, grad_ln_out: torch.Tensor, grad_link_pass: torch.Tensor): # type: ignore[override] + def backward(ctx, grad_ln_out: torch.Tensor, grad_link_pass: torch.Tensor): x_view, w, mean, rstd = ctx.saved_tensors grad_y = grad_ln_out.contiguous().view(ctx.M, ctx.D) @@ -394,4 +394,4 @@ def fused_ln_with_residual_link( "fused_ln_with_residual_link requires x and residual_link to be the " "same tensor instance (the caller must wire pair → both)." ) - return _LayerNormWithResidualLinkFn.apply(x, w, b, residual_link, eps, layout) # type: ignore[return-value] + return _LayerNormWithResidualLinkFn.apply(x, w, b, residual_link, eps, layout) diff --git a/src/transformers/models/esmfold2/kernels/fused_lnlin_swiglu.py b/src/transformers/models/esmfold2/kernels/fused_lnlin_swiglu.py index ff507d84fe..b163c5aa94 100644 --- a/src/transformers/models/esmfold2/kernels/fused_lnlin_swiglu.py +++ b/src/transformers/models/esmfold2/kernels/fused_lnlin_swiglu.py @@ -276,8 +276,8 @@ def _lnlin_swiglu_fwd( Rstd.stride(0), K, 1e-5, - BLOCK_SIZE=block, - num_warps=num_warps, # pyright: ignore[reportCallIssue] + BLOCK_SIZE=block, # ty:ignore[invalid-argument-type] + num_warps=num_warps, # ty:ignore[unknown-argument] ) cfg = _pick_fwd_config(K) @@ -304,13 +304,13 @@ def _lnlin_swiglu_fwd( lin.stride(1), out.stride(0), out.stride(1), - HAS_LN_BIAS=(LN_B is not None), + HAS_LN_BIAS=(LN_B is not None), # ty:ignore[invalid-argument-type] BLOCK_SIZE_M=cfg["BLOCK_SIZE_M"], BLOCK_SIZE_N=cfg["BLOCK_SIZE_N"], BLOCK_SIZE_K=cfg["BLOCK_SIZE_K"], GROUP_SIZE_M=cfg["GROUP_SIZE_M"], - num_stages=cfg["num_stages"], # pyright: ignore[reportCallIssue] - num_warps=cfg["num_warps"], # pyright: ignore[reportCallIssue] + num_stages=cfg["num_stages"], # ty:ignore[unknown-argument] + num_warps=cfg["num_warps"], # ty:ignore[unknown-argument] ) return out, lin, Mean, Rstd @@ -328,8 +328,8 @@ def _swiglu_bwd_inplace(dout: torch.Tensor, lin: torch.Tensor) -> torch.Tensor: N, dout.stride(0), lin.stride(0), - BLOCK=BLOCK, - num_warps=4, # pyright: ignore[reportCallIssue] + BLOCK=BLOCK, # ty:ignore[invalid-argument-type] + num_warps=4, # ty:ignore[unknown-argument] ) return lin @@ -410,6 +410,4 @@ def reset_parameters(self): nn.init.uniform_(self.W12, -bound, bound) def forward(self, X: torch.Tensor) -> torch.Tensor: - return FusedLNLinearSwiGLUFunction.apply( # type: ignore[return-value] - X, self.W12, self.LN_W, self.LN_B - ) + return FusedLNLinearSwiGLUFunction.apply(X, self.W12, self.LN_W, self.LN_B) diff --git a/src/transformers/models/esmfold2/kernels/trimul_einsum_triton.py b/src/transformers/models/esmfold2/kernels/trimul_einsum_triton.py index 98c1b757c5..238fba9132 100644 --- a/src/transformers/models/esmfold2/kernels/trimul_einsum_triton.py +++ b/src/transformers/models/esmfold2/kernels/trimul_einsum_triton.py @@ -153,9 +153,9 @@ def _batched_einsum( assert a.shape == b.shape, f"a {a.shape} ≠ b {b.shape}" assert a.ndim == 4 assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16 - assert ( - a.is_contiguous() and b.is_contiguous() - ), "trimul einsum kernel requires contiguous (D,B,L,L) inputs" + assert a.is_contiguous() and b.is_contiguous(), ( + "trimul einsum kernel requires contiguous (D,B,L,L) inputs" + ) D, B, L_row, L_col = a.shape assert L_row == L_col, "L_row must equal L_col for stage-3 einsum" diff --git a/src/transformers/models/esmfold2/kernels/trimul_with_residual.py b/src/transformers/models/esmfold2/kernels/trimul_with_residual.py index 45ec395c61..ca83035383 100644 --- a/src/transformers/models/esmfold2/kernels/trimul_with_residual.py +++ b/src/transformers/models/esmfold2/kernels/trimul_with_residual.py @@ -84,7 +84,7 @@ def _gated_gemm_with_residual_kernel( GROUP_M: tl.constexpr, PRECISION: tl.constexpr, HAS_DROP_MASK: tl.constexpr, - NEEDS_INT64: tl.constexpr = True, # type: ignore[assignment] + NEEDS_INT64: tl.constexpr = True, # ty:ignore[invalid-parameter-default] ): pid_m_raw = tl.program_id(axis=0) pid_n_raw = tl.program_id(axis=1) @@ -222,7 +222,7 @@ def _gated_gemm_with_residual_backward_kernel( TILE_K: tl.constexpr, GROUP_M: tl.constexpr, HAS_DROP_MASK: tl.constexpr, - NEEDS_INT64: tl.constexpr = True, # type: ignore[assignment] + NEEDS_INT64: tl.constexpr = True, # ty:ignore[invalid-parameter-default] ): pid_m_raw = tl.program_id(axis=0) pid_n_raw = tl.program_id(axis=1) @@ -342,9 +342,9 @@ def _gated_gemm_with_residual_bwd( M, K = x1.shape N = w1.shape[0] assert x2.shape == x1.shape - assert ( - w1.dtype == w2.dtype == torch.bfloat16 - ), f"weights must be bf16; got {w1.dtype}/{w2.dtype}" + assert w1.dtype == w2.dtype == torch.bfloat16, ( + f"weights must be bf16; got {w1.dtype}/{w2.dtype}" + ) assert x1.dtype == torch.bfloat16 and x2.dtype == torch.bfloat16 # Don't call .contiguous() unconditionally (avoids a full-tensor clone). @@ -440,7 +440,7 @@ def forward( return out @staticmethod - def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + def backward(ctx, grad_out: torch.Tensor): if ctx.has_drop_mask: x1, x2, w1, w2, drop_mask = ctx.saved_tensors else: @@ -523,7 +523,7 @@ def _gated_gemm_with_residual( or w2.requires_grad or residual.requires_grad ): - return GatedGEMMWithResidual.apply( # type: ignore[return-value] + return GatedGEMMWithResidual.apply( x1, x2, w1, w2, residual, drop_mask, n_row, n_col, precision ) return _gated_gemm_with_residual_fwd( @@ -586,11 +586,11 @@ def triangle_multiplicative_update_with_residual( # then a view/permute reaches the dbij layout downstream. a, b_t = fused_gated_dual_gemm_split(x, g_in_weight, p_in_weight, mask=mask) - # Stage 3: triangular einsum. - # Training: native (D, B, L, L) Triton kernel — its bwd reads the layout - # natively, eliminating the contiguous copy torch.einsum's autograd does. - # Inference: torch.einsum (cuBLAS bgemm) is faster fwd-only. - if torch.is_grad_enabled(): + # Stage 3: triangular einsum. The TILE_M=TILE_N=128 Triton kernel wins + # fwd+bwd on aligned lengths, but loses off-grid from wave quantization; + # inference stays on cuBLAS because the Triton kernel's edge is in backward. + _use_triton_einsum = torch.is_grad_enabled() and L_col % 128 == 0 + if _use_triton_einsum: x = trimul_batched_einsum(a, b_t, direction) elif direction == "outgoing": x = torch.einsum("dbik,dbjk->dbij", a, b_t) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2.py b/src/transformers/models/esmfold2/modeling_esmfold2.py index 06b590e770..561950ae33 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2.py @@ -21,20 +21,17 @@ from torch import Tensor try: - import transformer_engine.pytorch as te # type: ignore[import] - from transformer_engine.common.recipe import ( # type: ignore[import] - DelayedScaling, - Format, - ) + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import DelayedScaling, Format TE_AVAILABLE = True except ImportError: - te = None # type: ignore[assignment] - DelayedScaling = None # type: ignore[assignment] - Format = None # type: ignore[assignment] + te = None # ty:ignore[invalid-assignment] + DelayedScaling = None # ty:ignore[invalid-assignment] + Format = None # ty:ignore[invalid-assignment] TE_AVAILABLE = False -from ...modeling_utils import PreTrainedModel # type: ignore[import] +from ...modeling_utils import PreTrainedModel # ty:ignore[unresolved-import] from .configuration_esmfold2 import ESMFold2Config from .modeling_esmfold2_common import ( CHAR_VOCAB_SIZE, @@ -320,17 +317,21 @@ def forward( pair_chains_iptm = torch.zeros( Bm, n_chains, n_chains, device=tm_expected.device, dtype=tm_expected.dtype ) + # pair_chains_iptm[c1, c2] = max over rows i in chain c2 of the mean over + # columns j in chain c1 of tm_expected[i, j] (max-of-row-mean, as in the + # global iptm above), so iptm equals the max off-diagonal entry. for c1 in range(n_chains): chain_c1 = (expanded_asym == c1).float() * mask_f if chain_c1.sum() == 0: continue + col_mask = chain_c1.unsqueeze(-2) + avg_tm = (tm_expected * col_mask).sum(dim=-1) / ( + col_mask.sum(dim=-1) + _EPS + ) for c2 in range(n_chains): chain_c2 = (expanded_asym == c2).float() * mask_f - pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2) - denom = pair_m.sum(dim=(-1, -2)) + _EPS - pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum( - dim=(-1, -2) - ) / denom + row_vals = avg_tm.masked_fill(chain_c2 == 0, float("-inf")) + pair_chains_iptm[:, c1, c2] = row_vals.max(dim=-1).values.clamp(min=0.0) return { "plddt_logits": plddt_logits, @@ -362,7 +363,7 @@ def _convert_te_modules_to_fp8_inplace(module: nn.Module) -> None: """ if not TE_AVAILABLE: raise RuntimeError("transformer_engine is not available; cannot use fp8.") - from transformer_engine.pytorch import quantized_model_init # type: ignore[import] + from transformer_engine.pytorch import quantized_model_init def _walk(mod: nn.Module) -> None: for name, child in list(mod.named_children()): @@ -378,16 +379,16 @@ def _walk(mod: nn.Module) -> None: del child torch.cuda.empty_cache() with quantized_model_init(enabled=True): - new_mod = te.Linear( # type: ignore[union-attr] + new_mod = te.Linear( in_f, out_f, bias=has_bias, params_dtype=dtype ).to(device) - new_mod.weight.quantize_(w) # type: ignore[attr-defined,operator] + new_mod.weight.quantize_(w) # ty:ignore[call-non-callable, unresolved-attribute] if has_bias: assert b is not None - new_mod.bias.data.copy_(b) # type: ignore[union-attr] + new_mod.bias.data.copy_(b) # ty:ignore[call-non-callable] del w, b replaced = True - elif isinstance(child, te.Linear): # type: ignore[union-attr] + elif isinstance(child, te.Linear): # te.Linear with bf16 weight → re-init inside quantized_model_init for fp8. in_f, out_f = child.in_features, child.out_features has_bias = child.bias is not None @@ -402,16 +403,13 @@ def _walk(mod: nn.Module) -> None: del child torch.cuda.empty_cache() with quantized_model_init(enabled=True): - new_mod = te.Linear( # type: ignore[union-attr] - in_f, - out_f, - bias=has_bias, - params_dtype=dtype, # type: ignore[arg-type] - ).to(device) # type: ignore[arg-type] + new_mod = te.Linear( + in_f, out_f, bias=has_bias, params_dtype=dtype + ).to(device) # ty:ignore[no-matching-overload] new_mod.load_state_dict(state, strict=False) replaced = True - elif ( - hasattr(te, "LayerNormLinear") and isinstance(child, te.LayerNormLinear) # type: ignore[union-attr] + elif hasattr(te, "LayerNormLinear") and isinstance( + child, te.LayerNormLinear ): state = {k: v.detach().clone() for k, v in child.state_dict().items()} hidden_size = child.in_features @@ -422,7 +420,7 @@ def _walk(mod: nn.Module) -> None: del child torch.cuda.empty_cache() with quantized_model_init(enabled=True): - new_mod = te.LayerNormLinear( # type: ignore[union-attr] + new_mod = te.LayerNormLinear( hidden_size, out_features, bias=has_bias, @@ -430,30 +428,28 @@ def _walk(mod: nn.Module) -> None: ).to(device) new_mod.load_state_dict(state, strict=False) replaced = True - elif ( - hasattr(te, "LayerNormMLP") and isinstance(child, te.LayerNormMLP) # type: ignore[union-attr] - ): + elif hasattr(te, "LayerNormMLP") and isinstance(child, te.LayerNormMLP): state = {k: v.detach().clone() for k, v in child.state_dict().items()} - fc1_weight: Tensor = child.fc1_weight # type: ignore[attr-defined] + fc1_weight: Tensor = child.fc1_weight # ty:ignore[invalid-assignment] hidden_size = int(fc1_weight.shape[1]) # fc1 packed as (2*ffn_hidden_size, hidden_size) for swiglu. ffn_hidden_size = int(fc1_weight.shape[0]) // 2 has_bias = ( getattr(child, "fc1_bias", None) is not None - and child.fc1_bias is not None # type: ignore[attr-defined] + and child.fc1_bias is not None ) device = fc1_weight.device setattr(mod, name, nn.Identity()) del child torch.cuda.empty_cache() with quantized_model_init(enabled=True): - new_mod = te.LayerNormMLP( # type: ignore[union-attr] + new_mod = te.LayerNormMLP( hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size, bias=has_bias, activation="swiglu", params_dtype=torch.bfloat16, - ).to(device) # type: ignore[arg-type] + ).to(device) new_mod.load_state_dict(state, strict=False) replaced = True @@ -478,12 +474,12 @@ def _lm_precision_context(fp8: bool): """ with torch.autocast(device_type="cuda", dtype=torch.bfloat16): if fp8 and TE_AVAILABLE: - fp8_recipe = DelayedScaling( # type: ignore[misc] - fp8_format=Format.HYBRID, # type: ignore[union-attr] + fp8_recipe = DelayedScaling( + fp8_format=Format.HYBRID, amax_history_len=1, amax_compute_algo="most_recent", ) - with te.autocast(enabled=True, recipe=fp8_recipe): # type: ignore[union-attr] + with te.autocast(enabled=True, recipe=fp8_recipe): yield else: yield @@ -678,10 +674,10 @@ def apply_torch_compile( """ import torch._dynamo - torch._dynamo.config.cache_size_limit = 512 # type: ignore[attr-defined] - torch._dynamo.config.accumulated_cache_size_limit = 512 # type: ignore[attr-defined] + torch._dynamo.config.cache_size_limit = 512 + torch._dynamo.config.accumulated_cache_size_limit = 512 # capture_scalar_outputs avoids graph breaks at .item() in atom-attention path. - torch._dynamo.config.capture_scalar_outputs = True # type: ignore[attr-defined] + torch._dynamo.config.capture_scalar_outputs = True if dynamic is None: dynamic = mode == "dynamic_seqlen" @@ -702,7 +698,7 @@ def apply_torch_compile( def _maybe_compile(module: nn.Module) -> None: if isinstance(module, compile_targets): - module.forward = torch.compile(module.forward, **kwargs) # type: ignore[assignment] + module.forward = torch.compile(module.forward, **kwargs) # ty:ignore[invalid-assignment] self.apply(_maybe_compile) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2_common.py b/src/transformers/models/esmfold2/modeling_esmfold2_common.py index 4c27da1823..1b97cba771 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2_common.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2_common.py @@ -23,56 +23,44 @@ from torch.utils.checkpoint import checkpoint try: - from flash_attn import ( # type: ignore[import] - flash_attn_func, - flash_attn_varlen_func, - ) - from flash_attn.bert_padding import ( # type: ignore[import] - index_first_axis, - pad_input, - ) + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input FLASH_ATTN_AVAILABLE = True except ImportError: - flash_attn_func = None # type: ignore[assignment] - flash_attn_varlen_func = None # type: ignore[assignment] - index_first_axis = None # type: ignore[assignment] - pad_input = None # type: ignore[assignment] + flash_attn_func = None # ty:ignore[invalid-assignment] + flash_attn_varlen_func = None # ty:ignore[invalid-assignment] + index_first_axis = None # ty:ignore[invalid-assignment] + pad_input = None # ty:ignore[invalid-assignment] FLASH_ATTN_AVAILABLE = False try: - from cuequivariance_torch import ( # type: ignore[import] - attention_pair_bias as _cue_attn_pair_bias, - ) - from cuequivariance_torch.primitives.triangle import ( # type: ignore[import] + from cuequivariance_torch import attention_pair_bias as _cue_attn_pair_bias + from cuequivariance_torch.primitives.triangle import ( triangle_multiplicative_update as _cue_tri_mul, ) CUE_AVAILABLE = True except ImportError: - _cue_attn_pair_bias = None # type: ignore[assignment] - _cue_tri_mul = None # type: ignore[assignment] + _cue_attn_pair_bias = None # ty:ignore[invalid-assignment] + _cue_tri_mul = None # ty:ignore[invalid-assignment] CUE_AVAILABLE = False # Vendored inference-only Triton kernels. try: + from .kernels import FusedDropoutResidual as _FusedDropoutResidual + from .kernels import FusedLNLinearSwiGLU as _FusedLNLinearSwiGLU + from .kernels import fused_pair_bias as _fused_pair_bias from .kernels import ( - FusedDropoutResidual as _FusedDropoutResidual, # type: ignore[import] - ) - from .kernels import ( - FusedLNLinearSwiGLU as _FusedLNLinearSwiGLU, # type: ignore[import] - ) - from .kernels import fused_pair_bias as _fused_pair_bias # type: ignore[import] - from .kernels import ( # type: ignore[import] triangle_multiplicative_update_with_residual as _fused_trimul_with_residual, ) TRITON_KERNELS_AVAILABLE = True except ImportError: - _fused_pair_bias = None # type: ignore[assignment] - _fused_trimul_with_residual = None # type: ignore[assignment] - _FusedLNLinearSwiGLU = None # type: ignore[assignment] - _FusedDropoutResidual = None # type: ignore[assignment] + _fused_pair_bias = None + _fused_trimul_with_residual = None + _FusedLNLinearSwiGLU = None # ty:ignore[invalid-assignment] + _FusedDropoutResidual = None # ty:ignore[invalid-assignment] TRITON_KERNELS_AVAILABLE = False from .configuration_esmfold2 import ESMFold2Config @@ -81,6 +69,11 @@ BACKEND_CUEQ = "cuequivariance" _VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ) +# The vendored fused Triton kernels (LN+SwiGLU, trimul-with-residual) operate in +# bfloat16 only. Single-source that dtype here so every fused-path buffer and +# cast references one named constant instead of scattering ``torch.bfloat16``. +_FUSED_KERNEL_DTYPE = torch.bfloat16 + def _fused_active(module: nn.Module, tensor: Tensor) -> bool: """Common preconditions for the vendored fused Triton inference kernels.""" @@ -92,6 +85,20 @@ def _fused_active(module: nn.Module, tensor: Tensor) -> bool: ) +def _fused_pair_stack_active(module: nn.Module, tensor: Tensor) -> bool: + """Fused pair-stack kernels that support autograd.""" + return ( + TRITON_KERNELS_AVAILABLE + and getattr(module, "_kernel_backend", None) == BACKEND_FUSED + and tensor.is_cuda + ) + + +def _to_fused_kernel_dtype(t: Tensor) -> Tensor: + """Cast ``t`` to the fused-kernel dtype (no-op if already that dtype).""" + return t if t.dtype == _FUSED_KERNEL_DTYPE else t.to(_FUSED_KERNEL_DTYPE) + + def _cueq_active(module: nn.Module) -> bool: return CUE_AVAILABLE and getattr(module, "_kernel_backend", None) == BACKEND_CUEQ @@ -581,16 +588,16 @@ def forward(self, x: Tensor, attention_params: tuple) -> Tensor: attention_params[3], attention_params[4], ) - q_unpad = index_first_axis( # type: ignore[misc] + q_unpad = index_first_axis( q.reshape(-1, self.n_heads, self.head_dim), indices ) - k_unpad = index_first_axis( # type: ignore[misc] + k_unpad = index_first_axis( k.reshape(-1, self.n_heads, self.head_dim), indices ) - v_unpad = index_first_axis( # type: ignore[misc] + v_unpad = index_first_axis( v.reshape(-1, self.n_heads, self.head_dim), indices ) - out_unpad = flash_attn_varlen_func( # type: ignore[misc] + out_unpad = flash_attn_varlen_func( q_unpad, k_unpad, v_unpad, @@ -601,9 +608,9 @@ def forward(self, x: Tensor, attention_params: tuple) -> Tensor: softmax_scale=self.scale, window_size=(self.half_window, self.half_window), ) - out = pad_input(out_unpad, indices, B, N) # type: ignore[misc] + out = pad_input(out_unpad, indices, B, N) elif FLASH_ATTN_AVAILABLE: - out = flash_attn_func( # type: ignore[misc] + out = flash_attn_func( q, k, v, @@ -630,7 +637,7 @@ def forward(self, x: Tensor, attention_params: tuple) -> Tensor: ).transpose(1, 2) out = out * valid.unsqueeze(-1).unsqueeze(-1) - out = out.to(input_dtype).reshape(B, N, -1) # type: ignore[union-attr] + out = out.to(input_dtype).reshape(B, N, -1) out = out * torch.sigmoid(self.gate_proj(x_input)) return self.out_proj(out) @@ -1122,14 +1129,14 @@ def forward( else torch.zeros_like(pair_norm_w) ) z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16) - bias = _fused_pair_bias( # type: ignore[misc] + bias = _fused_pair_bias( z_bf, kernel_mask, self.pair_bias_proj.weight, num_heads=self.num_heads, pair_norm_w=pair_norm_w, pair_norm_b=pair_norm_b, - ) # (B, H, Q, K) + ) # (B, H, Q, K) # ty:ignore[call-non-callable] q_bhqd = q.transpose(1, 2) k_bhqd = k.transpose(1, 2) v_bhqd = v.transpose(1, 2) @@ -1151,7 +1158,7 @@ def forward( if attention_mask is not None else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) ) - out, _ = _cue_attn_pair_bias( # type: ignore[misc] + out, _ = _cue_attn_pair_bias( s=x, q=q.transpose(1, 2), k=k.transpose(1, 2), @@ -1794,11 +1801,11 @@ def sample( ) -> dict[str, Tensor | None]: """Diffusion sampling (Algorithm 18). - The Karras schedule is built with ``num_sampling_steps`` entries, then - clipped to ``max_inference_sigma`` (the high-σ tail above the cap is - dropped and the cap prepended). The number of denoising steps actually - run is therefore fewer than ``num_sampling_steps`` whenever the schedule - extends above the cap. + ``num_sampling_steps`` is the number of denoising steps actually run. + When ``max_inference_sigma`` is set, the Karras schedule built with + ``num_sampling_steps`` entries would lose its high-σ tail to the cap, + so we inflate the underlying schedule length here to land back at the + requested step count post-truncation. """ n_atoms = tok_idx.shape[1] device = s_inputs.device @@ -2418,7 +2425,7 @@ def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor p_in_weight, g_in_weight = self.split_kernel_weights() try: - return _cue_tri_mul( # type: ignore[misc] + return _cue_tri_mul( pair_grid, direction=self._kernel_flow_direction(), mask=visibility, @@ -2522,18 +2529,17 @@ def set_kernel_backend(self, backend: str | None) -> None: d_inner = self.ffn.hidden_features has_ln_bias = self.norm.bias is not None device = self.ffn.w12.weight.device - dtype = self.ffn.w12.weight.dtype fused = _FusedLNLinearSwiGLU( d_model=d_model, d_inner=d_inner, has_ln_bias=has_ln_bias, device=device, - dtype=dtype, + dtype=_FUSED_KERNEL_DTYPE, ) with torch.no_grad(): fused.LN_W.copy_(self.norm.weight) if has_ln_bias: - fused.LN_B.copy_(self.norm.bias) # type: ignore[union-attr] + fused.LN_B.copy_(self.norm.bias) # ty:ignore[unresolved-attribute] # FusedLNLinearSwiGLU.W12 is (d_model, 2*d_inner); transpose nn.Linear once. fused.W12.copy_(self.ffn.w12.weight.t().contiguous()) self._fused_swiglu = fused.eval().requires_grad_(False) @@ -2542,9 +2548,10 @@ def set_kernel_backend(self, backend: str | None) -> None: def _can_use_fused_path(self, x: Tensor) -> bool: return ( - _fused_active(self, x) + _fused_pair_stack_active(self, x) and self._fused_swiglu is not None - and x.dtype == torch.bfloat16 + and x.dtype == _FUSED_KERNEL_DTYPE + and self._fused_swiglu.W12.dtype == x.dtype ) def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor: @@ -2561,16 +2568,13 @@ def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor: out = torch.addmm( x.contiguous().view(-1, x_shape[-1]), hidden.view(-1, hidden.shape[-1]), - ffn.w3.weight.t(), + ffn.w3.weight.t().to(_FUSED_KERNEL_DTYPE), ) return out.view(x_shape) def forward(self, x: Tensor) -> Tensor: - # Inference-only fast path (addmm-fused residual + pre-alloc out) - # — diverges bit-exactly from ``x + ffn(norm(x))`` so we only use - # it when grad is disabled (binder-design / bit-exact tests run - # with grad on and need the reference path). - if not torch.is_grad_enabled() and self._can_use_fused_path(x): + # Fused fast path (addmm-fused residual). + if self._can_use_fused_path(x): fused = self._fused_swiglu assert fused is not None pre_w3 = fused @@ -2627,35 +2631,36 @@ def set_chunk_size(self, chunk_size: int | None) -> None: self.pair_transition.set_chunk_size(chunk_size) def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool: - return _fused_active(self, pair) and pair.dtype == torch.bfloat16 + return ( + _fused_pair_stack_active(self, pair) + and pair.dtype == _FUSED_KERNEL_DTYPE + and _fused_trimul_with_residual is not None + ) def _fused_trimul_with_residual( self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None ) -> Tensor: """Fused TriMul+residual call; weights from the corresponding engine.""" tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in - engine: TriangleMultiplicativeBlock = tri._engine # type: ignore[assignment] + engine: TriangleMultiplicativeBlock = tri._engine p_in_weight, g_in_weight = engine.split_kernel_weights() - def _bf16(t: Tensor) -> Tensor: - return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16) - - return _fused_trimul_with_residual( # type: ignore[misc] + return _fused_trimul_with_residual( pair, direction, residual=pair, drop_mask=None, # inference: no dropout, matches internal's eval path - norm_in_weight=_bf16(engine.norm_start.weight), - norm_in_bias=_bf16(engine.norm_start.bias), - p_in_weight=_bf16(p_in_weight), - g_in_weight=_bf16(g_in_weight), - norm_out_weight=_bf16(engine.norm_mix.weight), - norm_out_bias=_bf16(engine.norm_mix.bias), - p_out_weight=_bf16(engine.proj_emit.weight), - g_out_weight=_bf16(engine.proj_gate.weight), + norm_in_weight=_to_fused_kernel_dtype(engine.norm_start.weight), + norm_in_bias=_to_fused_kernel_dtype(engine.norm_start.bias), + p_in_weight=_to_fused_kernel_dtype(p_in_weight), + g_in_weight=_to_fused_kernel_dtype(g_in_weight), + norm_out_weight=_to_fused_kernel_dtype(engine.norm_mix.weight), + norm_out_bias=_to_fused_kernel_dtype(engine.norm_mix.bias), + p_out_weight=_to_fused_kernel_dtype(engine.proj_emit.weight), + g_out_weight=_to_fused_kernel_dtype(engine.proj_gate.weight), mask=pair_attention_mask, eps=_EPS, - ) + ) # ty:ignore[call-non-callable] def forward( self, pair: Tensor, pair_attention_mask: Tensor | None = None @@ -2711,7 +2716,7 @@ def forward( for block in self.blocks: fn = partial(block, pair_attention_mask=pair_attention_mask) if torch.is_grad_enabled(): - pair = checkpoint(fn, pair, use_reentrant=False) # pyright: ignore + pair = checkpoint(fn, pair, use_reentrant=False) else: pair = fn(pair) if pair.dtype != orig_dtype: diff --git a/src/transformers/models/esmfold2/modeling_esmfold2_experimental.py b/src/transformers/models/esmfold2/modeling_esmfold2_experimental.py index 122da498b6..2d2fb80fec 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2_experimental.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2_experimental.py @@ -29,7 +29,7 @@ import torch.nn.functional as F from torch import Tensor -from ...modeling_utils import PreTrainedModel # type: ignore[import] +from ...modeling_utils import PreTrainedModel # ty:ignore[unresolved-import] from .configuration_esmfold2 import ESMFold2Config from .modeling_esmfold2_common import ( CHAR_VOCAB_SIZE, @@ -306,17 +306,21 @@ def forward( pair_chains_iptm = torch.zeros( Bm, n_chains, n_chains, device=tm_expected.device, dtype=tm_expected.dtype ) + # pair_chains_iptm[c1, c2] = max over rows i in chain c2 of the mean over + # columns j in chain c1 of tm_expected[i, j] (max-of-row-mean, as in the + # global iptm above), so iptm equals the max off-diagonal entry. for c1 in range(n_chains): chain_c1 = (expanded_asym == c1).float() * mask_f if chain_c1.sum() == 0: continue + col_mask = chain_c1.unsqueeze(-2) + avg_tm = (tm_expected * col_mask).sum(dim=-1) / ( + col_mask.sum(dim=-1) + _EPS + ) for c2 in range(n_chains): chain_c2 = (expanded_asym == c2).float() * mask_f - pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2) - denom = pair_m.sum(dim=(-1, -2)) + _EPS - pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum( - dim=(-1, -2) - ) / denom + row_vals = avg_tm.masked_fill(chain_c2 == 0, float("-inf")) + pair_chains_iptm[:, c1, c2] = row_vals.max(dim=-1).values.clamp(min=0.0) return { "plddt_logits": plddt_logits, @@ -542,7 +546,7 @@ def __init__(self, config: ESMFold2Config) -> None: self.pair_loop_proj = nn.Sequential( nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False) ) - nn.init.zeros_(self.pair_loop_proj[1].weight) # type: ignore[arg-type] + nn.init.zeros_(self.pair_loop_proj[1].weight) # ty:ignore[invalid-argument-type] # Structure head self.structure_head = DiffusionStructureHead(config) @@ -598,7 +602,7 @@ def configure_lm_dropout( def load_esmc(self, esmc_model_path: str) -> None: """Load the ESMC LM backbone from a HuggingFace Hub repo ID or local directory.""" - from ..esmc.modeling_esmc import ESMCModel # type: ignore[import] + from ..esmc.modeling_esmc import ESMCModel esmc = ESMCModel.from_pretrained(esmc_model_path) self._esmc = esmc.bfloat16().to(self.device).eval() @@ -636,7 +640,7 @@ def infer_all_atom(self, structure_input, **forward_kwargs): "All-atom inference requires the `esm` companion package: " "`pip install esm`." ) from e - esmfold2 = esm.models.esmfold2 # type: ignore[attr-defined] + esmfold2 = esm.models.esmfold2 if isinstance(structure_input, esmfold2.ProteinInput): structure_input = esmfold2.StructurePredictionInput( @@ -656,7 +660,7 @@ def _output_to_molecular_complex(output: dict, features: dict, chain_infos: list import evolutionaryscale.opensource as esm # TODO: change to `import esm` when open sourcing import evolutionaryscale.opensource.models.esmfold2 # noqa: F401 # TODO: drop when open sourcing - esmfold2 = esm.models.esmfold2 # type: ignore[attr-defined] + esmfold2 = esm.models.esmfold2 ELEMENT_NUMBER_TO_SYMBOL = esmfold2.ELEMENT_NUMBER_TO_SYMBOL MolecularComplex = esmfold2.MolecularComplex From d509766508561910f0db8cb759634421088d5d19 Mon Sep 17 00:00:00 2001 From: recurse-agent Date: Fri, 14 Aug 2026 21:04:33 +0000 Subject: [PATCH 2/3] Allow ESMFold2 imports without a CUDA driver Co-Authored-By: OpenAI Codex --- src/transformers/models/esmfold2/modeling_esmfold2_common.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2_common.py b/src/transformers/models/esmfold2/modeling_esmfold2_common.py index 1b97cba771..17a309c7c0 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2_common.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2_common.py @@ -48,6 +48,8 @@ # Vendored inference-only Triton kernels. try: + if not torch.cuda.is_available(): + raise ImportError from .kernels import FusedDropoutResidual as _FusedDropoutResidual from .kernels import FusedLNLinearSwiGLU as _FusedLNLinearSwiGLU from .kernels import fused_pair_bias as _fused_pair_bias From 5e93c9a06dab35bb9f830e22fbb05d45b891ff2f Mon Sep 17 00:00:00 2001 From: recurse-agent Date: Sat, 15 Aug 2026 02:41:58 +0000 Subject: [PATCH 3/3] Chunk ESMFold2 confidence samples on demand Co-Authored-By: OpenAI Codex --- .../models/esmfold2/modeling_esmfold2.py | 77 +++++++++++++------ 1 file changed, 54 insertions(+), 23 deletions(-) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2.py b/src/transformers/models/esmfold2/modeling_esmfold2.py index 561950ae33..4325dbb94e 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2.py @@ -20,6 +20,7 @@ import torch.nn.functional as F from torch import Tensor + try: import transformer_engine.pytorch as te from transformer_engine.common.recipe import DelayedScaling, Format @@ -56,6 +57,7 @@ maybe_subsample_msa, ) + _EPS = 1e-6 _NONPOLYMER_ID = 4 @@ -877,6 +879,7 @@ def forward( msa_max_depth: int = 1024, msa_column_mask_rate: float = 0.1, msa_subsample_at_inference: bool = True, + confidence_chunk_size: int | None = None, **kwargs, ) -> dict[str, Tensor]: tok_mask = token_attention_mask @@ -995,15 +998,15 @@ def forward( msa_attention_mask = maybe_apply_msa_column_masking( msa_attention_mask, rate=msa_column_mask_rate ) - _msa_inputs = dict( - msa=msa, - msa_attention_mask=msa_attention_mask, - has_deletion=has_deletion, - deletion_value=deletion_value, - x_inputs=x_inputs, - max_depth=msa_max_depth, - subsample_enabled=msa_subsample_at_inference, - ) + _msa_inputs = { + "msa": msa, + "msa_attention_mask": msa_attention_mask, + "has_deletion": has_deletion, + "deletion_value": deletion_value, + "x_inputs": x_inputs, + "max_depth": msa_max_depth, + "subsample_enabled": msa_subsample_at_inference, + } # Method call (not inline loop) frees per-iter L²×c_z locals. z = self._run_one_loop( @@ -1054,20 +1057,48 @@ def forward( output: dict[str, Tensor] = {"distogram_logits": distogram_logits} output["sample_atom_coords"] = sample_coords - confidence_output = self.confidence_head( - s_inputs=x_inputs.detach(), - z=z.detach().float(), - x_pred=sample_coords.detach(), - distogram_atom_idx=disto_idx, - token_attention_mask=tok_mask, - atom_to_token=atom_to_token, - atom_attention_mask=atm_mask, - asym_id=asym_id, - mol_type=mol_type, - num_diffusion_samples=n_samples, - relative_position_encoding=relative_position_encoding.detach(), - token_bonds_encoding=token_bonds_encoding.detach(), - ) + confidence_kwargs = { + "s_inputs": x_inputs.detach(), + "z": z.detach().float(), + "distogram_atom_idx": disto_idx, + "token_attention_mask": tok_mask, + "atom_to_token": atom_to_token, + "atom_attention_mask": atm_mask, + "asym_id": asym_id, + "mol_type": mol_type, + "relative_position_encoding": relative_position_encoding.detach(), + "token_bonds_encoding": token_bonds_encoding.detach(), + } + if confidence_chunk_size is None or confidence_chunk_size >= n_samples: + confidence_output = self.confidence_head( + **confidence_kwargs, + x_pred=sample_coords.detach(), + num_diffusion_samples=n_samples, + ) + else: + if confidence_chunk_size < 1: + raise ValueError("confidence_chunk_size must be positive or None") + batch_size = x_inputs.shape[0] + sample_shape = sample_coords.shape[1:] + samples = sample_coords.reshape(batch_size, n_samples, *sample_shape) + confidence_chunks: dict[str, list[Tensor]] = {} + for start in range(0, n_samples, confidence_chunk_size): + stop = min(start + confidence_chunk_size, n_samples) + chunk_samples = samples[:, start:stop].reshape(-1, *sample_shape) + chunk_output = self.confidence_head( + **confidence_kwargs, + x_pred=chunk_samples.detach(), + num_diffusion_samples=stop - start, + ) + for name, value in chunk_output.items(): + value = value.reshape(batch_size, stop - start, *value.shape[1:]) + confidence_chunks.setdefault(name, []).append(value) + confidence_output = { + name: torch.cat(values, dim=1).reshape( + batch_size * n_samples, *values[0].shape[2:] + ) + for name, values in confidence_chunks.items() + } output.update(confidence_output) output["atom_pad_mask"] = ( atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask