diff --git a/.github/scripts/airtable_issue_sync.py b/.github/scripts/airtable_issue_sync.py new file mode 100644 index 00000000..431f9602 --- /dev/null +++ b/.github/scripts/airtable_issue_sync.py @@ -0,0 +1,380 @@ +#!/usr/bin/env python3 +"""Sync a GitHub issue into an Airtable base. + +Env: AIRTABLE_TOKEN, AIRTABLE_BASE, AIRTABLE_TABLE, and (from GitHub Actions) +GITHUB_EVENT_PATH. Optional: LABEL_TYPE_MAP, PRODUCT_KEYWORD_MAP, +DEFAULT_PRODUCTS, DRY_RUN. + +Everything written comes out of the webhook payload; the script makes no calls +to GitHub. + +An issue is matched to an existing row by URL and updated in place, otherwise a +row is created. Fields that already hold a value are left as they are. + +The Airtable token needs schema.bases:read, data.records:read and +data.records:write. +""" + +import json +import os +import re +import sys +import time +import urllib.error +import urllib.parse +import urllib.request + +API = "https://api.airtable.com/v0" + +F_TYPE = "Type of Request" +F_DATE = "Date" +F_PRODUCT = "Tool or Product" +F_SOURCE = "Issue Source" +F_DESCRIPTION = "Issue Description" +F_GITHUB_ID = "GitHub ID" +F_ORIGINAL_Q = "Link to Original Q" +F_COMPLETE = "Issue Complete?" +F_DETAILS = "Additional Details" +F_REFERENCE = "Original Record Reference" + +SOURCE_GITHUB = "Github" +STATE_OPEN = "Incomplete" +STATE_CLOSED = "Complete" +COMPLETE_STATES = {"complete", "out of scope"} + +# Airtable long text holds 100k characters; stay well under it. +MAX_TEXT_CHARS = 50_000 + +REQUEST_TIMEOUT_SECONDS = 30 + +# GitHub label (lowercased) -> "Type of Request" option. Extend with +# LABEL_TYPE_MAP rather than editing this. +DEFAULT_LABEL_TYPE_MAP = { + "bug": "Bug", + "defect": "Bug", + "error": "Error", + "crash": "Error", + "enhancement": "Feature Request", + "feature": "Feature Request", + "feature request": "Feature Request", + "question": "Support", + "support": "Support", + "help wanted": "Support", + "documentation": "Support", + "docs": "Support", + "feedback": "General Feedback", + "access": "Access issue", + "permissions": "Access issue", +} + +# Substring (lowercased) -> "Tool or Product" option, scanned over title, body +# and labels. +DEFAULT_PRODUCT_KEYWORD_MAP = { + "esmfold2": "ESMFold2", + "esmfold 2": "ESMFold2", + "esm atlas": "ESM Atlas", + "metagenomic atlas": "ESM Atlas", + "biohub platform": "Biohub Platform", + "esmc": "ESMC", + "esm3": "ESM3", + "esm 3": "ESM3", + "binder": "Binder", + "sae": "SAE", +} + + +def require_env(name): + v = os.environ.get(name) + if not v: + sys.exit(f"Missing required env var: {name}") + return v + + +def _lower_keys(raw): + return {k.lower(): v for k, v in json.loads(raw).items()} + + +AIRTABLE_TOKEN = require_env("AIRTABLE_TOKEN") +AIRTABLE_BASE = require_env("AIRTABLE_BASE") +AIRTABLE_TABLE = require_env("AIRTABLE_TABLE") +LABEL_TYPE_MAP = { + **DEFAULT_LABEL_TYPE_MAP, + **_lower_keys(os.environ.get("LABEL_TYPE_MAP", "{}")), +} +PRODUCT_KEYWORD_MAP = { + **DEFAULT_PRODUCT_KEYWORD_MAP, + **_lower_keys(os.environ.get("PRODUCT_KEYWORD_MAP", "{}")), +} +DEFAULT_PRODUCTS = json.loads(os.environ.get("DEFAULT_PRODUCTS", "[]")) +DRY_RUN = os.environ.get("DRY_RUN", "").lower() in ("1", "true", "yes") + +_warned = set() + + +def warn(msg): + if msg not in _warned: + print(f"WARN: {msg}") + _warned.add(msg) + + +def _request(req, attempts=4): + for attempt in range(attempts): + try: + with urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT_SECONDS) as resp: + return json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as err: + if err.code in (429, 500, 502, 503) and attempt < attempts - 1: + time.sleep(2**attempt) + continue + detail = err.read().decode("utf-8", "replace") + raise RuntimeError( + f"{req.get_method()} {req.full_url} -> {err.code}: {detail}" + ) + raise RuntimeError(f"{req.get_method()} {req.full_url}: retries exhausted") + + +def airtable(path, method="GET", body=None): + data = json.dumps(body).encode("utf-8") if body is not None else None + req = urllib.request.Request(API + path, data=data, method=method) + req.add_header("Authorization", f"Bearer {AIRTABLE_TOKEN}") + req.add_header("Content-Type", "application/json") + return _request(req) + + +def table_path(): + return f"/{AIRTABLE_BASE}/{urllib.parse.quote(AIRTABLE_TABLE, safe='')}" + + +def get_schema(): + """Field name -> field schema for the target table.""" + tables = airtable(f"/meta/bases/{AIRTABLE_BASE}/tables")["tables"] + for t in tables: + if AIRTABLE_TABLE in (t["id"], t["name"]): + return {f["name"]: f for f in t["fields"]} + names = ", ".join(f'"{t["name"]}"' for t in tables) + sys.exit( + f'Table "{AIRTABLE_TABLE}" not found in base {AIRTABLE_BASE}. Have: {names}' + ) + + +def truncate(text, url): + if len(text) <= MAX_TEXT_CHARS: + return text + return text[:MAX_TEXT_CHARS] + f"\n\n... (truncated; read the full issue at {url})" + + +def match_choices(fschema, values): + """Keep the values that already exist as options, in the schema's casing.""" + options = [c["name"] for c in fschema.get("options", {}).get("choices", [])] + by_lower = {o.lower(): o for o in options} + matched = [] + for v in values: + canonical = by_lower.get(str(v).strip().lower()) + if canonical is None: + warn(f'"{fschema["name"]}" has no option matching "{v}"; dropping it.') + elif canonical not in matched: + matched.append(canonical) + return matched + + +def coerce(fschema, value): + """Value in the shape Airtable wants, or None if it should not be written.""" + t = fschema["type"] + if t in ("multipleSelects", "singleSelect"): + values = value if isinstance(value, list) else [value] + matched = match_choices(fschema, [v for v in values if v]) + if not matched: + return None + return matched if t == "multipleSelects" else matched[0] + if t in ("multilineText", "singleLineText", "richText"): + return str(value) if value else None + if t == "date": + return str(value)[:10] if value else None + if t == "checkbox": + return bool(value) + warn(f'field "{fschema["name"]}" (type {t}) is not written by this sync; skipping.') + return None + + +def is_empty(current): + return current is None or current == "" or current == [] or current is False + + +class Row: + """Fields to write, validated against the live schema.""" + + def __init__(self, schema, existing): + self.schema = schema + self.existing = (existing or {}).get("fields", {}) + self.fields = {} + + def _prepare(self, name, value): + fschema = self.schema.get(name) + if not fschema: + warn(f'field "{name}" not found in the table; skipping.') + return None + return coerce(fschema, value) + + def own(self, name, value): + """Keep this field equal to the GitHub value.""" + prepared = self._prepare(name, value) + if prepared is not None and self.existing.get(name) != prepared: + self.fields[name] = prepared + + def fill(self, name, value): + """Write this field only while it is still empty.""" + if not is_empty(self.existing.get(name)): + return + prepared = self._prepare(name, value) + if prepared is not None: + self.fields[name] = prepared + + +def label_names(issue): + return [ + str(lb["name"] if isinstance(lb, dict) else lb) + for lb in issue.get("labels") or [] + ] + + +def request_types(issue): + types = [] + for label in label_names(issue): + mapped = LABEL_TYPE_MAP.get(label.strip().lower()) + if mapped and mapped not in types: + types.append(mapped) + return types + + +def products(issue): + haystack = "\n".join( + [ + issue.get("title") or "", + issue.get("body") or "", + " ".join(label_names(issue)), + ] + ).lower() + found = [] + # Longest keyword first so "esmfold2" is not shadowed by a shorter match. + for keyword in sorted(PRODUCT_KEYWORD_MAP, key=len, reverse=True): + product = PRODUCT_KEYWORD_MAP[keyword] + if product in found: + continue + if re.search(rf"(? 1: + ids = ", ".join(r["id"] for r in records) + warn(f"{len(records)} rows match {sync_key} ({ids}); updating the first.") + return records[0] if records else None + + +def desired_state(action, current): + """Follow the issue's open/closed state without discarding another value.""" + values = current if isinstance(current, list) else [current] if current else [] + is_complete = any(str(v).strip().lower() in COMPLETE_STATES for v in values) + if action == "closed": + return None if is_complete else STATE_CLOSED + if action == "reopened": + return STATE_OPEN if (not values or is_complete) else None + return None + + +def build_row(schema, existing, issue, repo_full_name, action): + url = issue["html_url"] + title = issue.get("title") or f"Issue #{issue['number']}" + body = (issue.get("body") or "").strip() or "(No description was provided.)" + login = (issue.get("user") or {}).get("login") + + row = Row(schema, existing) + row.own(F_ORIGINAL_Q, url) + row.own(F_REFERENCE, f"{repo_full_name}#{issue['number']}") + row.own(F_SOURCE, [SOURCE_GITHUB]) + row.own(F_GITHUB_ID, login) + + row.fill(F_DATE, issue.get("created_at")) + row.fill(F_DESCRIPTION, truncate(f"{title}\n\n{body}", url)) + row.fill(F_DETAILS, details_block(issue, repo_full_name)) + row.fill(F_TYPE, request_types(issue)) + row.fill(F_PRODUCT, products(issue)) + + state = ( + desired_state(action, row.existing.get(F_COMPLETE)) if existing else STATE_OPEN + ) + if state: + row.own(F_COMPLETE, [state]) + + return row.fields + + +def main(): + with open(require_env("GITHUB_EVENT_PATH"), encoding="utf-8") as fh: + event = json.load(fh) + action = event.get("action") + issue = event.get("issue") + if not issue: + print(f'No issue in payload (action="{action}"); nothing to do.') + return + + if (issue.get("user") or {}).get("type") == "Bot": + print(f"Issue #{issue['number']} was opened by a bot; skipping.") + return + + repo = event.get("repository") or {} + repo_full_name = repo.get("full_name") or os.environ.get("GITHUB_REPOSITORY", "") + sync_key = f"{repo_full_name}#{issue['number']}" + + schema = get_schema() + existing = find_record(sync_key, issue["html_url"]) + fields = build_row(schema, existing, issue, repo_full_name, action) + + if not fields: + print(f'{sync_key}: nothing to change (action="{action}").') + return + if DRY_RUN: + target = existing["id"] if existing else "(new record)" + print(f"DRY_RUN {sync_key} -> {target}\n{json.dumps(fields, indent=2)}") + return + + if existing: + airtable(f"{table_path()}/{existing['id']}", "PATCH", {"fields": fields}) + print( + f'Updated {sync_key} as {existing["id"]} (action="{action}"): ' + f"{', '.join(sorted(fields))}." + ) + else: + resp = airtable(table_path(), "POST", {"fields": fields}) + print(f'Created {sync_key} as {resp["id"]} (action="{action}").') + + +if __name__ == "__main__": + try: + main() + except Exception as err: # noqa: BLE001 + sys.exit(f"Sync failed: {err}") diff --git a/.github/workflows/airtable-issue-sync.yaml b/.github/workflows/airtable-issue-sync.yaml new file mode 100644 index 00000000..ce9d2bdb --- /dev/null +++ b/.github/workflows/airtable-issue-sync.yaml @@ -0,0 +1,30 @@ +name: Sync issues to Airtable + +on: + issues: + types: [opened, edited, reopened, closed, labeled, unlabeled] + +permissions: + contents: read + +concurrency: + group: airtable-sync-${{ github.event.issue.number }} + cancel-in-progress: false + +jobs: + sync: + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Sync issue to Airtable + run: python .github/scripts/airtable_issue_sync.py + env: + AIRTABLE_TOKEN: ${{ secrets.AIRTABLE_TOKEN }} + AIRTABLE_BASE: ${{ vars.AIRTABLE_BASE }} + AIRTABLE_TABLE: ${{ vars.AIRTABLE_TABLE }} diff --git a/cookbook/snippets/esmc.py b/cookbook/snippets/esmc.py index c83557a9..6b617860 100644 --- a/cookbook/snippets/esmc.py +++ b/cookbook/snippets/esmc.py @@ -3,7 +3,7 @@ import torch -from esm.models.esmc import ESMC +from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer from esm.sdk import batch_executor, esmc_client from esm.sdk.api import ( ESMCInferenceClient, @@ -56,18 +56,27 @@ def main(client: ESMCInferenceClient | ESM3ForgeInferenceClient): print(f"Client returned hidden states with shape {output.hidden_states.shape}") -def raw_forward(model: ESMC): +def raw_forward(model: EsmcForMaskedLM): protein = ESMProtein(sequence="AAAAA") assert protein.sequence is not None sequences = [protein.sequence, protein.sequence] # ================================================================ # Example usage: directly use the model # ================================================================ - input_ids = model._tokenize(sequences) - output = model(input_ids) + # EsmcModel / EsmcForMaskedLM are the canonical raw ESMC models. They expose + # a Hugging Face style ``forward`` (not the SDK ``encode``/``logits`` + # inference API - use ``esmc_client()`` for that). Tokenize with the ESMC + # tokenizer and pass ``input_ids`` directly. + tokenizer = EsmcTokenizer() + encoded = tokenizer(sequences, return_tensors="pt", padding=True) + output = model( + input_ids=encoded["input_ids"], + attention_mask=encoded["attention_mask"], + output_hidden_states=True, + ) logits, embeddings, hiddens = ( - output.sequence_logits, - output.embeddings, + output.logits, + output.last_hidden_state, output.hidden_states, ) print( @@ -132,11 +141,12 @@ def _get_logits(client: ESMCForgeInferenceClient, sequence: str) -> LogitsOutput print("ESM_API_KEY found. Trying to use model from Forge/Biohub Platform...") main(esmc_client(model="esmc-300m-2024-12")) else: - print("No ESM_API_KEY found. Trying to load model locally...") + print("No ESM_API_KEY found. Trying to load the model locally...") print( - "To try this script with a Forge/Biohub Platform API, please run ESM_API_KEY=your_api_key python esm3.py" + "To use the SDK inference API (encode/logits), run " + "ESM_API_KEY=your_api_key python esmc.py" ) - main(ESMC.from_pretrained("esm3_sm_open_v1")) - model = ESMC.from_pretrained("esmc_300m") - main(model) + # The local raw model uses the Hugging Face style forward API rather + # than the SDK inference client. + model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-300M") raw_forward(model) diff --git a/esm/models/esmc.py b/esm/models/esmc.py deleted file mode 100644 index 10a8ca7e..00000000 --- a/esm/models/esmc.py +++ /dev/null @@ -1,246 +0,0 @@ -from __future__ import annotations - -import contextlib - -import attr -import torch -import torch.nn as nn -from attr import dataclass - -try: - from flash_attn.bert_padding import pad_input, unpad_input - - is_flash_attn_available = True -except ImportError: - pad_input = None # ty:ignore[invalid-assignment] - unpad_input = None # ty:ignore[invalid-assignment] - is_flash_attn_available = False - -from esm.layers.regression_head import RegressionHead -from esm.layers.transformer_stack import TransformerStack -from esm.sdk.api import ( - ESMCInferenceClient, - ESMProtein, - ESMProteinTensor, - ForwardTrackData, - LogitsConfig, - LogitsOutput, -) -from esm.tokenization import EsmSequenceTokenizer -from esm.utils import encoding -from esm.utils.constants.models import ESMC_600M -from esm.utils.decoding import decode_sequence -from esm.utils.misc import stack_variable_length_tensors -from esm.utils.sampling import _BatchedESMProteinTensor - - -@dataclass -class ESMCOutput: - sequence_logits: torch.Tensor - embeddings: torch.Tensor | None - hidden_states: torch.Tensor | None - attentions: tuple[torch.FloatTensor, ...] | None = None - - -class ESMC(nn.Module, ESMCInferenceClient): - """ - ESMC model implementation. - - Args: - d_model (int): The dimensionality of the input and output feature vectors. - n_heads (int): The number of attention heads in the transformer layers. - n_layers (int): The number of transformer layers. - """ - - def __init__( - self, - d_model: int, - n_heads: int, - n_layers: int, - tokenizer: EsmSequenceTokenizer, - use_flash_attn: bool = True, - ): - super().__init__() - self.embed = nn.Embedding(64, d_model) - - self._use_flash_attn = is_flash_attn_available and use_flash_attn - self.transformer = TransformerStack( - d_model, - n_heads, - None, - n_layers, - n_layers_geom=0, - use_flash_attn=self._use_flash_attn, - ) - - self.sequence_head = RegressionHead(d_model, 64) - self.tokenizer = tokenizer - - @classmethod - def from_pretrained( - cls, - model_name: str = ESMC_600M, - device: torch.device | None = None, - use_flash_attn: bool = True, - ) -> ESMC: - from esm.pretrained import load_local_model - - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = load_local_model( - model_name, device=device, use_flash_attn=use_flash_attn - ) - if device.type != "cpu": - model = model.to(torch.bfloat16) - assert isinstance(model, ESMC) - return model - - @property - def device(self): - return next(self.parameters()).device - - @property - def raw_model(self): - return self - - def _tokenize(self, sequence: list[str]) -> torch.Tensor: - pad = self.tokenizer.pad_token_id - assert pad is not None - return stack_variable_length_tensors( - [ - encoding.tokenize_sequence(x, self.tokenizer, add_special_tokens=True) - for x in sequence - ], - constant_value=pad, - ).to(next(self.parameters()).device) - - def _detokenize(self, sequence: torch.Tensor) -> list[str]: - pad = self.tokenizer.pad_token_id - assert pad is not None - assert sequence.ndim == 2 - return [decode_sequence(x[x != pad][1:-1], self.tokenizer) for x in sequence] - - def forward( - self, - sequence_tokens: torch.Tensor | None = None, - sequence_id: torch.Tensor | None = None, - output_attentions: bool | None = None, - ) -> ESMCOutput: - """ - Performs forward pass through the ESMC model. Check utils to see how to tokenize inputs from raw data. - - Args: - sequence_tokens (torch.Tensor, optional): The amino acid tokens. - sequence_id (torch.Tensor, optional): The sequence ID. - output_attentions (bool, optional): Whether to return per-layer attention weights. - - Returns: - ESMCOutput: The output of the ESMC model. - - """ - output_attentions = bool(output_attentions) - if sequence_id is None: - # For EMSC, a boolean mask is created in place of sequence_id if not specified. - sequence_id = ( - sequence_tokens != self.tokenizer.pad_token_id - ) # ty:ignore[invalid-assignment] - - x = self.embed(sequence_tokens) - - B, L = x.shape[:2] - - # If sequence_id looks like a mask. - if self._use_flash_attn: - if output_attentions: - raise ValueError( - "output_attentions is not supported with flash attention." - ) - assert ( - sequence_id.dtype == torch.bool # ty:ignore[unresolved-attribute] - ), "sequence_id must be a boolean mask if Flash Attention is used" - assert sequence_id.shape == (B, L) # ty:ignore[unresolved-attribute] - assert unpad_input is not None - x, indices, *_ = unpad_input(x, sequence_id) - else: - indices = None - - x, _, hidden_states, attentions = self.transformer( - x, sequence_id=sequence_id, output_attentions=output_attentions - ) - - if self._use_flash_attn: - assert indices is not None - assert pad_input is not None - x = pad_input(x, indices, B, L) # Back to [B, L, D] - hidden_states = [ - # Back to [[B, L, D], ...] - pad_input(h, indices, B, L) - for h in hidden_states - ] - - # Stack hidden states into a [n_layers, B, L, D] matrix. - hidden_states = torch.stack(hidden_states, dim=0) - - sequence_logits = self.sequence_head(x) - output = ESMCOutput( - sequence_logits=sequence_logits, - embeddings=x, - hidden_states=hidden_states, - attentions=attentions, - ) - return output - - def encode(self, input: ESMProtein) -> ESMProteinTensor: - input = attr.evolve(input) # Make a copy - sequence_tokens = None - - if input.sequence is not None: - sequence_tokens = self._tokenize([input.sequence])[0] - return ESMProteinTensor( - sequence=sequence_tokens, - potential_sequence_of_concern=input.potential_sequence_of_concern, - ).to(next(self.parameters()).device) - - def decode(self, input: ESMProteinTensor) -> ESMProtein: - input = attr.evolve(input) # Make a copy - - assert input.sequence is not None - seq = input.sequence - if seq.ndim == 1: - seq = seq.unsqueeze(0) - sequence = self._detokenize(seq)[0] - - return ESMProtein(sequence=sequence) - - def logits( - self, - input: ESMProteinTensor | _BatchedESMProteinTensor, - config: LogitsConfig = LogitsConfig(), - ) -> LogitsOutput: - if not isinstance(input, _BatchedESMProteinTensor): - # Create batch dimension if necessary. - input = _BatchedESMProteinTensor.from_protein_tensor(input) - - device = torch.device(input.device) - - with ( - torch.no_grad(), - torch.autocast(enabled=True, device_type=device.type, dtype=torch.bfloat16) - if device.type == "cuda" - else contextlib.nullcontext(), - ): - output = self(sequence_tokens=input.sequence) - assert output.hidden_states is not None - output.hidden_states = ( - output.hidden_states[config.ith_hidden_layer : config.ith_hidden_layer + 1] - if config.ith_hidden_layer != -1 - else output.hidden_states - ) - - return LogitsOutput( - logits=ForwardTrackData( - sequence=output.sequence_logits if config.sequence else None - ), - embeddings=output.embeddings if config.return_embeddings else None, - hidden_states=output.hidden_states if config.return_hidden_states else None, - ) diff --git a/esm/models/esmc/__init__.py b/esm/models/esmc/__init__.py new file mode 100644 index 00000000..5a137dad --- /dev/null +++ b/esm/models/esmc/__init__.py @@ -0,0 +1,43 @@ +from esm.models.esmc.compatibility import ESMC, ESMCOutput +from esm.models.esmc.config import EsmcConfig +from esm.models.esmc.model import ( + EsmcForMaskedLM, + EsmcForSequenceClassification, + EsmcForTokenClassification, + EsmcMaskedLMOutput, + EsmcModel, + EsmcOutput, + EsmcPreTrainedModel, + EsmcSequenceClassifierOutput, + EsmcTokenClassifierOutput, +) +from esm.models.esmc.sae import ( + EsmcSaeConfig, + EsmcSaeLayer, + EsmcSaeModel, + EsmcSaeOutput, + EsmcSaeParams, +) +from esm.models.esmc.tokenizer import EsmcTokenizer + +__all__ = [ + # Deprecated: kept so pre-rename imports keep working. See compatibility.py. + "ESMC", + "ESMCOutput", + "EsmcConfig", + "EsmcForMaskedLM", + "EsmcForSequenceClassification", + "EsmcForTokenClassification", + "EsmcMaskedLMOutput", + "EsmcModel", + "EsmcOutput", + "EsmcPreTrainedModel", + "EsmcSaeConfig", + "EsmcSaeLayer", + "EsmcSaeModel", + "EsmcSaeOutput", + "EsmcSaeParams", + "EsmcSequenceClassifierOutput", + "EsmcTokenClassifierOutput", + "EsmcTokenizer", +] diff --git a/esm/models/esmc/checkpoint_layout.py b/esm/models/esmc/checkpoint_layout.py new file mode 100644 index 00000000..ced4bcaf --- /dev/null +++ b/esm/models/esmc/checkpoint_layout.py @@ -0,0 +1,180 @@ +"""Translation between the published ESMC checkpoint layout and this +implementation's in-memory layout. + +Published checkpoints use one projection per tensor — ``q_proj``, ``k_proj``, +``v_proj``, ``gate_proj``, ``up_proj`` — which is the layout ``transformers`` +reads and the only one its tensor-parallel plan can shard. + +This implementation instead keeps the projections fused, because +``transformer_engine``'s ``LayerNormLinear`` / ``LayerNormMLP`` fuse each +LayerNorm into the GEMM that follows it and take a single packed weight. That +fusion is what the fp8 path runs on, so the module tree keeps Transformer +Engine's parameter names and the two layouts are reconciled here instead. + +The conversion is exact in both directions: a fused weight is the row-wise +concatenation of its parts, so ``cat``/``chunk`` along dim 0 round-trips +bit-for-bit and no tensor is reshaped, transposed or reordered. + +The two key vocabularies are disjoint, so both functions convert what they +recognise and pass everything else through untouched. That keeps checkpoints +already saved in either layout loadable. +""" + +import re + +import torch + +_ENCODER_PREFIX = "esmc." + +# Published name -> in-memory name, for keys directly under the encoder root. +_ROOT = { + "embed_tokens.weight": "embed.weight", + "norm.weight": "transformer.norm.weight", +} + +# The masked-LM head is an ``nn.Sequential`` here, so its children are indices. +_LM_HEAD = { + "lm_head.dense.weight": "lm_head.0.weight", + "lm_head.dense.bias": "lm_head.0.bias", + "lm_head.layer_norm.weight": "lm_head.2.weight", + "lm_head.layer_norm.bias": "lm_head.2.bias", + "lm_head.decoder.weight": "lm_head.3.weight", + "lm_head.decoder.bias": "lm_head.3.bias", +} + +# Published name -> in-memory name, within one transformer block. The two +# LayerNorms live inside the fused modules they feed rather than beside them. +_BLOCK = { + "input_layernorm.weight": "attn.layernorm_qkv.layer_norm_weight", + "input_layernorm.bias": "attn.layernorm_qkv.layer_norm_bias", + "self_attn.q_norm.weight": "attn.q_ln.weight", + "self_attn.k_norm.weight": "attn.k_ln.weight", + "self_attn.o_proj.weight": "attn.out_proj.weight", + "post_attention_layernorm.weight": "ffn.layer_norm_weight", + "post_attention_layernorm.bias": "ffn.layer_norm_bias", + "mlp.down_proj.weight": "ffn.fc2_weight", +} + +# In-memory fused weight -> the published parts it concatenates, in row order. +# ``EsmcLayerNormMLP.forward`` splits fc1 as ``x1, x2 = chunk(2)`` and applies +# ``silu(x1) * x2``, so the gate rows come first; attention slices q, k, v in +# that order. +_FUSED = { + "attn.layernorm_qkv.weight": ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + ), + "ffn.fc1_weight": ("mlp.gate_proj.weight", "mlp.up_proj.weight"), +} + +_PUBLISHED_BLOCK = re.compile(r"^layers\.(\d+)\.(.+)$") +_NATIVE_BLOCK = re.compile(r"^transformer\.blocks\.(\d+)\.(.+)$") + + +def _split_encoder_prefix(key: str) -> tuple[str, str]: + """Return ``(prefix, remainder)``; the prefix is empty for a bare encoder.""" + if key.startswith(_ENCODER_PREFIX): + return _ENCODER_PREFIX, key[len(_ENCODER_PREFIX) :] + return "", key + + +def published_to_native_subtree( + raw: dict[str, torch.Tensor], prefix: str +) -> dict[str, torch.Tensor]: + """Translate only the keys under ``prefix``, leaving the rest untouched. + + Used when an ESMC encoder is bundled inside a larger checkpoint, so the host + model's own keys can never be caught by the ESMC key patterns. + """ + inner = {k[len(prefix) :]: v for k, v in raw.items() if k.startswith(prefix)} + out = {k: v for k, v in raw.items() if not k.startswith(prefix)} + out.update({prefix + k: v for k, v in published_to_native(inner).items()}) + return out + + +def published_to_native(raw: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Rewrite a published checkpoint onto this implementation's layout. + + Concatenates the separate q/k/v and gate/up projections into the fused + weights Transformer Engine expects. Keys already in the in-memory layout, + and keys belonging to neither vocabulary, are returned unchanged. + """ + out: dict[str, torch.Tensor] = {} + # (prefix, block index) -> {published leaf: tensor}, gathered so the parts of + # each fused weight can be concatenated once all of them have been seen. + parts: dict[tuple[str, str], dict[str, torch.Tensor]] = {} + part_leaves = {leaf for leaves in _FUSED.values() for leaf in leaves} + + for key, value in raw.items(): + prefix, rest = _split_encoder_prefix(key) + if rest in _ROOT: + out[prefix + _ROOT[rest]] = value + continue + if key in _LM_HEAD: + out[_LM_HEAD[key]] = value + continue + match = _PUBLISHED_BLOCK.match(rest) + if match is None: + out[key] = value + continue + index, leaf = match.group(1), match.group(2) + block = f"{prefix}transformer.blocks.{index}." + if leaf in _BLOCK: + out[block + _BLOCK[leaf]] = value + elif leaf in part_leaves: + parts.setdefault((prefix, index), {})[leaf] = value + else: + out[key] = value + + for (prefix, index), found in parts.items(): + block = f"{prefix}transformer.blocks.{index}." + for fused, leaves in _FUSED.items(): + present = [leaf for leaf in leaves if leaf in found] + if not present: + continue + if len(present) != len(leaves): + raise ValueError( + f"block {index} has an incomplete set of projections for " + f"{fused!r}: found {present}, need {list(leaves)}" + ) + out[block + fused] = torch.cat([found[leaf] for leaf in leaves], dim=0) + return out + + +def native_to_published(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Rewrite this implementation's state dict onto the published layout. + + The exact inverse of :func:`published_to_native`: splits each fused weight + back into its parts along dim 0. + """ + root = {native: pub for pub, native in _ROOT.items()} + lm_head = {native: pub for pub, native in _LM_HEAD.items()} + block_leaves = {native: pub for pub, native in _BLOCK.items()} + + out: dict[str, torch.Tensor] = {} + for key, value in state.items(): + prefix, rest = _split_encoder_prefix(key) + if rest in root: + out[prefix + root[rest]] = value + continue + if key in lm_head: + out[lm_head[key]] = value + continue + match = _NATIVE_BLOCK.match(rest) + if match is None: + out[key] = value + continue + index, leaf = match.group(1), match.group(2) + block = f"{prefix}layers.{index}." + if leaf in block_leaves: + out[block + block_leaves[leaf]] = value + elif leaf in _FUSED: + leaves = _FUSED[leaf] + # .clone(): safetensors refuses to write tensors that share storage. + chunks = torch.chunk(value, len(leaves), dim=0) + for name, chunk in zip(leaves, chunks, strict=True): + out[block + name] = chunk.clone() + else: + out[key] = value + return out diff --git a/esm/models/esmc/compatibility.py b/esm/models/esmc/compatibility.py new file mode 100644 index 00000000..e64d840d --- /dev/null +++ b/esm/models/esmc/compatibility.py @@ -0,0 +1,274 @@ +"""Deprecated ``ESMC`` surface, kept so existing code keeps running. + +``ESMC`` used to be one class holding an encoder, an LM head, a tokenizer and a +local implementation of ``ESMCInferenceClient``. Those are now four separate +things (:class:`EsmcModel`, :class:`EsmcForMaskedLM`, :class:`EsmcTokenizer` and +``esmc_client``). :class:`ESMC` here wraps a native +:class:`EsmcForMaskedLM` - the piece that owns the LM head, so +``sequence_logits`` is still available - and re-exposes the old method and +attribute names. + +``hidden_states`` is realigned back onto the old layout - one entry per block +output, before the final LayerNorm - because the native model follows the +``transformers`` convention instead; see :func:`_legacy_hidden_states`. + +""" + +import warnings +from dataclasses import dataclass + +import torch +import torch.nn as nn +from transformers import PreTrainedTokenizerFast + +from esm.models.esmc.config import ( + ESMC_6B_HF_REPO, + ESMC_300M_HF_REPO, + ESMC_600M_HF_REPO, + EsmcConfig, +) +from esm.models.esmc.model import EsmcForMaskedLM +from esm.models.esmc.tokenizer import EsmcTokenizer +from esm.sdk.api import ( + ESMProtein, + ESMProteinTensor, + ForwardTrackData, + LogitsConfig, + LogitsOutput, +) +from esm.utils.constants.models import ESMC_6B, ESMC_300M, ESMC_600M +from esm.utils.sampling import _BatchedESMProteinTensor + +_DEPRECATION_MESSAGE = ( + "ESMC is deprecated and now wraps the native EsmcForMaskedLM. Use " + "EsmcModel / EsmcForMaskedLM with EsmcTokenizer directly, or esmc_client(...) " + "for the remote inference API." +) + + +@dataclass +class ESMCOutput: + """Deprecated. Superseded by ``EsmcMaskedLMOutput``.""" + + sequence_logits: torch.Tensor + embeddings: torch.Tensor | None + hidden_states: torch.Tensor | None + attentions: tuple[torch.Tensor, ...] | None = None + + @property + def logits(self) -> torch.Tensor: + """Alias, so code written against ``EsmcMaskedLMOutput`` also works.""" + return self.sequence_logits + + def __getitem__(self, key: str): + """Mapping access, so ``out["logits"]`` works as it does on the native + output.""" + return self.logits if key == "logits" else getattr(self, key) + + +def _legacy_hidden_states( + stacked: torch.Tensor | None, last_prenorm: torch.Tensor | None +) -> torch.Tensor | None: + """Rebuild the old ``hidden_states`` stack from the native one. + + The old stack held one entry per block *output*, before the final + LayerNorm: ``(out_0, ..., out_{n-1})``, so ``hidden_states[i]`` was block + ``i``'s output. The native stack instead follows the ``transformers`` + convention of one entry per block *input* plus the normalised final state: + ``(embeddings, out_0, ..., out_{n-2}, norm(out_{n-1}))``. + + Dropping the leading embeddings realigns the indices, and the final + pre-LayerNorm state replaces the normalised one -- LayerNorm is not + invertible, so it has to be carried out of the model rather than recovered. + """ + if stacked is None: + return None + if last_prenorm is None: + raise ValueError( + "the native model returned hidden_states without " + "last_hidden_state_prenorm, so the legacy stack cannot be rebuilt" + ) + return torch.cat([stacked[1:-1], last_prenorm[None]], dim=0) + + +def _legacy_name_to_repo(model_name: str) -> str: + """Map a pre-rename local model name onto its HuggingFace repo id. + + Anything unrecognised is passed through, so a repo id works directly. + """ + return { + ESMC_300M: ESMC_300M_HF_REPO, + ESMC_600M: ESMC_600M_HF_REPO, + ESMC_6B: ESMC_6B_HF_REPO, + }.get(model_name, model_name) + + +class ESMC(nn.Module): + """Deprecated wrapper around :class:`EsmcForMaskedLM`. + + Accepts the old constructor arguments, or wraps an already-built native + model via ``model=``. + """ + + def __init__( + self, + d_model: int | None = None, + n_heads: int | None = None, + n_layers: int | None = None, + # Pre-port callers passed an EsmSequenceTokenizer, so accept any fast + # tokenizer: only __call__, pad_token_id and decode are used. + tokenizer: PreTrainedTokenizerFast | None = None, + use_flash_attn: bool = True, + *, + model: EsmcForMaskedLM | None = None, + ): + super().__init__() + warnings.warn(_DEPRECATION_MESSAGE, DeprecationWarning, stacklevel=2) + + if model is None: + defaults = EsmcConfig() + model = EsmcForMaskedLM( + EsmcConfig( + hidden_size=defaults.hidden_size if d_model is None else d_model, + num_attention_heads=defaults.num_attention_heads + if n_heads is None + else n_heads, + num_hidden_layers=defaults.num_hidden_layers + if n_layers is None + else n_layers, + attn_implementation="flash_attention_2" + if use_flash_attn + else "sdpa", + ) + ) + self.model = model + self.tokenizer = EsmcTokenizer() if tokenizer is None else tokenizer + + @classmethod + def from_pretrained( + cls, + model_name: str = ESMC_600M, + device: torch.device | str | None = None, + use_flash_attn: bool = True, + ) -> "ESMC": + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = torch.device(device) + model = EsmcForMaskedLM.from_pretrained( + _legacy_name_to_repo(model_name), + device=device, + # The old loader cast to bf16 whenever it left the CPU. + dtype=torch.bfloat16 if device.type != "cpu" else None, + attn_implementation="flash_attention_2" if use_flash_attn else "sdpa", + ) + return cls(model=model) + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + @property + def raw_model(self) -> "ESMC": + return self + + def state_dict(self, *args, **kwargs): + """Delegate so checkpoint keys match the unwrapped native model.""" + return self.model.state_dict(*args, **kwargs) + + def load_state_dict(self, state_dict, *args, **kwargs): + return self.model.load_state_dict(state_dict, *args, **kwargs) + + def _tokenize(self, sequence: list[str]) -> torch.Tensor: + encoded = self.tokenizer(sequence, return_tensors="pt", padding=True) + return encoded["input_ids"].to(self.device) + + def _detokenize(self, sequence: torch.Tensor) -> list[str]: + pad = self.tokenizer.pad_token_id + assert pad is not None + assert sequence.ndim == 2 + return [ + self.tokenizer.decode(row[row != pad][1:-1]).replace(" ", "") + for row in sequence + ] + + def forward( + self, + sequence_tokens: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + output_attentions: bool | None = None, + *, + input_ids: torch.Tensor | None = None, + ) -> ESMCOutput: + # ``input_ids`` is accepted so a caller written against the native model + # works here too, and no caller needs to branch on the model class. + if input_ids is not None: + if sequence_tokens is not None: + raise ValueError("pass either sequence_tokens or input_ids, not both") + sequence_tokens = input_ids + + # The old ``sequence_id`` defaulted to a boolean padding mask, which is + # what ``attention_mask`` means now; a non-bool tensor is the integer + # chain id the native model still calls ``sequence_id``. + attention_mask = None + if sequence_id is not None and sequence_id.dtype == torch.bool: + attention_mask, sequence_id = sequence_id, None + + out = self.model( + input_ids=sequence_tokens, + attention_mask=attention_mask, + sequence_id=sequence_id, + # The old forward always stacked and returned hidden states. + output_hidden_states=True, + output_attentions=bool(output_attentions), + return_dict=True, + ) + return ESMCOutput( + sequence_logits=out.logits, + embeddings=out.last_hidden_state, + hidden_states=_legacy_hidden_states( + out.hidden_states, out.last_hidden_state_prenorm + ), + attentions=out.attentions, + ) + + def encode(self, input: ESMProtein) -> ESMProteinTensor: + tokens = None if input.sequence is None else self._tokenize([input.sequence])[0] + return ESMProteinTensor( + sequence=tokens, + potential_sequence_of_concern=input.potential_sequence_of_concern, + ).to(self.device) + + def decode(self, input: ESMProteinTensor) -> ESMProtein: + tokens = input.sequence + assert tokens is not None + batched = tokens[None] if tokens.ndim == 1 else tokens + return ESMProtein(sequence=self._detokenize(batched)[0]) + + def logits( + self, + input: ESMProteinTensor | _BatchedESMProteinTensor, + config: LogitsConfig = LogitsConfig(), + ) -> LogitsOutput: + if not isinstance(input, _BatchedESMProteinTensor): + input = _BatchedESMProteinTensor.from_protein_tensor(input) + device = torch.device(input.device) + + with ( + torch.no_grad(), + torch.autocast( + device.type, dtype=torch.bfloat16, enabled=device.type == "cuda" + ), + ): + out = self(sequence_tokens=input.sequence) + + hidden = out.hidden_states + if hidden is not None and config.ith_hidden_layer != -1: + layer = config.ith_hidden_layer + hidden = hidden[layer : layer + 1] + return LogitsOutput( + logits=ForwardTrackData( + sequence=out.sequence_logits if config.sequence else None + ), + embeddings=out.embeddings if config.return_embeddings else None, + hidden_states=hidden if config.return_hidden_states else None, + ) diff --git a/esm/models/esmc/compatibility_test.py b/esm/models/esmc/compatibility_test.py new file mode 100644 index 00000000..b5becd46 --- /dev/null +++ b/esm/models/esmc/compatibility_test.py @@ -0,0 +1,162 @@ +"""Pins the deprecated ``ESMC`` surface so the compatibility wrapper can't rot.""" + +import warnings + +import pytest +import torch + +from esm.models.esmc import ESMC, ESMCOutput +from esm.models.esmc.compatibility import _legacy_name_to_repo +from esm.models.esmc.model import EsmcForMaskedLM +from esm.utils.constants.models import ESMC_6B, ESMC_300M, ESMC_600M + +D_MODEL, N_HEADS, N_LAYERS = 32, 4, 2 +SEQUENCES = ["AAAA", "MKV"] + + +def make_legacy_model() -> ESMC: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + return ESMC( + d_model=D_MODEL, n_heads=N_HEADS, n_layers=N_LAYERS, use_flash_attn=False + ).eval() + + +def test_legacy_constructor_warns_and_wraps_native_model(): + with pytest.warns(DeprecationWarning): + model = ESMC( + d_model=D_MODEL, n_heads=N_HEADS, n_layers=N_LAYERS, use_flash_attn=False + ) + assert isinstance(model.model, EsmcForMaskedLM) + + +def test_legacy_forward_returns_old_field_names(): + model = make_legacy_model() + tokens = model._tokenize(SEQUENCES) + + with torch.no_grad(): + out = model(sequence_tokens=tokens) + + assert isinstance(out, ESMCOutput) + b, length = tokens.shape + assert out.sequence_logits.shape == (b, length, 64) + assert out.embeddings is not None + assert out.embeddings.shape == (b, length, D_MODEL) + # One entry per block output, as before the port. + assert out.hidden_states is not None + assert out.hidden_states.shape == (N_LAYERS, b, length, D_MODEL) + + +def test_tokenize_detokenize_round_trip(): + model = make_legacy_model() + assert model._detokenize(model._tokenize(SEQUENCES)) == SEQUENCES + + +def test_boolean_sequence_id_is_treated_as_padding_mask(): + """The old ``sequence_id`` was a bool mask; that is ``attention_mask`` now. + + Passing it straight through as the native ``sequence_id`` would silently + mark every position valid, since ``False >= 0`` is true. + """ + model = make_legacy_model() + tokens = model._tokenize(SEQUENCES) + mask = tokens != model.tokenizer.pad_token_id + + with torch.no_grad(): + via_legacy = model(sequence_tokens=tokens, sequence_id=mask) + via_native = model.model( + input_ids=tokens, attention_mask=mask, output_hidden_states=True + ) + + assert via_legacy.embeddings is not None + torch.testing.assert_close(via_legacy.embeddings, via_native.last_hidden_state) + + +def test_output_attentions_is_forwarded(): + model = make_legacy_model() + tokens = model._tokenize(SEQUENCES) + + with torch.no_grad(): + assert model(sequence_tokens=tokens).attentions is None + attns = model(sequence_tokens=tokens, output_attentions=True).attentions + + assert attns is not None + assert len(attns) == N_LAYERS + + +def test_state_dict_keys_match_the_unwrapped_model(): + model = make_legacy_model() + assert list(model.state_dict()) == list(model.model.state_dict()) + model.load_state_dict(model.state_dict()) + + +def test_module_plumbing_reaches_the_wrapped_model(): + model = make_legacy_model() + assert sum(p.numel() for p in model.parameters()) > 0 + assert model.device == torch.device("cpu") + assert model.raw_model is model + + +@pytest.mark.parametrize("name", [ESMC_300M, ESMC_600M, ESMC_6B]) +def test_legacy_model_names_map_to_hub_repos(name: str): + assert _legacy_name_to_repo(name).startswith("biohub/ESMC-") + + +def test_unknown_name_passes_through_so_repo_ids_still_work(): + assert _legacy_name_to_repo("biohub/ESMC-300M") == "biohub/ESMC-300M" + + +def test_legacy_hidden_states_match_the_old_block_output_stack(): + """``main``'s stack held one entry per block *output*, before the final norm. + + The native stack follows the ``transformers`` convention instead - one entry + per block *input* plus the normalised final state - so the wrapper has to + realign it or every ``hidden_states[i]`` silently returns a different layer. + """ + model = make_legacy_model() + tokens = model._tokenize(SEQUENCES) + with torch.no_grad(): + legacy = model(sequence_tokens=tokens) + native = model.model( + input_ids=tokens, output_hidden_states=True, return_dict=True + ) + + assert legacy.hidden_states is not None and native.hidden_states is not None + assert legacy.hidden_states.shape[0] == N_LAYERS + assert native.hidden_states.shape[0] == N_LAYERS + 1 + # Index alignment: block i's output is native[i + 1]. + for i in range(N_LAYERS - 1): + assert torch.equal(legacy.hidden_states[i], native.hidden_states[i + 1]) + # The last entry is pre-LayerNorm, so it must differ from the normed output. + assert not torch.allclose(legacy.hidden_states[-1], legacy.embeddings) + + +def test_legacy_output_supports_native_access_patterns(): + model = make_legacy_model() + tokens = model._tokenize(SEQUENCES) + with torch.no_grad(): + out = model(sequence_tokens=tokens) + assert torch.equal(out.logits, out.sequence_logits) + assert torch.equal(out["logits"], out.sequence_logits) + + +def test_legacy_forward_accepts_input_ids(): + """So callers need not branch on the model class to pick a calling style.""" + model = make_legacy_model() + tokens = model._tokenize(SEQUENCES) + with torch.no_grad(): + by_alias = model(input_ids=tokens) + by_legacy_name = model(sequence_tokens=tokens) + assert torch.equal(by_alias["logits"], by_legacy_name.sequence_logits) + with pytest.raises(ValueError, match="not both"): + model(sequence_tokens=tokens, input_ids=tokens) + + +def test_retired_date_stamped_factories_are_still_importable(): + """They shipped publicly as ``esm.pretrained.ESMC_*_202412``.""" + from esm import pretrained + + for name in ("ESMC_300M_202412", "ESMC_600M_202412", "ESMC_6B_202412"): + assert callable(getattr(pretrained, name)) + for model_name in (ESMC_300M, ESMC_600M, ESMC_6B): + assert model_name in pretrained.LOCAL_MODEL_REGISTRY diff --git a/esm/models/esmc/config.py b/esm/models/esmc/config.py new file mode 100644 index 00000000..3a545928 --- /dev/null +++ b/esm/models/esmc/config.py @@ -0,0 +1,217 @@ +"""Configuration for the ESMC models.""" + +import json +import os +import warnings +from dataclasses import asdict, dataclass +from pathlib import Path + +CONFIG_NAME = "config.json" + +# Published ESMC checkpoints on the HuggingFace Hub. +ESMC_300M_HF_REPO = "biohub/ESMC-300M" +ESMC_600M_HF_REPO = "biohub/ESMC-600M" +ESMC_6B_HF_REPO = "biohub/ESMC-6B" + +# Every published ESMC uses this SwiGLU expansion. It stays a constant rather +# than a config field because ``intermediate_size`` is what actually fixes the +# shapes; this only derives a default for checkpoints that omit it. +ESMC_EXPANSION_RATIO = 8 / 3 + +# Field names as spelled by checkpoints published before the field names were +# aligned with ``transformers``. Delete once every ``biohub/ESMC-*`` repo has +# been re-published; until then a strict read would reject every live checkpoint. +_LEGACY_KEYS = { + "d_model": "hidden_size", + "n_heads": "num_attention_heads", + "n_layers": "num_hidden_layers", +} + + +def esmc_intermediate_size(hidden_size: int) -> int: + """Round the SwiGLU hidden dim up to a multiple of 256 after expansion.""" + return int(((ESMC_EXPANSION_RATIO * hidden_size) + 255) // 256 * 256) + + +def _resolve_legacy_keys(raw: dict) -> dict: + """Map pre-alignment field names onto their current spelling.""" + stale = {old: new for old, new in _LEGACY_KEYS.items() if old in raw} + if not stale: + return raw + resolved = dict(raw) + for old, new in stale.items(): + resolved.setdefault(new, resolved.pop(old)) + warnings.warn( + f"{CONFIG_NAME} uses pre-alignment ESMC field names {sorted(stale)}; " + f"the current names are {sorted(stale.values())}.", + FutureWarning, + stacklevel=3, + ) + return resolved + + +@dataclass +class EsmcConfig: + """Configuration for the ESMC models. + + Field names match ``transformers``' ``EsmcConfig`` so that one ``config.json`` + describes the model to both implementations. + + Parameters + ---------- + vocab_size : int + Number of amino acid tokens representable by ``input_ids``. + hidden_size : int + Dimensionality of the encoder layers. + num_attention_heads : int + Number of attention heads per layer. + num_hidden_layers : int + Number of transformer blocks. + intermediate_size : int | None + Width of the SwiGLU feed-forward hidden layer. Derived from + ``hidden_size`` via :func:`esmc_intermediate_size` when ``None``. + pad_token_id : int + Index of the padding token (``""``). + mask_token_id : int + Index of the mask token (``""``). + initializer_range : float + Standard deviation of the normal weight initialiser. + classifier_dropout : float + Dropout ratio for the classification heads. + num_labels : int + Number of labels for the classification heads. + problem_type : str | None + One of ``"regression"``, ``"single_label_classification"`` or + ``"multi_label_classification"``; inferred at loss time when ``None``. + output_hidden_states : bool + Default value for ``output_hidden_states`` at forward time. + output_attentions : bool + Default value for ``output_attentions`` at forward time. + return_dict : bool + Default value for ``return_dict`` at forward time. + attn_implementation : str + Attention backend: ``"sdpa"``, ``"eager"`` or ``"flash_attention_2"``. + A property of the run rather than of the weights, so it is honoured when + a ``config.json`` carries it but never written back — matching + ``transformers``, which keeps it off the serialised config entirely. + """ + + model_type = "esmc" + + vocab_size: int = 64 + hidden_size: int = 2560 + num_attention_heads: int = 40 + num_hidden_layers: int = 80 + intermediate_size: int | None = None + pad_token_id: int = 1 + mask_token_id: int = 32 + initializer_range: float = 0.02 + classifier_dropout: float = 0.1 + num_labels: int = 2 + problem_type: str | None = None + output_hidden_states: bool = False + output_attentions: bool = False + return_dict: bool = True + attn_implementation: str = "sdpa" + + def __post_init__(self) -> None: + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"hidden_size ({self.hidden_size}) is not a multiple of " + f"num_attention_heads ({self.num_attention_heads})." + ) + derived = esmc_intermediate_size(self.hidden_size) + if self.intermediate_size is None: + self.intermediate_size = derived + elif self.intermediate_size != derived: + # The feed-forward width is not a free parameter here: the layers + # derive it from hidden_size, so honouring a different value would + # need a different FFN. Refuse rather than build the wrong shapes. + raise ValueError( + f"intermediate_size ({self.intermediate_size}) does not match the " + f"value ESMC derives from hidden_size ({derived})." + ) + + @property + def use_return_dict(self) -> bool: + return self.return_dict + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_attention_heads + + @classmethod + def from_dict(cls, raw: dict) -> "EsmcConfig": + """Build a config from a parsed ``config.json``. + + Keys this class does not read are ignored. ``transformers`` serialises + knobs ESMC never varies — ``expansion_ratio``, ``scale_residue``, + ``hidden_act``, ``rope_parameters``, ``attention_bias``, ``mlp_bias``, + ``head_dim``, ``num_key_value_heads``, ``max_position_embeddings``, the + parallelism plans — and this implementation fixes all of them. + + Raises + ------ + KeyError + If a field the architecture depends on is absent. Substituting a + default there would silently build a different model and load + garbage into it. + """ + raw = _resolve_legacy_keys(raw) + + def required(key: str): + if key not in raw: + raise KeyError( + f"{CONFIG_NAME} is missing required ESMC field {key!r}; " + f"present keys are {sorted(raw)}" + ) + return raw[key] + + if "num_labels" in raw: + num_labels = raw["num_labels"] + elif "id2label" in raw: + num_labels = len(raw["id2label"]) + else: + num_labels = 2 + # hidden_size=1 keeps __post_init__'s divisibility check happy; only the + # non-architectural defaults are read off this instance. + defaults = cls(hidden_size=1, num_attention_heads=1) + return cls( + vocab_size=required("vocab_size"), + hidden_size=required("hidden_size"), + num_attention_heads=required("num_attention_heads"), + num_hidden_layers=required("num_hidden_layers"), + intermediate_size=raw.get("intermediate_size"), + pad_token_id=required("pad_token_id"), + mask_token_id=required("mask_token_id"), + initializer_range=raw.get("initializer_range", defaults.initializer_range), + classifier_dropout=raw.get( + "classifier_dropout", defaults.classifier_dropout + ), + num_labels=num_labels, + problem_type=raw.get("problem_type"), + output_hidden_states=raw.get( + "output_hidden_states", defaults.output_hidden_states + ), + output_attentions=raw.get("output_attentions", defaults.output_attentions), + return_dict=raw.get("return_dict", defaults.return_dict), + attn_implementation=raw.get( + "attn_implementation", defaults.attn_implementation + ), + ) + + @classmethod + def from_pretrained(cls, directory: str | os.PathLike) -> "EsmcConfig": + with open(Path(directory) / CONFIG_NAME) as f: + return cls.from_dict(json.load(f)) + + def to_dict(self) -> dict: + raw = asdict(self) + del raw["attn_implementation"] + return {"model_type": self.model_type, **raw} + + def save_pretrained(self, save_directory: str | os.PathLike) -> None: + directory = Path(save_directory) + directory.mkdir(parents=True, exist_ok=True) + with open(directory / CONFIG_NAME, "w") as f: + json.dump(self.to_dict(), f, indent=2, sort_keys=True) diff --git a/esm/models/esmc/kernels.py b/esm/models/esmc/kernels.py new file mode 100644 index 00000000..4cf84097 --- /dev/null +++ b/esm/models/esmc/kernels.py @@ -0,0 +1,89 @@ +"""Optional accelerated kernels used by the ESMC layers. + +ESMC runs on pure PyTorch by default, but selects fused kernels when the +matching package is installed and the model is being built for CUDA. These flags +report only what is importable; the target device decides what is usable. + +The kernels, in the order ESMC prefers them: + +* ``transformer_engine`` — fused LayerNorm+Linear / LayerNorm+MLP with an + fp32 reduction inside the LayerNorm. Recommended on GPU for accurate bf16 + inference; without it the pure-PyTorch fallback drifts ~O(10) in fp32 and + ~O(100) in bf16 on the unnormalized residual stream (perplexity stays + within rounding noise). +* ``xformers`` — preferred fused attention kernel with a deterministic + reduction order. +* ``flash-attn`` — secondary fused attention kernel (fp16 / bf16 only), plus + the Triton RoPE kernel and the varlen packing helpers. +""" + +import logging + +import torch + +logger = logging.getLogger(__name__) + +try: + import transformer_engine.pytorch as te + + TE_INSTALLED = True +except ImportError: + te = None # ty:ignore[invalid-assignment] + TE_INSTALLED = False + +try: + import xformers.ops as xops + + XFORMERS_INSTALLED = True +except ImportError: + xops = None # ty:ignore[invalid-assignment] + XFORMERS_INSTALLED = False + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_qkvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + + FLASH_ATTN_INSTALLED = True +except ImportError: + flash_attn_func = None # ty:ignore[invalid-assignment] + flash_attn_varlen_qkvpacked_func = None # ty:ignore[invalid-assignment] + pad_input = unpad_input = None # ty:ignore[invalid-assignment] + FLASH_ATTN_INSTALLED = False + +try: + from flash_attn.ops.triton.rotary import apply_rotary as apply_triton_rotary + + FLASH_ATTN_ROTARY_INSTALLED = True +except ImportError: + apply_triton_rotary = None # ty:ignore[invalid-assignment] + FLASH_ATTN_ROTARY_INSTALLED = False + + +if not TE_INSTALLED: + logger.warning( + "ESMC: Transformer Engine is not installed; falling back to " + "pure-PyTorch LayerNorm+Linear / LayerNorm+MLP. Outputs will differ " + "numerically - measured on the unnormalized residual stream (before " + "the final LayerNorm), ~O(10) max-diff in fp32 and ~O(100) in bf16; " + "after the final LayerNorm these shrink to a few ULP and perplexity " + "stays within rounding noise. Install with " + "`pip install transformer-engine[pytorch]` to enable fused fp32-" + "reduction LayerNorm." + ) + +if not XFORMERS_INSTALLED and not FLASH_ATTN_INSTALLED: + logger.warning( + "ESMC: neither xformers nor flash-attn is installed; falling back " + "to PyTorch ``F.scaled_dot_product_attention``. The attention " + "reduction order in bf16 differs from a fused kernel by ~1 bf16 " + "ULP per attention block; compounded across the 80-block stack " + "this reaches ~O(100) max-diff on the unnormalized residual stream. " + "Install xformers (preferred) with `pip install xformers` for a " + "fused attention kernel." + ) + +if torch.cuda.is_available() and not FLASH_ATTN_ROTARY_INSTALLED: + logger.warning( + "ESMC: flash-attn rotary kernel not installed; falling back to " + "pure-PyTorch RoPE. For faster GPU inference run `pip install flash-attn`." + ) diff --git a/esm/models/esmc/layers.py b/esm/models/esmc/layers.py new file mode 100644 index 00000000..fa12dd44 --- /dev/null +++ b/esm/models/esmc/layers.py @@ -0,0 +1,753 @@ +"""ESMC-specific transformer layers. + +RoPE, QK-LayerNorm attention, SwiGLU feed-forward and the transformer stack used +by the ESMC encoder. Each fused module (Transformer Engine LayerNorm+Linear / +LayerNorm+MLP, xformers / flash-attn attention, Triton RoPE) has a pure-PyTorch +fallback with matching parameter names so state dicts are interchangeable. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from esm.layers.rotary import apply_rotary_emb_torch +from esm.models.esmc.kernels import ( + FLASH_ATTN_INSTALLED, + FLASH_ATTN_ROTARY_INSTALLED, + XFORMERS_INSTALLED, + apply_triton_rotary, + flash_attn_func, + flash_attn_varlen_qkvpacked_func, + te, + xops, +) + + +class EsmcRotaryEmbedding(nn.Module): + """Rotary position embeddings (RoPE) as used by ESMC. + + Behaviourally this matches :class:`RotaryEmbedding`, with two differences + that are load-bearing for ESMC's numerics: + + * The ``inv_freq`` buffer is recomputed on every device move and kept in + fp32 even when the module is cast to bf16/fp16. CPU and CUDA ``pow`` + differ by ~1 fp32 ULP, which compounds across the deep attention stack. + * The ``forward`` selects the Flash-Attention Triton RoPE kernel when it is + available and the tensors live on CUDA, otherwise the pure-PyTorch path. + + Parameters + ---------- + dim : int + Size of a single attention head. + base : float + Frequency base for the sinusoidal positions. + interleaved : bool + If ``True`` rotate adjacent pairs (GPT-J style) instead of splitting + the head dimension in half (GPT-NeoX style). + scale_base : float | None + Enables XPos scaling when set. ESMC does not use it; the ``forward`` + raises if it is configured. + scaling_factor : float + Linear scaling factor applied to position indices. + pos_idx_in_fp32 : bool + Compute position indices in fp32 to avoid bf16 rounding at long + sequence lengths. + """ + + def __init__( + self, + dim: int, + base: float = 10000.0, + interleaved: bool = False, + scale_base: float | None = None, + scaling_factor: float = 1.0, + pos_idx_in_fp32: bool = True, + device=None, + ): + super().__init__() + self.dim = dim + self.base = base + self.interleaved = interleaved + self.scale_base = scale_base + self.scaling_factor = scaling_factor + self.pos_idx_in_fp32 = pos_idx_in_fp32 + + self._seq_len_cached = 0 + self._cos_cached: torch.Tensor | None = None + self._sin_cached: torch.Tensor | None = None + self._cos_k_cached: torch.Tensor | None = None + self._sin_k_cached: torch.Tensor | None = None + + self.reset_parameters(device=device) + + def reset_parameters(self, device=None) -> None: + inv_freq = self._compute_inv_freq(device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + arange = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + scale = ( + (arange + 0.4 * self.dim) / (1.4 * self.dim) + if self.scale_base is not None + else None + ) + self.register_buffer("scale", scale, persistent=False) + + def _compute_inv_freq(self, device=None) -> torch.Tensor: + return 1.0 / ( + self.base + ** ( + torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + / self.dim + ) + ) + + def _update_cos_sin_cache(self, seqlen: int, device=None, dtype=None) -> None: + if self.inv_freq.is_meta: + self.reset_parameters(device=device) + if ( + seqlen > self._seq_len_cached + or self._cos_cached is None + or self._cos_cached.device != device + or self._cos_cached.dtype != dtype + or (self.training and self._cos_cached.is_inference()) + ): + self._seq_len_cached = seqlen + if self.pos_idx_in_fp32: + t = ( + torch.arange(seqlen, device=device, dtype=torch.float32) + / self.scaling_factor + ) + inv_freq = ( + self.inv_freq.to(torch.float32) + if self.inv_freq.dtype != torch.float32 + else self.inv_freq + ) + else: + t = ( + 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) # 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 # ty:ignore[invalid-assignment] + power = ( + torch.arange(seqlen, dtype=_scale.dtype, device=_scale.device) + - seqlen // 2 + ) / 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) + self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) + self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) + + def _apply(self, fn, recurse=True): + if self.inv_freq.is_meta: + self.reset_parameters(device="cpu") + result = super()._apply(fn, recurse=recurse) + new_inv_freq = self._compute_inv_freq(device=self.inv_freq.device) + self.register_buffer("inv_freq", new_inv_freq, persistent=False) + self._seq_len_cached = 0 + self._cos_cached = None + self._sin_cached = None + self._cos_k_cached = None + self._sin_k_cached = None + return result + + def forward( + self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0 + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply RoPE to query and key tensors. + + Parameters + ---------- + q, k : torch.Tensor + Tensors of shape ``(batch, seqlen, n_heads, head_dim)``. + seqlen_offset : int + Offset used in incremental decoding. + """ + self._update_cos_sin_cache( + q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype + ) + assert self._cos_cached is not None and self._sin_cached is not None + + if self.scale is not None: + raise NotImplementedError("XPos scaling is not supported in this path.") + + cos = self._cos_cached[seqlen_offset:] + sin = self._sin_cached[seqlen_offset:] + + if FLASH_ATTN_ROTARY_INSTALLED and q.device.type == "cuda": + assert apply_triton_rotary is not None + 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) + return q_rot, k_rot + + +class EsmcTritonRotaryEmbedding(EsmcRotaryEmbedding): + """RoPE variant that delegates to the Flash-Attention Triton kernel. + + Used inside :class:`EsmcFlashMultiHeadAttention` when Flash Attention 2 is + available. The ``forward`` signature differs from + :class:`EsmcRotaryEmbedding` because Flash Attention packs Q, K, V together. + """ + + def forward( + self, qkv: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int + ) -> 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 + assert apply_triton_rotary is not None + + apply_triton_rotary( + qkv[:, 0], + self._cos_cached, + self._sin_cached, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + inplace=True, + ) + apply_triton_rotary( + qkv[:, 1], + self._cos_cached, + self._sin_cached, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + inplace=True, + ) + return qkv + + +def esmc_swiglu_hidden_dim(expansion_ratio: float, d_model: int) -> int: + """Round the hidden dim to the nearest multiple of 256 after expansion.""" + return int(((expansion_ratio * d_model) + 255) // 256 * 256) + + +class EsmcLayerNormMLP(nn.Module): + """LayerNorm + SwiGLU MLP (no bias), as used by ESMC. + + Shares the parameter names ``layer_norm_weight``, ``layer_norm_bias``, + ``fc1_weight`` and ``fc2_weight`` with Transformer Engine's fused + ``LayerNormMLP`` so the two are state-dict compatible: the fused kernel is + used when available and this pure-PyTorch module is the fallback. + """ + + def __init__( + self, hidden_size: int, ffn_hidden_size: int, eps: float = 1e-5 + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + self.eps = eps + self.layer_norm_weight = nn.Parameter(torch.ones(hidden_size)) + self.layer_norm_bias = nn.Parameter(torch.zeros(hidden_size)) + self.fc1_weight = nn.Parameter(torch.empty(2 * ffn_hidden_size, hidden_size)) + self.fc2_weight = nn.Parameter(torch.empty(hidden_size, ffn_hidden_size)) + nn.init.normal_(self.fc1_weight, std=0.02) + nn.init.normal_(self.fc2_weight, std=0.02) + + def forward(self, x: Tensor) -> Tensor: + x = F.layer_norm( + x, + (self.hidden_size,), + self.layer_norm_weight, + self.layer_norm_bias, + self.eps, + ) + x = F.linear(x, self.fc1_weight) + x1, x2 = x.chunk(2, dim=-1) + x = F.silu(x1) * x2 + return F.linear(x, self.fc2_weight) + + +def esmc_swiglu_ln_ffn( + d_model: int, expansion_ratio: float, bias: bool, use_te: bool +) -> nn.Module: + """LayerNorm + SwiGLU MLP, fused via Transformer Engine when ``use_te``.""" + assert not bias, "ESMC was trained with bias=False; bias=True not supported" + hidden = esmc_swiglu_hidden_dim(expansion_ratio, d_model) + if use_te: + assert te is not None + return te.LayerNormMLP( + hidden_size=d_model, + ffn_hidden_size=hidden, + bias=bias, + activation="swiglu", + init_method=None, + output_layer_init_method=None, + ) + return EsmcLayerNormMLP(hidden_size=d_model, ffn_hidden_size=hidden) + + +def esmc_gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool) -> nn.Sequential: + """LayerNorm + GELU MLP.""" + hidden = int(expansion_ratio * d_model) + return nn.Sequential( + nn.LayerNorm(d_model), + nn.Linear(d_model, hidden, bias=bias), + nn.GELU(), + nn.Linear(hidden, d_model, bias=bias), + ) + + +class EsmcLayerNormLinear(nn.Module): + """LayerNorm followed by a Linear projection (no bias). + + Shares the parameter names ``layer_norm_weight``, ``layer_norm_bias`` and + ``weight`` with Transformer Engine's fused ``LayerNormLinear`` so the two + are state-dict compatible: the fused kernel is used when available and this + pure-PyTorch module is the fallback. + """ + + def __init__(self, d_in: int, d_out: int, eps: float = 1e-5) -> None: + super().__init__() + self.d_in = d_in + self.eps = eps + self.layer_norm_weight = nn.Parameter(torch.ones(d_in)) + self.layer_norm_bias = nn.Parameter(torch.zeros(d_in)) + self.weight = nn.Parameter(torch.empty(d_out, d_in)) + nn.init.normal_(self.weight, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.layer_norm( + x, (self.d_in,), self.layer_norm_weight, self.layer_norm_bias, self.eps + ) + return F.linear(x, self.weight) + + +def make_esmc_attn_layernorm_qkv(d_model: int, bias: bool, use_te: bool) -> nn.Module: + """LayerNorm + fused QKV projection, fused via Transformer Engine when + ``use_te``.""" + assert not bias, "ESMC was trained with bias=False; bias=True not supported" + if use_te: + assert te is not None + return te.LayerNormLinear(d_model, d_model * 3, bias=bias, init_method=None) + return EsmcLayerNormLinear(d_model, d_model * 3) + + +def make_esmc_attn_out_proj(d_model: int, bias: bool, use_te: bool) -> nn.Module: + """Attention output projection, fused via Transformer Engine when ``use_te``.""" + if use_te: + assert te is not None + return te.Linear(d_model, d_model, bias=bias, init_method=None) + return nn.Linear(d_model, d_model, bias=bias) + + +def esmc_scaled_dot_product_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + n_heads: int, + d_head: int, + seq_id: torch.Tensor | None, +) -> torch.Tensor: + """Scaled dot-product attention with an optional boolean mask. + + ``seq_id`` is the mask itself, broadcastable to ``(B, heads, L, L)``, with + ``True`` meaning attend; :meth:`EsmcModel.forward` builds it. ``None`` means + every position attends to every other, which is what lets the fused kernels + take over. + + Dispatches in order of preference: + + 1. xformers ``memory_efficient_attention`` - preferred fused kernel, + requires ``xformers``, unmasked only. + 2. Flash Attention 2 (``flash_attn_func``) - secondary fused kernel, + requires ``flash-attn``, unmasked only, fp16 / bf16 only. + 3. PyTorch's ``F.scaled_dot_product_attention`` - last-resort path; the only + one that takes a mask, and the only fp32 one. + """ + if seq_id is None and XFORMERS_INSTALLED and q.is_cuda: + assert xops is not None + b, s, _ = q.shape + 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( + q4, k4, v4, attn_bias=None, scale=d_head**-0.5 + ) + return context.reshape(b, s, n_heads * d_head) + if ( + seq_id is None + and FLASH_ATTN_INSTALLED + and q.is_cuda + and q.dtype in (torch.float16, torch.bfloat16) + ): + assert flash_attn_func is not None + b, s, _ = q.shape + 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(q4, k4, v4, dropout_p=0.0, softmax_scale=d_head**-0.5) + 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) + v = v.view(b, s, n_heads, -1).transpose(1, 2) + context = F.scaled_dot_product_attention(q, k, v, seq_id) + _, h, _, d_out = context.shape + return context.transpose(1, 2).reshape(b, s, h * d_out) + + +class EsmcMultiHeadAttention(nn.Module): + """Multi-head self-attention with QK LayerNorm and RoPE, as used by ESMC. + + Parameters + ---------- + d_model : int + Model hidden dimension. + n_heads : int + Number of attention heads. + bias : bool + Whether to use bias in the linear layers (ESMC uses ``False``). + qk_layernorm : bool + Whether to apply LayerNorm to queries and keys before computing scores. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + bias: bool = False, + qk_layernorm: bool = True, + use_te: bool = False, + ): + super().__init__() + self.d_model = d_model + self.n_heads = n_heads + self.d_head = d_model // n_heads + + assert not bias, "ESMC was trained with bias=False; bias=True not supported" + self.layernorm_qkv = make_esmc_attn_layernorm_qkv(d_model, bias, use_te) + self.out_proj = make_esmc_attn_out_proj(d_model, bias, use_te) + + if qk_layernorm: + self.q_ln = nn.LayerNorm(d_model, bias=bias) + self.k_ln = nn.LayerNorm(d_model, bias=bias) + else: + self.q_ln = nn.Identity() + self.k_ln = nn.Identity() + + self.rotary = EsmcRotaryEmbedding(d_model // n_heads) + + def _apply_rotary( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + q = q.unflatten(-1, (self.n_heads, self.d_head)) + k = k.unflatten(-1, (self.n_heads, self.d_head)) + q, k = self.rotary(q, k) + q = q.flatten(-2, -1) + k = k.flatten(-2, -1) + return q, k + + def forward( + self, + x: torch.Tensor, + seq_id: torch.Tensor | None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return ``(context, attn_weights)``. + + ``attn_weights`` is ``None`` unless ``output_attentions=True`` - the + fused SDPA backends don't expose attention probabilities, so capturing + them forces a materialized ``softmax(Q @ K.T / sqrt(d)) @ V`` path with + shape ``(B, H, L, L)``. + """ + qkv = self.layernorm_qkv(x) + q, k, v = torch.chunk(qkv, 3, dim=-1) + q = self.q_ln(q).to(q.dtype) + k = self.k_ln(k).to(q.dtype) + q, k = self._apply_rotary(q, k) + + b, s, _ = q.shape + + if output_attentions: + # Manual SDPA so attention probabilities are observable. + q4 = q.view(b, s, self.n_heads, self.d_head).transpose(1, 2) + k4 = k.view(b, s, self.n_heads, self.d_head).transpose(1, 2) + v4 = v.view(b, s, self.n_heads, self.d_head).transpose(1, 2) + scale = self.d_head**-0.5 + attn_scores = (q4 @ k4.transpose(-2, -1)) * scale + if seq_id is not None: + attn_scores = attn_scores.masked_fill(~seq_id, float("-inf")) + attn_weights = torch.softmax(attn_scores, dim=-1) + context = (attn_weights @ v4).transpose(1, 2).reshape(b, s, -1) + return self.out_proj(context), attn_weights + + context = esmc_scaled_dot_product_attention( + q, k, v, n_heads=self.n_heads, d_head=self.d_head, seq_id=seq_id + ) + return self.out_proj(context), None + + +class EsmcFlashMultiHeadAttention(EsmcMultiHeadAttention): + """Flash-Attention 2 variant of :class:`EsmcMultiHeadAttention`.""" + + def __init__( + self, + d_model: int, + n_heads: int, + bias: bool = False, + qk_layernorm: bool = True, + use_te: bool = False, + ): + super().__init__( + d_model=d_model, + n_heads=n_heads, + bias=bias, + qk_layernorm=qk_layernorm, + use_te=use_te, + ) + self.rotary = EsmcTritonRotaryEmbedding(d_model // n_heads) + + def forward( + self, + x: torch.Tensor, + seq_id: torch.Tensor | None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if output_attentions: + raise ValueError( + "output_attentions=True is not supported with " + "attn_implementation='flash_attention_2'. " + "Re-load the model with attn_implementation='sdpa' (or 'eager')." + ) + assert seq_id is not None and seq_id.dtype == torch.bool + + seqlens = seq_id.sum(dim=-1, dtype=torch.int32) + cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + max_seqlen = int(seqlens.max().item()) + + qkv = self.layernorm_qkv(x) + q, k, v = torch.chunk(qkv, 3, dim=-1) + q = self.q_ln(q).to(q.dtype) + k = self.k_ln(k).to(q.dtype) + + # q/k/v are 2D (T, D) here: the parent model unpads the batch before the + # transformer stack to produce the varlen-flat layout that + # ``flash_attn_varlen_qkvpacked_func`` requires. + T = q.shape[0] + 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) + + assert flash_attn_varlen_qkvpacked_func is not None + 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 + return (self.out_proj(context.reshape(n_out, h_out * d_out)), None) + + +class EsmcUnifiedTransformerBlock(nn.Module): + """A single ESMC transformer block: pre-norm attention + pre-norm FFN with + residual scaling. + + Parameters + ---------- + d_model : int + Hidden dimension. + n_heads : int + Number of attention heads. + use_flash_attn : bool + Use the Flash Attention 2 kernel if available. + bias : bool + Whether linear layers include bias terms (ESMC uses ``False``). + expansion_ratio : float + Hidden-dim expansion ratio for the FFN. + residue_scaling_factor : float + Scales residual connections to stabilise deep networks + (``sqrt(n_layers / 36)`` is the ESM3 scheme). + qk_layernorm : bool + Whether to apply QK LayerNorm in attention. + ffn_type : str + Feed-forward activation: ``"swiglu"`` or ``"gelu"``. + use_te : bool + Build the Transformer Engine fused LayerNorm modules. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + use_flash_attn: bool = False, + bias: bool = False, + expansion_ratio: float = 4.0, + residue_scaling_factor: float = 1.0, + qk_layernorm: bool = True, + ffn_type: str = "swiglu", + use_te: bool = False, + ): + super().__init__() + + attn_cls = ( + EsmcFlashMultiHeadAttention if use_flash_attn else EsmcMultiHeadAttention + ) + self.attn = attn_cls( + d_model, n_heads, bias=bias, qk_layernorm=qk_layernorm, use_te=use_te + ) + + if ffn_type == "swiglu": + self.ffn = esmc_swiglu_ln_ffn(d_model, expansion_ratio, bias, use_te) + elif ffn_type == "gelu": + self.ffn = esmc_gelu_ln_ffn(d_model, expansion_ratio, bias) + else: + raise ValueError( + f"Unknown ffn_type: {ffn_type!r}. Choose 'swiglu' or 'gelu'." + ) + + self.scaling_factor = residue_scaling_factor + + def forward( + self, + x: torch.Tensor, + sequence_id: torch.Tensor | None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + Parameters + ---------- + x : torch.Tensor + Input of shape ``(batch, seq_len, d_model)``. + sequence_id : torch.Tensor | None + Chain-ID tensor used to restrict attention to tokens within the + same chain. SDPA blocks accept an integer tensor (``-1`` marks + padding); the flash-attn block takes a ``bool`` padding mask - the + caller selects which. ``None`` skips chain-aware masking (fast path). + output_attentions : bool + When ``True``, also returns the per-head attention weights. + """ + attn_out, attn_weights = self.attn( + x, sequence_id, output_attentions=output_attentions + ) + x = x + attn_out / self.scaling_factor + x = x + self.ffn(x) / self.scaling_factor + return x, attn_weights + + +class EsmcTransformerStack(nn.Module): + """Stack of :class:`EsmcUnifiedTransformerBlock` layers with a final + LayerNorm. + + Parameters + ---------- + d_model : int + Hidden dimension. + n_heads : int + Number of attention heads. + n_layers : int + Number of transformer blocks. + scale_residue : bool + When ``True`` apply ESM3 residue scaling ``sqrt(n_layers / 36)``. + bias : bool + Bias flag forwarded to every sub-module. + qk_layernorm : bool + QK LayerNorm flag forwarded to every block. + ffn_type : str + FFN activation type (``"swiglu"`` or ``"gelu"``). + expansion_ratio : float + FFN expansion ratio. + use_flash_attn : bool + Use the Flash Attention 2 kernel when available. + use_te : bool + Build the Transformer Engine fused LayerNorm modules. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + n_layers: int, + scale_residue: bool = True, + bias: bool = False, + qk_layernorm: bool = True, + ffn_type: str = "swiglu", + expansion_ratio: float = 8 / 3, + use_flash_attn: bool = False, + use_te: bool = False, + ): + super().__init__() + self.blocks = nn.ModuleList( + [ + EsmcUnifiedTransformerBlock( + d_model, + n_heads, + use_flash_attn=use_flash_attn, + residue_scaling_factor=math.sqrt(n_layers / 36) + if scale_residue + else 1.0, + expansion_ratio=expansion_ratio, + bias=bias, + qk_layernorm=qk_layernorm, + ffn_type=ffn_type, + use_te=use_te, + ) + for _ in range(n_layers) + ] + ) + self.norm = nn.LayerNorm(d_model, bias=False) + + def forward( + self, + x: torch.Tensor, + sequence_id: torch.Tensor | None = None, + layers_to_collect: list[int] | None = None, + output_attentions: bool = False, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + tuple[torch.Tensor, ...], + tuple[torch.Tensor, ...] | None, + ]: + """Run the full transformer stack. + + Parameters + ---------- + x : torch.Tensor + Input of shape ``(batch, seq_len, d_model)``. + sequence_id : torch.Tensor | None + Optional chain-id tensor forwarded to each block. + layers_to_collect : list[int] | None + Layer indices (0-based pre-block inputs plus ``n_layers`` for the + post-norm output) whose hidden states should be returned. + output_attentions : bool + When ``True``, collects the per-block attention weights. + + Returns + ------- + tuple + ``(post_norm, pre_norm, hidden_states, attentions)`` where + ``hidden_states`` is a (possibly empty) tuple of tensors and + ``attentions`` is a tuple of per-block ``(B, H, L, L)`` tensors or + ``None`` when ``output_attentions`` is ``False``. + """ + if layers_to_collect is None: + layers_to_collect = [] + + collected: list[torch.Tensor] = [] + all_attentions: list[torch.Tensor] = [] + for layer_idx, block in enumerate(self.blocks): + if layer_idx in layers_to_collect: + collected.append(x) + x, attn_weights = block(x, sequence_id, output_attentions=output_attentions) + if output_attentions and attn_weights is not None: + all_attentions.append(attn_weights) + + norm_x = self.norm(x) + if len(self.blocks) in layers_to_collect: + collected.append(norm_x) + + attentions = tuple(all_attentions) if output_attentions else None + return norm_x, x, tuple(collected), attentions diff --git a/esm/models/esmc/model.py b/esm/models/esmc/model.py new file mode 100644 index 00000000..48007e25 --- /dev/null +++ b/esm/models/esmc/model.py @@ -0,0 +1,835 @@ +"""ESMC encoder and task heads. + +ESMC is a protein language model trained by EvolutionaryScale with a +masked-token objective over amino acid sequences. The architecture is a +standard Transformer encoder with RoPE positional embeddings, QK LayerNorm and +SwiGLU feed-forward networks. +""" + +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Self + +import torch +import torch.nn as nn +from safetensors.torch import save_file +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from esm.models.esmc.checkpoint_layout import native_to_published, published_to_native +from esm.models.esmc.config import EsmcConfig +from esm.models.esmc.kernels import ( + FLASH_ATTN_INSTALLED, + TE_INSTALLED, + pad_input, + unpad_input, +) +from esm.models.esmc.layers import EsmcRotaryEmbedding, EsmcTransformerStack +from esm.models.esmc.sae import EsmcSaeLayer +from esm.models.hub import HubPreTrainedModel, resolve_model_dir + +_SAFETENSORS_INDEX = "model.safetensors.index.json" +_SAFETENSORS_SINGLE = "model.safetensors" + + +# --------------------------------------------------------------------------- +# Output dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class EsmcOutput: + """Output of :class:`EsmcModel`. + + Attributes + ---------- + last_hidden_state : torch.Tensor | None + Hidden states at the output of the last layer, after layer norm, of + shape ``(batch, seq_len, d_model)``. + hidden_states : torch.Tensor | None + Stacked hidden states of shape + ``(n_layers + 1, batch, seq_len, d_model)`` - one entry per block input + plus the final post-LayerNorm output; populated when + ``output_hidden_states=True``. + sae_outputs : dict[str, torch.Tensor] | None + SAE feature magnitudes keyed by SAE model name (sparse tensors). + attentions : tuple[torch.Tensor, ...] | None + Per-layer attention weights, populated when ``output_attentions=True``. + """ + + last_hidden_state: torch.Tensor | None = None + hidden_states: torch.Tensor | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + last_hidden_state_prenorm: torch.Tensor | None = None + + +@dataclass +class EsmcMaskedLMOutput: + """Output of :class:`EsmcForMaskedLM`.""" + + loss: torch.Tensor | None = None + logits: torch.Tensor | None = None + last_hidden_state: torch.Tensor | None = None + hidden_states: torch.Tensor | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + last_hidden_state_prenorm: torch.Tensor | None = None + + +@dataclass +class EsmcTokenClassifierOutput: + """Output of :class:`EsmcForTokenClassification`.""" + + loss: torch.Tensor | None = None + logits: torch.Tensor | None = None + last_hidden_state: torch.Tensor | None = None + hidden_states: torch.Tensor | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + + +@dataclass +class EsmcSequenceClassifierOutput: + """Output of :class:`EsmcForSequenceClassification`.""" + + loss: torch.Tensor | None = None + logits: torch.Tensor | None = None + last_hidden_state: torch.Tensor | None = None + hidden_states: torch.Tensor | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + attentions: tuple[torch.Tensor, ...] | None = None + + +# --------------------------------------------------------------------------- +# State-dict / weight loading helpers +# --------------------------------------------------------------------------- + + +def _drop_te_extra_state( + module: nn.Module, state_dict: dict, prefix: str, local_metadata: dict +) -> None: + """Strip Transformer Engine's non-tensor ``_extra_state`` entries; TE + returns these as ``io.BytesIO`` objects the checkpoint writer cannot + dtype-infer.""" + for key in list(state_dict): + if key.endswith("_extra_state"): + del state_dict[key] + + +def _ignore_te_extra_state(module: nn.Module, incompatible_keys) -> None: + """The load-side counterpart of :func:`_drop_te_extra_state`. + + Without it a model cannot even reload its own ``state_dict``, since the keys + stripped on the way out are still expected on the way in. They hold nothing + but fp8 scaling metadata, which is empty until fp8 runs and is recomputed + when it does. + """ + incompatible_keys.missing_keys[:] = [ + key + for key in incompatible_keys.missing_keys + if not key.endswith("_extra_state") + ] + + +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +class EsmcPreTrainedModel(HubPreTrainedModel): + """ESMC weight initialisation on top of the shared Hub loader.""" + + config_class = EsmcConfig + _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"] + + def _init_weights( + self, module: nn.Module, device: torch.device | str | None = None + ): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, EsmcRotaryEmbedding): + # `post_init` reaches here through `nn.Module.apply`, which passes no + # device, so fall back to where the buffers already are. Defaulting to + # the CPU moved them off a CUDA model and broke the next forward. + module.reset_parameters( + device=module.inv_freq.device if device is None else device + ) + + def _materialize_uninitialized(self, device: torch.device | str = "cpu") -> None: + """Materialize and initialize parameters/buffers still on the meta + device after a partial checkpoint load (e.g. a freshly-added head or + the non-persistent RoPE buffers).""" + for module in self.modules(): + direct = list(module._parameters.values()) + list(module._buffers.values()) + if any(t is not None and t.is_meta for t in direct): + module.to_empty(device=device, recurse=False) + self._init_weights(module, device=device) + + @classmethod + def _adapt_checkpoint_keys( + cls, raw: dict[str, torch.Tensor], model_keys: set[str] + ) -> dict[str, torch.Tensor]: + """Align published-checkpoint keys onto a target module tree. + + The published checkpoints are saved from ``EsmcForMaskedLM``, so + backbone keys carry an ``esmc.`` prefix and the head keys are + ``lm_head.*``. Keys are first translated out of the published layout + (see :mod:`.checkpoint_layout`), then matched against the target tree: + this keeps keys that already match, strips the ``esmc.`` prefix when + loading the bare encoder, drops Transformer-Engine ``_extra_state`` + entries, and ignores keys with no counterpart (e.g. ``lm_head.*`` when + loading :class:`EsmcModel`). + """ + adapted: dict[str, torch.Tensor] = {} + for key, value in raw.items(): + if key.endswith("_extra_state"): + continue + if key in model_keys: + adapted[key] = value + elif key.startswith("esmc.") and key[len("esmc.") :] in model_keys: + adapted[key[len("esmc.") :]] = value + return adapted + + @classmethod + def _normalize_checkpoint_layout( + cls, raw: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + return published_to_native(raw) + + def save_pretrained(self, save_directory: str | os.PathLike) -> None: + """Write ``config.json`` and ``model.safetensors`` in published layout. + + The in-memory state dict keeps Transformer Engine's fused parameters, so + it is translated on the way out; see :mod:`.checkpoint_layout`. + """ + directory = Path(save_directory) + directory.mkdir(parents=True, exist_ok=True) + self.config.save_pretrained(directory) + state = {k: v for k, v in self.state_dict().items() if v is not None} + save_file( + native_to_published(state), + str(directory / "model.safetensors"), + metadata={"format": "pt"}, + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | os.PathLike, + *, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + attn_implementation: str | None = None, + revision: str | None = None, + cache_dir: str | os.PathLike | None = None, + token: str | None = None, + local_files_only: bool = False, + force_download: bool = False, + ) -> Self: + """Load an ESMC model from a local directory or the HuggingFace Hub. + + ``device`` defaults to CUDA when one is present, and selects the fused + kernels as well as the placement: they are CUDA-only, so a model built + for the CPU must not use them. + """ + local_dir = resolve_model_dir( + pretrained_model_name_or_path, + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=local_files_only, + force_download=force_download, + ) + config = cls.config_class.from_pretrained(local_dir) + if attn_implementation is not None: + config.attn_implementation = attn_implementation + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + # Sets the default device the encoder reads to choose its kernels. + with torch.device(device): + return cls._load_pretrained(local_dir, config, device=device, dtype=dtype) + + +# --------------------------------------------------------------------------- +# Base encoder +# --------------------------------------------------------------------------- + + +class EsmcModel(EsmcPreTrainedModel): + """The bare ESMC encoder outputting raw hidden states. + + ``EsmcModel(config)`` is fully constructible from config alone with random + weights and performs no network or disk I/O, so it can be embedded as a + sub-module and populated later; checkpoint loading lives in + :meth:`from_pretrained`. Its ``state_dict`` has two top-level module groups, + ``embed.*`` (the token embedding) and ``transformer.*`` (the + :class:`EsmcTransformerStack`: ``transformer.blocks..*`` and + ``transformer.norm.weight``); the non-persistent RoPE buffers are omitted. + Per-layer hidden states are exposed via ``forward(..., output_hidden_states=True)`` + as :attr:`EsmcOutput.hidden_states`, shape + ``(n_layers + 1, batch, seq_len, d_model)``. + """ + + _SAE_KEY_RE = re.compile(r"layer(\d+)") + _keys_to_ignore_on_load_unexpected = [r"\._extra_state$", r"^lm_head\."] + + def __init__(self, config: EsmcConfig): + super().__init__(config) + # Every fused kernel is CUDA-only, so all of them hang off one condition: + # the device this model's parameters are being created on. Deriving them + # separately is what let the LayerNorm and attention gates disagree. + on_cuda = torch.get_default_device().type == "cuda" + self._use_flash_attn = ( + on_cuda + and FLASH_ATTN_INSTALLED + and config.attn_implementation == "flash_attention_2" + ) + self.embed = nn.Embedding(config.vocab_size, config.hidden_size) + self.transformer = EsmcTransformerStack( + config.hidden_size, + config.num_attention_heads, + config.num_hidden_layers, + use_flash_attn=self._use_flash_attn, + use_te=on_cuda and TE_INSTALLED, + ) + self._sae_models: nn.ModuleDict = nn.ModuleDict() + self._register_state_dict_hook(_drop_te_extra_state) + self.register_load_state_dict_post_hook(_ignore_te_extra_state) + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed + + def set_input_embeddings(self, value: nn.Embedding) -> None: + self.embed = value + + def add_sae_models(self, sae_models: list[EsmcSaeLayer]) -> None: + """Register one or more SAEs obtained from an :class:`EsmcSaeModel`. + + Each is keyed by ``f"layer{N}"`` (the backbone-layer index ``N`` the + SAE is trained against). Attaching two SAEs for the same backbone layer + raises - only one SAE per layer can be active. + """ + for layer in sae_models: + assert isinstance(layer, EsmcSaeLayer), ( + f"Expected an SAE layer (model.layers['']), got " + f"{type(layer).__name__}." + ) + key = f"layer{int(layer.layer)}" + if key in self._sae_models: + raise ValueError( + f"An SAE is already registered at {key!r}. Only one SAE " + "per backbone layer can be active - pick a different layer " + "on one of them, or attach in a fresh model." + ) + self._sae_models[key] = layer + + def _get_sae_layer_num_requested(self, model_name: str) -> int: + match = self._SAE_KEY_RE.fullmatch(model_name) + 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: + assert torch.all(input_ids != self.config.mask_token_id), ( + "SAE inputs must not contain mask tokens. " + "SAEs were trained on unmasked sequences." + ) + + def _get_sae_outputs( + self, + hidden_states: torch.Tensor, + layers_to_collect: list[int], + token_mask: torch.Tensor, + normalize_sae: bool = False, + ) -> dict[str, torch.Tensor]: + """Run all registered SAEs and return their feature magnitudes.""" + layer_to_idx = {layer: idx for idx, layer in enumerate(layers_to_collect)} + sae_outputs: dict[str, torch.Tensor] = {} + + for model_name, sae_module in self._sae_models.items(): + assert isinstance(sae_module, EsmcSaeLayer) + layer = sae_module + requested_layer = self._get_sae_layer_num_requested(model_name) + layer_idx = layer_to_idx[requested_layer] + layer_states = hidden_states[layer_idx].clone().to(self.device) + + sae_out = layer.get_sae_output(layer_states, token_mask) + features = sae_out.feature_magnitudes.detach() + + if normalize_sae: + features = (features / layer.max) * layer.idf + + sae_outputs[model_name] = features.to_sparse() + + return sae_outputs + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ) -> tuple[torch.Tensor, ...] | EsmcOutput: + """Encode ``input_ids`` and return hidden states. + + Parameters + ---------- + input_ids : torch.Tensor + Token ids of shape ``(batch, seq_len)``. + attention_mask : torch.Tensor | None + Boolean/int mask of real tokens; inferred from ``pad_token_id`` when + omitted and ``sequence_id`` is not given. + sequence_id : torch.Tensor | None + Integer chain-ID tensor for chain-aware attention masking; tokens + sharing a non-negative value attend to each other, padding is + ``-1``. May be combined with ``attention_mask``, which is folded in + by marking the masked positions ``-1``. The ``flash_attention_2`` + backend only supports single-chain inputs. + output_hidden_states, output_attentions, return_dict : bool | None + Override the corresponding config defaults. + compute_sae : bool + Run any SAEs registered via :meth:`add_sae_models`. + normalize_sae : bool + Scale SAE feature magnitudes by ``idf / max``. + """ + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + output_sae = compute_sae and len(self._sae_models) > 0 + + # Determine which intermediate layers to collect. When SAEs are + # registered we must collect at least the layers they target. + if output_hidden_states: + layers_to_collect: list[int] = list( + range(self.config.num_hidden_layers + 1) + ) + elif output_sae: + layers_to_collect = sorted( + {self._get_sae_layer_num_requested(name) for name in self._sae_models} + ) + else: + layers_to_collect = [] + + user_supplied_sequence_id = sequence_id is not None + if sequence_id is not None: + # Padding is spelled -1 in a sequence_id, so fold any mask the caller + # also passed into it rather than dropping one of the two. Without + # this, padding keeps a real chain id and stays visible to queries in + # that chain, which corrupts their output rather than just the pad + # rows. + if attention_mask is not None: + sequence_id = sequence_id.masked_fill(~attention_mask.bool(), -1) + bool_mask = sequence_id >= 0 + else: + if attention_mask is None: + attention_mask = ( + input_ids != self.config.pad_token_id + ) # ty:ignore[invalid-assignment] + assert attention_mask is not None + bool_mask = attention_mask.bool() + + x = self.embed(input_ids) + b, l_ = x.shape[:2] + + if self._use_flash_attn: + if sequence_id is not None and (sequence_id > 0).any(): + raise ValueError( + "Multi-chain ``sequence_id`` (any value > 0) is not " + "supported with attn_implementation='flash_attention_2'. " + "Re-load the model with attn_implementation='sdpa' (or " + "'eager') for chain-aware attention masking." + ) + assert unpad_input is not None + x, indices, *_ = unpad_input(x, bool_mask) + else: + indices = None + + # One of chain ids or padding decides the mask, never both - the same + # split `transformers` makes. Chain ids compare query against key; + # padding hides a key from every query, so it broadcasts over the query + # axis. ``True`` means attend. + if self._use_flash_attn: + trans_seq_id = bool_mask + elif user_supplied_sequence_id: + assert sequence_id is not None + trans_seq_id = ( + sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2) + ).unsqueeze(1) + elif bool_mask.all() and not output_attentions: + # No padding to mask, so the fused kernels can take it; + # output_attentions forces the manual branch regardless. + trans_seq_id = None + else: + trans_seq_id = bool_mask[:, None, None, :] + + last_hidden_state, prenorm, collected, attentions = self.transformer( + x, + sequence_id=trans_seq_id, + layers_to_collect=layers_to_collect, + output_attentions=output_attentions, + ) + + if self._use_flash_attn: + assert indices is not None and pad_input is not None + last_hidden_state = pad_input(last_hidden_state, indices, b, l_) + prenorm = pad_input(prenorm, indices, b, l_) + collected = tuple(pad_input(h, indices, b, l_) for h in collected) + + collected_tensor: torch.Tensor | None = ( + torch.stack(collected, dim=0) if collected else None + ) + + sae_outputs: dict[str, torch.Tensor] | None = None + if output_sae and collected_tensor is not None: + assert input_ids is not None + self._validate_sae_inputs(input_ids) + sae_outputs = self._get_sae_outputs( + collected_tensor, layers_to_collect, bool_mask, normalize_sae + ) + + hidden_states_tensor = collected_tensor if output_hidden_states else None + + if not return_dict: + return tuple( + v + for v in [ + last_hidden_state, + hidden_states_tensor, + sae_outputs, + attentions, + ] + if v is not None + ) # ty:ignore[invalid-return-type] + + return EsmcOutput( + last_hidden_state=last_hidden_state, + hidden_states=hidden_states_tensor, + sae_outputs=sae_outputs, + attentions=attentions, + last_hidden_state_prenorm=prenorm if output_hidden_states else None, + ) + + +# --------------------------------------------------------------------------- +# Heads +# --------------------------------------------------------------------------- + + +def _esmc_lm_head( + d_model: int, output_dim: int, hidden_dim: int | None = None +) -> nn.Sequential: + """Linear -> GELU -> LayerNorm -> Linear projection head for masked LM.""" + hidden_dim = hidden_dim if hidden_dim is not None else d_model + return nn.Sequential( + nn.Linear(d_model, hidden_dim), + nn.GELU(), + nn.LayerNorm(hidden_dim), + nn.Linear(hidden_dim, output_dim), + ) + + +class EsmcForMaskedLM(EsmcPreTrainedModel): + """ESMC with a masked language modelling head.""" + + def __init__(self, config: EsmcConfig): + super().__init__(config) + self.esmc = EsmcModel(config) + self.lm_head = _esmc_lm_head(config.hidden_size, config.vocab_size) + + def get_output_embeddings(self) -> nn.Linear: + return self.lm_head[-1] # ty:ignore[invalid-return-type] + + def set_output_embeddings(self, new_embeddings: nn.Linear) -> None: + self.lm_head[-1] = new_embeddings + + def add_sae_models(self, sae_models: list[EsmcSaeLayer]) -> None: + """Proxy to :meth:`EsmcModel.add_sae_models`.""" + self.esmc.add_sae_models(sae_models) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ) -> tuple[torch.Tensor, ...] | EsmcMaskedLMOutput: + """Run the encoder and project to per-token vocabulary logits. + + Parameters + ---------- + labels : torch.Tensor | None + Target token ids for the masked-LM loss; positions with label + ``-100`` are ignored. + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + encoder_outputs = self.esmc( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + compute_sae=compute_sae, + normalize_sae=normalize_sae, + ) + + logits = self.lm_head(encoder_outputs.last_hidden_state) + + loss: torch.Tensor | None = None + if labels is not None: + loss = CrossEntropyLoss(ignore_index=-100)( + logits.view(-1, self.config.vocab_size), labels.view(-1) + ) + + if not return_dict: + return tuple( + v + for v in [ + loss, + logits, + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + encoder_outputs.sae_outputs, + encoder_outputs.attentions, + ] + if v is not None + ) + + return EsmcMaskedLMOutput( + loss=loss, + logits=logits, + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + sae_outputs=encoder_outputs.sae_outputs, + attentions=encoder_outputs.attentions, + last_hidden_state_prenorm=encoder_outputs.last_hidden_state_prenorm, + ) + + +class _EsmcClassificationHead(nn.Module): + """Dense classification head applied to the ```` token.""" + + def __init__(self, config: EsmcConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.classifier_dropout) + self.out_proj = nn.Linear(config.hidden_size, config.num_labels) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + x = hidden_states[:, 0, :] # token + x = self.dropout(x) + x = torch.tanh(self.dense(x)) + x = self.dropout(x) + return self.out_proj(x) + + +class EsmcForSequenceClassification(EsmcPreTrainedModel): + """ESMC with a sequence-level classification head over the ```` + token. Supports regression, single-label and multi-label classification. + """ + + _keys_to_ignore_on_load_unexpected = [r"\._extra_state$", r"^lm_head\."] + _keys_to_ignore_on_load_missing = [r"^classifier\."] + + def __init__(self, config: EsmcConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.esmc = EsmcModel(config) + self.classifier = _EsmcClassificationHead(config) + + def add_sae_models(self, sae_models: list[EsmcSaeLayer]) -> None: + """Proxy to :meth:`EsmcModel.add_sae_models`.""" + self.esmc.add_sae_models(sae_models) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ) -> tuple[torch.Tensor, ...] | EsmcSequenceClassifierOutput: + """Run the encoder and classify each sequence from its ```` + representation. + + Parameters + ---------- + labels : torch.Tensor | None + Class indices (or float targets for regression). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + encoder_outputs = self.esmc( + input_ids, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + compute_sae=compute_sae, + normalize_sae=normalize_sae, + ) + logits = self.classifier(encoder_outputs.last_hidden_state) + + loss: torch.Tensor | None = None + if labels is not None: + labels = labels.to(logits.device) + + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + loss = loss_fct( + logits.squeeze() if self.num_labels == 1 else logits, + labels.squeeze() if self.num_labels == 1 else labels, + ) + elif self.config.problem_type == "single_label_classification": + loss = CrossEntropyLoss()( + logits.view(-1, self.num_labels), labels.view(-1) + ) + elif self.config.problem_type == "multi_label_classification": + loss = BCEWithLogitsLoss()(logits, labels) + + if not return_dict: + return tuple( + v + for v in [ + loss, + logits, + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + encoder_outputs.sae_outputs, + encoder_outputs.attentions, + ] + if v is not None + ) + + return EsmcSequenceClassifierOutput( + loss=loss, + logits=logits, + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + sae_outputs=encoder_outputs.sae_outputs, + attentions=encoder_outputs.attentions, + ) + + +class EsmcForTokenClassification(EsmcPreTrainedModel): + """ESMC with a per-token classification head.""" + + _keys_to_ignore_on_load_unexpected = [r"\._extra_state$", r"^lm_head\."] + _keys_to_ignore_on_load_missing = [r"^classifier\."] + + def __init__(self, config: EsmcConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.esmc = EsmcModel(config) + self.dropout = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + def add_sae_models(self, sae_models: list[EsmcSaeLayer]) -> None: + """Proxy to :meth:`EsmcModel.add_sae_models`.""" + self.esmc.add_sae_models(sae_models) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + labels: torch.Tensor | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ) -> tuple[torch.Tensor, ...] | EsmcTokenClassifierOutput: + """Run the encoder and classify every token. + + Parameters + ---------- + labels : torch.Tensor | None + Per-token class indices; positions with ``-100`` are ignored. + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + encoder_outputs = self.esmc( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + compute_sae=compute_sae, + normalize_sae=normalize_sae, + ) + + sequence_output = self.dropout(encoder_outputs.last_hidden_state) + logits = self.classifier(sequence_output) + + loss: torch.Tensor | None = None + if labels is not None: + loss = CrossEntropyLoss(ignore_index=-100)( + logits.view(-1, self.num_labels), labels.to(logits.device).view(-1) + ) + + if not return_dict: + return tuple( + v + for v in [ + loss, + logits, + encoder_outputs.last_hidden_state, + encoder_outputs.hidden_states, + encoder_outputs.sae_outputs, + encoder_outputs.attentions, + ] + if v is not None + ) + + return EsmcTokenClassifierOutput( + loss=loss, + logits=logits, + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + sae_outputs=encoder_outputs.sae_outputs, + attentions=encoder_outputs.attentions, + ) diff --git a/esm/models/esmc/sae.py b/esm/models/esmc/sae.py new file mode 100644 index 00000000..f22c0ddd --- /dev/null +++ b/esm/models/esmc/sae.py @@ -0,0 +1,385 @@ +"""ESMC sparse autoencoder (SAE) model. + +* :class:`EsmcSaeModel` - the published container, one repo per + ``(backbone, codebook_dim, k)`` group. Each backbone layer ships as a + ``layer_{i}.safetensors`` shard; ``from_pretrained`` downloads the whole + snapshot but loads no weights - callers materialize the layers they need via + :meth:`initialize_layers`. Single-layer repos auto-load so a bare + ``forward(x)`` works. +* :class:`EsmcSaeLayer` - the ``nn.Module`` that holds the weights for one + ``(backbone, codebook_dim, k, layer)`` SAE. Obtained via + ``model.layers[""]``. +""" + +import json +import os +from dataclasses import asdict, dataclass +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors.torch import load_file, save_file + +CONFIG_NAME = "config.json" + + +@dataclass +class EsmcSaeParams: + """Parameters for one backbone layer's SAE inside :class:`EsmcSaeModel`.""" + + d_model: int = 2560 + codebook_dim: int = 65536 + k: int = 64 + layer: int = 0 + + +@dataclass +class EsmcSaeConfig: + """Configuration for :class:`EsmcSaeModel`. + + A container holds one SAE per backbone layer for a fixed + ``(model, codebook_dim, k)`` group. All SAEs in a container share + ``d_model``, ``codebook_dim`` and ``k``; they differ only in the backbone + layer they were trained on. ``available_layers`` lists the backbone-layer + indices the repo ships; each entry ``i`` is stored on disk as + ``layer_{i}.safetensors``. + + Parameters + ---------- + d_model : int + Dimensionality of the ESMC hidden states fed into the SAEs. + codebook_dim : int + Number of sparse features in each SAE's codebook. + k : int + Top-k sparsity per SAE. + available_layers : list[int] + Which backbone-layer indices the repo ships. + """ + + model_type = "esmc_sae" + + d_model: int = 2560 + codebook_dim: int = 65536 + k: int = 64 + available_layers: list[int] | None = None + + def __post_init__(self) -> None: + self.available_layers = ( + list(self.available_layers) if self.available_layers is not None else [0] + ) + + @classmethod + def from_pretrained(cls, directory: str | os.PathLike) -> "EsmcSaeConfig": + with open(Path(directory) / CONFIG_NAME) as f: + raw = json.load(f) + return cls( + d_model=raw["d_model"], + codebook_dim=raw["codebook_dim"], + k=raw["k"], + available_layers=raw.get("available_layers"), + ) + + def save_pretrained(self, directory: str | os.PathLike) -> None: + payload = asdict(self) + payload["model_type"] = self.model_type + with open(Path(directory) / CONFIG_NAME, "w") as f: + json.dump(payload, f, indent=2) + + +@dataclass +class EsmcSaeOutput: + """Output of :class:`EsmcSaeModel` and :class:`EsmcSaeLayer`.""" + + feature_magnitudes: torch.Tensor + reconstruction_loss: torch.Tensor | None = None + + def to_sparse(self) -> None: + self.feature_magnitudes = self.feature_magnitudes.to_sparse() + + +class EsmcSaeLayer(nn.Module): + """One backbone layer's SAE - building block of :class:`EsmcSaeModel`. + + Obtained via ``model.layers[""]`` after calling + ``initialize_layers``. + """ + + idf: torch.Tensor + max: torch.Tensor + + def __init__(self, params: EsmcSaeParams): + super().__init__() + self.params = params + + self.W_enc = nn.Parameter(torch.empty(params.d_model, params.codebook_dim)) + self.W_dec = nn.Parameter(torch.empty(params.codebook_dim, params.d_model)) + self.b_dec = nn.Parameter(torch.zeros(params.d_model)) + # Per-feature normalization stats. Trained alongside the SAE for some + # variants; leaving these as ones makes the ``features / max * idf`` + # scaling in the backbone a no-op for variants that don't ship them. + self.register_buffer("idf", torch.ones(params.codebook_dim)) + self.register_buffer("max", torch.ones(params.codebook_dim)) + + @property + def layer(self) -> int: + """Backbone-layer index this SAE is trained against.""" + return self.params.layer + + def forward(self, x: torch.Tensor, **_kwargs: object) -> EsmcSaeOutput: + del _kwargs + x = self._zscore_normalize_representation(x) + + x_with_pre_encoder_bias = x - self.b_dec + preactivations = F.relu(x_with_pre_encoder_bias @ self.W_enc) + + topk = torch.topk(preactivations, self.params.k, dim=-1) + feature_magnitudes = torch.zeros_like(preactivations).scatter( + -1, topk.indices, topk.values + ) + + reconstructed = feature_magnitudes @ self.W_dec + self.b_dec + reconstruction_loss = (reconstructed - x).pow(2).mean(dim=-1) + + return EsmcSaeOutput( + feature_magnitudes=feature_magnitudes, + reconstruction_loss=reconstruction_loss, + ) + + def get_sae_output( + self, layer_states: torch.Tensor, token_mask: torch.Tensor + ) -> EsmcSaeOutput: + _, _, v_len = layer_states.shape + nonpad_states = layer_states[token_mask].view(-1, v_len) + return self(nonpad_states) + + def _zscore_normalize_representation(self, x: torch.Tensor) -> torch.Tensor: + x_mean = x.mean(dim=-1, keepdim=True) + x = x - x_mean + x_std = x.std(dim=-1, keepdim=True) + return x / (x_std + 1e-5) + + +class EsmcSaeModel(nn.Module): + """Container holding one SAE per backbone layer, all sharing the same + ``(d_model, codebook_dim, k)``. + + ``from_pretrained`` downloads the entire repo (every + ``layer_{i}.safetensors``) into the local cache but does **not** load any + weights into memory. Callers materialize the layers they actually need with + :meth:`initialize_layers`. The full set is available on disk after the + first call, so subsequent layer switches read from the local cache without + re-downloading. + + Examples + -------- + >>> model = EsmcSaeModel.from_pretrained( + ... "biohub/esmc-6b-2024-12-sae-k64-codebook16384" + ... ) + >>> model.initialize_layers([60]) # into memory + >>> out = model(layer_states, layer=60) # forward through layer 60 + >>> model.initialize_layers([45]) # add layer 45 (cached locally) + >>> model.release_layer(60) # free layer 60 + """ + + _device_marker: torch.Tensor + + def __init__(self, config: EsmcSaeConfig): + super().__init__() + self.config = config + # Layers are populated lazily by ``initialize_layers``; the container + # starts empty so ``from_pretrained`` doesn't materialize hundreds of + # GB of unused parameters. + self.layers = nn.ModuleDict() + # Zero-element buffer that rides along with ``.to(device/dtype)``. + # ``initialize_layers`` reads its current device/dtype so SAEs added + # after ``model.to("cuda")`` land on CUDA without re-passing ``device=``. + self.register_buffer("_device_marker", torch.empty(0), persistent=False) + self._snapshot_dir: str | None = None + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str | os.PathLike, + *, + revision: str | None = None, + cache_dir: str | os.PathLike | None = None, + token: str | None = None, + allow_patterns: list[str] | None = None, + local_files_only: bool = False, + force_download: bool = False, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> "EsmcSaeModel": + """Download (or reuse cached) the full repo and return the model. + + By default no weights are read into memory and the caller must invoke + :meth:`initialize_layers` before running :meth:`forward`. The single + exception is a repo that ships exactly one layer: that layer is + auto-loaded (honoring ``device`` / ``dtype``) so the bare + ``forward(x)`` call just works. + """ + local_dir = _resolve_snapshot_dir( + pretrained_model_name_or_path, + revision=revision, + cache_dir=cache_dir, + token=token, + allow_patterns=allow_patterns, + local_files_only=local_files_only, + force_download=force_download, + ) + config = EsmcSaeConfig.from_pretrained(local_dir) + model = cls(config) + model._snapshot_dir = str(local_dir) + if device is not None: + model.to(device) + if dtype is not None: + model.to(dtype) + assert config.available_layers is not None + if len(config.available_layers) == 1: + model.initialize_layers(list(config.available_layers)) + return model + + def initialize_layers( + self, + layers: list[int], + *, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + """Load the requested layers from the local snapshot into memory. + + Layers already present in :attr:`layers` are skipped - calling + ``initialize_layers([23])`` twice is idempotent. ``device`` / ``dtype`` + default to wherever the model itself lives (via the ``_device_marker`` + buffer that moves with ``.to(...)``). + """ + assert self._snapshot_dir is not None, ( + "EsmcSaeModel has no snapshot directory - call from_pretrained " + "first, or set _snapshot_dir manually." + ) + if device is None: + device = self._device_marker.device + if dtype is None: + dtype = self._device_marker.dtype + snapshot_dir = Path(self._snapshot_dir) + assert self.config.available_layers is not None + available = set(self.config.available_layers) + for layer_idx in layers: + key = str(layer_idx) + if key in self.layers: + continue + if layer_idx not in available: + raise KeyError( + f"Layer {layer_idx} is not in this repo. " + f"available_layers={sorted(available)}" + ) + shard = snapshot_dir / f"layer_{layer_idx}.safetensors" + if not shard.exists(): + raise FileNotFoundError( + f"Missing layer file {shard} - config lists layer " + f"{layer_idx} as available but the shard is not on disk." + ) + params = EsmcSaeParams( + d_model=self.config.d_model, + codebook_dim=self.config.codebook_dim, + k=self.config.k, + layer=layer_idx, + ) + # Build on the meta device so we don't allocate weights that + # ``load_state_dict`` would immediately overwrite. + with torch.device("meta"): + layer = EsmcSaeLayer(params) + layer.to_empty(device=device) + layer.load_state_dict(load_file(str(shard))) + layer.to(dtype=dtype) + self.layers[key] = layer + + def release_layer(self, layer: int) -> None: + """Drop the named layer from memory. No-op if not loaded.""" + key = str(layer) + if key in self.layers: + del self.layers[key] + + def loaded_layers(self) -> list[int]: + """Sorted list of layer indices currently materialized in memory.""" + return sorted(int(k) for k in self.layers.keys()) + + def forward( + self, x: torch.Tensor, layer: int | None = None, **kwargs: object + ) -> EsmcSaeOutput: + if layer is None: + if len(self.layers) == 1: + ((_only_key, only_layer),) = self.layers.items() + return only_layer(x, **kwargs) + if len(self.layers) == 0: + raise RuntimeError( + "No layers loaded - call initialize_layers([...]) first. " + f"available_layers={self.config.available_layers}" + ) + raise RuntimeError( + "Multiple layers are loaded - please select one via " + f"forward(x, layer=). Loaded layers: {self.loaded_layers()}" + ) + key = str(layer) + if key not in self.layers: + raise KeyError( + f"Layer {layer} is not loaded. Call initialize_layers([{layer}]) " + f"first. Loaded layers: {self.loaded_layers()}" + ) + return self.layers[key](x, **kwargs) + + def save_pretrained(self, save_directory: str | os.PathLike) -> None: + """Write ``config.json`` plus one ``layer_{i}.safetensors`` per loaded + layer. + + Only layers currently in :attr:`layers` are written. + ``available_layers`` in the saved config is synced to what's actually + on disk so a ``release_layer`` + ``save_pretrained`` round-trip never + advertises a layer whose shard is missing. + """ + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + self.config.available_layers = self.loaded_layers() + self.config.save_pretrained(str(save_directory)) + for key, layer in self.layers.items(): + shard = save_directory / f"layer_{key}.safetensors" + save_file( + { + k: v.detach().cpu().contiguous() + for k, v in layer.state_dict().items() + }, + str(shard), + ) + + +def _resolve_snapshot_dir( + pretrained_model_name_or_path: str | os.PathLike, + *, + revision: str | None, + cache_dir: str | os.PathLike | None, + token: str | None, + allow_patterns: list[str] | None, + local_files_only: bool, + force_download: bool, +) -> str: + """Local dir -> return as-is; hub id -> ``snapshot_download`` it. + + A directory only counts as "local" if it actually contains ``config.json``, + so a stale subdir named like a hub id doesn't accidentally shadow the hub + fetch. + """ + path = Path(pretrained_model_name_or_path) + if path.is_dir() and (path / CONFIG_NAME).exists(): + return str(path) + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=str(pretrained_model_name_or_path), + revision=revision, + cache_dir=None if cache_dir is None else str(cache_dir), + token=token, + allow_patterns=allow_patterns, + local_files_only=local_files_only, + force_download=force_download, + ) diff --git a/esm/models/esmc/tokenizer.py b/esm/models/esmc/tokenizer.py new file mode 100644 index 00000000..6a4b03b5 --- /dev/null +++ b/esm/models/esmc/tokenizer.py @@ -0,0 +1,177 @@ +"""Tokenization for ESMC.""" + +from tokenizers import AddedToken, Tokenizer +from tokenizers.models import BPE +from tokenizers.processors import TemplateProcessing +from transformers import PreTrainedTokenizerFast + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"} + +# Canonical amino acid vocabulary used by all ESMC checkpoints. +# Indices must be kept stable - they are hard-coded into the model weights. +SEQUENCE_VOCAB = [ + "", # 0 + "", # 1 + "", # 2 + "", # 3 + "L", # 4 + "A", # 5 + "G", # 6 + "V", # 7 + "S", # 8 + "E", # 9 + "R", # 10 + "T", # 11 + "I", # 12 + "D", # 13 + "P", # 14 + "K", # 15 + "Q", # 16 + "N", # 17 + "F", # 18 + "Y", # 19 + "M", # 20 + "H", # 21 + "W", # 22 + "C", # 23 + "X", # 24 ambiguous amino acid + "B", # 25 Asp/Asn ambiguity + "U", # 26 selenocysteine + "Z", # 27 Glu/Gln ambiguity + "O", # 28 pyrrolysine + ".", # 29 gap + "-", # 30 insertion + "|", # 31 chain-break + "", # 32 +] + + +class EsmcTokenizer(PreTrainedTokenizerFast): + r"""Construct an ESMC tokenizer. + + A character-level tokenizer backed by the HuggingFace ``tokenizers`` + library. It wraps every sequence with ```` and ```` tokens and + supports a ``|`` chain-break token for multi-chain inputs. + + Parameters + ---------- + unk_token : str + The unknown token. + cls_token : str + The classification token (prepended to every sequence). + pad_token : str + The padding token. + mask_token : str + The mask token, used for masked language modelling. + eos_token : str + The end-of-sequence token (appended to every sequence). + chain_break_token : str + Token inserted between chains in multi-chain protein inputs. + + Examples + -------- + >>> tokenizer = EsmcTokenizer() + >>> tokenizer("ACDEFGHIKLMNPQRSTVWY")["input_ids"] + [0, 5, 23, 13, 18, 9, 6, 21, 12, 15, 20, 17, 14, 16, 10, 8, 11, 7, 22, 19, 2] + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + unk_token="", + cls_token="", + pad_token="", + mask_token="", + eos_token="", + chain_break_token="|", + **kwargs, + ): + all_tokens = SEQUENCE_VOCAB + token_to_id = {tok: ind for ind, tok in enumerate(all_tokens)} + + # Normalise: always work with plain strings. + unk_token = self._ensure_str(unk_token) + cls_token = self._ensure_str(cls_token) + pad_token = self._ensure_str(pad_token) + mask_token = self._ensure_str(mask_token) + eos_token = self._ensure_str(eos_token) + chain_break_token = self._ensure_str(chain_break_token) + + # A character-level tokenizer is equivalent to BPE with no merges. + bpe = BPE(token_to_id, merges=[], unk_token=unk_token) + tokenizer = Tokenizer(bpe) + + special_tokens = [ + cls_token, + pad_token, + mask_token, + eos_token, + chain_break_token, + ] + tokenizer.add_special_tokens(special_tokens) + + # Automatically wrap every encoded sequence with ... . + tokenizer.post_processor = TemplateProcessing( + single=f"{cls_token} $A {eos_token}", + special_tokens=[ + (cls_token, tokenizer.token_to_id(cls_token)), + (eos_token, tokenizer.token_to_id(eos_token)), + ], + ) + + # Expose the chain-break token as an additional special token so it is + # preserved during encode/decode and can be looked up easily. + kwargs.setdefault("additional_special_tokens", []) + if chain_break_token not in kwargs["additional_special_tokens"]: + kwargs["additional_special_tokens"] = list( + kwargs["additional_special_tokens"] + ) + [chain_break_token] + + # Keep a reference before super().__init__ so the properties below work. + self._chain_break_token = chain_break_token + + super().__init__( + tokenizer_object=tokenizer, + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + **kwargs, + ) + + # The model uses as the sequence-start token; the base class expects + # ``bos_token``. Alias them to avoid confusion. + @property + def bos_token(self): + return self.cls_token + + @property + def bos_token_id(self): + return self.cls_token_id + + @property + def chain_break_token(self) -> str: + return self._chain_break_token + + @property + def chain_break_token_id(self) -> int: + token_id = self.convert_tokens_to_ids(self._chain_break_token) + assert isinstance(token_id, int) + return token_id + + @property + def all_token_ids(self): + return list(range(self.vocab_size)) + + @property + def special_token_ids(self): + return self.all_special_ids + + @staticmethod + def _ensure_str(token) -> str: + if isinstance(token, AddedToken): + return token.content + return str(token) diff --git a/esm/models/esmfold2/__init__.py b/esm/models/esmfold2/__init__.py index 972db0b9..b95b8267 100644 --- a/esm/models/esmfold2/__init__.py +++ b/esm/models/esmfold2/__init__.py @@ -1,3 +1,10 @@ +from typing import TYPE_CHECKING + +from esm.models.esmfold2.config import ( + ESMFOLD2_EXPERIMENTAL_HF_REPO, + ESMFOLD2_HF_REPO, + EsmFold2Config, +) from esm.models.esmfold2.conformers import load_ccd from esm.models.esmfold2.constants import ELEMENT_NUMBER_TO_SYMBOL from esm.models.esmfold2.prepare_input import ChainInfo, prepare_esmfold2_input @@ -19,9 +26,35 @@ MolecularComplexResult, ) +if TYPE_CHECKING: + from esm.models.esmfold2.experimental import EsmFold2ExperimentalModel + from esm.models.esmfold2.model import EsmFold2Model + + +def __getattr__(name: str) -> object: + if name == "EsmFold2ExperimentalModel": + from esm.models.esmfold2.experimental import EsmFold2ExperimentalModel + + return EsmFold2ExperimentalModel + if name == "EsmFold2Model": + from esm.models.esmfold2.model import EsmFold2Model + + return EsmFold2Model + raise AttributeError(name) + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) + + __all__ = [ + "ESMFOLD2_EXPERIMENTAL_HF_REPO", + "ESMFOLD2_HF_REPO", "ChainInfo", "CovalentBond", + "EsmFold2Config", + "EsmFold2ExperimentalModel", + "EsmFold2Model", "DistogramConditioning", "DNAInput", "ELEMENT_NUMBER_TO_SYMBOL", diff --git a/esm/models/esmfold2/config.py b/esm/models/esmfold2/config.py new file mode 100644 index 00000000..fcb09aa0 --- /dev/null +++ b/esm/models/esmfold2/config.py @@ -0,0 +1,517 @@ +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ESMFold2 model configuration.""" + +from __future__ import annotations + +import copy +import json +import os +import warnings +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, cast + +from esm.models.esmc.config import EsmcConfig + +CONFIG_NAME = "config.json" + +# Published ESMFold2 checkpoints on the HuggingFace Hub. +ESMFOLD2_HF_REPO = "biohub/ESMFold2" +ESMFOLD2_EXPERIMENTAL_HF_REPO = "biohub/ESMFold2-Experimental" + +# --------------------------------------------------------------------------- +# Nested dataclass configs +# --------------------------------------------------------------------------- + +_DEFAULT_ESMC_HF_REPO = "biohub/ESMC-6B" + +# Dotted source path in a pre-alignment ``config.json`` -> dotted path now. Field +# names follow the HuggingFace port so one config.json describes the model to +# both implementations. Delete once every ``biohub/ESMFold2*`` repo has been +# re-published; until then a strict read would reject every live checkpoint. +_LEGACY_PATHS = { + "d_single": "hidden_size", + "d_pair": "pairwise_hidden_size", + "inputs.d_inputs": "single_inputs_size", + "inputs.atom_encoder.swa_window_size": "sliding_window", + "inputs.atom_encoder.d_atom": "atom_encoder.hidden_size", + "inputs.atom_encoder.d_token": "atom_encoder.output_dim", + "inputs.atom_encoder.n_blocks": "atom_encoder.num_hidden_layers", + "inputs.atom_encoder.n_heads": "atom_encoder.num_attention_heads", + "inputs.atom_encoder.expansion_ratio": "atom_encoder.expansion_ratio", + "inputs.atom_encoder.spatial_rope_base_frequency": ( + "atom_encoder.spatial_rope_base_frequency" + ), + "inputs.atom_encoder.n_spatial_rope_pairs_per_axis": ( + "atom_encoder.n_spatial_rope_pairs_per_axis" + ), + "inputs.atom_encoder.n_uid_rope_pairs": "atom_encoder.n_uid_rope_pairs", + "inputs.atom_encoder.uid_rope_base_frequency": ( + "atom_encoder.uid_rope_base_frequency" + ), + "folding_trunk.n_layers": "folding_trunk_num_hidden_layers", + "folding_trunk.n_heads": "folding_trunk_num_attention_heads", + "folding_trunk.dropout": "folding_trunk_dropout", + "parcae.coda_n_layers": "parcae_num_coda_layers", + "msa_encoder_overwrite": "msa_encoder.overwrite", + "msa_encoder.d_msa": "msa_encoder.hidden_size", + "msa_encoder.d_hidden": "msa_encoder.outer_hidden_size", + "msa_encoder.n_layers": "msa_encoder.num_hidden_layers", + "msa_encoder.n_heads_msa": "msa_encoder.num_attention_heads", + "msa_encoder.msa_head_width": "msa_encoder.head_width", + "lm_encoder.n_layers": "lm_encoder.num_hidden_layers", + "confidence_head.folding_trunk.n_layers": "confidence_head.num_hidden_layers", + "confidence_head.folding_trunk.n_heads": "confidence_head.num_attention_heads", + "confidence_head.folding_trunk.dropout": "confidence_head.dropout", + "structure_head.diffusion_module.c_atom": ( + "structure_head.diffusion_module.atom_encoder.hidden_size" + ), + "structure_head.diffusion_module.atom_num_blocks": ( + "structure_head.diffusion_module.atom_encoder.num_hidden_layers" + ), + "structure_head.diffusion_module.atom_num_heads": ( + "structure_head.diffusion_module.atom_encoder.num_attention_heads" + ), + "structure_head.diffusion_module.c_token": ( + "structure_head.diffusion_module.token_hidden_size" + ), +} + + +def _pop_path(tree: dict, path: str): + """Remove and return ``path`` from a nested dict; ``_MISSING`` when absent.""" + parts = path.split(".") + node = tree + for part in parts[:-1]: + if not isinstance(node, dict) or part not in node: + return _MISSING + node = node[part] + if not isinstance(node, dict) or parts[-1] not in node: + return _MISSING + return node.pop(parts[-1]) + + +def _set_path(tree: dict, path: str, value) -> None: + parts = path.split(".") + node = tree + for part in parts[:-1]: + node = node.setdefault(part, {}) + node.setdefault(parts[-1], value) + + +_MISSING = object() + + +def _prune_empty(tree: dict) -> None: + """Drop dicts left empty by remapping, innermost first. + + The pre-alignment wrappers (``inputs``, ``folding_trunk``, + ``confidence_head.folding_trunk``) held nothing but remapped leaves, so once + those move out the wrapper is noise that would otherwise be kept as a stray + attribute and re-serialised. + """ + for key in list(tree): + value = tree[key] + if not isinstance(value, dict): + continue + _prune_empty(value) + if not value: + del tree[key] + + +def _resolve_legacy_keys(kwargs: dict) -> dict: + """Map a pre-alignment ``config.json`` onto the current field names.""" + if "num_recycles" in kwargs: + raise ValueError( + "config.json uses 'num_recycles'; ESMFold2 calls this 'num_loops'. " + "Re-export the checkpoint's config with the current field name " + "rather than letting the trunk silently fall back to the default " + "loop count." + ) + tree = copy.deepcopy(kwargs) + moved = [] + for old, new in _LEGACY_PATHS.items(): + value = _pop_path(tree, old) + if value is _MISSING: + continue + _set_path(tree, new, value) + moved.append(old) + _prune_empty(tree) + if moved: + warnings.warn( + f"config.json uses {len(moved)} pre-alignment ESMFold2 field " + f"path(s), e.g. {sorted(moved)[:3]}; see EsmFold2Config for the " + f"current names.", + FutureWarning, + stacklevel=3, + ) + return tree + + +@dataclass +class EsmFold2MsaEncoderConfig: + """Config for the optional MSA encoder module (Large MSA models only).""" + + enabled: bool = False + # If True, MSA encoder output replaces the pair stream; if False, it is added. + overwrite: bool = True + hidden_size: int = 128 + outer_hidden_size: int = 32 + num_hidden_layers: int = 4 + num_attention_heads: int = 8 + head_width: int = 32 + # Inference-time MSA diversity: in config, so a checkpoint ships the values + # it was tuned for rather than a call signature deciding them. + max_depth: int | None = 1024 + column_mask_rate: float = 0.1 + + +@dataclass +class EsmFold2ParcaeConfig: + """Release-only config for the parcae diffusion-loop scheduler.""" + + enabled: bool = True + poisson_mean: float = 3.0 + min_steps: int = 1 + max_steps: int | None = 6 + + +@dataclass +class EsmFold2LmEncoderConfig: + """Release-only config for the LM-side pair encoder.""" + + enabled: bool = True + num_hidden_layers: int = 4 + lm_dropout: float = 0.25 + per_loop_lm_dropout: bool = True + + +@dataclass +class EsmFold2AtomEncoderConfig: + """Config for SWA atom encoder/decoder with 3D RoPE. + + The attention window lives on the parent as ``sliding_window``, shared with + the structure head's atom encoder. + """ + + hidden_size: int = 128 + output_dim: int = 768 + num_hidden_layers: int = 3 + num_attention_heads: int = 4 + expansion_ratio: int = 2 + # 3D RoPE config + spatial_rope_base_frequency: float = 20.0 + n_spatial_rope_pairs_per_axis: int = 2 + n_uid_rope_pairs: int = 10 + uid_rope_base_frequency: float = 10000.0 + + +@dataclass +class EsmFold2DiffusionAtomEncoderConfig: + """Widths of the diffusion module's own atom encoder.""" + + hidden_size: int = 128 + num_hidden_layers: int = 3 + num_attention_heads: int = 4 + + +@dataclass +class EsmFold2DiffusionModuleConfig: + """Config for the DiffusionModule.""" + + sigma_data: float = 16.0 + token_hidden_size: int = 768 + fourier_dim: int = 256 + token_num_blocks: int = 12 + token_num_heads: int = 16 + transition_multiplier: int = 2 + atom_encoder: EsmFold2DiffusionAtomEncoderConfig = field( + default_factory=EsmFold2DiffusionAtomEncoderConfig + ) + # Re-derived by the port from the parent's widths; kept explicit here. + c_z: int = 256 + c_s_inputs: int = 451 + relpos_r_max: int = 32 + relpos_s_max: int = 2 + + def __post_init__(self): + if isinstance(self.atom_encoder, dict): + self.atom_encoder = EsmFold2DiffusionAtomEncoderConfig( + **cast(dict[str, Any], self.atom_encoder) + ) + + +@dataclass +class EsmFold2StructureHeadConfig: + """Config for the diffusion-based structure prediction head.""" + + diffusion_module: EsmFold2DiffusionModuleConfig = field( + default_factory=EsmFold2DiffusionModuleConfig + ) + distogram_bins: int = 128 + + # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1)) + train_noise_log_mean: float = -1.2 + train_noise_log_std: float = 1.5 + + # Sampling defaults (ODE) + gamma_0: float = 0.605 + gamma_min: float = 1.107 + noise_scale: float = 0.0 + step_scale: float = 1.0 + + # Inference schedule defaults + inference_s_max: float = 160.0 + inference_s_min: float = 4e-4 + inference_p: float = 8.0 + inference_num_steps: int = 68 + + def __post_init__(self): + if isinstance(self.diffusion_module, dict): + self.diffusion_module = EsmFold2DiffusionModuleConfig( + **self.diffusion_module # ty:ignore[invalid-argument-type] + ) + + +@dataclass +class EsmFold2ConfidenceHeadConfig: + enabled: bool = True + num_hidden_layers: int = 4 + num_plddt_bins: int = 50 + num_pde_bins: int = 64 + num_pae_bins: int = 64 + min_dist: float = 2.0 + max_dist: float = 52.0 + distogram_bins: int = 128 + # Not carried by the port, which fixes both; kept so a checkpoint can set them. + num_attention_heads: int = 8 + dropout: float = 0.0 + + +# --------------------------------------------------------------------------- +# Top-level config +# --------------------------------------------------------------------------- + + +class EsmFold2Config: + """ + Configuration for the ESMFold2 structure prediction model. + + Uses SWA atom encoders with 3D RoPE, a diffusion transformer, + a folding trunk, and an ESMC 6B PLM backbone. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control + the model outputs. Read the documentation from [`PretrainedConfig`] for more + information. + + Args: + d_single (`int`, defaults to 384): + Dimensionality of single (per-residue) representations. + d_pair (`int`, defaults to 256): + Dimensionality of pair (residue-residue) representations. + n_relative_residx_bins (`int`, defaults to 32): + Number of bins for relative residue index encoding. + n_relative_chain_bins (`int`, defaults to 2): + Number of bins for relative chain encoding. + num_loops (`int`, defaults to 20): + Number of trunk loops for iterative refinement. + num_diffusion_samples (`int`, defaults to 8): + Number of parallel structure predictions to generate. + lm_dropout (`float`, defaults to 0.0): + Dropout probability on LM pair embeddings. When > 0, dropout is + applied with ``training=True`` (including at inference) to match + the experimental training recipe used by binder design. + force_lm_dropout_during_inference (`bool`, defaults to False): + When True, apply ``lm_dropout`` even when ``model.eval()`` and + ``lm_dropout`` > 0. Binder-design loads set this to True. + lm_mask_pct (`float`, defaults to 0.0): + Fraction of sequence residues randomly replaced with the LM mask + token before running the PLM backbone, matching the training-time + input corruption. Single-sequence checkpoints set this to 0.1. + disable_msa_features (`bool`, defaults to False): + When True, zero out MSA-derived ``profile`` and ``deletion_mean`` + before the inputs embedder (experimental medium/large checkpoints). + single_inputs_size (`int`, defaults to 451): + Width of the per-residue input features. + sliding_window (`int`, defaults to 128): + Atom-attention window, shared by the inputs embedder and the + structure head's atom encoder. + folding_trunk_num_hidden_layers (`int`, defaults to 24): + Number of pairformer layers in the folding trunk. + parcae_num_coda_layers (`int`, defaults to 2): + Number of coda layers in the release-only parcae scheduler. + atom_encoder (`EsmFold2AtomEncoderConfig`): + Configuration for the inputs embedder's atom encoder. + structure_head (`EsmFold2StructureHeadConfig`): + Configuration for the diffusion-based structure prediction head. + confidence_head (`EsmFold2ConfidenceHeadConfig`): + Configuration for the confidence prediction head. + + Examples: + + ```python + >>> from transformers import EsmFold2Config, EsmFold2ExperimentalModel + + >>> # Initializing an ESMFold2 configuration + >>> configuration = EsmFold2Config(type="experimental") + + >>> # Initializing a model (with random weights) from the configuration + >>> model = EsmFold2ExperimentalModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "esmfold2" + has_no_defaults_at_init = True + + def __init__(self, **kwargs): + kwargs = _resolve_legacy_keys(kwargs) + # PretrainedConfig kept unrecognised kwargs as attributes; preserve that + # so nothing in a published config.json is silently discarded. The + # explicit assignments below then take precedence. + for key, value in kwargs.items(): + setattr(self, key, value) + + # Default "release" so a bare ``EsmFold2Config()`` (used internally by + # HF's ``save_pretrained``) doesn't raise ``KeyError('type')``. + self.type: str = kwargs.get("type", "release") + if self.type not in ("release", "experimental"): + raise ValueError( + f"EsmFold2Config.type must be 'release' or 'experimental', " + f"got {self.type!r}" + ) + + # Top-level scalar fields + self.hidden_size: int = kwargs.get("hidden_size", 384) + self.pairwise_hidden_size: int = kwargs.get("pairwise_hidden_size", 256) + self.single_inputs_size: int = kwargs.get("single_inputs_size", 451) + # Shared by the inputs embedder and the structure head's atom encoder. + self.sliding_window: int = kwargs.get("sliding_window", 128) + self.folding_trunk_num_hidden_layers: int = kwargs.get( + "folding_trunk_num_hidden_layers", 24 + ) + self.parcae_num_coda_layers: int = kwargs.get("parcae_num_coda_layers", 2) + # Not carried by the port, which fixes both; kept so a checkpoint can set them. + self.folding_trunk_num_attention_heads: int = kwargs.get( + "folding_trunk_num_attention_heads", 8 + ) + self.folding_trunk_dropout: float = kwargs.get("folding_trunk_dropout", 0.0) + self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32) + self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2) + self.num_loops: int = kwargs.get("num_loops", 20) + self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8) + # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs + # embedder. + self.disable_msa_features: bool = kwargs.get("disable_msa_features", False) + self.lm_dropout: float = kwargs.get("lm_dropout", 0.0) + self.force_lm_dropout_during_inference: bool = kwargs.get( + "force_lm_dropout_during_inference", False + ) + self.lm_mask_pct: float = kwargs.get("lm_mask_pct", 0.0) + + self.lm_d_model: int = kwargs.get("lm_d_model", 2560) + self.lm_num_layers: int = kwargs.get("lm_num_layers", 80) + # ``esmc_config`` describes a backbone bundled into the same checkpoint, + # so it loads in one pass. ``esmc_id`` is the older arrangement, where the + # backbone is a separate repo fetched afterwards by ``load_esmc``. + esmc_config = kwargs.get("esmc_config") + if isinstance(esmc_config, dict): + esmc_config = EsmcConfig.from_dict(esmc_config) + self.esmc_config: EsmcConfig | None = esmc_config + self.esmc_id: str = kwargs.get("esmc_id", _DEFAULT_ESMC_HF_REPO) + + def _init_nested(cls, val): + if isinstance(val, cls): + return val + if isinstance(val, dict): + return cls(**val) + return cls() + + self.atom_encoder = _init_nested( + EsmFold2AtomEncoderConfig, kwargs.get("atom_encoder") + ) + self.structure_head = _init_nested( + EsmFold2StructureHeadConfig, kwargs.get("structure_head") + ) + self.confidence_head = _init_nested( + EsmFold2ConfidenceHeadConfig, kwargs.get("confidence_head") + ) + self.msa_encoder = _init_nested( + EsmFold2MsaEncoderConfig, kwargs.get("msa_encoder") + ) + # Release-only modules — ignored when ``type == "experimental"``. + self.parcae = _init_nested(EsmFold2ParcaeConfig, kwargs.get("parcae")) + self.lm_encoder = _init_nested( + EsmFold2LmEncoderConfig, kwargs.get("lm_encoder") + ) + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str | os.PathLike, **kwargs: Any + ) -> "EsmFold2Config": + """Read ``config.json`` from a local directory or the HuggingFace Hub.""" + from esm.models.hub import resolve_model_dir + + local_dir = resolve_model_dir( + pretrained_model_name_or_path, + revision=kwargs.pop("revision", None), + cache_dir=kwargs.pop("cache_dir", None), + token=kwargs.pop("token", None), + local_files_only=kwargs.pop("local_files_only", False), + force_download=kwargs.pop("force_download", False), + ) + with open(Path(local_dir) / CONFIG_NAME) as f: + raw = json.load(f) + raw.pop("model_type", None) + # Caller overrides win, matching PretrainedConfig.from_pretrained. + raw.update(kwargs) + return cls(**raw) + + def save_pretrained(self, save_directory: str | os.PathLike) -> None: + directory = Path(save_directory) + directory.mkdir(parents=True, exist_ok=True) + with open(directory / CONFIG_NAME, "w") as f: + json.dump(self.to_dict(), f, indent=2, sort_keys=True) + + def to_dict(self): + output = {k: v for k, v in self.__dict__.items()} + output["model_type"] = self.model_type + if self.esmc_config is not None: + output["esmc_config"] = self.esmc_config.to_dict() + for name in ( + "atom_encoder", + "structure_head", + "confidence_head", + "msa_encoder", + "parcae", + "lm_encoder", + ): + output[name] = asdict(getattr(self, name)) + return output + + +__all__ = [ + "EsmFold2AtomEncoderConfig", + "EsmFold2Config", + "EsmFold2ConfidenceHeadConfig", + "EsmFold2DiffusionAtomEncoderConfig", + "EsmFold2DiffusionModuleConfig", + "EsmFold2LmEncoderConfig", + "EsmFold2MsaEncoderConfig", + "EsmFold2ParcaeConfig", + "EsmFold2StructureHeadConfig", +] diff --git a/esm/models/esmfold2/experimental.py b/esm/models/esmfold2/experimental.py new file mode 100644 index 00000000..0936a2e1 --- /dev/null +++ b/esm/models/esmfold2/experimental.py @@ -0,0 +1,1180 @@ +# coding=utf-8 +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ESMFold2 experimental variant — older architecture from during the +development of ESMFold2. + +Most users want :class:`EsmFold2Model` (in ``modeling_esmfold2``) instead; +this module retains an explicit refinement loop that re-injects the +previous pair representation through ``pair_loop_proj`` each iteration, +and exists to load checkpoints predating the standard architecture. +""" + +from __future__ import annotations + +from typing import cast + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from esm.models.esmc import EsmcModel +from esm.models.esmc.checkpoint_layout import published_to_native_subtree +from esm.models.esmfold2.config import EsmFold2Config +from esm.models.esmfold2.layers import ( + CHAR_VOCAB_SIZE, + MAX_ATOMIC_NUMBER, + NUM_RES_TYPES, + DiffusionStructureHead, + FoldingTrunk, + InputsEmbedder, + LanguageModelShim, + MSAPairWeightedAveraging, + OuterProductMean, + ResIdxAsymIdSymIdEntityIdEncoding, + RowAttentionPooling, + SwiGLUMLP, + TriangleMultiplicativeUpdate, + _categorical_mean, + _compute_intra_token_idx, + _seed_context, + compute_lm_hidden_states, + gather_rep_atom_coords, + gather_token_to_atom, + maybe_apply_msa_column_masking, + maybe_subsample_msa, +) + +# Not circular: model.py's reference to this module is inside ``from_pretrained``. +from esm.models.esmfold2.model import _IGNORED_FEATURE_KEYS +from esm.models.hub import HubPreTrainedModel, resolve_model_dir + +_EPS = 1e-5 +_NONPOLYMER_ID: int = 3 + + +# =========================================================================== +# ConfidenceHead +# =========================================================================== + + +class ConfidenceHead(nn.Module): + """Confidence head predicting per-atom pLDDT and pairwise PAE.""" + + boundaries: Tensor + + def __init__(self, config: EsmFold2Config) -> None: + super().__init__() + ch = config.confidence_head + d_single = config.hidden_size + d_pair = config.pairwise_hidden_size + d_inputs = config.single_inputs_size + + # Distogram bins boundary buffer + boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1) + self.register_buffer("boundaries", boundaries) + + self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair) + + self.s_norm = nn.LayerNorm(d_single) + + self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False) + + self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False) + + # s_to_z_prod + self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False) + + # s_input_to_s + self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False) + + self.s_inputs_norm = nn.LayerNorm(d_inputs) + self.z_norm = nn.LayerNorm(d_pair) + + # Row attention pooling + self.row_attention_pooling = RowAttentionPooling( + d_pair=d_pair, d_single=d_single + ) + + # Confidence folding trunk (4 blocks) + self.folding_trunk = FoldingTrunk( + n_layers=ch.num_hidden_layers, d_pair=d_pair, expansion_ratio=4 + ) + + # pLDDT head + self.plddt_ln = nn.LayerNorm(d_single) + max_atoms_per_token = 23 + self.plddt_weight = nn.Parameter( + torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins) + ) + + # PAE head + self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False) + + # ------------------------------------------------------------------ + # Kernel / chunking configuration + # ------------------------------------------------------------------ + + def set_kernel_backend(self, backend: str | None) -> None: + self.folding_trunk.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor: + if num_diffusion_samples == 1: + return x + return x.repeat_interleave(num_diffusion_samples, 0) + + @staticmethod + def _flatten_sample_axis(x: Tensor) -> Tensor: + if x.ndim == 4: + b, mult, n, c = x.shape + return x.reshape(b * mult, n, c) + return x + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + s_inputs: Tensor, + z: Tensor, + x_pred: Tensor, + distogram_atom_idx: Tensor, + token_attention_mask: Tensor, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + asym_id: Tensor, + mol_type: Tensor, + num_diffusion_samples: int = 1, + relative_position_encoding: Tensor | None = None, + token_bonds_encoding: Tensor | None = None, + ) -> dict[str, Tensor]: + """Run confidence head.""" + # Shared computation (batch-scale, before num_diffusion_samples expansion) + s_inputs_normed = self.s_inputs_norm(s_inputs) + + z_base = self.z_norm(z) + if relative_position_encoding is not None: + z_base = z_base + relative_position_encoding + if token_bonds_encoding is not None: + z_base = z_base + token_bonds_encoding + z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2) + z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1) + z_base = z_base + self.s_to_z_prod_out( + self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :] + * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :] + ) + + # Expand to num_diffusion_samples + pair = self._repeat_batch(z_base, num_diffusion_samples) + x_pred_flat = self._flatten_sample_axis(x_pred) + atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples) + atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples) + rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long() + mask = self._repeat_batch(token_attention_mask, num_diffusion_samples) + Bm = pair.shape[0] + + # Distogram from predicted coords + rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m) + rep_distances = torch.cdist( + rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" + ) + distogram_bins = ( + (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() + ) + pair = pair + self.dist_bin_pairwise_embed(distogram_bins) + + # Expand 1-D token mask → 2-D pair mask for folding trunk + pair_mask = mask[:, :, None].float() * mask[:, None, :].float() # [B*m, L, L] + + # FoldingTrunk + row attention pooling -> single + pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask) + single = self.row_attention_pooling(pair, mask) + + # Per-atom pLDDT + atom_mask_f = atom_mask_m.float() + s_at_atoms = gather_token_to_atom(single, atom_to_token_m) + s_at_atoms = self.plddt_ln(s_at_atoms) + + intra_idx = _compute_intra_token_idx(atom_to_token_m) + intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1) + w = self.plddt_weight[intra_idx] # [B*m, A, d_single, num_bins] + plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, w) + + plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0) + + # Per-token pLDDT (scatter mean) + L = single.shape[1] + plddt_sum = torch.zeros(Bm, L, device=single.device, dtype=plddt_per_atom.dtype) + atom_count = torch.zeros( + Bm, L, device=single.device, dtype=plddt_per_atom.dtype + ) + atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype) + plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t) + atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t) + plddt = plddt_sum / atom_count.clamp(min=1e-6) + + # Complex pLDDT (flat mean over all atoms) + complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / ( + atom_mask_f.sum(dim=-1) + _EPS + ) + + # Complex ipLDDT (interface-weighted) + expanded_type = self._repeat_batch(mol_type, num_diffusion_samples) + expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples) + is_ligand = (expanded_type == _NONPOLYMER_ID).float() + inter_chain = ( + expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) + ).float() + near_contact = (rep_distances < 8).float() + interface_per_token = ( + near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1) + ).amax(dim=-1) + iplddt_weight = torch.where( + is_ligand.bool(), + torch.full_like(interface_per_token, 2.0), + interface_per_token, + ) + iplddt_weight_atoms = gather_token_to_atom( + iplddt_weight.unsqueeze(-1), atom_to_token_m + ).squeeze(-1) + atom_iplddt_w = atom_mask_f * iplddt_weight_atoms + complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / ( + atom_iplddt_w.sum(dim=-1) + _EPS + ) + + # pLDDT at CA / representative atom + plddt_ca = plddt_per_atom.gather(1, rep_idx_m) + + # PAE + pae_logits = self.pae_head(pair) + pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach() + + # pTM / ipTM / per-chain-pair ipTM derived from pae_logits. + n_bins = pae_logits.shape[-1] + bin_width = 32.0 / n_bins + bin_centers = torch.arange( + 0.5 * bin_width, 32.0, bin_width, device=pae_logits.device + ) + mask_f = mask.float() + N_res = mask_f.sum(dim=-1, keepdim=True) + d0 = 1.24 * (N_res.clamp(min=19) - 15) ** (1 / 3) - 1.8 # [Bm, 1] + tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2) # [Bm, n_bins] + pae_probs = F.softmax(pae_logits, dim=-1) + tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum( + dim=-1 + ) # [Bm, L, L] + + pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2) # [Bm, L, L] + + # pTM: avg over all valid pairs per row, max over rows. + ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / ( + pair_mask_2d.sum(dim=-1) + _EPS + ) + ptm = ptm_per_row.max(dim=-1).values # [Bm] + + # ipTM: avg over inter-chain valid pairs per row, max over rows. + inter_chain_mask = ( + expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) + ).float() * pair_mask_2d + iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / ( + inter_chain_mask.sum(dim=-1) + _EPS + ) + iptm = iptm_per_row.max(dim=-1).values # [Bm] + + # Per-chain-pair ipTM: dense [Bm, N_chains, N_chains] padded to max chain id + 1. + max_chain_id = int(expanded_asym.max().item()) if Bm > 0 else 0 + n_chains = max_chain_id + 1 + 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 + 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, + "plddt": plddt.detach(), + "plddt_per_atom": plddt_per_atom.detach(), + "plddt_ca": plddt_ca.detach(), + "complex_plddt": complex_plddt.detach(), + "complex_iplddt": complex_iplddt.detach(), + "pae_logits": pae_logits, + "pae": pae, + "ptm": ptm.detach(), + "iptm": iptm.detach(), + "pair_chains_iptm": pair_chains_iptm.detach(), + } + + +# =========================================================================== +# MSA Encoder +# =========================================================================== + + +class _TransitionFFN(nn.Module): + """LayerNorm + SwiGLU FFN without residual (used inside MSAEncoderBlock).""" + + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + + def forward(self, x: Tensor) -> Tensor: + return self.ffn(self.norm(x)) + + +class MSAEncoderBlock(nn.Module): + """One block of the MSA encoder: MSA update + pair update.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_hidden: int = 32, + n_heads_msa: int = 8, + msa_head_width: int = 32, + ) -> None: + super().__init__() + self.outer_product_mean = OuterProductMean( + d_msa, d_hidden, d_pair, divide_outer_before_proj=True + ) + self.msa_pair_weighted_averaging = MSAPairWeightedAveraging( + d_msa, d_pair, n_heads_msa, msa_head_width + ) + self.msa_transition = _TransitionFFN(d_msa, expansion_ratio=4) + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = _TransitionFFN(d_pair, expansion_ratio=4) + + # _TransitionFFN is a plain norm+SwiGLU with neither knob. + def set_kernel_backend(self, backend: str | None) -> None: + self.tri_mul_out.set_kernel_backend(backend) + self.tri_mul_in.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.outer_product_mean.set_chunk_size(chunk_size) + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + + def forward( + self, + msa_repr: Tensor, + pair_repr: Tensor, + msa_attention_mask: Tensor, + pair_attention_mask: Tensor, + msa_track_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + """ + Args: + msa_repr: [B, L, M, d_msa] + pair_repr: [B, L, L, d_pair] + msa_attention_mask: [B, L, M] + pair_attention_mask:[B, L, L] + msa_track_mask: [B] bool — if False for a sample, zero out its contribution + Returns: + (msa_repr, pair_repr) + """ + mask4d = ( + msa_track_mask[:, None, None, None].to(dtype=msa_repr.dtype) + if msa_track_mask is not None + else None + ) + + def _maybe_mask(x: Tensor) -> Tensor: + return x * mask4d if mask4d is not None else x + + msa_repr = msa_repr + _maybe_mask( + self.msa_pair_weighted_averaging(msa_repr, pair_repr, pair_attention_mask) + ) + msa_repr = msa_repr + _maybe_mask(self.msa_transition(msa_repr)) + + pair_repr = pair_repr + _maybe_mask( + self.outer_product_mean(msa_repr, msa_attention_mask) + ) + pair_repr = pair_repr + _maybe_mask( + self.tri_mul_out(pair_repr, mask=pair_attention_mask) + ) + pair_repr = pair_repr + _maybe_mask( + self.tri_mul_in(pair_repr, mask=pair_attention_mask) + ) + pair_repr = pair_repr + _maybe_mask(self.pair_transition(pair_repr)) + + return msa_repr, pair_repr + + +class MSAEncoder(nn.Module): + """Embeds MSA features and runs encoder blocks to update the pair representation.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_inputs: int, + d_hidden: int = 32, + n_layers: int = 4, + n_heads_msa: int = 8, + msa_head_width: int = 32, + ) -> None: + super().__init__() + # 33 aa one-hot + has_deletion + deletion_value = 35 + self.embed = nn.Linear(35, d_msa, bias=False) + self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False) + self.blocks = nn.ModuleList( + [ + MSAEncoderBlock( + d_msa=d_msa, + d_pair=d_pair, + d_hidden=d_hidden, + n_heads_msa=n_heads_msa, + msa_head_width=msa_head_width, + ) + for _ in range(n_layers) + ] + ) + + def set_kernel_backend(self, backend: str | None) -> None: + for block in self.blocks: + cast(MSAEncoderBlock, block).set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(MSAEncoderBlock, block).set_chunk_size(chunk_size) + + def forward( + self, + x_pair: Tensor, + x_inputs: Tensor, + msa_oh: Tensor, + has_deletion: Tensor, + deletion_value: Tensor, + msa_attention_mask: Tensor, + ) -> Tensor: + """ + Args: + x_pair: [B, L, L, d_pair] current pair representation + x_inputs: [B, L, d_inputs] per-token input features + msa_oh: [B, L, M, 33] one-hot MSA (already transposed) + has_deletion: [B, L, M] + deletion_value: [B, L, M] + msa_attention_mask:[B, L, M] + Returns: + [B, L, L, d_pair] pair update + """ + B, L, M = msa_attention_mask.shape + + m_feat = torch.cat( + [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1 + ) + m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2) + + # Mask out the full update for samples with no real non-query MSA rows. + if M > 1: + msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2)) + else: + msa_track_mask = torch.zeros(B, dtype=torch.bool, device=x_pair.device) + + tok_mask = msa_attention_mask[:, :, 0] + pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1) + + for block in self.blocks: + m, x_pair = block( + m, x_pair, msa_attention_mask, pair_attention_mask, msa_track_mask + ) + + x_pair = x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype) + return x_pair + + +# =========================================================================== +# EsmFold2ExperimentalModel — the top-level model +# =========================================================================== + + +class EsmFold2ExperimentalModel(HubPreTrainedModel): + """ESMFold2 v2 structure prediction model.""" + + config_class = EsmFold2Config + _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"] + + @classmethod + def _normalize_checkpoint_layout( + cls, raw: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Translate a bundled ESMC encoder out of the published tensor layout. + + Scoped to the ``esmc.`` subtree so the trunk's own keys can never be + caught by the encoder's key patterns. + """ + return published_to_native_subtree(raw, "esmc.") + + def __init__(self, config: EsmFold2Config) -> None: + super().__init__(config) + + # InputsEmbedder + self.inputs_embedder = InputsEmbedder(config) + + # z_init projections + d_inputs = config.single_inputs_size + d_pair = config.pairwise_hidden_size + + self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False) + self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False) + + # Trunk relative position encoding + self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding( + n_relative_residx_bins=config.n_relative_residx_bins, + n_relative_chain_bins=config.n_relative_chain_bins, + d_pair=d_pair, + ) + + # Token bonds + self.token_bonds = nn.Linear(1, d_pair, bias=False) + + self.language_model = LanguageModelShim( + d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers + ) + # A bundled backbone is described by ``esmc_config`` and arrives in the + # same checkpoint, so it is built here and populated by the single + # ``from_pretrained`` pass. Otherwise ``load_esmc`` fetches it later. + self.esmc: nn.Module | None = ( + EsmcModel(config.esmc_config) if config.esmc_config is not None else None + ) + + # FoldingTrunk + self.folding_trunk = FoldingTrunk( + n_layers=config.folding_trunk_num_hidden_layers, + d_pair=d_pair, + expansion_ratio=4, + ) + + # Per-loop pair re-injection projection + 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 + ) # ty:ignore[invalid-argument-type] + + # Structure head + self.structure_head = DiffusionStructureHead(config) + + # Distogram head + self.distogram_head = nn.Linear( + d_pair, config.structure_head.distogram_bins, bias=True + ) + + if config.confidence_head.enabled: + self.confidence_head: ConfidenceHead | None = ConfidenceHead(config) + else: + self.confidence_head = None + + # MSA encoder (Large MSA models only) + msa_cfg = config.msa_encoder + if msa_cfg.enabled: + self.msa_encoder: MSAEncoder | None = MSAEncoder( + d_msa=msa_cfg.hidden_size, + d_pair=d_pair, + d_inputs=d_inputs, + d_hidden=msa_cfg.outer_hidden_size, + n_layers=msa_cfg.num_hidden_layers, + n_heads_msa=msa_cfg.num_attention_heads, + msa_head_width=msa_cfg.head_width, + ) + else: + self.msa_encoder = None + + self.post_init() + + def set_kernel_backend(self, backend: str | None) -> None: + """Select kernel backend (None / "fused" / "cuequivariance").""" + self.folding_trunk.set_kernel_backend(backend) + if self.confidence_head is not None: + self.confidence_head.set_kernel_backend(backend) + self.structure_head.set_kernel_backend(backend) + if self.msa_encoder is not None: + self.msa_encoder.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + """Set chunk size for memory-efficient triangle multiplicative updates.""" + self.folding_trunk.set_chunk_size(chunk_size) + if self.confidence_head is not None: + self.confidence_head.set_chunk_size(chunk_size) + if self.msa_encoder is not None: + self.msa_encoder.set_chunk_size(chunk_size) + + def configure_lm_dropout( + self, lm_dropout: float, *, force_lm_dropout_during_inference: bool = True + ) -> None: + """Configure LM embedding dropout (binder-design / critic scoring).""" + self.config.lm_dropout = lm_dropout + self.config.force_lm_dropout_during_inference = ( + force_lm_dropout_during_inference + ) + + def load_esmc(self, esmc_model_path: str) -> None: + """Fetch the ESMC LM backbone from a Hub repo ID or local directory. + + Only needed when the backbone is *not* bundled into this checkpoint; see + ``EsmFold2Config.esmc_config``. + """ + from esm.models.esmc import EsmcModel + + self.esmc = EsmcModel.from_pretrained(esmc_model_path) + self.finalize_esmc() + + def finalize_esmc(self) -> None: + """Cast the attached ESMC LM to bf16 and put it in eval mode.""" + if self.esmc is None: + raise RuntimeError("no ESMC LM is attached; nothing to cast.") + self.esmc = self.esmc.bfloat16().to(self.device).eval() + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path, + *, + load_esmc: bool = True, + esmc_precision: str = "bf16", + config: EsmFold2Config | None = None, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, + **kwargs, + ): + del esmc_precision # the experimental path always loads the LM in bf16 + local_dir = resolve_model_dir(pretrained_model_name_or_path, **kwargs) + if config is None: + config = EsmFold2Config.from_pretrained(local_dir) + model = cls._load_pretrained(local_dir, config, device=device, dtype=dtype) + # A bundled backbone came in with the trunk; only the separate-repo + # arrangement needs a second fetch. + if load_esmc and config.esmc_config is None: + model.load_esmc(model.config.esmc_id) + elif config.esmc_config is not None: + model.finalize_esmc() + return model + + @torch.no_grad() + def infer_protein(self, seq: str, **forward_kwargs) -> dict: + from esm.models.esmfold2.protein_utils import ( + OUTPUT_TO_PDB_FEATURE_KEYS, + prepare_protein_features, + ) + + features = prepare_protein_features(seq) + features = {k: v.to(self.device) for k, v in features.items()} + output = self(**features, **forward_kwargs) + for k in OUTPUT_TO_PDB_FEATURE_KEYS: + output[k] = features[k] + return output + + def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str: + return self.output_to_pdb(self.infer_protein(seq, **forward_kwargs)) + + @torch.no_grad() + def infer_all_atom(self, structure_input, **forward_kwargs): + try: + import esm as esm # TODO: change to `import esm` when open sourcing + import esm.models.esmfold2 # noqa: F401 # TODO: drop when open sourcing + except ImportError as e: + raise NotImplementedError( + "All-atom inference requires the `esm` companion package: " + "`pip install esm`." + ) from e + esmfold2 = esm.models.esmfold2 + + if isinstance(structure_input, esmfold2.ProteinInput): + structure_input = esmfold2.StructurePredictionInput( + sequences=[structure_input] + ) + processor = esmfold2.ESMFold2InputBuilder() + features, chain_infos = processor.prepare_input(structure_input) + features = { + k: v.to(self.device) if isinstance(v, Tensor) else v + for k, v in features.items() + } + output = self(**features, **forward_kwargs) + return self._output_to_molecular_complex(output, features, chain_infos) + + @staticmethod + def _output_to_molecular_complex(output: dict, features: dict, chain_infos: list): + import esm as esm # TODO: change to `import esm` when open sourcing + import esm.models.esmfold2 # noqa: F401 # TODO: drop when open sourcing + + esmfold2 = esm.models.esmfold2 + + ELEMENT_NUMBER_TO_SYMBOL = esmfold2.ELEMENT_NUMBER_TO_SYMBOL + MolecularComplex = esmfold2.MolecularComplex + MolecularComplexMetadata = esmfold2.MolecularComplexMetadata + + coords = output["sample_atom_coords"] + if coords.dim() == 4: + coords = coords[:, 0] + coords_np = coords.detach().cpu().numpy() + + plddt = output["plddt"].detach().cpu().numpy() + atom_to_token = features["atom_to_token"].cpu().numpy() + ref_chars = features["ref_atom_name_chars"].cpu().numpy() + ref_element = features["ref_element"].cpu().numpy() + atom_mask = features["atom_attention_mask"].cpu().numpy().astype(bool) + + if atom_to_token.ndim == 1: + atom_to_token = atom_to_token[None] + ref_chars = ref_chars[None] + ref_element = ref_element[None] + atom_mask = atom_mask[None] + + b = 0 + atoms_by_token: dict[int, list[int]] = {} + for a in range(atom_to_token.shape[1]): + if not atom_mask[b, a]: + continue + atoms_by_token.setdefault(int(atom_to_token[b, a]), []).append(a) + + flat_positions: list[np.ndarray] = [] + flat_elements: list[str] = [] + flat_names: list[str] = [] + flat_hetero: list[bool] = [] + sequence_tokens: list[str] = [] + token_to_atoms: list[list[int]] = [] + chain_ids_per_token: list[int] = [] + confidence_scores: list[float] = [] + chain_lookup: dict[int, str] = {} + entity_info: dict[int, str] = {} + + cursor = 0 + for chain in chain_infos: + chain_lookup[chain.asym_id] = chain.chain_id + entity_info[chain.entity_id] = ( + "polymer" if chain.mol_type != 3 else "non-polymer" + ) + for tok in chain.tokens: + sequence_tokens.append(tok.residue_name) + chain_ids_per_token.append(chain.asym_id) + confidence_scores.append(float(plddt[b, tok.token_index])) + start = cursor + for a in atoms_by_token.get(tok.token_index, []): + flat_positions.append(coords_np[b, a]) + name = "".join( + chr(int(c) + 32) if int(c) != 0 else " " + for c in ref_chars[b, a] + ).strip() + flat_names.append(name) + flat_elements.append( + ELEMENT_NUMBER_TO_SYMBOL.get(int(ref_element[b, a]), "X") + ) + flat_hetero.append(chain.mol_type == 3) + cursor += 1 + token_to_atoms.append([start, cursor]) + + return MolecularComplex( + id="prediction", + sequence=sequence_tokens, + atom_positions=np.array(flat_positions, dtype=np.float32), + atom_elements=np.array(flat_elements, dtype=object), + token_to_atoms=np.array(token_to_atoms, dtype=np.int32), + chain_id=np.array(chain_ids_per_token, dtype=np.int64), + plddt=np.array(confidence_scores, dtype=np.float32), + atom_names=np.array(flat_names, dtype=object), + atom_hetero=np.array(flat_hetero, dtype=bool), + metadata=MolecularComplexMetadata( + entity_lookup={k: str(v) for k, v in entity_info.items()}, + chain_lookup=chain_lookup, + assembly_composition=None, + ), + ) + + @staticmethod + def output_to_pdb(output: dict) -> str: + from esm.models.esmfold2.protein_utils import output_to_pdb as _output_to_pdb + + return _output_to_pdb(output) + + def _compute_lm_hidden_states( + self, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + token_mask: Tensor, + lm_mask_pct: float = 0.0, + ) -> Tensor: + """Run ESMC with BOS/EOS wrapping, return hidden states [B, L, N, D] with N=81 layers.""" + assert self.esmc is not None + return compute_lm_hidden_states( + self.esmc, + input_ids, + asym_id, + residue_index, + mol_type, + token_mask, + lm_mask_pct=lm_mask_pct, + ) + + def forward( + self, + # Token features + token_index: Tensor, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + mol_type: Tensor, + res_type: Tensor, + token_bonds: Tensor, + token_attention_mask: Tensor, + # Atom features + ref_pos: Tensor, + ref_element: Tensor, + ref_charge: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + atom_attention_mask: Tensor, + atom_to_token: Tensor, + distogram_atom_idx: Tensor, + # MSA features + deletion_mean: Tensor | None = None, + msa: Tensor | None = None, + has_deletion: Tensor | None = None, + deletion_value: Tensor | None = None, + msa_attention_mask: Tensor | None = None, + # LM features (auto-computed from input_ids if ESMC loaded) + input_ids: Tensor | None = None, + lm_hidden_states: Tensor | None = None, + # Used in design to provide a soft sequence input. + res_type_soft: Tensor | None = None, + # Inference config + num_loops: int | None = None, + num_diffusion_samples: int | None = None, + num_sampling_steps: int | None = None, + lm_mask_pct: float | None = None, + msa_max_depth: int | None = None, + msa_column_mask_rate: float | None = None, + noise_scale: float | None = None, + step_scale: float | None = None, + # Restates ``sample``'s own default; None there means "no cap". + max_inference_sigma: float | None = 256.0, + seed: int | None = None, + provide_soft_sequence_to_msa_and_profile: bool = True, + **unused_features: Tensor, + ) -> dict[str, Tensor]: + """Full ESMFold2 inference pipeline. + + Accepts tensors directly from ESMFold2InputBuilder.prepare_input(). + + Returns: + dict with sample_atom_coords, plddt, pae, distogram_logits, etc. + """ + unexpected = sorted(set(unused_features) - _IGNORED_FEATURE_KEYS) + if unexpected: + raise TypeError( + f"{type(self).__name__}.forward() got unexpected keyword " + f"argument(s) {unexpected}." + ) + + tok_mask = token_attention_mask + atm_mask = atom_attention_mask + disto_idx = distogram_atom_idx + + n_loops: int = num_loops if num_loops is not None else self.config.num_loops + n_samples: int = ( + num_diffusion_samples + if num_diffusion_samples is not None + else self.config.num_diffusion_samples + ) + + # One-hot res_type for input embedder concatenation + if res_type.dim() == 2: + res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float() + res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float() + else: + res_type_oh = res_type.float() + + # Profile: masked mean over MSA depth, with res_type fallback when no MSA. + if msa is not None: + msa_oh_profile = F.one_hot( + msa.long(), num_classes=NUM_RES_TYPES + ).float() # [B, M, L, V] + if msa_attention_mask is not None: + mask_f = msa_attention_mask.float().unsqueeze(-1) # [B, M, L, 1] + msa_oh_profile = msa_oh_profile * mask_f + valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1) + profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1) + else: + profile = msa_oh_profile.mean(dim=1) + else: + profile = res_type_oh + + # Used in design to provide a soft sequence input. + if res_type_soft is not None: + res_type_oh = res_type_soft.float() + if ( + not getattr(self.config, "disable_msa_features", False) + and provide_soft_sequence_to_msa_and_profile + ): + profile = res_type_oh + msa = res_type_oh.unsqueeze(1) + msa_attention_mask = tok_mask.unsqueeze(1) + + if deletion_mean is None: + deletion_mean = torch.zeros( + res_type.shape[0], res_type.shape[1], device=res_type.device + ) + + if getattr(self.config, "disable_msa_features", False): + profile = torch.zeros_like(profile) + deletion_mean = torch.zeros_like(deletion_mean) + + ref_element = F.one_hot( + ref_element.long(), num_classes=MAX_ATOMIC_NUMBER + ).float() + ref_atom_name_chars = F.one_hot( + ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE + ).float() + # Bias-free downstream Linears require zeroed padding. + atm_mask_f = atm_mask.float() + ref_element = ref_element * atm_mask_f.unsqueeze(-1) + ref_atom_name_chars = ref_atom_name_chars * atm_mask_f.unsqueeze(-1).unsqueeze( + -1 + ) + + atom_to_token = atom_to_token * atm_mask.long() + + use_amp = ref_pos.device.type == "cuda" + with ( + torch.set_grad_enabled(res_type_soft is not None), + torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16), + ): + # 1. Input embeddings + x_inputs = self.inputs_embedder( + aatype=res_type_oh, + profile=profile.float(), + deletion_mean=deletion_mean.float(), + ref_pos=ref_pos, + atom_attention_mask=atm_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=atom_to_token, + ) + + # 2. Initialize pair representation + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2( + x_inputs + ).unsqueeze(1) + + # 3. Positional encodings + relative_position_encoding = self.rel_pos( + residue_index=residue_index, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + token_index=token_index, + ) + token_bonds_encoding = self.token_bonds(token_bonds.float()) + z_init = z_init + relative_position_encoding + token_bonds_encoding + + # 4. Language model integration + if ( + lm_hidden_states is None + and input_ids is not None + and self.esmc is not None + ): + lm_hidden_states = self._compute_lm_hidden_states( + input_ids, + asym_id, + residue_index, + mol_type, + tok_mask, + lm_mask_pct=( + lm_mask_pct + if lm_mask_pct is not None + else self.config.lm_mask_pct + ), + ) + if lm_hidden_states is not None: + lm_z = self.language_model( + lm_hidden_states.detach(), lm_dropout=self.config.lm_dropout + ) + z_init = z_init + lm_z.to(z_init.dtype) + + # Inference-time MSA diversity: column mask is applied once here + # (shared across recycling loops); row subsampling is deferred to + # per-iter inside the loop below (fresh subset per loop). + _msa_inputs: dict | None = None + if self.msa_encoder is not None and msa is not None: + # ``None`` means "use the checkpoint's value". + msa_cfg = self.config.msa_encoder + depth = msa_cfg.max_depth if msa_max_depth is None else msa_max_depth + rate = ( + msa_cfg.column_mask_rate + if msa_column_mask_rate is None + else msa_column_mask_rate + ) + msa_attention_mask = maybe_apply_msa_column_masking( + msa_attention_mask, rate=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=depth, + subsample_enabled=depth is not None, + ) + + # Expand 1-D token mask → 2-D pair mask for folding trunk + pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() + + # Loop: pair-only folding trunk, MSA encoder runs inside each iteration. + z = torch.zeros_like(z_init) + for loop_num in range(n_loops + 1): + z = z_init + self.pair_loop_proj(z) + if _msa_inputs is not None and self.msa_encoder is not None: + # Fresh row subsample each iteration (column mask was applied + # once in forward, before this loop). + msa_i, mask_i, hd_i, dv_i = maybe_subsample_msa( + _msa_inputs["msa"], + _msa_inputs["msa_attention_mask"], + _msa_inputs["has_deletion"], + _msa_inputs["deletion_value"], + max_depth=_msa_inputs["max_depth"], + enabled=_msa_inputs["subsample_enabled"], + ) + if msa_i.dim() == 4: + B_msa, M, L_msa, _ = msa_i.shape + msa_oh = msa_i.permute(0, 2, 1, 3).float() + else: + B_msa, M, L_msa = msa_i.shape + msa_oh = F.one_hot( + msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES + ).float() # [B, L, M, 33] + msa_attn = ( + mask_i.permute(0, 2, 1).float() + if mask_i is not None + else tok_mask[:, :, None].expand(-1, -1, M).float() + ) + # Bias-free MSAEncoder.embed requires zeroed padding. + msa_oh = msa_oh * msa_attn.unsqueeze(-1) + hd = ( + hd_i.permute(0, 2, 1).float() + if hd_i is not None + else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + ) + dv = ( + dv_i.permute(0, 2, 1).float() + if dv_i is not None + else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + ) + z = z + self.msa_encoder( + x_pair=z, + x_inputs=_msa_inputs["x_inputs"], + msa_oh=msa_oh, + has_deletion=hd, + deletion_value=dv, + msa_attention_mask=msa_attn, + ).to(z.dtype) + z = self.folding_trunk(z, pair_attention_mask=pair_mask) + + # 6. Distogram (inside the trunk autocast so z stays bf16) + distogram_logits = self.distogram_head(z + z.transpose(-2, -3)) + + # 7. Diffusion sampling (always no_grad; optional seed for parity) + with torch.no_grad(), _seed_context(seed): + structure_output = self.structure_head.sample( + z_trunk=z.float(), + s_inputs=x_inputs, + s_trunk=None, + relative_position_encoding=relative_position_encoding, + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=atm_mask, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + ref_space_uid=ref_space_uid, + tok_idx=atom_to_token, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=tok_mask, + num_diffusion_samples=n_samples, + num_sampling_steps=num_sampling_steps, + noise_scale=noise_scale, + step_scale=step_scale, + max_inference_sigma=max_inference_sigma, + return_atom_repr=False, + ) + + sample_coords = structure_output["sample_atom_coords"] + assert sample_coords is not None + output: dict[str, Tensor] = {"distogram_logits": distogram_logits} + output["sample_atom_coords"] = sample_coords + + if self.confidence_head is not None: + 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(), + ) + output.update(confidence_output) + + # Pass-through tensors used by output decoders. + output["atom_pad_mask"] = ( + atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask + ) + output["residue_index"] = residue_index + output["entity_id"] = entity_id + + return output + + +__all__ = ["EsmFold2ExperimentalModel"] diff --git a/esm/models/esmfold2/kernels/__init__.py b/esm/models/esmfold2/kernels/__init__.py new file mode 100644 index 00000000..ef99caeb --- /dev/null +++ b/esm/models/esmfold2/kernels/__init__.py @@ -0,0 +1,23 @@ +# coding=utf-8 +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Triton inference kernels for ESMFold2.""" + +from esm.models.esmfold2.kernels.fused_attention_pair_bias import fused_pair_bias +from esm.models.esmfold2.kernels.fused_dropout_residual import FusedDropoutResidual +from esm.models.esmfold2.kernels.fused_lnlin_swiglu import FusedLNLinearSwiGLU +from esm.models.esmfold2.kernels.trimul_with_residual import ( + triangle_multiplicative_update_with_residual, +) + +__all__ = [ + "fused_pair_bias", + "FusedDropoutResidual", + "FusedLNLinearSwiGLU", + "triangle_multiplicative_update_with_residual", +] diff --git a/esm/models/esmfold2/kernels/fused_attention_pair_bias.py b/esm/models/esmfold2/kernels/fused_attention_pair_bias.py new file mode 100644 index 00000000..98b72239 --- /dev/null +++ b/esm/models/esmfold2/kernels/fused_attention_pair_bias.py @@ -0,0 +1,697 @@ +"""Fused Triton kernel for AttentionPairBias (forward + backward). + +Inspired by cuequivariance's ``attention_pair_bias`` +(https://docs.nvidia.com/cuda/cuequivariance/index.html); independently +re-implemented in Triton with a backward pass and no sequence-length gate. + +Fuses ``LayerNorm(z) -> z @ w_proj_z.T -> (1-mask)*(-INF)`` into a single +kernel that emits a ``bias[B, H, Q, K]`` tensor, then dispatches the +attention itself to ``torch.nn.functional.scaled_dot_product_attention``. +""" + +# ruff: noqa: E402 + +import os +import warnings + +warnings.filterwarnings("ignore", category=FutureWarning) +warnings.filterwarnings("ignore", category=UserWarning) +warnings.filterwarnings("ignore", category=DeprecationWarning) +os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1") +os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1") + +import torch +import triton +import triton.language as tl + +# Static config — runtime autotune cold-start is unshippable for inference. +_BIAS_AUTOTUNE_CONFIGS = [ + triton.Config({"TILE_K": 64, "TILE_C": 64}, num_stages=3, num_warps=4) +] + + +@triton.autotune( + configs=_BIAS_AUTOTUNE_CONFIGS, + key=["Q", "K", "DIM_Z", "NUM_HEADS", "HEADS_PER_BLK", "HAS_MASK", "AFFINE"], +) +@triton.jit +def _pair_bias_kernel( + z_ptr, # [B, Q, K, DIM_Z] pair tensor, bf16/fp16/fp32 + mask_ptr, # [B, K] key padding mask (bool/uint8) + w_proj_z_ptr, # [NUM_HEADS, DIM_Z] bias-projection weight + w_ln_ptr, # [DIM_Z] LN gamma (or unused) + b_ln_ptr, # [DIM_Z] LN beta (or unused) + out_ptr, # [B, NUM_HEADS, Q, K] fused output bias, same dtype as z + mean_ptr, # [B, Q, K] saved LN mean (fp32) or dummy + rstd_ptr, # [B, Q, K] saved LN rstd (fp32) or dummy + B, + Q, + K, + NEG_INF: tl.constexpr, + EPS: tl.constexpr, + DIM_Z: tl.constexpr, + DIM_Z_PAD: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEADS_PER_BLK: tl.constexpr, + TILE_K: tl.constexpr, + TILE_C: tl.constexpr, # tile over the DIM_Z reduction axis + HAS_MASK: tl.constexpr, + AFFINE: tl.constexpr, + SAVE_STATS: tl.constexpr, +): + """One CTA owns (one Q-row × TILE_K K-positions × HEADS_PER_BLK heads). + + Grid: (cdiv(K, TILE_K), Q, B * cdiv(NUM_HEADS, HEADS_PER_BLK)). + + Each thread block computes a tile of the output bias + ``bias[b, h_blk, q, k_tile]`` and stores it to ``out_ptr``. Layout of + ``out_ptr`` is ``[B, NUM_HEADS, Q, K]`` so the downstream SDPA call can + pass it directly as ``attn_mask`` (broadcast-compatible with the standard + BHQK attention layout). + + When ``SAVE_STATS`` is set (training path), the per-row ``mean`` and + ``rstd`` are stored into ``mean_ptr``/``rstd_ptr`` of shape ``[B, Q, K]`` + so the backward kernel can avoid recomputing them. Only the first head-block + in each (B, Q, K) tile writes the stats to avoid races. + """ + pid_k = tl.program_id(0) + pid_q = tl.program_id(1) + pid_bh = tl.program_id(2) + NUM_HEAD_BLKS: tl.constexpr = (NUM_HEADS + HEADS_PER_BLK - 1) // HEADS_PER_BLK + pid_b = pid_bh // NUM_HEAD_BLKS + pid_hblk = pid_bh % NUM_HEAD_BLKS + + offs_k = pid_k * TILE_K + tl.arange(0, TILE_K) + offs_z_full = tl.arange(0, DIM_Z_PAD) + offs_h = pid_hblk * HEADS_PER_BLK + tl.arange(0, HEADS_PER_BLK) + mask_k = offs_k < K + mask_h = offs_h < NUM_HEADS + + z_full_ptrs = ( + z_ptr + + pid_b * Q * K * DIM_Z + + pid_q * K * DIM_Z + + offs_k[:, None] * DIM_Z + + offs_z_full[None, :] + ) + mask_z_full = offs_z_full < DIM_Z + z_full = tl.load( + z_full_ptrs, mask=mask_k[:, None] & mask_z_full[None, :], other=0.0 + ).to(tl.float32) + + mean = tl.sum(z_full, axis=1) / DIM_Z + z_centered = z_full - mean[:, None] + z_centered = tl.where(mask_z_full[None, :], z_centered, 0.0) + var = tl.sum(z_centered * z_centered, axis=1) / DIM_Z + rstd = 1.0 / tl.sqrt(var + EPS) + + # Save mean/rstd for backward — only the first head-block writes, all + # head-blocks have the same value so this avoids racy duplicate writes. + if SAVE_STATS: + if pid_hblk == 0: + stats_ptrs = mean_ptr + pid_b * Q * K + pid_q * K + offs_k + tl.store(stats_ptrs, mean, mask=mask_k) + stats_ptrs2 = rstd_ptr + pid_b * Q * K + pid_q * K + offs_k + tl.store(stats_ptrs2, rstd, mask=mask_k) + + acc = tl.zeros([TILE_K, HEADS_PER_BLK], dtype=tl.float32) + num_tiles_c = tl.cdiv(DIM_Z, TILE_C) + for tc in range(0, num_tiles_c): + offs_c = tc * TILE_C + tl.arange(0, TILE_C) + mask_c = offs_c < DIM_Z + + z_slice_ptrs = ( + z_ptr + + pid_b * Q * K * DIM_Z + + pid_q * K * DIM_Z + + offs_k[:, None] * DIM_Z + + offs_c[None, :] + ) + z_slice = tl.load( + z_slice_ptrs, mask=mask_k[:, None] & mask_c[None, :], other=0.0 + ).to(tl.float32) + + z_norm = (z_slice - mean[:, None]) * rstd[:, None] + if AFFINE: + gamma = tl.load(w_ln_ptr + offs_c, mask=mask_c, other=1.0).to(tl.float32) + beta = tl.load(b_ln_ptr + offs_c, mask=mask_c, other=0.0).to(tl.float32) + z_norm = z_norm * gamma[None, :] + beta[None, :] + z_norm = tl.where(mask_c[None, :], z_norm, 0.0) + + w_ptrs = w_proj_z_ptr + (offs_h[None, :] * DIM_Z + offs_c[:, None]) + w_tile = tl.load(w_ptrs, mask=mask_h[None, :] & mask_c[:, None], other=0.0).to( + tl.float32 + ) + + acc = tl.dot(z_norm.to(tl.float32), w_tile, acc, input_precision="tf32x3") + + if HAS_MASK: + m_tile = tl.load(mask_ptr + pid_b * K + offs_k, mask=mask_k, other=0).to( + tl.int32 + ) + acc = acc + tl.where(m_tile == 0, NEG_INF, 0.0)[:, None] + # Mask out-of-bounds K positions (we pad to TILE_K). + acc = tl.where(mask_k[:, None], acc, NEG_INF) + + out_ptrs = ( + out_ptr + + pid_b * NUM_HEADS * Q * K + + offs_h[None, :] * Q * K + + pid_q * K + + offs_k[:, None] + ) + tl.store( + out_ptrs, + acc.to(out_ptr.type.element_ty), + mask=mask_k[:, None] & mask_h[None, :], + ) + + +# Backward is reduction-heavy (atomics into d_w_proj_z / d_pair_norm_*); modest +# TILE_K=32 balances ILP and register pressure on Hopper. +_BIAS_BWD_CONFIGS = [triton.Config({"TILE_K": 32}, num_stages=3, num_warps=4)] + + +@triton.autotune( + configs=_BIAS_BWD_CONFIGS, key=["Q", "K", "DIM_Z", "NUM_HEADS", "AFFINE"] +) +@triton.jit +def _pair_bias_backward_kernel( + z_ptr, # [B, Q, K, DIM_Z] input pair tensor (bf16) + w_proj_z_ptr, # [NUM_HEADS, DIM_Z] bias-projection weight (bf16) + w_ln_ptr, # [DIM_Z] LN gamma (bf16, or dummy) + b_ln_ptr, # [DIM_Z] LN beta (bf16, or dummy) + mean_ptr, # [B, Q, K] saved LN mean (fp32) + rstd_ptr, # [B, Q, K] saved LN rstd (fp32) + d_bias_ptr, # [B, NUM_HEADS, Q, K] upstream gradient (bf16) + d_z_ptr, # [B, Q, K, DIM_Z] output d_z (bf16) + d_w_proj_z_ptr, # [NUM_HEADS, DIM_Z] output d_w_proj_z (fp32 accum) + d_ln_w_ptr, # [DIM_Z] output d_pair_norm_w (fp32 accum) + d_ln_b_ptr, # [DIM_Z] output d_pair_norm_b (fp32 accum) + Q, + K, + EPS: tl.constexpr, + DIM_Z: tl.constexpr, + DIM_Z_PAD: tl.constexpr, + NUM_HEADS: tl.constexpr, + NUM_HEADS_PAD: tl.constexpr, + TILE_K: tl.constexpr, + AFFINE: tl.constexpr, +): + """Backward for ``fused_pair_bias``. + + Grid: (cdiv(K, TILE_K), Q, B). + + Each CTA processes a (b, q, k_tile) slab and: + + 1. Loads ``z[b,q,k,:]``, ``mean``, ``rstd``, ``d_bias[b,:,q,k]``, + ``w_proj_z[:,:]``, ``gamma``, ``beta``. + 2. Computes ``z_hat = (z - mean) * rstd``, ``z_norm = z_hat * gamma + beta``. + 3. ``d_z_norm[k,c] = sum_h d_bias[h,k] * w_proj_z[h,c]`` (matmul, k×c). + 4. Atomic-add ``d_w_proj_z[h,c] += d_bias[k,h]^T @ z_norm[k,c]``. + 5. Atomic-add ``d_pair_norm_b[c] += sum_k d_z_norm[k,c]``. + 6. Atomic-add ``d_pair_norm_w[c] += sum_k d_z_norm[k,c] * z_hat[k,c]``. + 7. ``d_z_hat = d_z_norm * gamma``. + 8. LN bwd: ``d_z = (d_z_hat - mean_c(d_z_hat) - z_hat * mean_c(d_z_hat * z_hat)) * rstd``. + 9. Store ``d_z``. + + All math is fp32 internally; only loads/stores are bf16. Atomic adds + target fp32 buffers (bf16 atomics are not supported on Hopper). + """ + pid_k = tl.program_id(0) + pid_q = tl.program_id(1) + pid_b = tl.program_id(2) + + offs_k = pid_k * TILE_K + tl.arange(0, TILE_K) + offs_z = tl.arange(0, DIM_Z_PAD) + offs_h = tl.arange(0, NUM_HEADS_PAD) + mask_k = offs_k < K + mask_z = offs_z < DIM_Z + mask_h = offs_h < NUM_HEADS + + z_ptrs = ( + z_ptr + + pid_b * Q * K * DIM_Z + + pid_q * K * DIM_Z + + offs_k[:, None] * DIM_Z + + offs_z[None, :] + ) + z = tl.load(z_ptrs, mask=mask_k[:, None] & mask_z[None, :], other=0.0).to( + tl.float32 + ) + mean_ptrs = mean_ptr + pid_b * Q * K + pid_q * K + offs_k + rstd_ptrs = rstd_ptr + pid_b * Q * K + pid_q * K + offs_k + mean = tl.load(mean_ptrs, mask=mask_k, other=0.0) + rstd = tl.load(rstd_ptrs, mask=mask_k, other=0.0) + + z_hat = (z - mean[:, None]) * rstd[:, None] + z_hat = tl.where(mask_k[:, None] & mask_z[None, :], z_hat, 0.0) + + if AFFINE: + gamma = tl.load(w_ln_ptr + offs_z, mask=mask_z, other=1.0).to(tl.float32) + beta = tl.load(b_ln_ptr + offs_z, mask=mask_z, other=0.0).to(tl.float32) + else: + gamma = tl.full([DIM_Z_PAD], 1.0, dtype=tl.float32) + beta = tl.full([DIM_Z_PAD], 0.0, dtype=tl.float32) + # Recompute normalized output (needed for d_w_proj_z). + z_norm = z_hat * gamma[None, :] + beta[None, :] + z_norm = tl.where(mask_k[:, None] & mask_z[None, :], z_norm, 0.0) + + d_bias_ptrs = ( + d_bias_ptr + + pid_b * NUM_HEADS * Q * K + + offs_h[None, :] * Q * K + + pid_q * K + + offs_k[:, None] + ) + d_bias = tl.load(d_bias_ptrs, mask=mask_k[:, None] & mask_h[None, :], other=0.0).to( + tl.float32 + ) + + w_proj_ptrs = w_proj_z_ptr + offs_h[:, None] * DIM_Z + offs_z[None, :] + w_proj = tl.load(w_proj_ptrs, mask=mask_h[:, None] & mask_z[None, :], other=0.0).to( + tl.float32 + ) + + d_z_norm = tl.dot(d_bias, w_proj, input_precision="tf32x3") + d_z_norm = tl.where(mask_k[:, None] & mask_z[None, :], d_z_norm, 0.0) + + d_w = tl.dot(tl.trans(d_bias), z_norm, input_precision="tf32x3") + d_w_ptrs = d_w_proj_z_ptr + offs_h[:, None] * DIM_Z + offs_z[None, :] + tl.atomic_add(d_w_ptrs, d_w, mask=mask_h[:, None] & mask_z[None, :], sem="relaxed") + + if AFFINE: + d_b_tile = tl.sum(d_z_norm, axis=0) + tl.atomic_add(d_ln_b_ptr + offs_z, d_b_tile, mask=mask_z, sem="relaxed") + + d_w_tile = tl.sum(d_z_norm * z_hat, axis=0) + tl.atomic_add(d_ln_w_ptr + offs_z, d_w_tile, mask=mask_z, sem="relaxed") + + d_z_hat = d_z_norm * gamma[None, :] + d_z_hat = tl.where(mask_k[:, None] & mask_z[None, :], d_z_hat, 0.0) + sum_dzh = tl.sum(d_z_hat, axis=1) + sum_dzh_zhat = tl.sum(d_z_hat * z_hat, axis=1) + mean_dzh = sum_dzh / DIM_Z + mean_dzh_zhat = sum_dzh_zhat / DIM_Z + d_z = (d_z_hat - mean_dzh[:, None] - z_hat * mean_dzh_zhat[:, None]) * rstd[:, None] + d_z = tl.where(mask_k[:, None] & mask_z[None, :], d_z, 0.0) + + d_z_ptrs = ( + d_z_ptr + + pid_b * Q * K * DIM_Z + + pid_q * K * DIM_Z + + offs_k[:, None] * DIM_Z + + offs_z[None, :] + ) + tl.store( + d_z_ptrs, + d_z.to(d_z_ptr.type.element_ty), + mask=mask_k[:, None] & mask_z[None, :], + ) + + +def _next_pow2(x: int) -> int: + p = 1 + while p < x: + p *= 2 + return max(p, 16) + + +def _round_up_heads_per_blk(num_heads: int) -> int: + """Always return 16 — ``tl.dot`` requires M, N, K >= 16 on Hopper. + + For the AF3-style transformer ``num_heads=16``; ``HEADS_PER_BLK=16`` gives + a single CTA per (b, q, k_tile) and is the cheapest schedule. Head-counts + below 16 are zero-padded inside the kernel. + """ + del num_heads # head-tile is fixed at 16 by the tl.dot lower bound + return 16 + + +def _launch_forward( + z: torch.Tensor, + mask: torch.Tensor | None, + w_proj_z: torch.Tensor, + pair_norm_w: torch.Tensor | None, + pair_norm_b: torch.Tensor | None, + num_heads: int, + eps: float, + inf: float, + save_stats: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Forward kernel launch helper. Returns ``(bias, mean, rstd)``. + + ``mean``/``rstd`` are ``None`` unless ``save_stats=True``. + """ + 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})" + + z = z.contiguous() + w_proj_z = w_proj_z.contiguous() + affine = pair_norm_w is not None or pair_norm_b is not None + if affine: + if pair_norm_w is None: + pair_norm_w = torch.ones(DIM_Z, device=z.device, dtype=z.dtype) + if pair_norm_b is None: + pair_norm_b = torch.zeros(DIM_Z, device=z.device, dtype=z.dtype) + pair_norm_w = pair_norm_w.contiguous() + pair_norm_b = pair_norm_b.contiguous() + if mask is not None: + mask = mask.contiguous() + assert mask.shape == (B, K), f"mask {mask.shape} ≠ ({B}, {K})" + + out = torch.empty((B, num_heads, Q, K), device=z.device, dtype=z.dtype) + if save_stats: + mean = torch.empty((B, Q, K), device=z.device, dtype=torch.float32) + rstd = torch.empty((B, Q, K), device=z.device, dtype=torch.float32) + else: + mean = None + rstd = None + heads_per_blk = _round_up_heads_per_blk(num_heads) + DIM_Z_PAD = _next_pow2(DIM_Z) + _dummy = torch.empty(1, device=z.device, dtype=z.dtype) + _dummy_f32 = torch.empty(1, device=z.device, dtype=torch.float32) + num_head_blks = (num_heads + heads_per_blk - 1) // heads_per_blk + grid = lambda meta: (triton.cdiv(K, meta["TILE_K"]), Q, B * num_head_blks) + _pair_bias_kernel[grid]( + z, + mask if mask is not None else _dummy, + w_proj_z, + pair_norm_w if affine else _dummy, + pair_norm_b if affine else _dummy, + out, + mean if save_stats else _dummy_f32, + rstd if save_stats else _dummy_f32, + B, + Q, + K, + NEG_INF=-float(inf), + EPS=eps, + DIM_Z=DIM_Z, + DIM_Z_PAD=DIM_Z_PAD, + NUM_HEADS=num_heads, + HEADS_PER_BLK=heads_per_blk, + HAS_MASK=mask is not None, + AFFINE=affine, + SAVE_STATS=save_stats, + ) + return out, mean, rstd + + +def _launch_backward( + z: torch.Tensor, + w_proj_z: torch.Tensor, + pair_norm_w: torch.Tensor | None, + pair_norm_b: torch.Tensor | None, + mean: torch.Tensor, + rstd: torch.Tensor, + d_bias: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Backward kernel launch helper. Returns ``(d_z, d_w_proj_z, d_pair_norm_w, d_pair_norm_b)``. + + The latter two are ``None`` when affine is off. + """ + B, Q, K, DIM_Z = z.shape + NUM_HEADS = w_proj_z.shape[0] + affine = pair_norm_w is not None + if affine: + assert pair_norm_b is not None + + z = z.contiguous() + w_proj_z = w_proj_z.contiguous() + d_bias = d_bias.contiguous() + if affine: + assert pair_norm_w is not None and pair_norm_b is not None + pair_norm_w = pair_norm_w.contiguous() + pair_norm_b = pair_norm_b.contiguous() + + d_z = torch.empty_like(z) + # fp32 accumulators for atomic adds — bf16 atomic_add is not supported on Hopper + d_w_proj_z_f32 = torch.zeros( + (NUM_HEADS, DIM_Z), device=z.device, dtype=torch.float32 + ) + if affine: + d_pair_norm_w_f32 = torch.zeros(DIM_Z, device=z.device, dtype=torch.float32) + d_pair_norm_b_f32 = torch.zeros(DIM_Z, device=z.device, dtype=torch.float32) + else: + d_pair_norm_w_f32 = torch.zeros(1, device=z.device, dtype=torch.float32) + d_pair_norm_b_f32 = torch.zeros(1, device=z.device, dtype=torch.float32) + + DIM_Z_PAD = _next_pow2(DIM_Z) + # min 16 to satisfy Triton's tl.dot requirement (M, N, K >= 16 on Hopper). + NUM_HEADS_PAD = max(16, _next_pow2(NUM_HEADS)) + _dummy = torch.empty(1, device=z.device, dtype=z.dtype) + + grid = lambda meta: (triton.cdiv(K, meta["TILE_K"]), Q, B) + _pair_bias_backward_kernel[grid]( + z, + w_proj_z, + pair_norm_w if affine else _dummy, + pair_norm_b if affine else _dummy, + mean, + rstd, + d_bias, + d_z, + d_w_proj_z_f32, + d_pair_norm_w_f32, + d_pair_norm_b_f32, + Q, + K, + EPS=eps, + DIM_Z=DIM_Z, + DIM_Z_PAD=DIM_Z_PAD, + NUM_HEADS=NUM_HEADS, + NUM_HEADS_PAD=NUM_HEADS_PAD, + AFFINE=affine, + ) + + d_w_proj_z = d_w_proj_z_f32.to(w_proj_z.dtype) + d_pair_norm_w: torch.Tensor | None + d_pair_norm_b: torch.Tensor | None + if affine: + assert pair_norm_w is not None and pair_norm_b is not None + d_pair_norm_w = d_pair_norm_w_f32.to(pair_norm_w.dtype) + d_pair_norm_b = d_pair_norm_b_f32.to(pair_norm_b.dtype) + else: + d_pair_norm_w = None + d_pair_norm_b = None + + return d_z, d_w_proj_z, d_pair_norm_w, d_pair_norm_b + + +class FusedPairBias(torch.autograd.Function): + """Autograd wrapper around ``_pair_bias_kernel`` and ``_pair_bias_backward_kernel``. + + Forward saves ``(z, w_proj_z, pair_norm_w, pair_norm_b, mean, rstd)`` so that + the backward kernel can recompute ``z_hat`` without re-doing the full + LN reduction. ``mask`` is non-differentiable; we save it in ``ctx`` only as + a marker (the kernel applies it inside the forward, but mask positions get + -INF and so contribute 0 gradient through softmax → no special handling + needed for ``d_bias``). + """ + + @staticmethod + def forward(ctx, z, mask, w_proj_z, pair_norm_w, pair_norm_b, num_heads, eps, inf): + out, mean, rstd = _launch_forward( + z, + mask, + w_proj_z, + pair_norm_w, + pair_norm_b, + num_heads, + eps, + inf, + save_stats=True, + ) + affine = pair_norm_w is not None or pair_norm_b is not None + ctx.save_for_backward( + z, + w_proj_z, + pair_norm_w if affine else None, + pair_norm_b if affine else None, + mean, + rstd, + ) + ctx.affine = affine + ctx.eps = eps + return out + + @staticmethod + def backward(ctx, d_bias): + z, w_proj_z, pair_norm_w, pair_norm_b, mean, rstd = ctx.saved_tensors + d_z, d_w_proj_z, d_pair_norm_w, d_pair_norm_b = _launch_backward( + z, + w_proj_z, + pair_norm_w, + pair_norm_b, + mean, + rstd, + d_bias.contiguous(), + ctx.eps, + ) + # Order must match forward args: (z, mask, w_proj_z, pair_norm_w, pair_norm_b, num_heads, eps, inf) + return ( + d_z, + None, # mask + d_w_proj_z, + d_pair_norm_w, + d_pair_norm_b, + None, # num_heads + None, # eps + None, # inf + ) + + +@torch._dynamo.disable +def fused_pair_bias( + z: torch.Tensor, + mask: torch.Tensor | None, + w_proj_z: torch.Tensor, + pair_norm_w: torch.Tensor | None, + pair_norm_b: torch.Tensor | None, + *, + num_heads: int, + eps: float = 1e-5, + inf: float = 1e6, +) -> torch.Tensor: + """Compute ``bias[B, H, Q, K] = LN(z) @ w_proj_z.T + (1-mask)*-INF``. + + Dispatches to the autograd-aware ``FusedPairBias`` path when autograd is + enabled (i.e. any input requires_grad and we are not in a no_grad/inference + context). Otherwise falls back to the forward-only kernel. + + Parameters + ---------- + z : (B, Q, K, DIM_Z) + mask : (B, K) bool or None. True = keep, False = mask out (-INF added). + w_proj_z : (num_heads, DIM_Z) + pair_norm_w, pair_norm_b : (DIM_Z,) or None. Pass both or neither. + num_heads : int + eps, inf : LN epsilon and masking-infinity respectively. + + Returns + ------- + bias : (B, num_heads, Q, K) — same dtype as z. + """ + use_autograd = torch.is_grad_enabled() and ( + z.requires_grad + or w_proj_z.requires_grad + or (pair_norm_w is not None and pair_norm_w.requires_grad) + or (pair_norm_b is not None and pair_norm_b.requires_grad) + ) + if use_autograd: + out_t: torch.Tensor = FusedPairBias.apply( # type: ignore[assignment] + z, mask, w_proj_z, pair_norm_w, pair_norm_b, num_heads, eps, inf + ) + return out_t + out, _, _ = _launch_forward( + z, + mask, + w_proj_z, + pair_norm_w, + pair_norm_b, + num_heads, + eps, + inf, + save_stats=False, + ) + return out + + +def fused_attention_pair_bias( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + z: torch.Tensor | None, + mask: torch.Tensor | None, + x_for_gate: torch.Tensor, + *, + w_proj_z: torch.Tensor | None, + w_proj_g: torch.Tensor, + w_proj_o: torch.Tensor, + pair_norm_w: torch.Tensor | None = None, + pair_norm_b: torch.Tensor | None = None, + eps: float = 1e-5, + inf: float = 1e6, + precomputed_bias: torch.Tensor | None = None, +) -> torch.Tensor: + """End-to-end fused AttentionPairBias forward (no-conditioning path). + + Pipeline: + bias = fused_pair_bias(z, mask, w_proj_z, ln_w, ln_b) # Triton + attn = SDPA(q, k, v, attn_mask=bias) # cuDNN / Flash + gate = sigmoid(linear(x_for_gate, w_proj_g)) # cuBLAS + out = linear(gate * attn, w_proj_o) # cuBLAS + + Parameters + ---------- + q, k, v : (B, H, Q|K, D) — already-projected query/key/value, head-split. + z : (B, Q, K, DIM_Z) — raw (unnormed) pair tensor. Ignored when + ``precomputed_bias`` is supplied. + mask : (B, K) bool or None. Ignored when ``precomputed_bias`` is supplied. + x_for_gate : (B, Q, d_model) — pre-norm input for the gate + ``sigmoid(x @ w_proj_g.T)``. In the production module this is the + adaln/pre_norm output ``x`` (same tensor that feeds ``q``). + w_proj_z : (H, DIM_Z). Ignored when ``precomputed_bias`` is supplied. + w_proj_g, w_proj_o : (d_model, d_model) + pair_norm_w, pair_norm_b : (DIM_Z,). Ignored when ``precomputed_bias`` is + supplied. + precomputed_bias : (B, H, Q, K) optional cached bias tensor. When provided + the LN+proj+mask Triton kernel is skipped entirely — the bias is reused + as the SDPA ``attn_mask`` directly. This is the 50× diffusion-step + amortization path: ``z`` is constant within a loop, so the bias can + be computed once and reused across all 50 denoise steps. Compute it + with ``fused_pair_bias`` once, then pass it back here on every step. + + Returns + ------- + out : (B, Q, d_model) + """ + B, H, Q, D = q.shape + 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." + bias = precomputed_bias + else: + if z is None or w_proj_z is None: + raise ValueError( + "Either precomputed_bias OR (z, w_proj_z) must be supplied" + ) + bias = fused_pair_bias( + z, mask, w_proj_z, pair_norm_w, pair_norm_b, num_heads=H, eps=eps, inf=inf + ) # (B, H, Q, K) + + # Match the dtype expected by cuDNN attention. q/k/v are already bf16 in + # inference; ``bias`` inherits z's dtype. + if not torch.compiler.is_compiling(): + with torch.nn.attention.sdpa_kernel( + backends=[ + torch.nn.attention.SDPBackend.CUDNN_ATTENTION, + torch.nn.attention.SDPBackend.FLASH_ATTENTION, + torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, + ], + set_priority=True, + ): + attn = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=bias, is_causal=False + ) + else: + attn = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=bias, is_causal=False + ) + attn = attn.transpose(1, 2).contiguous().view(B, Q, d_model) + + gate = torch.sigmoid(torch.nn.functional.linear(x_for_gate, w_proj_g)) + out = torch.nn.functional.linear(gate * attn, w_proj_o) + return out + + +__all__ = ["FusedPairBias", "fused_attention_pair_bias", "fused_pair_bias"] diff --git a/esm/models/esmfold2/kernels/fused_dropout_residual.py b/esm/models/esmfold2/kernels/fused_dropout_residual.py new file mode 100644 index 00000000..02f7ed3c --- /dev/null +++ b/esm/models/esmfold2/kernels/fused_dropout_residual.py @@ -0,0 +1,247 @@ +"""Fused row-shared-dropout + residual-add kernel for the pair stream. + +For the pairformer pattern: + + pair_new = pair + Dropout(r, batch_dim=1)(delta) + +where `delta` is the output of e.g. a triangle multiplication, the row-shared +dropout multiplies `delta` by a Bernoulli mask of shape `[B, 1, N_col, D]` +(broadcast over the row dim) scaled by `1/(1-r)`. Naively this materializes +two intermediate `[B, N, N, D]` tensors (one for `delta * mask`, one for +`pair + ...`) — three full HBM round-trips of the pair tensor. + +This kernel reads `pair`, `delta`, and the small `[N_col, D]` shared mask +(via modulo-`N_col` indexing — no broadcast materialization) and writes the +combined `pair + delta * mask` once. + +Backward is also a single Triton kernel: it mutates the saved mask buffer in +place to produce `ddelta = dout * mask`, and the residual gradient passes +through unchanged. + +Convention: the mask is expected to already be scaled (i.e. it's the output +of `nn.Dropout(r)(ones_like(...))` so that retained entries have value +`1/(1-r)`). +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import triton +import triton.language as tl + + +@triton.jit +def _fused_dropout_residual_fwd_kernel( + pair_ptr, # [M, D] M = B*N_row*N_col + delta_ptr, # [M, D] + mask_ptr, # [B*N_col, D] per-batch row-shared mask + out_ptr, # [M, D] + M, + D: tl.constexpr, + N_COL: tl.constexpr, + STRIDE_B: tl.constexpr, # = N_row * N_col, used to recover batch index + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + pid_d = tl.program_id(1).to(tl.int64) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + mm_m = offs_m < M + mm_d = offs_d < D + + # m = b*N_row*N_col + i*N_col + j → b = m // (N_row*N_col), j = m % N_col + # mask is laid out as [B*N_col, D] so per-batch index is b*N_col + j. + b = offs_m // STRIDE_B + j = offs_m % N_COL + mask_row = b * N_COL + j + + pair_ptrs = pair_ptr + offs_m[:, None] * D + offs_d[None, :] + delta_ptrs = delta_ptr + offs_m[:, None] * D + offs_d[None, :] + out_ptrs = out_ptr + offs_m[:, None] * D + offs_d[None, :] + mask_ptrs = mask_ptr + mask_row[:, None] * D + offs_d[None, :] + + p = tl.load(pair_ptrs, mask=mm_m[:, None] & mm_d[None, :], other=0.0) + de = tl.load(delta_ptrs, mask=mm_m[:, None] & mm_d[None, :], other=0.0) + mk = tl.load(mask_ptrs, mask=mm_m[:, None] & mm_d[None, :], other=0.0) + + o = p + de * mk + tl.store( + out_ptrs, o.to(out_ptr.type.element_ty), mask=mm_m[:, None] & mm_d[None, :] + ) + + +@triton.jit +def _fused_dropout_residual_bwd_kernel( + dout_ptr, # [M, D] + mask_ptr, # [B*N_col, D] + ddelta_ptr, # [M, D] + M, + D: tl.constexpr, + N_COL: tl.constexpr, + STRIDE_B: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + pid_d = tl.program_id(1).to(tl.int64) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + mm_m = offs_m < M + mm_d = offs_d < D + + b = offs_m // STRIDE_B + j = offs_m % N_COL + mask_row = b * N_COL + j + mask_ptrs = mask_ptr + mask_row[:, None] * D + offs_d[None, :] + + do = tl.load( + dout_ptr + offs_m[:, None] * D + offs_d[None, :], + mask=mm_m[:, None] & mm_d[None, :], + other=0.0, + ) + mk = tl.load(mask_ptrs, mask=mm_m[:, None] & mm_d[None, :], other=0.0) + tl.store( + ddelta_ptr + offs_m[:, None] * D + offs_d[None, :], + (do * mk).to(ddelta_ptr.type.element_ty), + mask=mm_m[:, None] & mm_d[None, :], + ) + + +def _fused_dropout_residual_fwd( + pair_2d: torch.Tensor, + delta_2d: torch.Tensor, + mask_2d: torch.Tensor, + n_row: int, + n_col: int, +) -> torch.Tensor: + M, D = pair_2d.shape + out = torch.empty_like(pair_2d) + BLOCK_M = 64 + BLOCK_D = min(128, _next_pow2(D)) + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(D, BLOCK_D)) + _fused_dropout_residual_fwd_kernel[grid]( + pair_2d, + delta_2d, + mask_2d, + out, + M, + 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 + + +def _fused_dropout_residual_bwd( + dout_2d: torch.Tensor, mask_2d: torch.Tensor, n_row: int, n_col: int +) -> torch.Tensor: + M, D = dout_2d.shape + ddelta = torch.empty_like(dout_2d) + BLOCK_M = 64 + BLOCK_D = min(128, _next_pow2(D)) + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(D, BLOCK_D)) + _fused_dropout_residual_bwd_kernel[grid]( + dout_2d, + mask_2d, + ddelta, + M, + 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 + + +def _next_pow2(n: int) -> int: + p = 1 + while p < n: + p <<= 1 + return p + + +class FusedDropoutResidualFn(torch.autograd.Function): + """`pair_out = pair + delta * mask` fused into one kernel. + + Args: + pair: [B, N_row, N_col, D] + delta: [B, N_row, N_col, D] + mask: [B, 1, N_col, D] or [1, 1, N_col, D] — already scaled to 1/(1-r) for kept entries + + The mask is row-shared (size 1 along `N_row`) — we read it directly with + modular indexing rather than broadcasting and materializing a full [B, N, N, D] + copy. + """ + + @staticmethod + def forward(ctx, pair, delta, mask): + in_shape = pair.shape # [B, N_row, N_col, D] + N_row = in_shape[-3] + N_col = in_shape[-2] + D = in_shape[-1] + pair_2d = pair.contiguous().view(-1, D) + delta_2d = delta.contiguous().view(-1, D) + # [B, 1, N_col, D] → [B*N_col, D]; kernel indexes b*N_col + (m % N_col) + # so each batch sample uses its own draw. + mask_2d = mask.contiguous().view(-1, D) + out_2d = _fused_dropout_residual_fwd(pair_2d, delta_2d, mask_2d, N_row, N_col) + ctx.save_for_backward(mask_2d) + ctx.in_shape = in_shape + ctx.n_row = N_row + ctx.n_col = N_col + return out_2d.view(in_shape) + + @staticmethod + def backward(ctx, dout): + (mask_2d,) = ctx.saved_tensors + D = ctx.in_shape[-1] + dout_2d = dout.contiguous().view(-1, D) + ddelta_2d = _fused_dropout_residual_bwd(dout_2d, mask_2d, ctx.n_row, ctx.n_col) + # Residual: gradient passes through. Mask: not differentiable. + return dout, ddelta_2d.view(ctx.in_shape), None + + +class FusedDropoutResidual(nn.Module): + """Fused module for the pattern `pair + Dropout(r, batch_dim=1)(delta)`. + + The kernel only supports row-shared dropout — i.e. a `[B, 1, N_col, D]` + mask broadcast over the row axis (dim 1) — so this module hardcodes that + layout. The `j = m % N_col` indexing in the Triton kernel would read out + of bounds with any other sharing pattern, so we don't expose `batch_dim` + as a parameter. Use plain `Dropout(r, batch_dim=2)` for col-shared. + + Usage in a pairformer block: + + # Before + pair = pair + self.row_drop(self.tri_mul_out(pair, mask=...)) + + # After + pair = self.row_drop(pair, self.tri_mul_out(pair, mask=...)) + + where `self.row_drop` is `FusedDropoutResidual(r)`. + """ + + def __init__(self, r: float): + super().__init__() + self.r = r + + def forward(self, pair: torch.Tensor, delta: torch.Tensor) -> torch.Tensor: + if not self.training or self.r == 0.0: + return pair + delta + # Use delta.dtype: F.dropout's cuRAND draw depends on input dtype, so + # building from pair (fp32) vs delta (bf16 under autocast) breaks + # same-seed parity with the unfused `pair + Dropout(delta)` path. + shape = list(pair.shape) + 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) diff --git a/esm/models/esmfold2/kernels/fused_dual_gemm.py b/esm/models/esmfold2/kernels/fused_dual_gemm.py new file mode 100644 index 00000000..b6791385 --- /dev/null +++ b/esm/models/esmfold2/kernels/fused_dual_gemm.py @@ -0,0 +1,610 @@ +"""Native Triton ``fused_sigmoid_gated_dual_gemm`` for TriMul stage 2. + +Computes ``sigmoid(x @ w1.T) * (x @ w2.T)`` with optional row-shared mask, +bias, and transposed output. Forward + backward implemented in bf16. + +Inspired by cuequivariance's ``fused_sigmoid_gated_dual_gemm`` +(https://docs.nvidia.com/cuda/cuequivariance/index.html); independently +re-implemented in Triton. +""" + +# ruff: noqa: E402 + +from __future__ import annotations + +import os +import warnings + +warnings.filterwarnings("ignore", category=FutureWarning) +warnings.filterwarnings("ignore", category=UserWarning) +warnings.filterwarnings("ignore", category=DeprecationWarning) +os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1") +os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1") + +import torch +import triton +import triton.language as tl + +# Static config — runtime autotune cold-start is unshippable for inference. +_AUTOTUNE_CONFIGS = [ + triton.Config( + {"TILE_M": 128, "TILE_N": 64, "TILE_K": 32, "GROUP_M": 8}, + num_stages=4, + num_warps=4, + ) +] + + +@triton.autotune( + configs=_AUTOTUNE_CONFIGS, key=["M", "N", "K", "HAS_MASK", "TRANSPOSE_OUT"] +) +@triton.jit +def _gated_dual_gemm_kernel( + x_ptr, # [M, K] bf16 + w1_ptr, # [N, K] bf16 — gate weight + w2_ptr, # [N, K] bf16 — value weight + mask_ptr, # [M] bf16 — row-shared mask broadcast to (M, N) + out_ptr, # [M, N] or [N, M] bf16 — sigmoid(x@w1) * (x@w2) + M, + N, + K, + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + TILE_K: tl.constexpr, + GROUP_M: tl.constexpr, + HAS_MASK: tl.constexpr, + TRANSPOSE_OUT: tl.constexpr, # store (N, M) instead of (M, N) + 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 + val_acc = Σ_K (x[:, k] @ w2[:, k]) over k + delta = sigmoid(gate_acc) * val_acc + if mask: delta *= mask[m_tile] (broadcast over N) + store delta (transposed if TRANSPOSE_OUT) + """ + pid_m_raw = tl.program_id(0) + pid_n_raw = tl.program_id(1) + + # GROUP_M swizzle for L2 reuse of ``x`` across consecutive CTAs. + num_pid_m = tl.cdiv(M, TILE_M) + num_pid_n = tl.cdiv(N, TILE_N) + pid = pid_n_raw * num_pid_m + pid_m_raw # row-major program id + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + if NEEDS_INT64: + pid_m = tl.cast(pid_m, tl.int64) + pid_n = tl.cast(pid_n, tl.int64) + M = tl.cast(M, tl.int64) + N = tl.cast(N, tl.int64) + K = tl.cast(K, tl.int64) + + start_m = pid_m * TILE_M + start_n = pid_n * TILE_N + + offs_m = start_m + tl.arange(0, TILE_M) + offs_n = start_n + tl.arange(0, TILE_N) + offs_k = tl.arange(0, TILE_K) + if NEEDS_INT64: + offs_m = tl.cast(offs_m, tl.int64) + offs_n = tl.cast(offs_n, tl.int64) + offs_k = tl.cast(offs_k, tl.int64) + + x_ptrs = x_ptr + (offs_m[:, None] * K + offs_k[None, :]) + w1_base = w1_ptr + (offs_n[None, :] * K + offs_k[:, None]) + w2_base = w2_ptr + (offs_n[None, :] * K + offs_k[:, None]) + + gate_acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + val_acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + mask_m = offs_m < M + + for _ in range(0, tl.cdiv(K, TILE_K)): + x_raw = tl.load(x_ptrs, mask=mask_m[:, None], other=0.0) + x_op = x_raw.to(w1_ptr.type.element_ty) + w1_tile = tl.load(w1_base) + w2_tile = tl.load(w2_base) + gate_acc = tl.dot(x_op, w1_tile, gate_acc) + val_acc = tl.dot(x_op, w2_tile, val_acc) + x_ptrs += TILE_K + w1_base += TILE_K + w2_base += TILE_K + + delta = tl.sigmoid(gate_acc) * val_acc # fp32 + + if HAS_MASK: + mask_tile = tl.load(mask_ptr + offs_m, mask=mask_m, other=0.0).to(tl.float32) + delta = delta * mask_tile[:, None] + + if TRANSPOSE_OUT: + out_ptrs = out_ptr + (offs_n[:, None] * M + offs_m[None, :]) + tl.store( + out_ptrs, tl.trans(delta).to(out_ptr.type.element_ty), mask=mask_m[None, :] + ) + else: + out_ptrs = out_ptr + (offs_m[:, None] * N + offs_n[None, :]) + tl.store(out_ptrs, delta.to(out_ptr.type.element_ty), mask=mask_m[:, None]) + + +# Backward emits per-element grad_gate_logits and grad_val_acc by recomputing +# the two forward GEMMs. Weight grads (d_w1, d_w2, d_x) are done in cuBLAS in +# the autograd Function — cuBLAS beats Triton bf16 at these reduction shapes. +# Narrower TILE_N=64 vs forward's 128 (bwd writes two M*N tensors, doubling +# register pressure). +_BWD_AUTOTUNE_CONFIGS = [ + triton.Config( + {"TILE_M": 64, "TILE_N": 64, "TILE_K": 64, "GROUP_M": 8}, + num_stages=3, + num_warps=4, + ) +] + + +@triton.autotune( + configs=_BWD_AUTOTUNE_CONFIGS, + key=["M", "N", "K", "HAS_MASK", "GRAD_OUT_TRANSPOSED", "GRAD_OUT_SPLIT"], +) +@triton.jit +def _gated_dual_gemm_backward_kernel( + grad_out_ptr, # [M, N] (or [N, M] if GRAD_OUT_TRANSPOSED, or [N/2, M] if GRAD_OUT_SPLIT) + grad_out2_ptr, # GRAD_OUT_SPLIT only: second half of (N, M); else dummy + x_ptr, # [M, K] bf16 — saved input + w1_ptr, # [N, K] bf16 + w2_ptr, # [N, K] bf16 + mask_ptr, # [M] bf16 — row-shared (unused if HAS_MASK=0) + grad_gate_logits_ptr, # [M, N] bf16 — out: d (x @ w1.T) + grad_val_acc_ptr, # [M, N] bf16 — out: d (x @ w2.T) + grad_mask_partials_ptr, # [num_pid_n, M] fp32 — partial mask grads + M, + N, + K, + HALF_N: tl.constexpr, # = N // 2, only used when GRAD_OUT_SPLIT=1 + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + TILE_K: tl.constexpr, + GROUP_M: tl.constexpr, + 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, # ty:ignore[invalid-parameter-default] +): + """Per (TILE_M, TILE_N) output tile: + gate_acc = Σ_K x[:, k] @ w1[:, k] + val_acc = Σ_K x[:, k] @ w2[:, k] + g = sigmoid(gate_acc) + grad_o = grad_out_tile (post-mask: multiply by mask if HAS_MASK) + d_gate_logits = grad_o * val_acc * g * (1 - g) + d_val_acc = grad_o * g + d_mask_partial[pid_n, m] = sum_n (grad_out_tile * g * val_acc) [if HAS_MASK] + """ + pid_m_raw = tl.program_id(axis=0) + pid_n_raw = tl.program_id(axis=1) + + num_pid_m = tl.cdiv(M, TILE_M) + num_pid_n = tl.cdiv(N, TILE_N) + pid = pid_n_raw * num_pid_m + pid_m_raw + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + if NEEDS_INT64: + pid_m = tl.cast(pid_m, tl.int64) + pid_n = tl.cast(pid_n, tl.int64) + M = tl.cast(M, tl.int64) + N = tl.cast(N, tl.int64) + K = tl.cast(K, tl.int64) + + start_m = pid_m * TILE_M + start_n = pid_n * TILE_N + + offs_m = start_m + tl.arange(0, TILE_M) + offs_n = start_n + tl.arange(0, TILE_N) + offs_k = tl.arange(0, TILE_K) + if NEEDS_INT64: + offs_m = tl.cast(offs_m, tl.int64) + offs_n = tl.cast(offs_n, tl.int64) + offs_k = tl.cast(offs_k, tl.int64) + + x_ptrs = x_ptr + (offs_m[:, None] * K + offs_k[None, :]) + w1_base = w1_ptr + (offs_n[None, :] * K + offs_k[:, None]) + w2_base = w2_ptr + (offs_n[None, :] * K + offs_k[:, None]) + + gate_acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + val_acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + mask_m = offs_m < M + + # Recompute fwd GEMMs (cheaper than saving N*M activations). + for _ in range(0, tl.cdiv(K, TILE_K)): + x_tile = tl.load(x_ptrs, mask=mask_m[:, None], other=0.0) + x_op = x_tile.to(w1_ptr.type.element_ty) + w1_tile = tl.load(w1_base) + w2_tile = tl.load(w2_base) + gate_acc = tl.dot(x_op, w1_tile, gate_acc) + val_acc = tl.dot(x_op, w2_tile, val_acc) + x_ptrs += TILE_K + w1_base += TILE_K + w2_base += TILE_K + + g = tl.sigmoid(gate_acc) + + if GRAD_OUT_SPLIT: + # Two (N/2, M) grad tensors; caller guarantees TILE_N divides HALF_N. + if start_n < HALF_N: + offs_n_local = offs_n + grad_o_ptrs = grad_out_ptr + (offs_n_local[None, :] * M + offs_m[:, None]) + else: + offs_n_local = offs_n - HALF_N + grad_o_ptrs = grad_out2_ptr + (offs_n_local[None, :] * M + offs_m[:, None]) + elif GRAD_OUT_TRANSPOSED: + grad_o_ptrs = grad_out_ptr + (offs_n[None, :] * M + offs_m[:, None]) + else: + grad_o_ptrs = grad_out_ptr + (offs_m[:, None] * N + offs_n[None, :]) + grad_o = tl.load(grad_o_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + if HAS_MASK: + # d_mask = row-sum BEFORE mask multiply; per-pid_n partial, host reduces. + d_mask_partial = tl.sum(grad_o * g * val_acc, axis=1) + d_mask_partials_ptrs = grad_mask_partials_ptr + pid_n * M + offs_m + tl.store(d_mask_partials_ptrs, d_mask_partial, mask=mask_m) + mask_tile = tl.load(mask_ptr + offs_m, mask=mask_m, other=0.0).to(tl.float32) + grad_o = grad_o * mask_tile[:, None] + + # Both grad ptrs index into one (M, 2N) buffer (val_acc cols [0:N], + # gate_logits cols [N:2N]) — lets downstream d_x / d_w fold into single GEMMs. + d_val_acc = (grad_o * g).to(grad_val_acc_ptr.type.element_ty) + d_val_acc_ptrs = grad_val_acc_ptr + (offs_m[:, None] * (2 * N) + offs_n[None, :]) + tl.store(d_val_acc_ptrs, d_val_acc, mask=mask_m[:, None]) + + d_gate_logits = (grad_o * val_acc * g * (1.0 - g)).to( + grad_gate_logits_ptr.type.element_ty + ) + d_gate_logits_ptrs = grad_gate_logits_ptr + ( + offs_m[:, None] * (2 * N) + offs_n[None, :] + ) + tl.store(d_gate_logits_ptrs, d_gate_logits, mask=mask_m[:, None]) + + +# Backward TILE_N must match the kernel's autotuned tile in ``_BWD_AUTOTUNE_CONFIGS``; +# host code uses it to size the per-tile mask-grad partials buffer statically. +_BWD_TILE_N = 64 + + +def _fused_gated_dual_gemm_bwd( + grad_out: torch.Tensor, # (M, N) or (N, M) layout (see grad_out_transposed) + x: torch.Tensor, # [..., K] (un-flattened, original) + w1: torch.Tensor, + w2: torch.Tensor, + mask: torch.Tensor | None, + grad_out_transposed: bool = False, + grad_out_split: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Compute (d_x, d_w1, d_w2, d_mask) for ``fused_gated_dual_gemm``. + + Three grad_out layouts supported: + * default: ``grad_out`` of shape ``(..., N)`` flattened to (M, N). + * ``grad_out_transposed=True``: shape ``(N, ...)`` flattened to (N, M). + * ``grad_out_split=(g1, g2)``: two tensors, each of shape ``(N/2, ...)`` + flattened to ``(N/2, M)``; the kernel reads from the appropriate + pointer per tile. Avoids the autograd-introduced concat that + ``torch.chunk(dim=0)``'s backward inserts (saves 3-4 ms at prod + shape, B=5 L=768 c_z=128). + """ + in_shape = x.shape + K = in_shape[-1] + x_c = x if x.is_contiguous() else x.contiguous() + x_2d = x_c.view(-1, K) + 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 x_2d.dtype == torch.bfloat16, "bwd only supports bf16 x" + + if grad_out_split is not None: + g1, g2 = grad_out_split + half_n = N // 2 + # Materialize only if non-contig (avoids a full-tensor copy). + if not g1.is_contiguous(): + g1 = g1.contiguous() + if not g2.is_contiguous(): + g2 = g2.contiguous() + assert g1.numel() == half_n * M and g2.numel() == half_n * M, ( + f"split grad sizes mismatch: g1={g1.numel()} g2={g2.numel()} " + f"vs half_n*M={half_n * M}" + ) + grad_out_2d_a = g1.view(half_n, M) + grad_out_2d_b = g2.view(half_n, M) + elif grad_out_transposed: + grad_out_2d_a = grad_out.contiguous().view(N, M) + grad_out_2d_b = grad_out_2d_a # unused (dummy) + else: + grad_out_2d_a = grad_out.contiguous().view(M, N) + grad_out_2d_b = grad_out_2d_a # unused (dummy) + + mask_flat = mask.contiguous().view(-1) if mask is not None else None + if mask_flat is not None: + assert mask_flat.shape[0] == M + + # (M, 2N) combined-grad buffer; matched stacked_w order (w2 then w1) + # lets d_w and d_x each collapse to a single cuBLAS GEMM. + grad_combined = torch.empty((M, 2 * N), device=x.device, dtype=torch.bfloat16) + grad_val_acc = grad_combined[:, :N] + grad_gate_logits = grad_combined[:, N:] + + tiles_n_max = triton.cdiv(N, _BWD_TILE_N) + if mask is not None: + grad_mask_partials = torch.empty( + (tiles_n_max, M), device=x.device, dtype=torch.float32 + ) + else: + grad_mask_partials = torch.empty((1,), device=x.device, dtype=torch.float32) + + _dummy_mask = torch.zeros((), device=x.device, dtype=torch.bfloat16) + + # (M, 2N) layout → use stride 2N for the int32-overflow gate. + NEEDS_INT64 = (M * K >= 2**31 - 1) or (M * (2 * N) >= 2**31 - 1) + + def grid(meta): + assert N % meta["TILE_N"] == 0 + return (triton.cdiv(M, meta["TILE_M"]), N // meta["TILE_N"]) + + half_n = N // 2 if grad_out_split is not None else 1 + + _gated_dual_gemm_backward_kernel[grid]( + grad_out_2d_a, + grad_out_2d_b, + x_2d, + w1.contiguous(), + w2.contiguous(), + mask_flat if mask_flat is not None else _dummy_mask, + grad_gate_logits, + grad_val_acc, + grad_mask_partials, + M, + N, + K, + HALF_N=half_n, + HAS_MASK=mask is not None, + GRAD_OUT_TRANSPOSED=grad_out_transposed, + GRAD_OUT_SPLIT=grad_out_split is not None, + NEEDS_INT64=NEEDS_INT64, + ) + + # Stacked weights for the single-GEMM d_x path. + stacked_w = torch.cat([w2, w1], dim=0) # (2N, K), matches val|gate layout + + d_x_2d = grad_combined @ stacked_w + d_x = d_x_2d.view(in_shape) + + d_w_combined = grad_combined.t() @ x_2d # (2N, K) + d_w2 = d_w_combined[:N] + d_w1 = d_w_combined[N:] + + if mask is not None: + actual_tiles = triton.cdiv(N, _BWD_TILE_N) + d_mask_flat = grad_mask_partials[:actual_tiles].sum(dim=0).to(mask.dtype) + d_mask = d_mask_flat.view(mask.shape) + else: + d_mask = None + + return d_x, d_w1, d_w2, d_mask + + +class FusedGatedDualGEMM(torch.autograd.Function): + """Autograd wrapper for ``fused_gated_dual_gemm`` (bf16 path, + ``transpose_out=False``). + + Saves ``x``, ``w1``, ``w2``, ``mask`` for the backward kernel; the + backward kernel recomputes the two dual-GEMM activations (cheap vs + materializing them in fwd context). + """ + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + mask: torch.Tensor | None, + ) -> torch.Tensor: + out = _fused_gated_dual_gemm_fwd(x, w1, w2, mask=mask, transpose_out=False) + if mask is None: + ctx.save_for_backward(x, w1, w2) + ctx.has_mask = False + else: + ctx.save_for_backward(x, w1, w2, mask) + ctx.has_mask = True + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + if ctx.has_mask: + x, w1, w2, mask = ctx.saved_tensors + else: + x, w1, w2 = ctx.saved_tensors + mask = None + if not grad_out.is_contiguous(): + grad_out = grad_out.contiguous() + d_x, d_w1, d_w2, d_mask = _fused_gated_dual_gemm_bwd(grad_out, x, w1, w2, mask) + return d_x, d_w1, d_w2, d_mask + + +class FusedGatedDualGEMMSplit(torch.autograd.Function): + """Autograd wrapper that returns the dual-GEMM output as two split halves. + + Forward: runs the kernel with ``transpose_out=True``, producing a single + ``(N=2*c_z, *trailing)`` buffer where the gate half (first c_z) and the + value half (next c_z) live contiguously along dim 0. Returns the two + views ``(a, b_t)`` so downstream einsums consume them in + ``(c_z, B, L, L)`` layout — no chunk in the autograd graph. + + Backward: receives ``(grad_a, grad_b_t)`` (each ``(c_z, *trailing)``) + and passes them as the ``grad_out_split`` pair to the backward kernel. + Eliminates the ~3.5 ms ``torch.chunk(dim=0)`` backward concat at the + prod shape (B=5, L=768, c_z=128) by reading the two halves from + separate pointers inside the kernel. + """ + + @staticmethod + def forward( + ctx, + x: torch.Tensor, # (B, L, L, c_z) bf16 + w1: torch.Tensor, # (N=2*c_z, c_z) bf16 — gate + w2: torch.Tensor, # (N=2*c_z, c_z) bf16 — value + mask: torch.Tensor | None, # (B, L, L) bf16 or None + trailing_shape: tuple, # (B, L, L) for output view + ) -> tuple[torch.Tensor, torch.Tensor]: + out = _fused_gated_dual_gemm_fwd(x, w1, w2, mask=mask, transpose_out=True) + # Slicing on a leading dim of a contiguous tensor stays contiguous (no copy). + N = w1.shape[0] + half_n = N // 2 + out_view = out.view((N,) + trailing_shape) + a = out_view[:half_n] + b_t = out_view[half_n:] + if mask is None: + ctx.save_for_backward(x, w1, w2) + ctx.has_mask = False + else: + ctx.save_for_backward(x, w1, w2, mask) + ctx.has_mask = True + return a, b_t + + @staticmethod + def backward(ctx, grad_a: torch.Tensor, grad_b_t: torch.Tensor): + if ctx.has_mask: + x, w1, w2, mask = ctx.saved_tensors + else: + x, w1, w2 = ctx.saved_tensors + mask = None + # 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, # ty:ignore[invalid-argument-type] + x, + w1, + w2, + mask, + grad_out_split=(grad_a, grad_b_t), + ) + return d_x, d_w1, d_w2, d_mask, None + + +def fused_gated_dual_gemm_split( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Split-output variant of ``fused_gated_dual_gemm``: returns ``(a, b_t)`` + with layout ``(c_z, B, L, L)`` each, avoiding chunk in the autograd graph. + + Inference fallback (no grad) just calls the regular fwd with + ``transpose_out=True`` and chunks the result — same layout, no kernel + change required. + """ + trailing_shape = tuple(x.shape[:-1]) + 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) + 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) + half_n = N // 2 + return out_view[:half_n], out_view[half_n:] + + +def _fused_gated_dual_gemm_fwd( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + mask: torch.Tensor | None = None, + transpose_out: bool = False, +) -> torch.Tensor: + """Native Triton implementation of sigmoid-gated dual GEMM. + + Computes ``sigmoid(x @ w1.T) * (x @ w2.T)`` with optional row-shared mask. + + Shapes: + x: (..., K) bf16 + w1: (N, K) bf16 + w2: (N, K) bf16 + mask: (...,) bf16 — flattened to a per-row scalar + out: (..., N) or (N, ...) when ``transpose_out=True`` + """ + in_shape = x.shape + K = in_shape[-1] + x_2d = x.contiguous().view(-1, K) + M = x_2d.shape[0] + N = w1.shape[0] + + 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}" + + out_dtype = torch.bfloat16 + if transpose_out: + out = torch.empty((N, M), device=x.device, dtype=out_dtype) + out_shape = (N,) + in_shape[:-1] + else: + out = torch.empty((M, N), device=x.device, dtype=out_dtype) + out_shape = in_shape[:-1] + (N,) + + mask_flat = mask.contiguous().view(-1) if mask is not None else None + if mask_flat is not None: + assert mask_flat.shape[0] == M, f"mask len {mask_flat.shape[0]} ≠ M {M}" + + _dummy = torch.zeros((), device=x.device, dtype=torch.float32) + + NEEDS_INT64 = (M * K >= 2**31 - 1) or (M * N >= 2**31 - 1) + + def grid(meta): + assert N % meta["TILE_N"] == 0 + return (triton.cdiv(M, meta["TILE_M"]), N // meta["TILE_N"]) + + _gated_dual_gemm_kernel[grid]( + x_2d, + w1.contiguous(), + w2.contiguous(), + mask_flat if mask_flat is not None else _dummy, + out, + M, + N, + K, + HAS_MASK=mask is not None, + TRANSPOSE_OUT=transpose_out, + NEEDS_INT64=NEEDS_INT64, + ) + return out.view(out_shape) + + +def fused_gated_dual_gemm( + x: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + mask: torch.Tensor | None = None, + transpose_out: bool = False, +) -> torch.Tensor: + """Public wrapper: dispatches to autograd Function if grad is enabled. + + ``transpose_out=True`` is inference-only (bypasses autograd). + """ + if torch.is_grad_enabled() and ( + x.requires_grad or w1.requires_grad or w2.requires_grad + ): + assert not transpose_out, ( + "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) + return _fused_gated_dual_gemm_fwd(x, w1, w2, mask=mask, transpose_out=transpose_out) diff --git a/esm/models/esmfold2/kernels/fused_ln_residual.py b/esm/models/esmfold2/kernels/fused_ln_residual.py new file mode 100644 index 00000000..3ab8d4ad --- /dev/null +++ b/esm/models/esmfold2/kernels/fused_ln_residual.py @@ -0,0 +1,397 @@ +"""Triton LayerNorm (bf16 IO, fp32 stats) with optional fused residual-add in +the *backward* pass. + +Inspired by cuequivariance's ``layer_norm_transpose`` +(https://docs.nvidia.com/cuda/cuequivariance/index.html); independently +re-implemented in Triton to add the bwd residual-link fusion. + +Used by ``trimul_with_residual.triangle_multiplicative_update_with_residual`` +for two LN calls: + + * Stage 1 LN: ``x_in = LN(pair)`` with layout ``bijd->bijd``. The downstream + residual add ``out = residual + delta`` (where ``residual = pair``) flows + ``grad_out`` straight back to ``grad_pair``. Fusing that add into LN's + bwd kernel saves one full pair-tensor read + one full write (~250 MB at + B=5 L=768 c_z=128 bf16). Exposed via :func:`fused_ln_with_residual_link`. + + * Stage 4 LN: ``x_out = LN(einsum(...))`` with layout ``dbij->bijd``. No + residual to fuse — :func:`fused_ln_transpose` is the plain variant. + +The residual-link fusion in the bwd pass is the motivation for shipping this: +the bwd kernel computes ``grad_x = LN_bwd(grad_y, x, w, mean, rstd)`` AND +optionally adds an external ``grad_residual`` tensor into ``grad_x`` in the +same pass — saves one HBM round-trip of the (M, D) tensor. + +Forward + backward use a single static ``triton.Config`` each (runtime +autotune cold-start is unshippable for inference paths). The bwd pass +emits per-tile ``grad_w``/``grad_b`` fp32 partials reduced host-side; this +is the standard Triton LN-bwd idiom (see openai/triton tutorial 05). +""" + +import torch +import triton +import triton.language as tl + +# Layout enum: 0 == "bijd->bijd" (bnd contig), 1 == "dbij->bijd" (dbn contig in). +_LAYOUT_BND_BND = 0 +_LAYOUT_DBN_BND = 1 + +_FWD_TILE_M = 64 +_FWD_NUM_WARPS = 8 +_FWD_NUM_STAGES = 2 + +_BWD_TILE_M = 64 +_BWD_NUM_WARPS = 8 +_BWD_NUM_STAGES = 2 + + +@triton.jit +def _ln_fwd_kernel( + x_ptr, # input + w_ptr, # [D] + b_ptr, # [D] + out_ptr, # [M, D] contig output + mean_ptr, # [M] fp32 + rstd_ptr, # [M] fp32 + M, + D: tl.constexpr, + EPS: tl.constexpr, + LAYOUT: tl.constexpr, + TILE_M: tl.constexpr, +): + pid = tl.program_id(axis=0).to(tl.int64) + M64 = M.to(tl.int64) + + offs_m = pid * TILE_M + tl.arange(0, TILE_M).to(tl.int64) + offs_d = tl.arange(0, D).to(tl.int64) + mask_m = offs_m < M64 + + if LAYOUT == 0: + x_ptrs = x_ptr + offs_m[:, None] * D + offs_d[None, :] + x = tl.load(x_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + else: # LAYOUT == 1: (D, M) contig + x_ptrs = x_ptr + offs_d[None, :] * M64 + offs_m[:, None] + x = tl.load(x_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + mean = tl.sum(x, axis=1) / D + x_c = x - mean[:, None] + var = tl.sum(x_c * x_c, axis=1) / D + rstd = 1.0 / tl.sqrt(var + EPS) + x_hat = x_c * rstd[:, None] + + tl.store(mean_ptr + offs_m, mean, mask=mask_m) + tl.store(rstd_ptr + offs_m, rstd, mask=mask_m) + + w = tl.load(w_ptr + offs_d).to(tl.float32) + b = tl.load(b_ptr + offs_d).to(tl.float32) + y = x_hat * w[None, :] + b[None, :] + + out_ptrs = out_ptr + offs_m[:, None] * D + offs_d[None, :] + tl.store(out_ptrs, y.to(out_ptr.type.element_ty), mask=mask_m[:, None]) + + +@triton.jit +def _ln_bwd_kernel( + grad_y_ptr, # [M, D] (always bnd; LN out is bnd) + x_ptr, # input — layout LAYOUT + w_ptr, # [D] + mean_ptr, # [M] fp32 + rstd_ptr, # [M] fp32 + grad_x_ptr, # output — layout LAYOUT (same as x) + grad_w_partial_ptr, # [num_tiles, D] fp32 + grad_b_partial_ptr, # [num_tiles, D] fp32 + grad_residual_ptr, # [M, D] in bnd layout; ignored if HAS_RESIDUAL=0 + M, + D: tl.constexpr, + LAYOUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + TILE_M: tl.constexpr, +): + pid = tl.program_id(axis=0).to(tl.int64) + M64 = M.to(tl.int64) + + offs_m = pid * TILE_M + tl.arange(0, TILE_M).to(tl.int64) + offs_d = tl.arange(0, D).to(tl.int64) + mask_m = offs_m < M64 + + grad_y_ptrs = grad_y_ptr + offs_m[:, None] * D + offs_d[None, :] + grad_y = tl.load(grad_y_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + if LAYOUT == 0: + x_ptrs = x_ptr + offs_m[:, None] * D + offs_d[None, :] + x = tl.load(x_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + else: + x_ptrs = x_ptr + offs_d[None, :] * M64 + offs_m[:, None] + x = tl.load(x_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + mean = tl.load(mean_ptr + offs_m, mask=mask_m, other=0.0) + rstd = tl.load(rstd_ptr + offs_m, mask=mask_m, other=0.0) + w = tl.load(w_ptr + offs_d).to(tl.float32) + + x_hat = (x - mean[:, None]) * rstd[:, None] + wdy = grad_y * w[None, :] + c1 = tl.sum(wdy, axis=1) / D + c2 = tl.sum(wdy * x_hat, axis=1) / D + grad_x = (wdy - (c1[:, None] + x_hat * c2[:, None])) * rstd[:, None] + + if HAS_RESIDUAL: + gr_ptrs = grad_residual_ptr + offs_m[:, None] * D + offs_d[None, :] + gr = tl.load(gr_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + grad_x = grad_x + gr + + if LAYOUT == 0: + gx_ptrs = grad_x_ptr + offs_m[:, None] * D + offs_d[None, :] + tl.store(gx_ptrs, grad_x.to(grad_x_ptr.type.element_ty), mask=mask_m[:, None]) + else: + gx_ptrs = grad_x_ptr + offs_d[None, :] * M64 + offs_m[:, None] + tl.store(gx_ptrs, grad_x.to(grad_x_ptr.type.element_ty), mask=mask_m[:, None]) + + # Per-tile partial reduction → write to (num_tiles, D) fp32 buffer. + # Final reduction (sum over num_tiles) happens host-side as a torch.sum + # (standard Triton LN-bwd idiom; see openai/triton tutorial 05). + mask_f = mask_m[:, None].to(tl.float32) + dw_tile = tl.sum(grad_y * x_hat * mask_f, axis=0) + db_tile = tl.sum(grad_y * mask_f, axis=0) + dw_offs = pid * D + offs_d + db_offs = pid * D + offs_d + tl.store(grad_w_partial_ptr + dw_offs, dw_tile) + tl.store(grad_b_partial_ptr + db_offs, db_tile) + + +def _layout_to_int(layout: str) -> int: + if layout in ("bijd->bijd", "bnd->bnd"): + return _LAYOUT_BND_BND + if layout in ("dbij->bijd", "dbn->bnd"): + return _LAYOUT_DBN_BND + raise ValueError(f"unsupported layout {layout!r}") + + +def _reshape_for_layout( + x: torch.Tensor, layout: str +) -> tuple[tuple[int, ...], int, int, torch.Tensor]: + """Reshape x to 2D LN view. Returns (out_shape, M, D, x_view).""" + if layout == "bijd->bijd": + B, II, J, D = x.shape + M = B * II * J + return (B, II, J, D), M, D, x.contiguous().view(M, D) + if layout == "bnd->bnd": + B, N, D = x.shape + M = B * N + return (B, N, D), M, D, x.contiguous().view(M, D) + if layout == "dbij->bijd": + D, B, II, J = x.shape + M = B * II * J + return (B, II, J, D), M, D, x.contiguous().view(D, M) + if layout == "dbn->bnd": + D, B, N = x.shape + M = B * N + return (B, N, D), M, D, x.contiguous().view(D, M) + raise ValueError(f"unsupported layout {layout!r}") + + +def _ln_fwd( + x_view: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eps: float, + layout_int: int, + M: int, + D: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + out = torch.empty((M, D), device=x_view.device, dtype=x_view.dtype) + mean = torch.empty((M,), device=x_view.device, dtype=torch.float32) + rstd = torch.empty((M,), device=x_view.device, dtype=torch.float32) + grid = (triton.cdiv(M, _FWD_TILE_M),) + _ln_fwd_kernel[grid]( + x_view, + w, + b, + out, + mean, + rstd, + M, + 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 + + +def _ln_bwd( + grad_y: torch.Tensor, + x_view: torch.Tensor, + w: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + layout_int: int, + grad_residual: torch.Tensor | None, + M: int, + D: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if layout_int == _LAYOUT_BND_BND: + grad_x = torch.empty((M, D), device=x_view.device, dtype=x_view.dtype) + else: + grad_x = torch.empty((D, M), device=x_view.device, dtype=x_view.dtype) + + num_tiles = triton.cdiv(M, _BWD_TILE_M) + grad_w_partial = torch.empty( + (num_tiles, D), device=x_view.device, dtype=torch.float32 + ) + grad_b_partial = torch.empty( + (num_tiles, D), device=x_view.device, dtype=torch.float32 + ) + + has_residual = grad_residual is not None + _dummy = torch.empty((), device=x_view.device, dtype=grad_y.dtype) + grid = (num_tiles,) + _ln_bwd_kernel[grid]( + grad_y, + x_view, + w, + mean, + rstd, + grad_x, + grad_w_partial, + grad_b_partial, + grad_residual if has_residual else _dummy, + M, + 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) + grad_b = grad_b_partial.sum(dim=0).to(w.dtype) + return grad_x, grad_w, grad_b + + +class _LayerNormTransposeFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, x: torch.Tensor, w: torch.Tensor, b: torch.Tensor, eps: float, layout: str + ) -> torch.Tensor: + layout_int = _layout_to_int(layout) + out_shape, M, D, x_view = _reshape_for_layout(x, layout) + out_bnd, mean, rstd = _ln_fwd(x_view, w, b, eps, layout_int, M, D) + ctx.save_for_backward(x_view, w, mean, rstd) + ctx.layout_int = layout_int + ctx.M = M + ctx.D = D + ctx.x_orig_shape = x.shape + return out_bnd.view(*out_shape) + + @staticmethod + 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( + grad_y, x_view, w, mean, rstd, ctx.layout_int, None, ctx.M, ctx.D + ) + grad_x = grad_x.view(*ctx.x_orig_shape) + return grad_x, grad_w, grad_b, None, None + + +def fused_ln_transpose( + x: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + eps: float = 1e-5, + layout: str = "bijd->bijd", +) -> torch.Tensor: + """Plain LN replacement for ``layer_norm_transpose`` (no residual fusion).""" + return _LayerNormTransposeFn.apply(x, w, b, eps, layout) + + +# Stage-1 LN with residual-add folded into the bwd kernel. The Function returns +# (ln_out, residual_alias); downstream uses ln_out, stage-5 uses residual_alias +# as the residual input. In bwd we receive both grad tensors and fold +# grad_residual_alias into grad_x in-kernel (saves one HBM round-trip). +# residual_alias is a fresh tensor (not a view) so autograd reliably routes +# its grad back through this Function. + + +class _LayerNormWithResidualLinkFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + residual_link: torch.Tensor, + eps: float, + layout: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + layout_int = _layout_to_int(layout) + out_shape, M, D, x_view = _reshape_for_layout(x, layout) + out_bnd, mean, rstd = _ln_fwd(x_view, w, b, eps, layout_int, M, D) + ctx.save_for_backward(x_view, w, mean, rstd) + ctx.layout_int = layout_int + ctx.M = M + ctx.D = D + ctx.x_orig_shape = x.shape + ctx.residual_link_shape = residual_link.shape + + ln_out = out_bnd.view(*out_shape) + # residual_alias: returning ``residual_link`` itself works in + # custom Functions — autograd treats the input-as-output case + # correctly (sums grads back into the input). The downstream + # graph node ``out = residual_alias + delta`` sees a regular + # tensor and produces grad of shape == residual_link. + # Note: we use .view_as for safety so the AutogradMeta is fresh. + return ln_out, residual_link.view_as(residual_link) + + @staticmethod + 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) + + # grad_link_pass shape == residual_link shape. + if grad_link_pass is None: + grad_residual = None + else: + grad_residual = grad_link_pass.contiguous().view(ctx.M, ctx.D) + + grad_x, grad_w, grad_b = _ln_bwd( + grad_y, x_view, w, mean, rstd, ctx.layout_int, grad_residual, ctx.M, ctx.D + ) + grad_x = grad_x.view(*ctx.x_orig_shape) + + # We folded grad_link_pass into grad_x → return None for that input's + # grad slot, since we've already accounted for it via x's grad path + # (the caller is expected to pass x == residual_link, so grad_x += + # grad_residual is exactly the combined grad on the shared leaf). + # If caller passes DIFFERENT tensors for x and residual_link, the + # link's grad is *lost* — that's a usage error we accept by contract. + return grad_x, grad_w, grad_b, None, None, None + + +def fused_ln_with_residual_link( + x: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor, + residual_link: torch.Tensor, + eps: float = 1e-5, + layout: str = "bijd->bijd", +) -> tuple[torch.Tensor, torch.Tensor]: + """LN(x) with a residual-link passthrough that fuses + ``grad_residual_link`` into the LN backward kernel. + + Contract: ``x`` and ``residual_link`` MUST refer to the same leaf + tensor (same identity). The returned ``residual_alias`` MUST be used as + the residual input to the downstream add — otherwise the fusion routes + the grad incorrectly. + + Returns ``(ln_out, residual_alias)``. + """ + if x is not residual_link: + raise ValueError( + "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) diff --git a/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py b/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py new file mode 100644 index 00000000..b163c5aa --- /dev/null +++ b/esm/models/esmfold2/kernels/fused_lnlin_swiglu.py @@ -0,0 +1,413 @@ +"""Fused LayerNorm + Linear(d, 2*d_inner) + SwiGLU(silu(x1) * x2) kernel. + +Collapses the standard LayerNorm -> Linear -> chunk -> SiLU -> mul sequence +into a single Triton kernel: + + out[..., d_inner] = silu(x1) * x2 where (x1, x2) = chunk(LN(x) @ W12, 2) + +Compared to the unfused PyTorch sequence this kernel: + + 1. Eliminates the [M, 2*d_inner] HBM read between Linear and SwiGLU. + 2. Halves X reads in the matmul: each program produces both halves of the + linear output (against W12_a and W12_b) in one pass over X. + 3. Fuses LayerNorm into the matmul k-loop. + +The output ordering uses silu of the FIRST half of W12's output, so the +fused module matches the standard LN+SwiGLU MLP composition. + +Backward: SwiGLU bwd is a small Triton kernel that mutates the saved [M, 2N] +linear-output buffer in place to produce dlin (no extra alloc). LN+Linear bwd +uses cuBLAS gemm + ATen's `native_layer_norm_backward` — fast on H100 and +avoids needing per-shape autotuned Triton bwd configs. + +Tuned forward configs were autotuned on H100 80GB for the pair-transition +production shapes (d_pair=256, d_inner=1024, M=B*N*N at N=384 / 640). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + + +@triton.jit +def _ln_stats_kernel( + X_ptr, + X_row_stride, + Mean_ptr, + Mean_row_stride, + Rstd_ptr, + Rstd_row_stride, + K, + eps, + BLOCK_SIZE: tl.constexpr, +): + """Per-row LayerNorm reduction: writes mean and 1/sqrt(var+eps) for each row of X.""" + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK_SIZE) + mask = cols < K + + x = tl.load(X_ptr + row * X_row_stride + cols, mask=mask, other=0.0) + mean = tl.sum(x, axis=0) / K + centered = tl.where(mask, x - mean, 0.0) + var = tl.sum(centered * centered, axis=0) / K + rstd = 1.0 / tl.sqrt(var + eps) + + tl.store(Mean_ptr + row * Mean_row_stride, mean) + tl.store(Rstd_ptr + row * Rstd_row_stride, rstd) + + +@triton.jit +def _lnlin_swiglu_fwd_kernel( + X_ptr, # [M, K] + W_ptr, # [K, 2N] — first half (cols [0:N]) feeds the silu input; + # second half (cols [N:2N]) is the gate output + LN_W_ptr, # [K] + LN_B_ptr, # [K] + Lin_ptr, # [M, 2N] — full linear output (saved for backward) + Out_ptr, # [M, N] — silu(x1) * x2 (input to the next Linear) + Mean_ptr, + Rstd_ptr, + M, + N, + K, + stride_xm, + stride_xk, + stride_wk, + stride_wn, + stride_lin_m, + stride_lin_n, + stride_out_m, + stride_out_n, + HAS_LN_BIAS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """One program produces a [BLOCK_M, BLOCK_N] tile of `Out`. To avoid + re-reading X for the two halves of W12, the program does TWO matmul + accumulators (against W_a and W_b) in the same K-loop, producing both + halves of the linear output for the same M tile. Then SwiGLU is applied + in registers and both the [M, 2N] linear output and [M, N] swiglu output + are written.""" + pid = tl.program_id(axis=0).to(tl.int64) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + + mean = tl.load(Mean_ptr + offs_m, mask=offs_m < M, other=0.0) + rstd = tl.load(Rstd_ptr + offs_m, mask=offs_m < M, other=0.0) + + x_ptrs = X_ptr + (offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk) + wa_ptrs = W_ptr + (offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn) + wb_ptrs = W_ptr + (offs_k[:, None] * stride_wk + (N + offs_n[None, :]) * stride_wn) + + a_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + b_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)): + kk = k * BLOCK_SIZE_K + k_remaining = K - kk + k_mask = offs_k < k_remaining + + x = tl.load( + x_ptrs, + mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), + other=0.0, + ) + ln_w = tl.load(LN_W_ptr + kk + offs_k, mask=k_mask, other=0.0) + if HAS_LN_BIAS: + ln_b = tl.load(LN_B_ptr + kk + offs_k, mask=k_mask, other=0.0) + x_hat = ((x - mean[:, None]) * rstd[:, None]) * ln_w[None, :] + ln_b[ + None, : + ] + else: + x_hat = ((x - mean[:, None]) * rstd[:, None]) * ln_w[None, :] + + wa = tl.load( + wa_ptrs, + mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), + other=0.0, + ) + wb = tl.load( + wb_ptrs, + mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), + other=0.0, + ) + + a_acc = tl.dot(x_hat, wa, a_acc) + b_acc = tl.dot(x_hat, wb, b_acc) + + x_ptrs += BLOCK_SIZE_K * stride_xk + wa_ptrs += BLOCK_SIZE_K * stride_wk + wb_ptrs += BLOCK_SIZE_K * stride_wk + + a_bf = a_acc.to(Lin_ptr.type.element_ty) + b_bf = b_acc.to(Lin_ptr.type.element_ty) + + out_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + lin_a_ptrs = ( + Lin_ptr + offs_m[:, None] * stride_lin_m + offs_n[None, :] * stride_lin_n + ) + lin_b_ptrs = ( + Lin_ptr + offs_m[:, None] * stride_lin_m + (N + offs_n[None, :]) * stride_lin_n + ) + tl.store(lin_a_ptrs, a_bf, mask=out_mask) + tl.store(lin_b_ptrs, b_bf, mask=out_mask) + + # SwiGLU = silu(a) * b. SiLU computed in fp32 for numerical accuracy. + sig = tl.sigmoid(a_acc) + silu_a = a_acc * sig + swiglu = silu_a * b_acc + out_ptrs = Out_ptr + offs_m[:, None] * stride_out_m + offs_n[None, :] * stride_out_n + tl.store(out_ptrs, swiglu.to(Out_ptr.type.element_ty), mask=out_mask) + + +# SwiGLU backward: in-place into the saved linear-output buffer (no alloc). +@triton.jit +def _swiglu_bwd_inplace_kernel( + dout_ptr, lin_ptr, M, N, stride_d, stride_l, BLOCK: tl.constexpr +): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK) + mask = cols < N + + a_ptr = lin_ptr + row * stride_l + cols + b_ptr = lin_ptr + row * stride_l + N + cols + do_ptr = dout_ptr + row * stride_d + cols + + a = tl.load(a_ptr, mask=mask, other=0.0).to(tl.float32) + b = tl.load(b_ptr, mask=mask, other=0.0).to(tl.float32) + do = tl.load(do_ptr, mask=mask, other=0.0).to(tl.float32) + + sig = tl.sigmoid(a) + silu = a * sig + # d(silu)/da = sig + silu*(1 - sig) = sig*(1 + a*(1 - sig)) + silu_grad = sig * (1.0 + a * (1.0 - sig)) + + da = do * b * silu_grad + db = do * silu + + tl.store(a_ptr, da.to(lin_ptr.type.element_ty), mask=mask) + tl.store(b_ptr, db.to(lin_ptr.type.element_ty), mask=mask) + + +_FWD_CONFIG_K128 = dict( + BLOCK_SIZE_M=128, + BLOCK_SIZE_N=128, + BLOCK_SIZE_K=32, + GROUP_SIZE_M=8, + num_stages=4, + num_warps=8, +) +_FWD_CONFIG_K256 = dict( + BLOCK_SIZE_M=64, + BLOCK_SIZE_N=64, + BLOCK_SIZE_K=64, + GROUP_SIZE_M=8, + num_stages=3, + num_warps=4, +) + + +def _pick_fwd_config(K: int) -> dict: + if K == 128: + return _FWD_CONFIG_K128 + if K == 256: + return _FWD_CONFIG_K256 + # Reasonable default for other K. May be sub-optimal — autotune for new shapes. + return _FWD_CONFIG_K256 + + +def _next_pow2(n: int) -> int: + p = 1 + while p < n: + p <<= 1 + return p + + +def _ln_stats_settings(K: int) -> tuple[int, int]: + BLOCK = _next_pow2(K) + if BLOCK <= 256: + num_warps = 4 + elif BLOCK <= 1024: + num_warps = 8 + else: + num_warps = 16 + return BLOCK, num_warps + + +def _lnlin_swiglu_fwd( + x_2d: torch.Tensor, W12: torch.Tensor, LN_W: torch.Tensor, LN_B: torch.Tensor | None +): + assert x_2d.is_contiguous(), "X must be contiguous" + M, K = x_2d.shape + K2, two_N = W12.shape + assert K2 == K and two_N % 2 == 0, f"W12 shape mismatch: {W12.shape} vs K={K}" + N = two_N // 2 + + out = torch.empty((M, N), dtype=x_2d.dtype, device=x_2d.device) + lin = torch.empty((M, two_N), dtype=x_2d.dtype, device=x_2d.device) + Mean = torch.empty((M,), dtype=x_2d.dtype, device=x_2d.device) + Rstd = torch.empty((M,), dtype=x_2d.dtype, device=x_2d.device) + + block, num_warps = _ln_stats_settings(K) + _ln_stats_kernel[(M,)]( + x_2d, + x_2d.stride(0), + Mean, + Mean.stride(0), + Rstd, + Rstd.stride(0), + K, + 1e-5, + BLOCK_SIZE=block, # ty:ignore[invalid-argument-type] + num_warps=num_warps, # ty:ignore[unknown-argument] + ) + + cfg = _pick_fwd_config(K) + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + _lnlin_swiglu_fwd_kernel[grid]( + x_2d, + W12, + LN_W, + LN_B if LN_B is not None else LN_W, # ptr always non-null + lin, + out, + Mean, + Rstd, + M, + N, + K, + x_2d.stride(0), + x_2d.stride(1), + W12.stride(0), + W12.stride(1), + lin.stride(0), + lin.stride(1), + out.stride(0), + out.stride(1), + 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"], # ty:ignore[unknown-argument] + num_warps=cfg["num_warps"], # ty:ignore[unknown-argument] + ) + return out, lin, Mean, Rstd + + +def _swiglu_bwd_inplace(dout: torch.Tensor, lin: torch.Tensor) -> torch.Tensor: + """In-place SwiGLU backward: writes da, db into the [M, 2N] `lin` buffer. + + Returns the same tensor (now containing dlin values).""" + M, N = dout.shape + BLOCK = _next_pow2(N) + _swiglu_bwd_inplace_kernel[(M,)]( + dout, + lin, + M, + N, + dout.stride(0), + lin.stride(0), + BLOCK=BLOCK, # ty:ignore[invalid-argument-type] + num_warps=4, # ty:ignore[unknown-argument] + ) + return lin + + +class FusedLNLinearSwiGLUFunction(torch.autograd.Function): + """ + Forward: out = silu(x1) * x2 where (x1, x2) = chunk(LN(X) @ W12, 2) + Backward: standard chain via cuBLAS gemms + ATen native_layer_norm_backward. + """ + + @staticmethod + @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.bfloat16) + def forward(ctx, X, W12, LN_W, LN_B): + x_shape = X.shape + x_2d = X.contiguous().view(-1, x_shape[-1]) + out, lin, mean, rstd = _lnlin_swiglu_fwd(x_2d, W12, LN_W, LN_B) + ctx.save_for_backward(x_2d, W12, LN_W, LN_B, mean, rstd, lin) + ctx.x_shape = x_shape + ctx.has_ln_bias = LN_B is not None + return out.view(*x_shape[:-1], out.shape[-1]) + + @staticmethod + @torch.amp.custom_bwd(device_type="cuda") + def backward(ctx, dout): + x_2d, W12, LN_W, LN_B, mean, rstd, lin = ctx.saved_tensors + dout_2d = dout.contiguous().view(-1, dout.shape[-1]) + K = x_2d.shape[1] + + # SwiGLU backward, in-place into the saved `lin` buffer (no new alloc). + dlin = _swiglu_bwd_inplace(dout_2d, lin) + + # LN+Linear backward via cuBLAS + ATen LN bwd. + # x_norm needed for dW12; recomputed via F.layer_norm (cuDNN, ~100 µs at d=256). + x_norm = F.layer_norm(x_2d, (K,), LN_W, LN_B, eps=1e-5) + dW12 = x_norm.transpose(0, 1) @ dlin # [K, 2N] + dx_norm = dlin @ W12.transpose(0, 1) # [M, K] + del x_norm + + # native_layer_norm_backward output_mask = [need_dX, need_dGamma, need_dBeta]. + # We always need dX and dGamma; dBeta only when LN bias was used. + output_mask = [True, True, ctx.has_ln_bias] + dX, dLN_W, dLN_B = torch.ops.aten.native_layer_norm_backward( + dx_norm, x_2d, [K], mean.float(), rstd.float(), LN_W, LN_B, output_mask + ) + # If LN_B was None (no bias), ATen returns dLN_B=None and we just pass it through. + return dX.view(ctx.x_shape), dW12, dLN_W, dLN_B + + +class FusedLNLinearSwiGLU(nn.Module): + """Fused module for `LayerNorm(d) -> Linear(d, 2*hidden) -> silu(x1)*x2`. + + Convention: silu of FIRST half. Output: [..., hidden].""" + + def __init__( + self, + d_model: int, + d_inner: int, + has_ln_bias: bool = True, + device=None, + dtype=None, + ): + super().__init__() + factory = {"device": device, "dtype": dtype} + self.d_model = d_model + self.d_inner = d_inner + self.LN_W = nn.Parameter(torch.empty(d_model, **factory)) + self.LN_B = ( + nn.Parameter(torch.empty(d_model, **factory)) if has_ln_bias else None + ) + self.W12 = nn.Parameter(torch.empty(d_model, 2 * d_inner, **factory)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.ones_(self.LN_W) + if self.LN_B is not None: + nn.init.zeros_(self.LN_B) + bound = 1.0 / math.sqrt(self.d_model) + nn.init.uniform_(self.W12, -bound, bound) + + def forward(self, X: torch.Tensor) -> torch.Tensor: + return FusedLNLinearSwiGLUFunction.apply(X, self.W12, self.LN_W, self.LN_B) diff --git a/esm/models/esmfold2/kernels/trimul_einsum_triton.py b/esm/models/esmfold2/kernels/trimul_einsum_triton.py new file mode 100644 index 00000000..98c1b757 --- /dev/null +++ b/esm/models/esmfold2/kernels/trimul_einsum_triton.py @@ -0,0 +1,242 @@ +"""Triton kernel for trimul stage-3 batched einsum in native (D, B, L, L) layout. + + outgoing: ``out[d,b,i,j] = sum_k a[d,b,i,k] * b[d,b,j,k]`` (TRANSPOSE_B=True) + incoming: ``out[d,b,i,j] = sum_k a[d,b,k,i] * b[d,b,k,j]`` (TRANSPOSE_A=True) + +All three tensors share the dense ``(D, B, L, L)`` row-major layout produced +upstream by ``fused_gated_dual_gemm_split`` and consumed downstream by +``layer_norm_transpose(..., layout="dbij->bijd")``. Operating natively in +this layout avoids the contiguous copy that ``torch.einsum``'s autograd +inserts in its backward when saved tensors don't match its expected layout. + +A single kernel covers all 4 transpose patterns (fwd + 4 bwd grad patterns) +via ``TRANSPOSE_A`` / ``TRANSPOSE_B`` constexpr flags. +""" + +import torch +import triton +import triton.language as tl + +# Static config — runtime autotune cold-start is unshippable for inference. +# Triton fwd loses to cuBLAS bgemm; bwd wins by avoiding torch.einsum's +# autograd-internal contiguous copy on (D, B, L, L) saved tensors. +_AUTOTUNE_CONFIGS = [ + triton.Config( + {"TILE_M": 128, "TILE_N": 128, "TILE_K": 64, "GROUP_M": 8}, + num_stages=3, + num_warps=8, + ) +] + + +@triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["L", "TRANSPOSE_A", "TRANSPOSE_B"]) +@triton.jit +def _batched_einsum_kernel( + a_ptr, # (D, B, L, L) row-major + b_ptr, # (D, B, L, L) row-major + out_ptr, # (D, B, L, L) row-major + D, + B, + L, + stride_a_d, # = B*L*L + stride_a_b, # = L*L + stride_b_d, + stride_b_b, + stride_out_d, + stride_out_b, + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + TILE_K: tl.constexpr, + GROUP_M: tl.constexpr, + TRANSPOSE_A: tl.constexpr, + TRANSPOSE_B: tl.constexpr, +): + # Grid: (num_m * num_n, D * B); (m, n) swizzled for L2 reuse, (d, b) on axis-1. + pid_mn = tl.program_id(axis=0) + pid_db = tl.program_id(axis=1) + + pid_d = pid_db // B + pid_b = pid_db % B + + num_pid_m = tl.cdiv(L, TILE_M) + num_pid_n = tl.cdiv(L, TILE_N) + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid_mn // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid_mn % num_pid_in_group) % group_size_m) + pid_n = (pid_mn % num_pid_in_group) // group_size_m + + # int64 indexing — (D,B,L,L) at L=1024 D=128 is well over 2**31. + pid_d = tl.cast(pid_d, tl.int64) + pid_b = tl.cast(pid_b, tl.int64) + pid_m_i = tl.cast(pid_m, tl.int64) + pid_n_i = tl.cast(pid_n, tl.int64) + L_i = tl.cast(L, tl.int64) + + a_plane = a_ptr + pid_d * stride_a_d + pid_b * stride_a_b + b_plane = b_ptr + pid_d * stride_b_d + pid_b * stride_b_b + out_plane = out_ptr + pid_d * stride_out_d + pid_b * stride_out_b + + offs_m = pid_m_i * TILE_M + tl.arange(0, TILE_M).to(tl.int64) + offs_n = pid_n_i * TILE_N + tl.arange(0, TILE_N).to(tl.int64) + offs_k = tl.arange(0, TILE_K).to(tl.int64) + + mask_m = offs_m < L_i + mask_n = offs_n < L_i + + if TRANSPOSE_A: + a_ptrs = a_plane + (offs_k[:, None] * L_i + offs_m[None, :]) + a_step = TILE_K * L_i + else: + a_ptrs = a_plane + (offs_m[:, None] * L_i + offs_k[None, :]) + a_step = TILE_K + + if TRANSPOSE_B: + b_ptrs = b_plane + (offs_n[None, :] * L_i + offs_k[:, None]) + b_step = TILE_K + else: + b_ptrs = b_plane + (offs_k[:, None] * L_i + offs_n[None, :]) + b_step = TILE_K * L_i + + acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + n_k_tiles = tl.cdiv(L, TILE_K) + for kk in range(0, n_k_tiles): + k_remaining = L_i - kk * TILE_K + mask_k = offs_k < k_remaining + + if TRANSPOSE_A: + a_mask = mask_k[:, None] & mask_m[None, :] + a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0) + a_tile = tl.trans(a_tile) # tl.dot wants (M, K) + else: + a_mask = mask_m[:, None] & mask_k[None, :] + a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0) + + if TRANSPOSE_B: + b_mask = mask_k[:, None] & mask_n[None, :] + b_tile = tl.load(b_ptrs, mask=b_mask, other=0.0) + else: + b_mask = mask_k[:, None] & mask_n[None, :] + b_tile = tl.load(b_ptrs, mask=b_mask, other=0.0) + + acc = tl.dot(a_tile, b_tile, acc) + + a_ptrs += a_step + b_ptrs += b_step + + out_ptrs = out_plane + (offs_m[:, None] * L_i + offs_n[None, :]) + out_mask = mask_m[:, None] & mask_n[None, :] + tl.store(out_ptrs, acc.to(out_ptr.type.element_ty), mask=out_mask) + + +def _batched_einsum( + a: torch.Tensor, b: torch.Tensor, transpose_a: bool, transpose_b: bool +) -> torch.Tensor: + """Compute ``(A op_a) @ (B op_b)`` per (d, b) plane. + + Parameters + ---------- + a, b : torch.Tensor + Shape (D, B, L, L), bf16, contiguous. + transpose_a : bool + If True, contract over the row dim of A (k = leading), else over col dim. + transpose_b : bool + If True, B is read as (n, k) (i.e. B^T enters the matmul: out += A·B^T), + else as (k, n) (out += A·B). + + Returns + ------- + out : torch.Tensor + Shape (D, B, L, L), bf16. + """ + 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" + + D, B, L_row, L_col = a.shape + assert L_row == L_col, "L_row must equal L_col for stage-3 einsum" + L = L_row + + out = torch.empty_like(a) + + stride_a_d, stride_a_b = a.stride(0), a.stride(1) + stride_b_d, stride_b_b = b.stride(0), b.stride(1) + stride_out_d, stride_out_b = out.stride(0), out.stride(1) + + def grid(meta): + num_m = triton.cdiv(L, meta["TILE_M"]) + num_n = triton.cdiv(L, meta["TILE_N"]) + return (num_m * num_n, D * B) + + _batched_einsum_kernel[grid]( + a, + b, + out, + D, + B, + L, + stride_a_d, + stride_a_b, + stride_b_d, + stride_b_b, + stride_out_d, + stride_out_b, + TRANSPOSE_A=transpose_a, + TRANSPOSE_B=transpose_b, + ) + return out + + +class TrimulBatchedEinsum(torch.autograd.Function): + """Autograd-aware wrapper for trimul stage-3 batched einsum. + + Forward direction is one of: + "outgoing": out[d,b,i,j] = sum_k a[d,b,i,k] * b[d,b,j,k] (A · B^T) + "incoming": out[d,b,i,j] = sum_k a[d,b,k,i] * b[d,b,k,j] (A^T · B) + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, a: torch.Tensor, b: torch.Tensor, direction: str + ) -> torch.Tensor: + if direction == "outgoing": + out = _batched_einsum(a, b, transpose_a=False, transpose_b=True) + elif direction == "incoming": + out = _batched_einsum(a, b, transpose_a=True, transpose_b=False) + else: + raise ValueError(f"unknown direction {direction!r}") + ctx.save_for_backward(a, b) + ctx.direction = direction + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + a, b = ctx.saved_tensors + direction = ctx.direction + if not grad_out.is_contiguous(): + grad_out = grad_out.contiguous() + + if direction == "outgoing": + grad_a = _batched_einsum(grad_out, b, transpose_a=False, transpose_b=False) + grad_b = _batched_einsum(grad_out, a, transpose_a=True, transpose_b=False) + else: # incoming + grad_a = _batched_einsum(b, grad_out, transpose_a=False, transpose_b=True) + grad_b = _batched_einsum(a, grad_out, transpose_a=False, transpose_b=False) + return grad_a, grad_b, None + + +def trimul_batched_einsum( + a: torch.Tensor, b: torch.Tensor, direction: str +) -> torch.Tensor: + """Public entry: dispatches to autograd Function if grad is enabled.""" + if torch.is_grad_enabled() and (a.requires_grad or b.requires_grad): + return TrimulBatchedEinsum.apply(a, b, direction) # type: ignore[return-value] + if direction == "outgoing": + return _batched_einsum(a, b, transpose_a=False, transpose_b=True) + elif direction == "incoming": + return _batched_einsum(a, b, transpose_a=True, transpose_b=False) + raise ValueError(f"unknown direction {direction!r}") diff --git a/esm/models/esmfold2/kernels/trimul_with_residual.py b/esm/models/esmfold2/kernels/trimul_with_residual.py new file mode 100644 index 00000000..6f642c3f --- /dev/null +++ b/esm/models/esmfold2/kernels/trimul_with_residual.py @@ -0,0 +1,625 @@ +"""TriMul with output residual + dropout-mask epilogue fused into the final GEMM. + +Inspired by cuequivariance's ``triangle_multiplicative_update`` +(https://docs.nvidia.com/cuda/cuequivariance/index.html); independently +re-implemented in Triton with the output residual and dropout mask folded +into the final gated GEMM so the ``delta = TriMul(pair)`` intermediate is +never written to HBM. + +Replaces the standard pattern + + pair_new = pair + dropout_mask * triangle_multiplicative_update(pair) + +— where the two ops materialize the full ``delta = TriMul(pair)`` tensor +between them — with a single fused path. The final gated GEMM fuses three +operations end-to-end: + + 1. Computes ``delta = sigmoid(x_in @ Wg) * (x_out @ Wp)`` per output tile. + 2. Multiplies by the row-shared dropout mask in-register + (``mask[b, j, d]`` indexed by decoded ``b, j`` from the flat ``m`` index). + 3. Adds the prior pair residual loaded from HBM. + 4. Writes the new pair tensor. + +Saves: one full pair-tensor write (``delta``) + one full pair-tensor read +(``delta`` re-read by FusedDropoutResidual). Roughly ~250 MB per call at +L=768, B=5 in bf16; compounds across loops. + +Forward + backward both implemented for the bf16 path. fp8 IO paths and +``transpose_out`` remain inference-only. The backward kernel recomputes +``gate = sigmoid(x1 @ w1.T)`` and ``val = x2 @ w2.T`` (the two dual GEMMs) +to keep saved-context cost down, and emits per-element +``d_gate_logits``, ``d_val_acc``, ``d_drop_mask_partials``; outer cuBLAS +calls handle the weight/input GEMM reductions. The residual add is +identity in the backward: ``d_residual = d_out``. +""" +# ruff: noqa: E402 + +from __future__ import annotations + +import os +import warnings + +warnings.filterwarnings("ignore", category=FutureWarning) +warnings.filterwarnings("ignore", category=UserWarning) +warnings.filterwarnings("ignore", category=DeprecationWarning) +os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1") +os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1") + +import torch +import triton +import triton.language as tl + +from esm.models.esmfold2.kernels.fused_dual_gemm import fused_gated_dual_gemm_split +from esm.models.esmfold2.kernels.fused_ln_residual import ( + fused_ln_transpose, + fused_ln_with_residual_link, +) +from esm.models.esmfold2.kernels.trimul_einsum_triton import trimul_batched_einsum + +# Static config — runtime autotune cold-start is unshippable for inference. +_AUTOTUNE_CONFIGS = [ + triton.Config( + {"TILE_M": 64, "TILE_N": 64, "TILE_K": 64, "GROUP_M": 8}, + num_stages=3, + num_warps=4, + ) +] + + +@triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["M", "N", "K", "HAS_DROP_MASK"]) +@triton.jit +def _gated_gemm_with_residual_kernel( + x1_ptr, # [M, K] — gate input (pre-residual pair) + x2_ptr, # [M, K] — value input (post-LN_out) + w1_ptr, # [N, K] — gate weight + w2_ptr, # [N, K] — value weight + residual_ptr, # [M, N] pair tensor pre-TriMul + drop_mask_ptr, # [B*N_COL, N] row-shared dropout mask + o_ptr, # [M, N] output = residual + drop_mask * sigmoid(x1@w1) * (x2@w2) + M, + N, + K, + STRIDE_B: tl.constexpr, # = N_ROW * N_COL (for decoding b from m) + N_COL: tl.constexpr, # = L_col (for decoding j from m) + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + TILE_K: tl.constexpr, + GROUP_M: tl.constexpr, + PRECISION: tl.constexpr, + HAS_DROP_MASK: tl.constexpr, + 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) + + # GROUP_M swizzle for L2 reuse of x rows across consecutive CTAs. + num_pid_m = tl.cdiv(M, TILE_M) + num_pid_n = tl.cdiv(N, TILE_N) + pid = pid_n_raw * num_pid_m + pid_m_raw + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + if NEEDS_INT64: + pid_m = tl.cast(pid_m, tl.int64) + pid_n = tl.cast(pid_n, tl.int64) + M = tl.cast(M, tl.int64) + N = tl.cast(N, tl.int64) + K = tl.cast(K, tl.int64) + + start_m = pid_m * TILE_M + start_n = pid_n * TILE_N + + offs_xm = start_m + tl.arange(0, TILE_M) + offs_wn = start_n + tl.arange(0, TILE_N) + offs_k = tl.arange(0, TILE_K) + + if NEEDS_INT64: + offs_xm = tl.cast(offs_xm, tl.int64) + offs_wn = tl.cast(offs_wn, tl.int64) + offs_k = tl.cast(offs_k, tl.int64) + + x1_ptrs = x1_ptr + (offs_xm[:, None] * K + offs_k[None, :]) + x2_ptrs = x2_ptr + (offs_xm[:, None] * K + offs_k[None, :]) + w_tile_offs = offs_wn[None, :] * K + offs_k[:, None] + + acc_1 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + acc_2 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + mask_m = offs_xm < M + + for _ in range(0, tl.cdiv(K, TILE_K)): + x1 = tl.load(x1_ptrs, mask=mask_m[:, None], other=0.0).to( + w1_ptr.type.element_ty + ) + w1_ptrs = w1_ptr + w_tile_offs + w1 = tl.load(w1_ptrs) + if PRECISION == 0: + acc_1 = tl.dot(x1, w1, acc_1) + elif PRECISION == 2: + acc_1 = tl.dot(x1, w1, acc_1, input_precision="tf32x3") + else: + acc_1 = tl.dot(x1, w1, acc_1, input_precision="ieee") + x1_ptrs += TILE_K + w1_ptr += TILE_K + + for _ in range(0, tl.cdiv(K, TILE_K)): + x2 = tl.load(x2_ptrs, mask=mask_m[:, None], other=0.0).to( + w2_ptr.type.element_ty + ) + w2_ptrs = w2_ptr + w_tile_offs + w2 = tl.load(w2_ptrs) + if PRECISION == 0: + acc_2 = tl.dot(x2, w2, acc_2) + elif PRECISION == 2: + acc_2 = tl.dot(x2, w2, acc_2, input_precision="tf32x3") + else: + acc_2 = tl.dot(x2, w2, acc_2, input_precision="ieee") + x2_ptrs += TILE_K + w2_ptr += TILE_K + + acc_1 = 1.0 / (1.0 + tl.exp(-acc_1)) + delta = acc_1 * acc_2 + + offs_om = pid_m * TILE_M + tl.arange(0, TILE_M) + offs_on = pid_n * TILE_N + tl.arange(0, TILE_N) + if NEEDS_INT64: + offs_om = tl.cast(offs_om, tl.int64) + offs_on = tl.cast(offs_on, tl.int64) + + # Row-shared dropout: mask is [B*N_COL, N]; decode (b, j) from m. + if HAS_DROP_MASK: + b = offs_om // STRIDE_B + j = offs_om % N_COL + mask_row = b * N_COL + j + drop_ptrs = drop_mask_ptr + mask_row[:, None] * N + offs_on[None, :] + drop_tile = tl.load(drop_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + delta = delta * drop_tile + + resid_ptrs = residual_ptr + offs_om[:, None] * N + offs_on[None, :] + resid_tile = tl.load(resid_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + out_val = delta + resid_tile + + o_ptrs = o_ptr + offs_om[:, None] * N + offs_on[None, :] + o_mask = mask_m[:, None] + tl.store(o_ptrs, out_val.to(o_ptr.type.element_ty), mask=o_mask) + + +# Backward recomputes the two fwd GEMMs, emits per-element grads for the +# outer cuBLAS GEMMs (d_x1, d_x2, d_w1, d_w2). d_residual is identity through +# the residual add. d_drop_mask is row-shared → per-pid_n partial + host reduce. + + +# Narrow TILE_N=32 in bwd: register pressure is dominated by the two output +# gradient tensors (vs forward's single output). +_BWD_AUTOTUNE_CONFIGS = [ + triton.Config( + {"TILE_M": 64, "TILE_N": 32, "TILE_K": 64, "GROUP_M": 8}, + num_stages=3, + num_warps=4, + ) +] + + +@triton.autotune(configs=_BWD_AUTOTUNE_CONFIGS, key=["M", "N", "K", "HAS_DROP_MASK"]) +@triton.jit +def _gated_gemm_with_residual_backward_kernel( + grad_out_ptr, # [M, N] bf16 — incoming grad + x1_ptr, # [M, K] bf16 — saved gate input + x2_ptr, # [M, K] bf16 — saved value input + w1_ptr, # [N, K] bf16 + w2_ptr, # [N, K] bf16 + drop_mask_ptr, # [B*N_COL, N] bf16 — row-shared (unused if HAS_DROP_MASK=0) + grad_gate_logits_ptr, # [M, N] bf16 — out + grad_val_acc_ptr, # [M, N] bf16 — out + grad_drop_mask_buf_ptr, # [B*N_COL, N] fp32 — drop_mask grad accumulator (atomic_add) + M, + N, + K, + STRIDE_B: tl.constexpr, + N_COL: tl.constexpr, + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + TILE_K: tl.constexpr, + GROUP_M: tl.constexpr, + HAS_DROP_MASK: tl.constexpr, + 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) + + num_pid_m = tl.cdiv(M, TILE_M) + num_pid_n = tl.cdiv(N, TILE_N) + pid = pid_n_raw * num_pid_m + pid_m_raw + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + if NEEDS_INT64: + pid_m = tl.cast(pid_m, tl.int64) + pid_n = tl.cast(pid_n, tl.int64) + M = tl.cast(M, tl.int64) + N = tl.cast(N, tl.int64) + K = tl.cast(K, tl.int64) + + start_m = pid_m * TILE_M + start_n = pid_n * TILE_N + + offs_xm = start_m + tl.arange(0, TILE_M) + offs_wn = start_n + tl.arange(0, TILE_N) + offs_k = tl.arange(0, TILE_K) + if NEEDS_INT64: + offs_xm = tl.cast(offs_xm, tl.int64) + offs_wn = tl.cast(offs_wn, tl.int64) + offs_k = tl.cast(offs_k, tl.int64) + + x1_ptrs = x1_ptr + (offs_xm[:, None] * K + offs_k[None, :]) + x2_ptrs = x2_ptr + (offs_xm[:, None] * K + offs_k[None, :]) + w_tile_offs = offs_wn[None, :] * K + offs_k[:, None] + + acc_1 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + acc_2 = tl.zeros((TILE_M, TILE_N), dtype=tl.float32) + mask_m = offs_xm < M + + w1p = w1_ptr + for _ in range(0, tl.cdiv(K, TILE_K)): + x1 = tl.load(x1_ptrs, mask=mask_m[:, None], other=0.0).to( + w1_ptr.type.element_ty + ) + w1 = tl.load(w1p + w_tile_offs) + acc_1 = tl.dot(x1, w1, acc_1) + x1_ptrs += TILE_K + w1p += TILE_K + + w2p = w2_ptr + for _ in range(0, tl.cdiv(K, TILE_K)): + x2 = tl.load(x2_ptrs, mask=mask_m[:, None], other=0.0).to( + w2_ptr.type.element_ty + ) + w2 = tl.load(w2p + w_tile_offs) + acc_2 = tl.dot(x2, w2, acc_2) + x2_ptrs += TILE_K + w2p += TILE_K + + g = tl.sigmoid(acc_1) + delta = g * acc_2 # pre-mask delta (forward post-mask = delta * drop_mask) + + offs_om = pid_m * TILE_M + tl.arange(0, TILE_M) + offs_on = pid_n * TILE_N + tl.arange(0, TILE_N) + if NEEDS_INT64: + offs_om = tl.cast(offs_om, tl.int64) + offs_on = tl.cast(offs_on, tl.int64) + + grad_o_ptrs = grad_out_ptr + offs_om[:, None] * N + offs_on[None, :] + grad_o = tl.load(grad_o_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + if HAS_DROP_MASK: + # Multiple m's collide on the same mask_row, so use fp32 atomic_add to a + # (B*N_COL, N) accumulator (handles intra- and inter-tile collisions). + b = offs_om // STRIDE_B + j = offs_om % N_COL + mask_row = b * N_COL + j + drop_ptrs = drop_mask_ptr + mask_row[:, None] * N + offs_on[None, :] + drop_tile = tl.load(drop_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + d_drop_per_elem = grad_o * delta + acc_ptrs = grad_drop_mask_buf_ptr + mask_row[:, None] * N + offs_on[None, :] + tl.atomic_add(acc_ptrs, d_drop_per_elem, mask=mask_m[:, None]) + + grad_o = grad_o * drop_tile + + d_val_acc = (grad_o * g).to(grad_val_acc_ptr.type.element_ty) + d_val_acc_ptrs = grad_val_acc_ptr + offs_om[:, None] * N + offs_on[None, :] + tl.store(d_val_acc_ptrs, d_val_acc, mask=mask_m[:, None]) + + d_gate_logits = (grad_o * acc_2 * g * (1.0 - g)).to( + grad_gate_logits_ptr.type.element_ty + ) + d_gate_logits_ptrs = grad_gate_logits_ptr + offs_om[:, None] * N + offs_on[None, :] + tl.store(d_gate_logits_ptrs, d_gate_logits, mask=mask_m[:, None]) + + +def _gated_gemm_with_residual_bwd( + grad_out: torch.Tensor, # [M, N] bf16 + x1: torch.Tensor, # [M, K] bf16 + x2: torch.Tensor, # [M, K] bf16 + w1: torch.Tensor, + w2: torch.Tensor, + drop_mask: torch.Tensor | None, + n_row: int, + n_col: int, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, +]: + """Returns (d_x1, d_x2, d_w1, d_w2, d_residual, d_drop_mask).""" + 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 x1.dtype == torch.bfloat16 and x2.dtype == torch.bfloat16 + + # Don't call .contiguous() unconditionally (avoids a full-tensor clone). + if not grad_out.is_contiguous(): + grad_out = grad_out.contiguous() + grad_out_2d = grad_out.view(M, N) + x1_c = x1 if x1.is_contiguous() else x1.contiguous() + x2_c = x2 if x2.is_contiguous() else x2.contiguous() + w1_c = w1 if w1.is_contiguous() else w1.contiguous() + w2_c = w2 if w2.is_contiguous() else w2.contiguous() + + grad_gate_logits = torch.empty((M, N), device=x1.device, dtype=torch.bfloat16) + grad_val_acc = torch.empty((M, N), device=x1.device, dtype=torch.bfloat16) + + # fp32 accumulator for row-shared mask_row collisions across m. + if drop_mask is not None: + drop_mask = drop_mask.contiguous().view(-1, N) + n_mask_rows = drop_mask.shape[0] + grad_drop_buf = torch.zeros( + (n_mask_rows, N), device=x1.device, dtype=torch.float32 + ) + else: + grad_drop_buf = torch.empty((1,), device=x1.device, dtype=torch.float32) + + _dummy_mask = torch.zeros((), device=x1.device, dtype=torch.bfloat16) + + NEEDS_INT64 = (M * K >= 2**31 - 1) or (M * N >= 2**31 - 1) + + def grid(meta): + return (triton.cdiv(M, meta["TILE_M"]), triton.cdiv(N, meta["TILE_N"])) + + _gated_gemm_with_residual_backward_kernel[grid]( + grad_out_2d, + x1_c, + x2_c, + w1_c, + w2_c, + drop_mask if drop_mask is not None else _dummy_mask, + grad_gate_logits, + grad_val_acc, + grad_drop_buf, + M, + N, + K, + STRIDE_B=n_row * n_col, + N_COL=n_col, + HAS_DROP_MASK=drop_mask is not None, + NEEDS_INT64=NEEDS_INT64, + ) + + d_w1 = grad_gate_logits.t() @ x1_c # [N, K] + d_w2 = grad_val_acc.t() @ x2_c # [N, K] + d_x1 = grad_gate_logits @ w1_c # [M, K] + d_x2 = grad_val_acc @ w2_c # [M, K] + + d_residual = grad_out_2d + + if drop_mask is not None: + d_drop_mask = grad_drop_buf.to(drop_mask.dtype) + else: + d_drop_mask = None + + return d_x1, d_x2, d_w1, d_w2, d_residual, d_drop_mask + + +class GatedGEMMWithResidual(torch.autograd.Function): + """Autograd wrapper for ``_gated_gemm_with_residual`` (bf16 path).""" + + @staticmethod + def forward( + ctx, + x1: torch.Tensor, + x2: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + residual: torch.Tensor, + drop_mask: torch.Tensor | None, + n_row: int, + n_col: int, + precision: int, + ) -> torch.Tensor: + out = _gated_gemm_with_residual_fwd( + x1, x2, w1, w2, residual, drop_mask, n_row, n_col, precision=precision + ) + if drop_mask is None: + ctx.save_for_backward(x1, x2, w1, w2) + ctx.has_drop_mask = False + else: + ctx.save_for_backward(x1, x2, w1, w2, drop_mask) + ctx.has_drop_mask = True + ctx.n_row = n_row + ctx.n_col = n_col + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + if ctx.has_drop_mask: + x1, x2, w1, w2, drop_mask = ctx.saved_tensors + else: + x1, x2, w1, w2 = ctx.saved_tensors + drop_mask = None + d_x1, d_x2, d_w1, d_w2, d_res, d_drop = _gated_gemm_with_residual_bwd( + grad_out.contiguous(), x1, x2, w1, w2, drop_mask, ctx.n_row, ctx.n_col + ) + return d_x1, d_x2, d_w1, d_w2, d_res, d_drop, None, None, None + + +def _gated_gemm_with_residual_fwd( + x1: torch.Tensor, + x2: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + residual: torch.Tensor, + drop_mask: torch.Tensor | None, + n_row: int, + n_col: int, + precision: int = 0, +) -> torch.Tensor: + """Run the residual+dropout-aware final GEMM. + + Shapes: x1, x2: (M, K). w1, w2: (N, K). residual: (M, N). drop_mask: (B*n_col, N) or None. + Output: (M, N). + """ + M, K = x1.shape + N = w1.shape[0] + x1 = x1.contiguous() + x2 = x2.contiguous() + w1 = w1.contiguous() + w2 = w2.contiguous() + residual = residual.contiguous().view(M, N) + if drop_mask is not None: + drop_mask = drop_mask.contiguous().view(-1, N) + out = residual.new_empty((M, N)) + + NEEDS_INT64 = (M * K >= 2**31 - 1) or (M * N >= 2**31 - 1) + + def grid(meta): + return (triton.cdiv(M, meta["TILE_M"]), triton.cdiv(N, meta["TILE_N"])) + + _gated_gemm_with_residual_kernel[grid]( + x1, + x2, + w1, + w2, + residual, + drop_mask if drop_mask is not None else residual, # dummy pass-through + out, + M, + N, + K, + STRIDE_B=n_row * n_col, + N_COL=n_col, + PRECISION=precision, + HAS_DROP_MASK=drop_mask is not None, + NEEDS_INT64=NEEDS_INT64, + ) + return out + + +def _gated_gemm_with_residual( + x1: torch.Tensor, + x2: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + residual: torch.Tensor, + drop_mask: torch.Tensor | None, + n_row: int, + n_col: int, + precision: int = 0, +) -> torch.Tensor: + """Public wrapper: dispatches to autograd Function if grad is enabled.""" + if torch.is_grad_enabled() and ( + x1.requires_grad + or x2.requires_grad + or w1.requires_grad + or w2.requires_grad + or residual.requires_grad + ): + return GatedGEMMWithResidual.apply( + x1, x2, w1, w2, residual, drop_mask, n_row, n_col, precision + ) + return _gated_gemm_with_residual_fwd( + x1, x2, w1, w2, residual, drop_mask, n_row, n_col, precision=precision + ) + + +@torch._dynamo.disable +def triangle_multiplicative_update_with_residual( + pair: torch.Tensor, + direction: str, + residual: torch.Tensor, + drop_mask: torch.Tensor | None, + *, + norm_in_weight: torch.Tensor, + norm_in_bias: torch.Tensor, + p_in_weight: torch.Tensor, + g_in_weight: torch.Tensor, + norm_out_weight: torch.Tensor, + norm_out_bias: torch.Tensor, + p_out_weight: torch.Tensor, + g_out_weight: torch.Tensor, + mask: torch.Tensor | None = None, + eps: float = 1e-5, + precision: int = 0, +) -> torch.Tensor: + """Fused TriMul-and-residual: pair_new = residual + drop_mask * TriMul(pair). + + The intermediate ``delta = TriMul(pair)`` is never materialized in HBM — + the dropout-mask multiply and residual-add happen in-kernel at the final + GEMM. Saves one pair-tensor write + read vs the + ``TriMul → FusedDropoutResidual`` baseline. + + Shapes: + pair, residual: (B, L, L, c_z) + drop_mask: (B, 1, L, c_z) or (B*L, c_z) — row-shared + mask: (B, L, L) or None + weights: (N, K) bf16 — see arg list for each stage. + """ + assert pair.shape == residual.shape + B, L_row, L_col, c_z = pair.shape + + # Stage 1: input LayerNorm. When ``residual is pair`` (common case for trimul), + # use the fused LN-fwd + LN-bwd-with-residual variant so the bwd pass adds + # ``grad_residual`` into ``grad_x`` in one kernel (saves one pair-tensor + # HBM round-trip). + if residual is pair: + x, residual_alias = fused_ln_with_residual_link( + pair, norm_in_weight, norm_in_bias, pair, eps=eps, layout="bijd->bijd" + ) + else: + x = fused_ln_transpose( + pair, norm_in_weight, norm_in_bias, eps=eps, layout="bijd->bijd" + ) + residual_alias = residual + x_in = x + + # Stage 2: gated dual GEMM → ab (transposed for the einsum's dbij layout). + # Train path runs transpose_out=False (no bwd for in-kernel transpose), + # 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. 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) + else: + x = torch.einsum("dbki,dbkj->dbij", a, b_t) + + # Stage 4: output LayerNorm (back to bijd). + x_out = fused_ln_transpose( + x, norm_out_weight, norm_out_bias, eps=eps, layout="dbij->bijd" + ) + + # Stage 5: fused output gated GEMM + dropout mask + residual add. + x_in_2d = x_in.reshape(-1, c_z) + x_out_2d = x_out.reshape(-1, c_z) + # Stage-5 residual MUST be the LN1 alias when LN1-residual-fusion is on, + # otherwise the grad routing through the fused LN bwd won't fire. + residual_2d = residual_alias.reshape(-1, c_z) + out_2d = _gated_gemm_with_residual( + x_in_2d, + x_out_2d, + g_out_weight, + p_out_weight, + residual_2d, + drop_mask, + n_row=L_row, + n_col=L_col, + precision=precision, + ) + return out_2d.view(B, L_row, L_col, c_z) diff --git a/esm/models/esmfold2/layers.py b/esm/models/esmfold2/layers.py new file mode 100644 index 00000000..b3460955 --- /dev/null +++ b/esm/models/esmfold2/layers.py @@ -0,0 +1,2811 @@ +# coding=utf-8 +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Shared building blocks for ESMFold2 HuggingFace model variants.""" + +from __future__ import annotations + +import random +from contextlib import contextmanager +from functools import partial +from typing import cast + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.utils.checkpoint import checkpoint + +try: + 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 # 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 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 # ty:ignore[invalid-assignment] + _cue_tri_mul = None # ty:ignore[invalid-assignment] + CUE_AVAILABLE = False + +# Vendored inference-only Triton kernels. +try: + from esm.models.esmfold2.kernels import ( + FusedDropoutResidual as _FusedDropoutResidual, + ) + from esm.models.esmfold2.kernels import FusedLNLinearSwiGLU as _FusedLNLinearSwiGLU + from esm.models.esmfold2.kernels import fused_pair_bias as _fused_pair_bias + from esm.models.esmfold2.kernels import ( + triangle_multiplicative_update_with_residual as _fused_trimul_with_residual, + ) + + TRITON_KERNELS_AVAILABLE = True +except ImportError: + _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 esm.models.esmfold2.config import EsmFold2Config + +BACKEND_FUSED = "fused" +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.""" + return ( + TRITON_KERNELS_AVAILABLE + and getattr(module, "_kernel_backend", None) == BACKEND_FUSED + and not torch.is_grad_enabled() + and tensor.is_cuda + ) + + +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 + + +class DropoutResidual(nn.Module): + """``residual + dropout(delta)`` with row/col-shared dropout. + + Same signature on both paths. ``use_fused_kernels=True`` + ``batch_dim=1`` + routes through ``FusedDropoutResidual`` (single-pass over pair tensor, + in-place residual add). Falls back to unfused otherwise. + """ + + def __init__( + self, r: float, batch_dim: int, use_fused_kernels: bool = False + ) -> None: + super().__init__() + assert batch_dim in (1, 2), f"batch_dim must be 1 or 2, got {batch_dim}" + self._use_fused_kernels = ( + use_fused_kernels and batch_dim == 1 and _FusedDropoutResidual is not None + ) + self._batch_dim = batch_dim + self._r = r + if self._use_fused_kernels: + assert _FusedDropoutResidual is not None + self._impl: nn.Module = _FusedDropoutResidual(r) + else: + self._impl = nn.Dropout(r) + + def forward(self, residual: Tensor, delta: Tensor) -> Tensor: + if self._use_fused_kernels: + return self._impl(residual, delta) + # Unfused: row/col-shared dropout via [1, ...] mask broadcast. + if self._r == 0.0 or not self.training: + return residual + delta + shape = list(delta.shape) + shape[self._batch_dim] = 1 + mask = self._impl(delta.new_ones(shape)) + return residual + delta * mask + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +CHAR_VOCAB_SIZE: int = 64 +MAX_CHARS: int = 4 +XYZ_DIMS: int = 3 +MAX_ATOMIC_NUMBER: int = 128 + +# Input feature dim = 3 + 1 + 1 + 128 + 64*4 = 389 +ATOM_FEATURE_DIM: int = ( + XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS +) + + +NUM_RES_TYPES: int = 33 + +_EPS = 1e-5 + +# Default for the triangle / OPM / pair-transition L² ops. Caps peak memory +# so L≈2k folds on an 80 GB GPU (~76 GB peak at chunk=128 for L=1438; +# chunk=64 leaves headroom for the largest foldbench targets). Override via +# ``model.set_chunk_size(...)``; pass None to disable chunking (faster for +# short L but OOM-prone past ~600). +_DEFAULT_CHUNK_SIZE = 64 + + +# =========================================================================== +# MSA inference-time diversity augmentations +# =========================================================================== +# Operate on encoder-input tensors of shape [B, depth, L] (depth = dim 1; ``msa`` +# may also be one-hot [B, depth, L, C]). + + +def maybe_subsample_msa( + msa: Tensor, + msa_attention_mask: Tensor | None, + has_deletion: Tensor | None, + deletion_value: Tensor | None, + *, + max_depth: int | None, + enabled: bool, +) -> tuple[Tensor, Tensor | None, Tensor | None, Tensor | None]: + """Randomly subsample the MSA to ``max_depth`` rows, keeping query row 0. + + No-op when disabled, ``max_depth`` is None, or depth <= ``max_depth``. + """ + if not enabled or max_depth is None: + return msa, msa_attention_mask, has_deletion, deletion_value + depth = msa.size(1) + if depth <= 1 or depth <= max_depth: + return msa, msa_attention_mask, has_deletion, deletion_value + indices = torch.zeros(max_depth, dtype=torch.long, device=msa.device) + perm = torch.randperm(depth - 1, device=msa.device)[: max_depth - 1] + indices[1:] = perm + 1 + indices = indices.sort().values + msa = msa[:, indices] + if msa_attention_mask is not None: + msa_attention_mask = msa_attention_mask[:, indices] + if has_deletion is not None: + has_deletion = has_deletion[:, indices] + if deletion_value is not None: + deletion_value = deletion_value[:, indices] + return msa, msa_attention_mask, has_deletion, deletion_value + + +def maybe_apply_msa_column_masking( + msa_attention_mask: Tensor | None, rate: float +) -> Tensor | None: + """Mask fraction ``rate`` of MSA columns in non-query rows of + ``msa_attention_mask``, keeping query row 0. No-op when absent, ``rate <= 0``, + or depth <= 1. + """ + if msa_attention_mask is None or rate <= 0.0 or msa_attention_mask.size(1) <= 1: + return msa_attention_mask + B, _M, L = msa_attention_mask.shape + col_keep = torch.rand(B, L, device=msa_attention_mask.device) >= rate + col_keep = col_keep.unsqueeze(1).expand_as(msa_attention_mask).clone() + col_keep[:, 0, :] = True + return msa_attention_mask.bool() & col_keep + + +# =========================================================================== +# Atom-token utilities +# =========================================================================== + + +def gather_token_to_atom(token_features: Tensor, atom_to_token_idx: Tensor) -> Tensor: + """Broadcast per-token features to per-atom features using gather. + + Args: + token_features: [B, L, d] + atom_to_token_idx: [B, A] int64 + + Returns: + [B, A, d] + """ + idx = atom_to_token_idx.unsqueeze(-1).expand(-1, -1, token_features.size(-1)) + return torch.gather(token_features, 1, idx) + + +def scatter_atom_to_token( + atom_features: Tensor, + atom_to_token_idx: Tensor, + n_tokens: int, + atom_mask: Tensor | None = None, +) -> Tensor: + """Aggregate per-atom features to per-token features (mean). + + Args: + atom_features: [B, A, d] + atom_to_token_idx: [B, A] int64 + n_tokens: L + atom_mask: [B, A] bool + + Returns: + [B, L, d] + """ + B, A, d = atom_features.shape + n_out = n_tokens + idx = atom_to_token_idx + if atom_mask is not None: + idx = torch.where(atom_mask, atom_to_token_idx, n_tokens) + n_out = n_tokens + 1 + idx_expanded = idx.unsqueeze(-1).expand(B, A, d) + out = torch.zeros( + B, n_out, d, device=atom_features.device, dtype=atom_features.dtype + ) + out.scatter_reduce_( + 1, idx_expanded, atom_features, reduce="mean", include_self=False + ) + return out[:, :n_tokens, :] + + +def gather_rep_atom_coords(coords: Tensor, rep_atom_idx: Tensor) -> Tensor: + """Gather representative atom coordinates for each token. + + Args: + coords: [B, A, 3] + rep_atom_idx: [B, L] int64 + + Returns: + [B, L, 3] + """ + idx = rep_atom_idx.unsqueeze(-1).expand(-1, -1, coords.size(-1)) + return torch.gather(coords, 1, idx) + + +def _compute_intra_token_idx(atom_to_token: Tensor) -> Tensor: + """Compute local atom index within each token (vectorised). + + Atoms belonging to the same token are contiguous, so this computes a + running count that resets at each token boundary. + + Args: + atom_to_token: [B, A] flat index mapping each atom to its token. + + Returns: + [B, A] tensor with values in [0, max_atoms_per_token - 1]. + """ + same_as_prev = F.pad( + atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False + ) + ones = torch.ones_like(atom_to_token) + cumsum = torch.cumsum(ones, dim=-1) + group_start = cumsum.masked_fill(same_as_prev, 0) + group_start = torch.cummax(group_start, dim=-1).values + return cumsum - group_start + + +def _categorical_mean(logits: Tensor, start: float, end: float) -> Tensor: + """Expected value of a categorical distribution over evenly-spaced bins. + + Equivalent to ``CategoricalMixture(logits, bins=logits.shape[-1], start, end).mean()``. + + Args: + logits: [..., n_bins] + start: left boundary + end: right boundary + + Returns: + [...] expected value + """ + n_bins = logits.shape[-1] + edges = torch.linspace( + start, end, n_bins + 1, device=logits.device, dtype=torch.float32 + ) + v_bins = (edges[:-1] + edges[1:]) / 2 # [n_bins] + return (logits.float().softmax(-1) @ v_bins.unsqueeze(1)).squeeze(-1) + + +# =========================================================================== +# TransitionLayer (used in DiffusionConditioning) +# =========================================================================== + + +class TransitionLayer(nn.Module): + """SwiGLU transition: norm -> a_proj, b_proj -> silu(a)*b -> out_proj.""" + + def __init__(self, d_model: int, n: int, eps: float = 1e-5) -> None: + super().__init__() + hidden = n * d_model + self.norm = nn.LayerNorm(d_model, eps=eps) + self.a_proj = nn.Linear(d_model, hidden, bias=False) + self.b_proj = nn.Linear(d_model, hidden, bias=False) + self.out_proj = nn.Linear(hidden, d_model, bias=False) + + def forward(self, x: Tensor) -> Tensor: + x = self.norm(x) + a = self.a_proj(x) + b = self.b_proj(x) + return self.out_proj(F.silu(a) * b) + + +# =========================================================================== +# AdaptiveLayerNorm (used in DiffusionTransformer) +# =========================================================================== + + +class AdaptiveLayerNorm(nn.Module): + """Adaptive layer normalization (adaLN-Zero).""" + + def __init__(self, d_model: int, d_cond: int, eps: float = 1e-5) -> None: + super().__init__() + self.d_model = d_model + self.d_cond = d_cond + self.eps = eps + self.s_scale = nn.Parameter(torch.ones(d_cond)) + self.s_gate = nn.Linear(d_cond, d_model, bias=True) + self.s_shift = nn.Linear(d_cond, d_model, bias=False) + + def forward(self, a: Tensor, s: Tensor) -> Tensor: + a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps) + s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps) + return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm) + + +# =========================================================================== +# FourierEmbedding +# =========================================================================== + + +class FourierEmbedding(nn.Module): + """Fourier embedding: cos(2*pi*(t*w + b)).""" + + w: Tensor + b: Tensor + + def __init__(self, c: int) -> None: + super().__init__() + self.c = c + self.register_buffer("w", torch.randn(c)) + self.register_buffer("b", torch.randn(c)) + + def forward(self, t_hat: Tensor) -> Tensor: + t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1) + return torch.cos( + 2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :]) + ) + + +# =========================================================================== +# SwiGLU / SwiGLUMLP +# =========================================================================== + + +def _compute_swiglu_hidden_size(d_model: int, expansion_ratio: int) -> int: + return expansion_ratio * d_model + + +class SwiGLU(nn.Module): + """SwiGLU with packed w12 and output w3.""" + + def __init__( + self, + in_features: int, + hidden_features: int, + out_features: int | None = None, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + self.hidden_features = hidden_features + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.split(self.hidden_features, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +class SwiGLUMLP(SwiGLU): + """SwiGLU MLP with packed weights, no bias.""" + + def __init__( + self, d_model: int, expansion_ratio: int = 4, bias: bool = False + ) -> None: + hidden = _compute_swiglu_hidden_size(d_model, expansion_ratio) + super().__init__( + in_features=d_model, hidden_features=hidden, out_features=d_model, bias=bias + ) + + +# =========================================================================== +# SWA Atom Attention components +# =========================================================================== + + +def _rotate_half(x: Tensor) -> Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: + """Apply RoPE with batch-dependent cos/sin. + + Args: + x: [B, L, H, D] + cos: [B, L, D/2] + sin: [B, L, D/2] + """ + ro_dim = cos.shape[-1] * 2 + cos = cos.unsqueeze(2).repeat(1, 1, 1, 2) + sin = sin.unsqueeze(2).repeat(1, 1, 1, 2) + return torch.cat( + [x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim]) * sin, x[..., ro_dim:]], + dim=-1, + ) + + +@torch.compiler.disable +def build_3d_rope( + ref_pos: Tensor, + ref_space_uid: Tensor, + head_dim: int, + n_spatial_per_axis: int = 4, + n_uid_pairs: int = 2, + spatial_base_freq: float = 10000.0, + uid_base_freq: float = 10.0, +) -> tuple[Tensor, Tensor]: + """Build cos/sin for 3D RoPE + UID RoPE.""" + device = ref_pos.device + B, N = ref_pos.shape[:2] + half_dim = head_dim // 2 + n_spatial_total = 3 * n_spatial_per_axis + + spatial_inv_freq = 1.0 / ( + spatial_base_freq + ** ( + torch.arange(0, n_spatial_per_axis, dtype=torch.float32, device=device) + / n_spatial_per_axis + ) + ) + uid_inv_freq = 1.0 / ( + uid_base_freq + ** ( + torch.arange(0, n_uid_pairs, dtype=torch.float32, device=device) + / n_uid_pairs + ) + ) + + pos_f32 = ref_pos.float() + spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq) + spatial_freqs = spatial_freqs.reshape(B, N, n_spatial_total) + + uid_f32 = ref_space_uid.float() + uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq) + + n_active = n_spatial_total + n_uid_pairs + freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1) + + if n_active < half_dim: + padding = torch.zeros( + B, N, half_dim - n_active, device=device, dtype=torch.float32 + ) + freqs = torch.cat([freqs, padding], dim=-1) + + cos = freqs.cos().to(torch.bfloat16) + sin = freqs.sin().to(torch.bfloat16) + return cos, sin + + +def qk_norm(x: Tensor) -> Tensor: + return F.rms_norm(x, (x.size(-1),)).to(x.dtype) + + +# =========================================================================== +# SwiGLUFFN (atom transformer blocks) +# =========================================================================== + + +class SwiGLUFFN(nn.Module): + """SwiGLU FFN with rounded hidden size for hardware alignment.""" + + def __init__(self, d_model: int, expansion_ratio: int = 2) -> None: + super().__init__() + hidden_size = ((expansion_ratio * (d_model // 3) * 2) + 255) // 256 * 256 + self.w_up = nn.Linear(d_model, 2 * hidden_size, bias=False) + self.w_down = nn.Linear(hidden_size, d_model, bias=False) + + def forward(self, x: Tensor) -> Tensor: + x = x.to(self.w_up.weight.dtype) + x1, x2 = self.w_up(x).chunk(2, dim=-1) + return self.w_down(F.silu(x1) * x2) + + +# =========================================================================== +# SWA3DRoPEAttention +# =========================================================================== + + +class SWA3DRoPEAttention(nn.Module): + """Sliding window attention with 3D RoPE. Has Wqkv, gate_proj, out_proj.""" + + def __init__(self, d_model: int, n_heads: int, half_window: int = 64) -> None: + super().__init__() + self.n_heads = n_heads + self.head_dim = d_model // n_heads + self.scale = self.head_dim**-0.5 + self.half_window = half_window + + self.Wqkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + self.gate_proj = nn.Linear(d_model, d_model, bias=False) + + def forward(self, x: Tensor, attention_params: tuple) -> Tensor: + B, N = x.shape[:2] + cos, sin = attention_params[0], attention_params[1] + + x_input = x + qkv = self.Wqkv(x) + qkv = qkv.view(B, N, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4) + q, k, v = qkv.unbind(0) + q, k = qk_norm(q), qk_norm(k) + + q = apply_rotary_emb_3d(q, cos, sin) + k = apply_rotary_emb_3d(k, cos, sin) + + input_dtype = q.dtype + if q.dtype not in (torch.float16, torch.bfloat16): + q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16() + + if len(attention_params) > 2 and FLASH_ATTN_AVAILABLE: + assert index_first_axis is not None + assert flash_attn_varlen_func is not None + assert pad_input is not None + indices, cu_seqlens, max_seqlen = ( + attention_params[2], + attention_params[3], + attention_params[4], + ) + q_unpad = index_first_axis( + q.reshape(-1, self.n_heads, self.head_dim), indices + ) + k_unpad = index_first_axis( + k.reshape(-1, self.n_heads, self.head_dim), indices + ) + v_unpad = index_first_axis( + v.reshape(-1, self.n_heads, self.head_dim), indices + ) + out_unpad = flash_attn_varlen_func( + q_unpad, + k_unpad, + v_unpad, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + softmax_scale=self.scale, + window_size=(self.half_window, self.half_window), + ) + out = pad_input(out_unpad, indices, B, N) + elif FLASH_ATTN_AVAILABLE: + assert flash_attn_func is not None + out = flash_attn_func( + q, + k, + v, + softmax_scale=self.scale, + window_size=(self.half_window, self.half_window), + ) + else: + if len(attention_params) > 2: + valid = torch.zeros(B * N, dtype=torch.bool, device=q.device) + valid[attention_params[2]] = True + valid = valid.view(B, N) + else: + valid = torch.ones(B, N, dtype=torch.bool, device=q.device) + rank = torch.cumsum(valid, dim=1) - 1 + within = (rank.unsqueeze(2) - rank.unsqueeze(1)).abs() <= self.half_window + allowed = within & valid.unsqueeze(1) & valid.unsqueeze(2) + allowed |= torch.eye(N, dtype=torch.bool, device=q.device) + out = F.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + attn_mask=allowed.unsqueeze(1), + scale=self.scale, + ).transpose(1, 2) + out = out * valid.unsqueeze(-1).unsqueeze(-1) + + out = out.to(input_dtype).reshape(B, N, -1) + out = out * torch.sigmoid(self.gate_proj(x_input)) + return self.out_proj(out) + + +# =========================================================================== +# SWAAtomBlock, SWAAtomTransformer +# =========================================================================== + + +def _rms_adaln_raw(x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift + + +def _gated_residual_raw(x: Tensor, gate: Tensor, y: Tensor) -> Tensor: + return x + gate * y + + +class SWAAtomBlock(nn.Module): + """adaLN-Zero + SWA attention + SwiGLU FFN. + + Creates adaln_modulation = Sequential(SiLU(), Linear) -> keys like adaln_modulation.1.weight + """ + + def __init__( + self, + d_atom: int, + n_heads: int, + half_window: int = 64, + expansion_ratio: int = 2, + use_compile_fusions: bool = False, + ) -> None: + super().__init__() + self.attn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) + self.ffn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) + + adaln_linear = nn.Linear(d_atom, 6 * d_atom, bias=False) + nn.init.zeros_(adaln_linear.weight) + self.adaln_modulation = nn.Sequential(nn.SiLU(), adaln_linear) + + self.attn = SWA3DRoPEAttention(d_atom, n_heads, half_window=half_window) + self.ffn = SwiGLUFFN(d_atom, expansion_ratio) + + self._rms_adaln = ( + torch.compile(_rms_adaln_raw) if use_compile_fusions else _rms_adaln_raw + ) + self._gated_residual = ( + torch.compile(_gated_residual_raw) + if use_compile_fusions + else _gated_residual_raw + ) + + def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor: + mod = self.adaln_modulation(c_l) + if mod.dim() == 2: + mod = mod.unsqueeze(1) + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1) + + attn_input = self._rms_adaln(x, scale_a, shift_a) + attn_out = self.attn(attn_input, attention_params) + x = self._gated_residual(x, gate_a, attn_out) + + ffn_input = self._rms_adaln(x, scale_f, shift_f) + ffn_out = self.ffn(ffn_input) + x = self._gated_residual(x, gate_f, ffn_out) + return x + + +class SWAAtomTransformer(nn.Module): + """Stack of SWAAtomBlocks.""" + + def __init__( + self, + d_atom: int = 128, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.swa_window_size = swa_window_size + self.head_dim = d_atom // n_heads + self.spatial_rope_base_frequency = spatial_rope_base_frequency + self.n_spatial_rope_pairs_per_axis = n_spatial_rope_pairs_per_axis + self.n_uid_rope_pairs = n_uid_rope_pairs + self.uid_rope_base_frequency = uid_rope_base_frequency + + self.blocks = nn.ModuleList( + [ + SWAAtomBlock( + d_atom=d_atom, + n_heads=n_heads, + half_window=swa_window_size // 2, + expansion_ratio=expansion_ratio, + ) + for _ in range(n_blocks) + ] + ) + + def _build_3d_rope( + self, ref_pos: Tensor, ref_space_uid: Tensor + ) -> tuple[Tensor, Tensor]: + return build_3d_rope( + ref_pos=ref_pos, + ref_space_uid=ref_space_uid, + head_dim=self.head_dim, + n_spatial_per_axis=self.n_spatial_rope_pairs_per_axis, + n_uid_pairs=self.n_uid_rope_pairs, + spatial_base_freq=self.spatial_rope_base_frequency, + uid_base_freq=self.uid_rope_base_frequency, + ) + + def forward( + self, + q_l: Tensor, + c_l: Tensor, + attention_params: tuple, + return_intermediates: bool = False, + ) -> Tensor | tuple[Tensor, list[Tensor]]: + intermediates: list[Tensor] = [] + for block in self.blocks: + q_l = block(q_l, c_l, attention_params) + if return_intermediates: + intermediates.append(q_l) + if return_intermediates: + return q_l, intermediates + return q_l + + +# =========================================================================== +# EsmFold2AtomEncoder (for both inputs_embedder and diffusion_module) +# =========================================================================== + + +class EsmFold2AtomEncoder(nn.Module): + """SWA atom encoder with atom_linear, atom_norm, atom_to_token_linear, [coords_linear], atom_transformer. + + Args: + d_atom: atom hidden dim + d_token: token dim for atom_to_token aggregation + n_blocks, n_heads, swa_window_size, expansion_ratio: transformer params + structure_prediction: if True, creates coords_linear and uses full d_token + spatial_rope_base_frequency, n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs, uid_rope_base_frequency: 3D RoPE config + """ + + def __init__( + self, + d_atom: int = 128, + d_token: int = 768, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + structure_prediction: bool = True, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.d_atom = d_atom + self.d_token = d_token + self.structure_prediction = structure_prediction + + self.atom_linear = nn.Linear(ATOM_FEATURE_DIM, d_atom, bias=False) + self.atom_norm = nn.LayerNorm(d_atom) + + if structure_prediction: + self.coords_linear = nn.Linear(6, d_atom, bias=False) + + self.atom_transformer = SWAAtomTransformer( + d_atom=d_atom, + n_blocks=n_blocks, + n_heads=n_heads, + swa_window_size=swa_window_size, + expansion_ratio=expansion_ratio, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + # Output aggregation: d_token for structure prediction, d_token//2 for inputs + out_dim = d_token if structure_prediction else d_token // 2 + self.atom_to_token_linear = nn.Linear(d_atom, out_dim, bias=False) + + def forward( + self, + ref_pos: Tensor, + atom_attention_mask: Tensor, + ref_space_uid: Tensor, + ref_charge: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + atom_to_token: Tensor, + r_l: Tensor | None = None, + pred_r1: Tensor | None = None, + s_i: Tensor | None = None, + z_ij: Tensor | None = None, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + inference_cache: dict | None = None, + ) -> tuple[Tensor, Tensor, Tensor, tuple, list[Tensor]]: + """Returns (a, q, c, attention_params, intermediates). + + ``inference_cache`` caches step-invariant tensors (c_base, 3D RoPE, + attention indices, n_tokens) across diffusion steps. + """ + B, N = ref_pos.shape[:2] + + layer_cache = None + if inference_cache is not None: + layer_cache = inference_cache.setdefault("atomencoder", {}) + + if layer_cache is None or len(layer_cache) == 0: + atom_feats = torch.cat( + [ + ref_pos, + ref_charge.unsqueeze(-1), + atom_attention_mask.unsqueeze(-1), + ref_element, + ref_atom_name_chars.reshape(B, N, MAX_CHARS * CHAR_VOCAB_SIZE), + ], + dim=-1, + ) + c_base = self.atom_norm(self.atom_linear(atom_feats)) + cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid) + cos = cos.repeat_interleave(num_diffusion_samples, 0) + sin = sin.repeat_interleave(num_diffusion_samples, 0) + mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0) + seqlens = mask_exp.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten() + max_seqlen = int(seqlens.max().item()) + cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + attention_params = (cos, sin, indices, cu_seqlens, max_seqlen) + n_tokens = int(atom_to_token.max().item()) + 1 + if layer_cache is not None: + layer_cache["c_base"] = c_base + layer_cache["attention_params"] = attention_params + layer_cache["mask_exp"] = mask_exp + layer_cache["n_tokens"] = n_tokens + layer_cache["atom_to_token_exp"] = atom_to_token.repeat_interleave( + num_diffusion_samples, 0 + ) + else: + c_base = layer_cache["c_base"] + attention_params = layer_cache["attention_params"] + mask_exp = layer_cache["mask_exp"] + n_tokens = layer_cache["n_tokens"] + + c = c_base + + q = c + + if self.structure_prediction and r_l is not None: + q = q.repeat_interleave(num_diffusion_samples, 0) + if pred_r1 is None: + pred_r1 = torch.zeros_like(r_l) + r_input = torch.cat([r_l, pred_r1], dim=-1) + r_to_q = self.coords_linear(r_input) + q = q + r_to_q + + c = c.repeat_interleave(num_diffusion_samples, 0) + + result = self.atom_transformer( + q_l=q, + c_l=c, + attention_params=attention_params, + return_intermediates=return_intermediates, + ) + if return_intermediates: + q, intermediates = result + else: + q = result + intermediates = [] + + q_to_a = F.relu(self.atom_to_token_linear(q)) + if layer_cache is not None and "atom_to_token_exp" in layer_cache: + atom_to_token_exp = layer_cache["atom_to_token_exp"] + else: + atom_to_token_exp = atom_to_token.repeat_interleave( + num_diffusion_samples, 0 + ) + a = scatter_atom_to_token( + q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool() + ) + + return a, q, c, attention_params, intermediates + + +# =========================================================================== +# EsmFold2AtomDecoder +# =========================================================================== + + +class EsmFold2AtomDecoder(nn.Module): + """SWA atom decoder with token_to_atom_linear, atom_transformer, norm, output_linear.""" + + def __init__( + self, + d_atom: int = 128, + d_token: int = 768, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.token_to_atom_linear = nn.Linear(d_token, d_atom, bias=False) + + self.atom_transformer = SWAAtomTransformer( + d_atom=d_atom, + n_blocks=n_blocks, + n_heads=n_heads, + swa_window_size=swa_window_size, + expansion_ratio=expansion_ratio, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + self.norm = nn.LayerNorm(d_atom) + self.output_linear = nn.Linear(d_atom, XYZ_DIMS, bias=False) + + def forward( + self, + a_i: Tensor, + q_l: Tensor, + c_l: Tensor, + p_lm: tuple, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + ) -> tuple[Tensor, list[Tensor]]: + """Returns (r_update, intermediates).""" + atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0) + a_to_q = self.token_to_atom_linear(a_i) + a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp) + q_l = q_l + a_to_q + + result = self.atom_transformer( + q_l=q_l, + c_l=c_l, + attention_params=p_lm, + return_intermediates=return_intermediates, + ) + if return_intermediates: + q_l, intermediates = result + else: + q_l = result + intermediates = [] + + r_l = self.output_linear(self.norm(q_l)) + return r_l, intermediates + + +# =========================================================================== +# AttentionPairBias (DiffusionTransformer attention block) +# =========================================================================== + + +class AttentionPairBias(nn.Module): + """Gated multi-head attention with pair bias conditioning.""" + + def __init__( + self, + d_model: int, + d_pair: int, + num_heads: int, + d_cond: int | None = None, + use_conditioning: bool = True, + ) -> None: + super().__init__() + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + d_cond = d_cond or d_model + + if use_conditioning: + self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) + self.out_gate = nn.Linear(d_cond, d_model, bias=True) + # adaln init: weight=0, bias=-2 + nn.init.zeros_(self.out_gate.weight) + nn.init.constant_(self.out_gate.bias, -2.0) + else: + self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) + + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False) + self.g_proj = nn.Linear(d_model, d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + + if d_pair > 0: + self.pair_norm = nn.LayerNorm(d_pair, eps=1e-5) + self.pair_bias_proj = nn.Linear(d_pair, num_heads, bias=False) + + self._kernel_backend: str | None = None + + def set_kernel_backend(self, backend: str | None) -> None: + if backend not in _VALID_BACKENDS: + raise ValueError( + f"backend must be one of {_VALID_BACKENDS}, got {backend!r}" + ) + self._kernel_backend = backend + + def _is_zero_beta(self, beta: Tensor | float) -> bool: + if isinstance(beta, (int, float)): + return beta == 0.0 + return bool((beta == 0).all()) + + def _can_use_fused_pair_bias( + self, z: Tensor, n_queries: int, beta: Tensor | float + ) -> bool: + return ( + _fused_active(self, z) + and z.dim() == 4 + and self._is_zero_beta(beta) + and hasattr(self, "pair_bias_proj") + and hasattr(self, "pair_norm") + ) + + def _can_use_cueq_pair_bias( + self, z: Tensor, n_queries: int, beta: Tensor | float + ) -> bool: + return ( + _cueq_active(self) + and n_queries > 750 + and z.dim() == 4 + and self._is_zero_beta(beta) + and hasattr(self, "pair_bias_proj") + ) + + def forward( + self, + a: Tensor, + s: Tensor | None, + z: Tensor, + beta: Tensor | float = 0.0, + attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + ) -> Tensor: + bsz, n_queries, d_model = a.shape + + if s is not None: + x = self.adaln(a, s) + else: + x = self.pre_norm(a) + + n_keys = x.shape[1] + q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim) + kv = self.kv_proj(x) + k, v = kv.chunk(2, dim=-1) + k = k.view(bsz, n_keys, self.num_heads, self.head_dim) + v = v.view(bsz, n_keys, self.num_heads, self.head_dim) + + # Expand z for num_diffusion_samples + if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1: + z = z.repeat_interleave(num_diffusion_samples, dim=0) + if ( + attention_mask is not None + and attention_mask.shape[0] != bsz + and num_diffusion_samples > 1 + ): + attention_mask = attention_mask.repeat_interleave( + num_diffusion_samples, dim=0 + ) + + if self._can_use_fused_pair_bias(z, n_queries, beta): + kernel_mask = ( + attention_mask + if attention_mask is not None + else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) + ) + pair_norm_w = self.pair_norm.weight + pair_norm_b = ( + self.pair_norm.bias + if self.pair_norm.bias is not None + else torch.zeros_like(pair_norm_w) + ) + z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16) + 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) # ty:ignore[call-non-callable] + q_bhqd = q.transpose(1, 2) + k_bhqd = k.transpose(1, 2) + v_bhqd = v.transpose(1, 2) + attn_out = F.scaled_dot_product_attention( + q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype) + ) + g = torch.sigmoid(self.g_proj(x)).view( + bsz, n_queries, self.num_heads, self.head_dim + ) + ctx = g * attn_out.transpose(1, 2) + out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + if s is not None: + out = torch.sigmoid(self.out_gate(s)) * out + return out + + if self._can_use_cueq_pair_bias(z, n_queries, beta): + assert _cue_attn_pair_bias is not None + kernel_mask = ( + attention_mask + if attention_mask is not None + else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) + ) + out, _ = _cue_attn_pair_bias( + s=x, + q=q.transpose(1, 2), + k=k.transpose(1, 2), + v=v.transpose(1, 2), + z=z, + mask=kernel_mask, + num_heads=self.num_heads, + w_proj_z=self.pair_bias_proj.weight, + w_proj_g=self.g_proj.weight, + w_proj_o=self.out_proj.weight, + w_ln_z=self.pair_norm.weight, + b_ln_z=self.pair_norm.bias, + return_z_proj=False, + is_cached_z_proj=False, + ) + else: + # Standard attention with pair bias + g = torch.sigmoid(self.g_proj(x)).view( + bsz, n_queries, self.num_heads, self.head_dim + ) + + logits = ( + torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale + ) + + if z.dim() == 4: + pair_bias = self.pair_bias_proj(self.pair_norm(z)) + else: + pair_bias = z.unsqueeze(-1) + logits = logits + pair_bias.to(dtype=logits.dtype) + + if attention_mask is not None: + min_val = torch.finfo(logits.dtype).min + mask_bias = torch.where( + attention_mask.bool()[:, None, :, None], 0.0, min_val + ) + logits = logits + mask_bias.to(dtype=logits.dtype) + + attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype) + ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v) + ctx = g * ctx + out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + + if s is not None: + out = torch.sigmoid(self.out_gate(s)) * out + return out + + +# =========================================================================== +# ConditionedTransitionBlock +# =========================================================================== + + +class ConditionedTransitionBlock(nn.Module): + """Conditioned SwiGLU transition with adaptive layer norm.""" + + def __init__( + self, + d_model: int, + d_cond: int | None = None, + transition_multiplier: int = 2, + use_conditioning: bool = True, + ) -> None: + super().__init__() + d_cond = d_cond or d_model + hidden = transition_multiplier * d_model + + if use_conditioning: + self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) + self.output_gate = nn.Linear(d_cond, d_model, bias=True) + nn.init.zeros_(self.output_gate.weight) + nn.init.constant_(self.output_gate.bias, -2.0) + else: + self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) + + self.lin_swish = nn.Linear(d_model, 2 * hidden, bias=False) + self.lin_out = nn.Linear(hidden, d_model, bias=False) + + def forward(self, a: Tensor, s: Tensor | None) -> Tensor: + if s is not None: + x = self.adaln(a, s) + else: + x = self.pre_norm(a) + + swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1) + b = F.silu(swish_a) * swish_b + out = self.lin_out(b) + + if s is not None: + out = torch.sigmoid(self.output_gate(s)) * out + return out + + +# =========================================================================== +# DiffusionTransformer (token transformer) +# =========================================================================== + + +class DiffusionTransformer(nn.Module): + """Diffusion denoising transformer with attention pair bias.""" + + def __init__( + self, + d_model: int, + d_pair: int, + num_heads: int, + num_blocks: int, + d_cond: int | None = None, + transition_multiplier: int = 2, + use_conditioning: bool = True, + ) -> None: + super().__init__() + d_cond = d_cond or d_model + + self.attn_blocks = nn.ModuleList( + [ + AttentionPairBias( + d_model=d_model, + d_pair=d_pair, + num_heads=num_heads, + d_cond=d_cond, + use_conditioning=use_conditioning, + ) + for _ in range(num_blocks) + ] + ) + self.transition_blocks = nn.ModuleList( + [ + ConditionedTransitionBlock( + d_model=d_model, + d_cond=d_cond, + transition_multiplier=transition_multiplier, + use_conditioning=use_conditioning, + ) + for _ in range(num_blocks) + ] + ) + + def set_kernel_backend(self, backend: str | None) -> None: + for attn in self.attn_blocks: + cast(AttentionPairBias, attn).set_kernel_backend(backend) + + def forward( + self, + a: Tensor, + s: Tensor | None, + z: Tensor, + beta: Tensor | float = 0.0, + attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + ) -> tuple[Tensor, list[Tensor]]: + intermediates: list[Tensor] = [] + x = a + for attn, transition in zip(self.attn_blocks, self.transition_blocks): + x = x + attn( + x, + s, + z, + beta, + attention_mask=attention_mask, + num_diffusion_samples=num_diffusion_samples, + ) + x = x + transition(x, s) + if return_intermediates: + intermediates.append(x) + return x, intermediates + + +# =========================================================================== +# DiffusionConditioning +# =========================================================================== + + +class DiffusionConditioning(nn.Module): + """Conditions pair and single representations on noise timestep.""" + + def __init__( + self, + c_z: int = 256, + c_s: int = 768, + c_s_inputs: int = 451, + sigma_data: float = 16.0, + fourier_dim: int = 256, + transition_multiplier: int = 2, + layer_norm_eps: float = 1e-5, + ) -> None: + super().__init__() + self.sigma_data = float(sigma_data) + self.c_z = c_z + self.c_s = c_s + self.c_s_inputs = c_s_inputs + + self.z_input_norm = nn.LayerNorm(2 * c_z, eps=layer_norm_eps) + self.z_proj = nn.Linear(2 * c_z, c_z, bias=False) + self.z_transitions = nn.ModuleList( + [ + TransitionLayer(c_z, n=transition_multiplier, eps=layer_norm_eps) + for _ in range(2) + ] + ) + + self.s_input_norm = nn.LayerNorm(c_s_inputs, eps=layer_norm_eps) + self.s_proj = nn.Linear(c_s_inputs, c_s, bias=False) + self.fourier = FourierEmbedding(fourier_dim) + self.noise_norm = nn.LayerNorm(fourier_dim, eps=layer_norm_eps) + self.noise_proj = nn.Linear(fourier_dim, c_s, bias=False) + self.s_transitions = nn.ModuleList( + [ + TransitionLayer(c_s, n=transition_multiplier, eps=layer_norm_eps) + for _ in range(2) + ] + ) + + def forward( + self, + t_hat: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + z_trunk: Tensor, + relative_position_encoding: Tensor, + sigma_data: float | None = None, + num_diffusion_samples: int = 1, + inference_cache: dict[str, Tensor] | None = None, + ) -> tuple[Tensor, Tensor]: + sigma = self.sigma_data if sigma_data is None else float(sigma_data) + base_batch = z_trunk.shape[0] + target_batch = base_batch * num_diffusion_samples + + # z conditioning (cached across diffusion steps — independent of t_hat) + if inference_cache is not None and "z" in inference_cache: + z = inference_cache["z"] + else: + z_rel = relative_position_encoding.to(dtype=torch.float32) + z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1) + z = self.z_proj(self.z_input_norm(z)) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for block in self.z_transitions: + z = z + block(z) + if inference_cache is not None: + inference_cache["z"] = z + + # s conditioning + s_inputs_eff = s_inputs + if s_inputs_eff.shape[0] != target_batch: + s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0) + + s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32))) + + # Noise embedding + t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1) + if t.numel() == 1: + t = t.expand(target_batch) + elif t.shape[0] != target_batch: + t = t.repeat_interleave(num_diffusion_samples, 0) + t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20)) + n = self.fourier(t_noise) + n = self.noise_proj(self.noise_norm(n)) + s = s + n.unsqueeze(1) + + for block in self.s_transitions: + s = s + block(s) + + return s, z + + +# =========================================================================== +# DiffusionModule +# =========================================================================== + + +class DiffusionModule(nn.Module): + """Diffusion denoising module for structure prediction.""" + + def __init__( + self, + c_atom: int = 128, + c_token: int = 768, + c_z: int = 256, + c_s_inputs: int = 451, + sigma_data: float = 16.0, + fourier_dim: int = 256, + atom_num_blocks: int = 3, + atom_num_heads: int = 4, + token_num_blocks: int = 12, + token_num_heads: int = 16, + transition_multiplier: int = 2, + swa_window_size: int = 128, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.sigma_data = float(sigma_data) + + self.conditioning = DiffusionConditioning( + c_z=c_z, + c_s=c_token, # conditioning s output is c_token + c_s_inputs=c_s_inputs, + sigma_data=sigma_data, + fourier_dim=fourier_dim, + transition_multiplier=transition_multiplier, + ) + + # Atom encoder (structure_prediction=True, with coords_linear) + self.atom_encoder = EsmFold2AtomEncoder( + d_atom=c_atom, + d_token=c_token, + n_blocks=atom_num_blocks, + n_heads=atom_num_heads, + swa_window_size=swa_window_size, + expansion_ratio=2, + structure_prediction=True, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + # Atom decoder + self.atom_decoder = EsmFold2AtomDecoder( + d_atom=c_atom, + d_token=c_token, + n_blocks=atom_num_blocks, + n_heads=atom_num_heads, + swa_window_size=swa_window_size, + expansion_ratio=2, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + self.s_to_token = nn.Linear(c_token, c_token, bias=False) + nn.init.zeros_(self.s_to_token.weight) + + # Token transformer (DiffusionTransformer with pair bias) + self.token_transformer = DiffusionTransformer( + d_model=c_token, + d_pair=c_z, + num_heads=token_num_heads, + num_blocks=token_num_blocks, + d_cond=c_token, + transition_multiplier=transition_multiplier, + use_conditioning=True, + ) + + self.s_step_norm = nn.LayerNorm(c_token) + self.token_norm = nn.LayerNorm(c_token) + + def set_kernel_backend(self, backend: str | None) -> None: + self.token_transformer.set_kernel_backend(backend) + + def forward( + self, + x_noisy: Tensor, + t_hat: Tensor, + ref_pos: Tensor, + ref_charge: Tensor, + ref_mask: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + tok_idx: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + z_trunk: Tensor, + relative_position_encoding: Tensor, + asym_id: Tensor, + residue_index: Tensor, + entity_id: Tensor, + token_index: Tensor, + sym_id: Tensor, + sigma_data: float | None = None, + token_attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + return_token_repr: bool = False, + return_atom_repr: bool = False, + inference_cache: dict[str, Tensor] | None = None, + ) -> dict[str, Tensor | None]: + bsz = x_noisy.shape[0] + sigma = self.sigma_data if sigma_data is None else float(sigma_data) + t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape( + -1 + ) + if t.numel() == 1: + t = t.expand(bsz) + + # Step 1: conditioning (pair z is cached across diffusion steps) + s, z = self.conditioning( + t_hat=t, + s_inputs=s_inputs, + s_trunk=s_trunk, + z_trunk=z_trunk, + relative_position_encoding=relative_position_encoding, + sigma_data=sigma, + num_diffusion_samples=num_diffusion_samples, + inference_cache=inference_cache, + ) + + # Step 2: normalize noisy coords + denom = torch.sqrt(t * t + sigma * sigma) + r_noisy = x_noisy / denom[:, None, None] + + # Step 3: atom encoder + a, q_skip, c_skip, p_skip, enc_intermediates = self.atom_encoder( + ref_pos=ref_pos, + atom_attention_mask=ref_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=tok_idx, + r_l=r_noisy, + s_i=s_trunk, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + inference_cache=inference_cache, + ) + + # Step 4: add conditioned s + a = a + self.s_to_token(self.s_step_norm(s)) + + # Step 5: token transformer + a, _ = self.token_transformer( + a, + s, + z, + beta=0.0, + attention_mask=token_attention_mask, + num_diffusion_samples=num_diffusion_samples, + ) + + # Step 6: token norm + a = self.token_norm(a) + + # Step 7: atom decoder + r_update, dec_intermediates = self.atom_decoder( + a_i=a, + q_l=q_skip, + c_l=c_skip, + p_lm=p_skip, + atom_to_token=tok_idx, + atom_attention_mask=ref_mask, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + ) + + # Step 8: compute denoised output + sigma2 = sigma * sigma + t2 = t * t + out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy + out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update + + # Collect atom intermediates from encoder + decoder + atom_intermediates: Tensor | None = None + if return_atom_repr: + all_ints = enc_intermediates + dec_intermediates + if all_ints: + atom_intermediates = torch.stack(all_ints, dim=2) + + return { + "x_denoised": out, + "token_repr": a if return_token_repr else None, + "atom_intermediates": atom_intermediates, + } + + +# =========================================================================== +# DiffusionStructureHead +# =========================================================================== + + +class DiffusionStructureHead(nn.Module): + """Wrapper around DiffusionModule with diffusion sampling.""" + + def __init__(self, config: EsmFold2Config) -> None: + super().__init__() + dm = config.structure_head.diffusion_module + swa_cfg = config.atom_encoder + sh = config.structure_head + + self.diffusion_module = DiffusionModule( + c_atom=dm.atom_encoder.hidden_size, + c_token=dm.token_hidden_size, + c_z=dm.c_z, + c_s_inputs=dm.c_s_inputs, + sigma_data=dm.sigma_data, + fourier_dim=dm.fourier_dim, + atom_num_blocks=dm.atom_encoder.num_hidden_layers, + atom_num_heads=dm.atom_encoder.num_attention_heads, + token_num_blocks=dm.token_num_blocks, + token_num_heads=dm.token_num_heads, + transition_multiplier=dm.transition_multiplier, + swa_window_size=config.sliding_window, + spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, + uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + ) + + # Sampling hyperparameters + self.sigma_data = dm.sigma_data + self.gamma_0 = sh.gamma_0 + self.gamma_min = sh.gamma_min + self.noise_scale = sh.noise_scale + self.step_scale = sh.step_scale + self.inference_s_max = sh.inference_s_max + self.inference_s_min = sh.inference_s_min + self.inference_p = sh.inference_p + self.inference_num_steps = sh.inference_num_steps + + def set_kernel_backend(self, backend: str | None) -> None: + self.diffusion_module.set_kernel_backend(backend) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def inference_noise_schedule( + self, num_steps: int | None = None, device: torch.device | None = None + ) -> Tensor: + """Karras power-law noise schedule.""" + steps = self.inference_num_steps if num_steps is None else int(num_steps) + if steps == 1: + return torch.tensor( + [self.inference_s_max * self.sigma_data, 0.0], + device=device, + dtype=torch.float32, + ) + p = float(self.inference_p) + inv_p = 1.0 / p + k = torch.arange(steps, device=device, dtype=torch.float32) + base = self.inference_s_max**inv_p + (k / (steps - 1)) * ( + self.inference_s_min**inv_p - self.inference_s_max**inv_p + ) + schedule = self.sigma_data * base.pow(p) + return F.pad(schedule, (0, 1), value=0.0) + + @staticmethod + def _random_rotations(n: int, dtype: torch.dtype, device: torch.device) -> Tensor: + q = torch.randn((n, 4), dtype=dtype, device=device) + scale = torch.sqrt((q * q).sum(dim=1)) + signs = torch.where(q[:, 0] < 0, -scale, scale) + q = q / signs[:, None] + r, i, j, k = torch.unbind(q, dim=-1) + two_s = 2.0 / (q * q).sum(dim=-1) + return torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + dim=-1, + ).reshape(n, 3, 3) + + def _center_random_augmentation( + self, x: Tensor, atom_mask: Tensor, second_coords: Tensor | None = None + ) -> tuple[Tensor, Tensor | None]: + """Algorithm 19: center + random rotation + translation.""" + bsz = x.shape[0] + mask = atom_mask.unsqueeze(-1) # [B, A, 1] + denom = mask.sum(dim=1, keepdim=True).clamp(min=1) + mean = (x * mask).sum(dim=1, keepdim=True) / denom + x = x - mean + if second_coords is not None: + second_coords = second_coords - mean + + r = self._random_rotations(bsz, x.dtype, x.device) + x = torch.einsum("bmd,bds->bms", x, r) + if second_coords is not None: + second_coords = torch.einsum("bmd,bds->bms", second_coords, r) + + t = torch.randn_like(x[:, 0:1, :]) + x = x + t + if second_coords is not None: + second_coords = second_coords + t + return x, second_coords + + @staticmethod + def _weighted_rigid_align( + x: Tensor, x_gt: Tensor, w: Tensor, mask: Tensor + ) -> Tensor: + """Kabsch alignment: align x to x_gt with weights w.""" + w = (mask * w).unsqueeze(-1) # [B, N, 1] + denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8) + mu = (x * w).sum(dim=-2, keepdim=True) / denom + mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom + x_c = x - mu + xgt_c = x_gt - mu_gt + H = torch.einsum("bni,bnj->bij", w * xgt_c, x_c) + H32 = H.float() + U, _, Vh = torch.linalg.svd(H32, driver="gesvd" if H32.is_cuda else None) + det = torch.linalg.det(U @ Vh) + ones = torch.ones_like(det) + R = (U @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ Vh).to( + H.dtype + ) + return x_c @ R.transpose(-1, -2) + mu_gt + + # ------------------------------------------------------------------ + # Sampling + # ------------------------------------------------------------------ + + @torch.inference_mode() + def sample( + self, + z_trunk: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + relative_position_encoding: Tensor, + ref_pos: Tensor, + ref_charge: Tensor, + ref_mask: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + tok_idx: Tensor, + asym_id: Tensor, + residue_index: Tensor, + entity_id: Tensor, + token_index: Tensor, + sym_id: Tensor, + token_attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + num_sampling_steps: int | None = None, + max_inference_sigma: float | None = 256.0, + noise_scale: float | None = None, + step_scale: float | None = None, + return_atom_repr: bool = False, + use_inference_cache: bool = True, + ) -> dict[str, Tensor | None]: + """Diffusion sampling (Algorithm 18). + + ``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 + target_batch = s_inputs.shape[0] * num_diffusion_samples + + inference_cache: dict[str, Tensor] | None = {} if use_inference_cache else None + + steps = ( + self.inference_num_steps + if num_sampling_steps is None + else int(num_sampling_steps) + ) + + schedule = self.inference_noise_schedule(steps, device) + if max_inference_sigma is not None: + schedule = schedule[schedule <= float(max_inference_sigma)] + schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma)) + + lam = self.noise_scale if noise_scale is None else float(noise_scale) + eta = self.step_scale if step_scale is None else float(step_scale) + + x = schedule[0] * torch.randn( + target_batch, n_atoms, 3, device=device, dtype=torch.float32 + ) + atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float() + + gammas = torch.where( + schedule > self.gamma_min, + torch.full_like(schedule, self.gamma_0), + torch.zeros_like(schedule), + ) + + x_denoised_prev: Tensor | None = None + token_repr: Tensor | None = None + diff_atom_intermediates: Tensor | None = None + + step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:])) + num_steps = len(step_pairs) + + for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_pairs): + x, x_denoised_prev = self._center_random_augmentation( + x, atom_mask, second_coords=x_denoised_prev + ) + + sigma_tm_val = float(sigma_tm.item()) + t_hat_val = sigma_tm_val * (1.0 + float(gamma.item())) + eps_std = lam * max(t_hat_val**2 - sigma_tm_val**2, 0.0) ** 0.5 + x_noisy = x + eps_std * torch.randn_like(x) + + request_atom_repr = return_atom_repr and step_idx == num_steps - 1 + + dm_out = self.diffusion_module( + x_noisy=x_noisy, + t_hat=torch.full( + (target_batch,), t_hat_val, device=device, dtype=torch.float32 + ), + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=ref_mask, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + ref_space_uid=ref_space_uid, + tok_idx=tok_idx, + s_inputs=s_inputs, + s_trunk=s_trunk, + z_trunk=z_trunk, + relative_position_encoding=relative_position_encoding, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=token_attention_mask, + num_diffusion_samples=num_diffusion_samples, + return_token_repr=True, + return_atom_repr=request_atom_repr, + inference_cache=inference_cache, + ) + + x_denoised = dm_out["x_denoised"] + token_repr = dm_out["token_repr"] + if request_atom_repr: + diff_atom_intermediates = dm_out.get("atom_intermediates") + + # Reverse diffusion alignment (Kabsch) + with torch.autocast(device_type="cuda", enabled=False): + x_noisy = self._weighted_rigid_align( + x_noisy.float(), x_denoised.float(), atom_mask, atom_mask + ) + x_noisy = x_noisy.to(dtype=x_denoised.dtype) + + # ODE/SDE step + sigma_t_val = float(sigma_t.item()) + denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val + x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma + + x_denoised_prev = x_denoised + + result: dict[str, Tensor | None] = { + "sample_atom_coords": x, + "diff_token_repr": token_repr, + } + if return_atom_repr: + result["diff_atom_intermediates"] = diff_atom_intermediates + return result + + +# =========================================================================== +# RowAttentionPooling +# =========================================================================== + + +class RowAttentionPooling(nn.Module): + """Row-wise attention pooling: attn_proj, out_proj.""" + + def __init__(self, d_pair: int, d_single: int) -> None: + super().__init__() + self.attn_proj = nn.Linear(d_pair, 1, bias=False) + self.out_proj = nn.Linear(d_pair, d_single, bias=False) + + def forward(self, z: Tensor, mask: Tensor) -> Tensor: + scores = self.attn_proj(z).squeeze(-1) + mask_bias = torch.where( + mask[:, None, :].bool(), + torch.zeros_like(scores), + torch.full_like(scores, -1e9), + ) + scores = scores + mask_bias + weights = F.softmax(scores, dim=-1) + pooled = torch.einsum("bnm,bnmd->bnd", weights, z) + return self.out_proj(pooled) + + +# =========================================================================== +# InputsEmbedder +# =========================================================================== + + +class InputsEmbedder(nn.Module): + """Embeds input features including atom-level encoding via SWA attention.""" + + def __init__(self, config: EsmFold2Config) -> None: + super().__init__() + swa_cfg = config.atom_encoder + + self.atom_attention_encoder = EsmFold2AtomEncoder( + d_atom=swa_cfg.hidden_size, + d_token=swa_cfg.output_dim, + n_blocks=swa_cfg.num_hidden_layers, + n_heads=swa_cfg.num_attention_heads, + swa_window_size=config.sliding_window, + expansion_ratio=swa_cfg.expansion_ratio, + structure_prediction=False, # no coords_linear + spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, + uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + ) + + def forward( + self, + aatype: Tensor, + profile: Tensor, + deletion_mean: Tensor, + ref_pos: Tensor, + atom_attention_mask: Tensor, + ref_space_uid: Tensor, + ref_charge: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + atom_to_token: Tensor, + ) -> Tensor: + """Embed inputs into per-token features. + + Returns: + [B, L, d_inputs] concatenation of atom encoding, aatype, profile, + and deletion_mean. + """ + a, _q, _c, _attn_params, _intermediates = self.atom_attention_encoder( + ref_pos=ref_pos, + atom_attention_mask=atom_attention_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=atom_to_token, + ) + return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1) + + +# =========================================================================== +# ResIdxAsymIdSymIdEntityIdEncoding (trunk relative position) +# =========================================================================== + + +class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module): + """embed.weight [d_pair, n_features] where n_features = 2*(2*r_bins+2) + 1 + (2*c_bins+2). + + For default r_bins=32, c_bins=2: 2*66 + 1 + 6 = 139. + """ + + def __init__( + self, + n_relative_residx_bins: int = 32, + n_relative_chain_bins: int = 2, + d_pair: int = 256, + ) -> None: + super().__init__() + self.n_relative_residx_bins = n_relative_residx_bins + self.n_relative_chain_bins = n_relative_chain_bins + self.d_pair = d_pair + + n_feats_residue = 2 * n_relative_residx_bins + 2 + n_feats_token = 2 * n_relative_residx_bins + 2 + n_feats_chain = 2 * n_relative_chain_bins + 2 + n_feats_same_entity = 1 + total_feats = ( + n_feats_residue + n_feats_token + n_feats_chain + n_feats_same_entity + ) + self.embed = nn.Linear(total_feats, d_pair, bias=False) + + def forward( + self, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + token_index: Tensor, + ) -> Tensor: + bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1) + bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1) + bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1) + + dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1) + dij_residue = torch.clip( + dij_residue + self.n_relative_residx_bins, + 0, + 2 * self.n_relative_residx_bins, + ) + dij_residue = torch.where( + bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1 + ) + aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2) + + dij_token = torch.clip( + token_index.unsqueeze(2) + - token_index.unsqueeze(1) + + self.n_relative_residx_bins, + 0, + 2 * self.n_relative_residx_bins, + ) + dij_token = torch.where( + bij_same_chain & bij_same_residue, + dij_token, + 2 * self.n_relative_residx_bins + 1, + ) + aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2) + + dij_chain = torch.clip( + sym_id.unsqueeze(2) - sym_id.unsqueeze(1) + self.n_relative_chain_bins, + 0, + 2 * self.n_relative_chain_bins, + ) + dij_chain = torch.where( + bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain + ) + aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2) + + feats = torch.cat( + [ + aij_rel_pos.float(), + aij_rel_token.float(), + bij_same_entity.float().unsqueeze(-1), + aij_rel_chain.float(), + ], + dim=-1, + ) + + return self.embed(feats) + + +# =========================================================================== +# SingleToPair (for LanguageModelShim) +# =========================================================================== + + +class SingleToPair(nn.Module): + """downproject, output_mlp (Sequential of Linear, GELU, Linear).""" + + def __init__(self, input_dim: int, downproject_dim: int, output_dim: int) -> None: + super().__init__() + self.downproject = nn.Linear(input_dim, downproject_dim) + self.output_mlp = nn.Sequential( + nn.Linear(2 * downproject_dim, output_dim), + nn.GELU(), + nn.Linear(output_dim, output_dim), + ) + + def forward(self, x: Tensor) -> Tensor: + x = self.downproject(x) + x = torch.cat( + [(x.unsqueeze(2) * x.unsqueeze(1)), (x.unsqueeze(2) - x.unsqueeze(1))], + dim=3, + ) + return self.output_mlp(x) + + +# =========================================================================== +# LanguageModelShim +# =========================================================================== + + +class LanguageModelShim(nn.Module): + """Shim holding the trainable projection weights for LM integration. + + Contains: + - base_z_combine: nn.Parameter [num_layers+1] + - base_z_linear: Sequential(LayerNorm(d_model), Linear(d_model, d_z, bias=False)) + - base_z_mlp: Sequential(SingleToPair(d_z, d_z, d_z), LayerNorm(d_z)) + """ + + def __init__( + self, d_z: int = 256, d_model: int = 2560, num_layers: int = 80 + ) -> None: + super().__init__() + + self.base_z_mlp = nn.Sequential(SingleToPair(d_z, d_z, d_z), nn.LayerNorm(d_z)) + self.base_z_linear = nn.Sequential( + nn.LayerNorm(d_model), nn.Linear(d_model, d_z, bias=False) + ) + self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1)) + + def forward(self, hidden_states: Tensor, *, lm_dropout: float = 0.0) -> Tensor: + """Project pre-computed ESMC hidden states to pair representation. + + Args: + hidden_states: [B, L, num_layers+1, d_model] from ESMC 6B. + lm_dropout: Dropout probability applied to the pair + representation after ``base_z_mlp``. + + Returns: + [B, L, L, d_pair] pair representation. + """ + lm_z = self.base_z_linear(hidden_states) # [B, L, 81, d_z] + weights = self.base_z_combine.softmax(0) # [81] + lm_z = (weights @ lm_z).squeeze(-2) # [B, L, d_z] + lm_z = self.base_z_mlp(lm_z) # [B, L, L, d_z] + if lm_dropout > 0: + lm_z = F.dropout(lm_z, p=lm_dropout, training=True) + return lm_z + + +# =========================================================================== +# Reproducibility helper (mirrors evolutionaryscale.utils.reproducibility) +# =========================================================================== + + +@contextmanager +def _seed_context(seed: int | None, *, cuda: bool = True): + """Temporarily seed Python, NumPy, and PyTorch RNGs.""" + if seed is None: + yield + return + py_state = random.getstate() + np_state = np.random.get_state() + torch_state = torch.get_rng_state() + cuda_states = ( + torch.cuda.get_rng_state_all() if cuda and torch.cuda.is_available() else None + ) + seed = int(seed) % (2**32) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if cuda and torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + try: + yield + finally: + random.setstate(py_state) + np.random.set_state(np_state) + torch.set_rng_state(torch_state) + if cuda_states is not None: + torch.cuda.set_rng_state_all(cuda_states) + + +# =========================================================================== +# EsmFold2ExperimentalModel — the top-level model +# =========================================================================== + + +def compute_lm_hidden_states( + esmc: nn.Module, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + token_mask: Tensor, + pad_to_multiple: int | None = None, + lm_mask_pct: float = 0.0, + mask_token_id: int = 32, +) -> Tensor: + """Run ESMC with BOS/EOS wrapping, return hidden states [B, L, N, D] with N=81 layers. + + Atom-tokenized modified residues (HYP, MSE, ACE, NH2, ...) span multiple + structure tokens but share a single ``(asym_id, residue_index)`` key — + collapse them to one LM token per residue before running the LM (the LM + was trained on per-residue inputs, not per-atom), then scatter the + hidden states back to the per-token layout. + """ + B, L = input_ids.shape + device = input_ids.device + protein_mask = (mol_type == 0) & token_mask + + lm_input_list = [] + lm_lengths = [] + # Per-batch maps from (original protein-token index) to (LM input position). + expand_maps: list[Tensor] = [] + for b in range(B): + mask_b = protein_mask[b] + ids_b = input_ids[b][mask_b] + asym_b = asym_id[b][mask_b] + res_b = residue_index[b][mask_b] + + # Collapse: keep first token per (asym_id, residue_index) key, in + # input order. ``inverse`` maps each original protein-token to its + # collapsed residue index. + keys = torch.stack((asym_b, res_b), dim=1) + unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True) + n_unique = unique_keys.size(0) + token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long) + first_pos = torch.full( + (n_unique,), keys.size(0), device=device, dtype=torch.long + ) + first_pos.scatter_reduce_( + 0, inverse, token_positions, reduce="amin", include_self=True + ) + ordered = torch.argsort(first_pos) + first_pos_ordered = first_pos[ordered] + ids_collapsed = ids_b[first_pos_ordered] + asym_collapsed = asym_b[first_pos_ordered] + remap = torch.empty_like(ordered) + remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long) + inverse_ordered = remap[inverse] + + chain_ids = asym_collapsed.unique(sorted=True) + # [BOS] chain1 [EOS BOS] chain2 ... [EOS] + parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)] + # Per-chain LM positions accumulate; track them for the expand map. + per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long) + cursor = 1 # position 0 is the leading BOS + for i, cid in enumerate(chain_ids): + in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0] + parts.append(ids_collapsed[in_chain]) + per_token_lm_pos[in_chain] = torch.arange( + cursor, cursor + in_chain.shape[0], device=device, dtype=torch.long + ) + cursor += in_chain.shape[0] + if i < len(chain_ids) - 1: + parts.append(torch.tensor([2, 0], device=device, dtype=ids_b.dtype)) + cursor += 2 # EOS + BOS + parts.append(torch.tensor([2], device=device, dtype=ids_b.dtype)) + lm_seq = torch.cat(parts) + lm_input_list.append(lm_seq) + lm_lengths.append(lm_seq.shape[0]) + + # Original protein-token position → LM input position. + prot_pos_b = mask_b.nonzero(as_tuple=True)[0] + expand_map = torch.full((L,), -1, device=device, dtype=torch.long) + expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered] + expand_maps.append(expand_map) + + # Pad to longest LM input; round to ``pad_to_multiple`` when fp8 is on + # (TE fp8 kernels assert prod(shape[:-1]) % 8 == 0). + max_len = max(lm_lengths) + if pad_to_multiple is not None and pad_to_multiple > 1: + max_len = ((max_len + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple + lm_input_ids = torch.full( + (B, max_len), + 1, + device=device, + dtype=input_ids.dtype, # PAD=1 + ) + for b in range(B): + lm_input_ids[b, : lm_lengths[b]] = lm_input_list[b] + + # sequence_id for chain-aware attention; PAD tokens get -1 (no attention). + sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1 # BOS=0 + sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1) # PAD=1 + + if lm_mask_pct > 0.0: + # Randomly corrupt residues with the LM mask token, matching training; + # BOS/EOS/PAD (ids 0/2/1) are never masked so chain structure is kept. + special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2) + do_mask = ( + torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct + ) & ~special + lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id) + + with torch.inference_mode(): + esmc_out = esmc( + input_ids=lm_input_ids, sequence_id=sequence_id, output_hidden_states=True + ) + + hs = esmc_out.hidden_states # [n_layers+1, B, max_len, D] + n_layers_plus_1, _, _, D = hs.shape + result = torch.zeros(B, L, n_layers_plus_1, D, device=device, dtype=hs.dtype) + for b in range(B): + mb = protein_mask[b] + em = expand_maps[b][mb] # [n_protein_tokens] LM positions + # hs[:, b, em, :] -> [n_layers+1, n_protein_tokens, D] + gathered = hs[:, b, em, :].permute(1, 0, 2) + result[b, mb.nonzero(as_tuple=True)[0]] = gathered + + return result.detach() + + +# =========================================================================== +# TriangleMultiplicativeUpdate +# =========================================================================== +class TriangleMultiplicativeBlock(nn.Module): + """Triangle multiplicative update block with gated signal routing.""" + + _FLOW_TO_EINSUM = {"outgoing": "bikd,bjkd->bijd", "incoming": "bkid,bkjd->bijd"} + _VALID_FLOWS = ("outgoing", "incoming") + + def __init__(self, input_channels: int, latent_channels: int, flow: str) -> None: + super().__init__() + if flow not in self._FLOW_TO_EINSUM: + raise ValueError( + f"Invalid flow={flow!r}. Expected one of {self._VALID_FLOWS}." + ) + + self.input_channels = input_channels + self.latent_channels = latent_channels + self.flow = flow + self._einsum_equation = self._FLOW_TO_EINSUM[flow] + self.norm_start = nn.LayerNorm(self.input_channels, eps=_EPS) + self.norm_mix = nn.LayerNorm(self.latent_channels, eps=_EPS) + self.proj_bundle = nn.Linear( + self.input_channels, 4 * self.latent_channels, bias=False + ) + self.proj_emit = nn.Linear( + self.latent_channels, self.input_channels, bias=False + ) + self.proj_gate = nn.Linear(self.input_channels, self.input_channels, bias=False) + + self._use_kernels: bool = False + # Default chunked for memory on long sequences; tests override with + # ``set_chunk_size(None)`` for the unchunked path under bit-exact bf16 + # parity checks. + self._chunk_size: int | None = 64 + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def split_kernel_weights(self) -> tuple[Tensor, Tensor]: + return ( + self.proj_bundle.weight[: 2 * self.latent_channels, :], + self.proj_bundle.weight[2 * self.latent_channels :, :], + ) + + def _kernel_flow_direction(self) -> str: + return self.flow + + def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor: + return torch.einsum(self._einsum_equation, left_stream, right_stream) + + def _triangular_contract_chunked( + self, left_stream: Tensor, right_stream: Tensor, chunk_size: int + ) -> Tensor: + """Compute the triangular einsum in chunks along the output i-dimension.""" + L = left_stream.shape[1] if self.flow == "outgoing" else left_stream.shape[2] + chunks = [] + for start in range(0, L, chunk_size): + end = min(start + chunk_size, L) + if self.flow == "outgoing": + chunk = torch.einsum( + self._einsum_equation, left_stream[:, start:end], right_stream + ) + else: + chunk = torch.einsum( + self._einsum_equation, left_stream[:, :, start:end], right_stream + ) + chunks.append(chunk) + return torch.cat(chunks, dim=1) + + def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor: + if visibility is None: + visibility = pair_grid.new_ones(pair_grid.shape[:-1]) + + if self._use_kernels: + assert _cue_tri_mul is not None + p_in_weight, g_in_weight = self.split_kernel_weights() + + try: + return _cue_tri_mul( + pair_grid, + direction=self._kernel_flow_direction(), + mask=visibility, + norm_in_weight=self.norm_start.weight, + norm_in_bias=self.norm_start.bias, + p_in_weight=p_in_weight, + g_in_weight=g_in_weight, + norm_out_weight=self.norm_mix.weight, + norm_out_bias=self.norm_mix.bias, + p_out_weight=self.proj_emit.weight, + g_out_weight=self.proj_gate.weight, + eps=_EPS, + ) + except Exception as e: + import logging as _logging + + _logging.getLogger(__name__).warning( + "cuequivariance triangle_multiplicative_update kernel failed " + "(flow=%s, shape=%s, dtype=%s); falling back to chunked einsum. " + "Error: %s", + self.flow, + tuple(pair_grid.shape), + pair_grid.dtype, + e, + ) + + normalized_grid = self.norm_start(pair_grid) + bundled = self.proj_bundle(normalized_grid) + signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1) + routed = signal * torch.sigmoid(gate_logits) + routed = routed * visibility.unsqueeze(-1) + + left_stream, right_stream = routed.float().chunk(2, dim=-1) + if self._chunk_size is not None: + contracted = self._triangular_contract_chunked( + left_stream, right_stream, self._chunk_size + ) + else: + contracted = self._triangular_contract(left_stream, right_stream) + mixed = self.proj_emit(self.norm_mix(contracted)) + output_gate = torch.sigmoid(self.proj_gate(normalized_grid)) + return mixed * output_gate + + +class TriangleMultiplicativeUpdate(nn.Module): + """Thin wrapper exposing the triangular mixer with explicit orientation (v3).""" + + def __init__(self, dim: int = 128, _outgoing: bool = True) -> None: + super().__init__() + flow = "outgoing" if _outgoing else "incoming" + self._engine = TriangleMultiplicativeBlock( + input_channels=dim, latent_channels=dim, flow=flow + ) + + def set_kernel_backend(self, backend: str | None) -> None: + # Engine uses cueq when backend=="cuequivariance"; the "fused" backend + # routes through the parent PairUpdateBlock's fused path (bypassing this). + self._engine._use_kernels = backend == BACKEND_CUEQ + if backend == BACKEND_CUEQ and not CUE_AVAILABLE: + raise RuntimeError( + "backend='cuequivariance' but cuequivariance_torch is not installed." + ) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._engine.set_chunk_size(chunk_size) + + def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor: + return self._engine(z, visibility=mask) + + +# =========================================================================== +# FoldingTrunk: Transition, PairUpdateBlock, FoldingTrunk +# =========================================================================== + + +class Transition(nn.Module): + """LN + SwiGLU FFN with addmm-fused residual; optional Triton LN+w12+SwiGLU kernel.""" + + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + # Default chunked; set_chunk_size(None) disables for bit-exact parity tests. + self._chunk_size: int | None = 64 + self._fused_swiglu: nn.Module | None = None + self._kernel_backend: str | None = None + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def set_kernel_backend(self, backend: str | None) -> None: + """Install / uninstall FusedLNLinearSwiGLU (no cueq equivalent).""" + if backend not in _VALID_BACKENDS: + raise ValueError( + f"backend must be one of {_VALID_BACKENDS}, got {backend!r}" + ) + self._kernel_backend = backend + if backend == BACKEND_FUSED and TRITON_KERNELS_AVAILABLE: + assert _FusedLNLinearSwiGLU is not None + d_model = self.norm.normalized_shape[0] + d_inner = self.ffn.hidden_features + has_ln_bias = self.norm.bias is not None + device = self.ffn.w12.weight.device + fused = _FusedLNLinearSwiGLU( + d_model=d_model, + d_inner=d_inner, + has_ln_bias=has_ln_bias, + device=device, + 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) # 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) + else: + self._fused_swiglu = None + + def _can_use_fused_path(self, x: Tensor) -> bool: + return ( + _fused_pair_stack_active(self, x) + and self._fused_swiglu is not None + and x.dtype == _FUSED_KERNEL_DTYPE + and self._fused_swiglu.W12.dtype == x.dtype + ) + + def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor: + """SwiGLU through silu(x1)*x2, before the final w3.""" + ffn = self.ffn + x12 = ffn.w12(x_normed) + x1, x2 = x12.split(ffn.hidden_features, dim=-1) + return F.silu(x1) * x2 + + def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor: + """x + w3(hidden) via single cuBLAS addmm — avoids transition-output allocation.""" + ffn = self.ffn + x_shape = x.shape + out = torch.addmm( + x.contiguous().view(-1, x_shape[-1]), + hidden.view(-1, hidden.shape[-1]), + ffn.w3.weight.t().to(_FUSED_KERNEL_DTYPE), + ) + return out.view(x_shape) + + def forward(self, x: Tensor) -> Tensor: + # 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 + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + hidden = pre_w3(x) + return self._addmm_residual(x, hidden) + out = torch.empty_like(x) + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + hidden = pre_w3(sl) + out[:, s:e] = self._addmm_residual(sl, hidden) + return out + # Reference path — bit-exact with main: x + ffn(norm(x)). + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + return x + self.ffn(self.norm(x)) + out_list: list[Tensor] = [] + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + out_list.append(sl + self.ffn(self.norm(sl))) + return torch.cat(out_list, dim=1) + + +class PairUpdateBlock(nn.Module): + """tri_mul_out, tri_mul_in, pair_transition.""" + + def __init__(self, d_pair: int = 256, expansion_ratio: int = 4) -> None: + super().__init__() + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = Transition(d_pair, expansion_ratio=expansion_ratio) + self._kernel_backend: str | None = None + # Row-shared dropout-residual; r=0 for inference (HF model is inference-only). + # backend='fused' swaps in the FusedDropoutResidual Triton kernel. + self.row_drop = DropoutResidual(0.0, batch_dim=1, use_fused_kernels=False) + + def set_kernel_backend(self, backend: str | None) -> None: + if backend not in _VALID_BACKENDS: + raise ValueError( + f"backend must be one of {_VALID_BACKENDS}, got {backend!r}" + ) + self.tri_mul_out.set_kernel_backend(backend) + self.tri_mul_in.set_kernel_backend(backend) + self.pair_transition.set_kernel_backend(backend) + self._kernel_backend = backend + self.row_drop = DropoutResidual( + 0.0, batch_dim=1, use_fused_kernels=(backend == BACKEND_FUSED) + ) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + self.pair_transition.set_chunk_size(chunk_size) + + def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool: + 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 + p_in_weight, g_in_weight = engine.split_kernel_weights() + + return _fused_trimul_with_residual( + pair, + direction, + residual=pair, + drop_mask=None, # inference: no dropout, matches internal's eval path + 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 + ) -> Tensor: + if self._can_use_fused_trimul_with_residual(pair): + pair = self._fused_trimul_with_residual( + pair, "outgoing", pair_attention_mask + ) + pair = self._fused_trimul_with_residual( + pair, "incoming", pair_attention_mask + ) + else: + pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask)) + pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask)) + pair = self.pair_transition(pair) + return pair + + +class FoldingTrunk(nn.Module): + """ModuleList of PairUpdateBlocks.""" + + def __init__( + self, n_layers: int = 24, d_pair: int = 256, expansion_ratio: int = 4 + ) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [ + PairUpdateBlock(d_pair=d_pair, expansion_ratio=expansion_ratio) + for _ in range(n_layers) + ] + ) + + def set_kernel_backend(self, backend: str | None) -> None: + for block in self.blocks: + cast(PairUpdateBlock, block).set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(PairUpdateBlock, block).set_chunk_size(chunk_size) + + def forward( + self, pair: Tensor, pair_attention_mask: Tensor | None = None + ) -> Tensor: + # Cast pair → bf16 internally when the fused trimul backend is enabled + # (its bwd kernel requires bf16). Other backends keep the input dtype. + orig_dtype = pair.dtype + fused_on = ( + len(self.blocks) > 0 + and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED + ) + if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16: + pair = pair.to(torch.bfloat16) + 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) + else: + pair = fn(pair) + if pair.dtype != orig_dtype: + pair = pair.to(orig_dtype) + return pair + + +# =========================================================================== +# MSA Encoder +# =========================================================================== + + +class OuterProductMean(nn.Module): + """Outer-product mean: maps an MSA representation into a pair update. + + The order of the ``/ n_valid`` divide vs. the ``Wout`` projection is + selectable via ``divide_outer_before_proj`` because different ESMFold2 + checkpoints were trained with different orderings: + + * ``False`` (default): ``Wout(outer) / n_valid`` — the projection bias + is scaled by 1/n_valid alongside the outer product. + * ``True``: ``Wout(outer / n_valid)`` — the projection bias is added + unscaled, post-divide. + """ + + def __init__( + self, + d_msa: int, + d_hidden: int, + d_pair: int, + divide_outer_before_proj: bool = False, + ) -> None: + super().__init__() + self.d_hidden = d_hidden + self.divide_outer_before_proj = divide_outer_before_proj + self.norm = nn.LayerNorm(d_msa) + self.W = nn.Linear(d_msa, 2 * d_hidden, bias=False) + self.Wout = nn.Linear(d_hidden * d_hidden, d_pair, bias=True) + # Off for bit-exact bf16; ``set_chunk_size(64)`` for long sequences. + self._chunk_size: int | None = None + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def forward(self, m: Tensor, msa_attention_mask: Tensor) -> Tensor: + m_norm = self.norm(m) + x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype) + a, b = x.chunk(2, dim=-1) + mask_f = msa_attention_mask.to(a.dtype) + n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0) + if self._chunk_size is None: + outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2) + if self.divide_outer_before_proj: + return self.Wout(outer / n_valid) + return self.Wout(outer) / n_valid + # Chunk along the left (i) axis so the peak einsum intermediate is + # [B, chunk, L, c, d] instead of [B, L, L, c, d]. + L = a.shape[1] + out_chunks: list[Tensor] = [] + for s in range(0, L, self._chunk_size): + e = min(s + self._chunk_size, L) + outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, s:e], b).flatten(-2) + if self.divide_outer_before_proj: + out_chunks.append(self.Wout(outer_chunk / n_valid[:, s:e])) + else: + out_chunks.append(self.Wout(outer_chunk) / n_valid[:, s:e]) + return torch.cat(out_chunks, dim=1) + + +class MSAPairWeightedAveraging(nn.Module): + """Pair-biased MSA row update (AF3 Supplement Algorithm 10).""" + + def __init__( + self, d_msa: int, d_pair: int, n_heads: int = 8, head_width: int = 32 + ) -> None: + super().__init__() + self.n_heads = n_heads + self.head_width = head_width + self.norm_single = nn.LayerNorm(d_msa) + self.compute_bias = nn.Sequential( + nn.LayerNorm(d_pair), nn.Linear(d_pair, n_heads, bias=False) + ) + self.Wv = nn.Linear(d_msa, n_heads * head_width, bias=False) + self.Wgate = nn.Linear(d_msa, n_heads * head_width, bias=False) + self.Wout = nn.Linear(n_heads * head_width, d_msa, bias=False) + + def forward( + self, msa_repr: Tensor, pair_repr: Tensor, pair_attention_mask: Tensor + ) -> Tensor: + """ + Args: + msa_repr: [B, L, M, d_msa] + pair_repr: [B, L, L, d_pair] + pair_attention_mask:[B, L, L] + Returns: + [B, L, M, d_msa] + """ + B, L, M, _ = msa_repr.shape + h, dh = self.n_heads, self.head_width + + msa_normed = self.norm_single(msa_repr) + bias = self.compute_bias(pair_repr) # [B, L, L, n_heads] + bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5) + attn = torch.softmax(bias, dim=-2) # softmax over j + + v = self.Wv(msa_normed).reshape(B, L, M, h, dh) + gate = torch.sigmoid(self.Wgate(msa_normed)).reshape(B, L, M, h, dh) + + output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate) + return self.Wout(output.reshape(B, L, M, h * dh)) diff --git a/esm/models/esmfold2/model.py b/esm/models/esmfold2/model.py new file mode 100644 index 00000000..42b67d65 --- /dev/null +++ b/esm/models/esmfold2/model.py @@ -0,0 +1,1304 @@ +"""PyTorch ESMFold2 model — the standard released architecture. + +Quickstart:: + + from transformers import EsmFold2Model + + model = EsmFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval() + open("ubq.pdb", "w").write(model.infer_protein_as_pdb("MQIFVKTLTGKT...")) + +For multi-chain / ligand / MSA inputs see ``ESMFold2InputBuilder`` in the +companion ``esm`` package. +""" + +import math +from contextlib import contextmanager +from typing import Any, cast + +import torch +import torch.nn as nn +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 + + TE_AVAILABLE = True +except ImportError: + te = None # ty:ignore[invalid-assignment] + DelayedScaling = None # ty:ignore[invalid-assignment] + Format = None # ty:ignore[invalid-assignment] + TE_AVAILABLE = False + +from esm.models.esmc import EsmcModel +from esm.models.esmc.checkpoint_layout import published_to_native_subtree +from esm.models.esmfold2.config import EsmFold2Config +from esm.models.esmfold2.layers import ( + CHAR_VOCAB_SIZE, + MAX_ATOMIC_NUMBER, + NUM_RES_TYPES, + DiffusionStructureHead, + FoldingTrunk, + InputsEmbedder, + LanguageModelShim, + MSAPairWeightedAveraging, + OuterProductMean, + ResIdxAsymIdSymIdEntityIdEncoding, + RowAttentionPooling, + SwiGLUMLP, + TriangleMultiplicativeUpdate, + _categorical_mean, + _compute_intra_token_idx, + compute_lm_hidden_states, + gather_rep_atom_coords, + gather_token_to_atom, + maybe_apply_msa_column_masking, + maybe_subsample_msa, +) +from esm.models.hub import HubPreTrainedModel, resolve_model_dir + +_EPS = 1e-6 +_NONPOLYMER_ID = 4 + +# Default for the triangle / OPM / pair-transition L² ops. Caps peak memory +# so L≈2k folds on an 80 GB GPU (~76 GB peak at chunk=128 for L=1438; +# chunk=64 leaves headroom for the largest foldbench targets). Override via +# ``model.set_chunk_size(...)``; pass None to disable chunking (faster for +# short L but OOM-prone past ~600). +_DEFAULT_CHUNK_SIZE = 64 + +# Keys ``prepare_esmfold2_input`` emits that inference has no consumer for. +# ``forward`` accepts and drops these and rejects every other unknown keyword; it +# cannot drop the catch-all entirely because ``fold`` splats the whole feature +# dict, which is a strict superset of the signature. +_IGNORED_FEATURE_KEYS = frozenset( + {"pocket_feature", "gt_coords", "is_resolved", "frames_idx"} +) + + +class PairTransition(nn.Module): + """LayerNorm + SwiGLU feed-forward residual block on the pair representation.""" + + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + self._chunk_size: int | None = _DEFAULT_CHUNK_SIZE + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def forward(self, x: Tensor) -> Tensor: + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + return self.ffn(self.norm(x)) + out: list[Tensor] = [] + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + out.append(self.ffn(self.norm(sl))) + return torch.cat(out, dim=1) + + +class ConfidenceHead(nn.Module): + """Predicts pLDDT, PAE, PDE, resolved-atom probability and distogram bins.""" + + boundaries: Tensor + + def __init__(self, config: "EsmFold2Config") -> None: + super().__init__() + ch = config.confidence_head + d_single = config.hidden_size + d_pair = config.pairwise_hidden_size + d_inputs = config.single_inputs_size + + boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1) + self.register_buffer("boundaries", boundaries) + self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair) + + self.s_norm = nn.LayerNorm(d_single) + self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False) + self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False) + self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False) + self.s_inputs_norm = nn.LayerNorm(d_inputs) + self.z_norm = nn.LayerNorm(d_pair) + + self.row_attention_pooling = RowAttentionPooling( + d_pair=d_pair, d_single=d_single + ) + + self.folding_trunk = FoldingTrunk( + n_layers=ch.num_hidden_layers, d_pair=d_pair, expansion_ratio=4 + ) + + # Heads. + self.plddt_ln = nn.LayerNorm(d_single) + max_atoms_per_token = 23 + self.plddt_weight = nn.Parameter( + torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins) + ) + + self.pae_ln = nn.LayerNorm(d_pair) + self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False) + + self.pde_ln = nn.LayerNorm(d_pair) + self.pde_head = nn.Linear(d_pair, ch.num_pde_bins, bias=False) + + self.resolved_ln = nn.LayerNorm(d_single) + # 2 = resolved logits ([unresolved, resolved]). + self.resolved_weight = nn.Parameter( + torch.zeros(max_atoms_per_token, d_single, 2) + ) + + def set_kernel_backend(self, backend: str | None) -> None: + self.folding_trunk.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + + @staticmethod + def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor: + return ( + x + if num_diffusion_samples == 1 + else x.repeat_interleave(num_diffusion_samples, 0) + ) + + @staticmethod + def _flatten_sample_axis(x: Tensor) -> Tensor: + if x.ndim == 4: + b, mult, n, c = x.shape + return x.reshape(b * mult, n, c) + return x + + def forward( + self, + s_inputs: Tensor, + z: Tensor, + x_pred: Tensor, + distogram_atom_idx: Tensor, + token_attention_mask: Tensor, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + asym_id: Tensor, + mol_type: Tensor, + num_diffusion_samples: int = 1, + relative_position_encoding: Tensor | None = None, + token_bonds_encoding: Tensor | None = None, + ) -> dict[str, Tensor]: + s_inputs_normed = self.s_inputs_norm(s_inputs) + + z_base = self.z_norm(z) + if relative_position_encoding is not None: + z_base = z_base + relative_position_encoding + if token_bonds_encoding is not None: + z_base = z_base + token_bonds_encoding + z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2) + z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1) + z_base = z_base + self.s_to_z_prod_out( + self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :] + * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :] + ) + + pair = self._repeat_batch(z_base, num_diffusion_samples) + x_pred_flat = self._flatten_sample_axis(x_pred) + atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples) + atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples) + rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long() + mask = self._repeat_batch(token_attention_mask, num_diffusion_samples) + Bm = pair.shape[0] + + rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m) + rep_distances = torch.cdist( + rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" + ) + distogram_bins = ( + (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() + ) + pair = pair + self.dist_bin_pairwise_embed(distogram_bins) + + pair_mask = mask[:, :, None].float() * mask[:, None, :].float() + + # FoldingTrunk handles the bf16 cast internally during inference so + # each block's fused trimul engages. In-place residual avoids an + # extra fp32 pair allocation. + with torch.amp.autocast("cuda", enabled=pair.is_cuda, dtype=torch.bfloat16): + pair_delta = self.folding_trunk(pair, pair_attention_mask=pair_mask) + pair.add_(pair_delta.float()) + del pair_delta + single = self.row_attention_pooling(pair, mask) + + atom_mask_f = atom_mask_m.float() + s_at_atoms = gather_token_to_atom(single, atom_to_token_m) + s_at_atoms_ln = self.plddt_ln(s_at_atoms) + + intra_idx = _compute_intra_token_idx(atom_to_token_m) + intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1) + w_plddt = self.plddt_weight[intra_idx] + plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms_ln, w_plddt) + plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0) + + L = single.shape[1] + plddt_sum = torch.zeros(Bm, L, device=single.device, dtype=plddt_per_atom.dtype) + atom_count = torch.zeros( + Bm, L, device=single.device, dtype=plddt_per_atom.dtype + ) + atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype) + plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t) + atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t) + plddt = plddt_sum / atom_count.clamp(min=1e-6) + + complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / ( + atom_mask_f.sum(dim=-1) + _EPS + ) + + expanded_type = self._repeat_batch(mol_type, num_diffusion_samples) + expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples) + is_ligand = (expanded_type == _NONPOLYMER_ID).float() + inter_chain = ( + expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) + ).float() + near_contact = (rep_distances < 8).float() + interface_per_token = ( + near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1) + ).amax(dim=-1) + iplddt_weight = torch.where( + is_ligand.bool(), + torch.full_like(interface_per_token, 2.0), + interface_per_token, + ) + iplddt_weight_atoms = gather_token_to_atom( + iplddt_weight.unsqueeze(-1), atom_to_token_m + ).squeeze(-1) + atom_iplddt_w = atom_mask_f * iplddt_weight_atoms + complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / ( + atom_iplddt_w.sum(dim=-1) + _EPS + ) + + plddt_ca = plddt_per_atom.gather(1, rep_idx_m) + + # PAE + pae_logits = self.pae_head(self.pae_ln(pair)) + pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach() + + # PDE + pde_logits = self.pde_head(self.pde_ln(pair)) + pde = _categorical_mean(pde_logits, start=0.0, end=32.0).detach() + + # Resolved (per-atom binary). + s_at_atoms_res = self.resolved_ln(s_at_atoms) + w_res = self.resolved_weight[intra_idx] + resolved_logits = torch.einsum("...c,...cb->...b", s_at_atoms_res, w_res) + + # pTM / ipTM from pae_logits. + n_bins = pae_logits.shape[-1] + bin_width = 32.0 / n_bins + bin_centers = torch.arange( + 0.5 * bin_width, 32.0, bin_width, device=pae_logits.device + ) + mask_f = mask.float() + N_res = mask_f.sum(dim=-1, keepdim=True) + d0 = 1.24 * (N_res.clamp(min=19) - 15) ** (1 / 3) - 1.8 + tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2) + pae_probs = F.softmax(pae_logits, dim=-1) + tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1) + + pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2) + ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / ( + pair_mask_2d.sum(dim=-1) + _EPS + ) + ptm = ptm_per_row.max(dim=-1).values + + inter_chain_mask = ( + expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) + ).float() * pair_mask_2d + iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / ( + inter_chain_mask.sum(dim=-1) + _EPS + ) + iptm = iptm_per_row.max(dim=-1).values + + max_chain_id = int(expanded_asym.max().item()) if Bm > 0 else 0 + n_chains = max_chain_id + 1 + 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 + 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, + "plddt": plddt.detach(), + "plddt_per_atom": plddt_per_atom.detach(), + "plddt_ca": plddt_ca.detach(), + "complex_plddt": complex_plddt.detach(), + "complex_iplddt": complex_iplddt.detach(), + "pae_logits": pae_logits, + "pae": pae, + "pde_logits": pde_logits, + "pde": pde, + "resolved_logits": resolved_logits, + "ptm": ptm.detach(), + "iptm": iptm.detach(), + "pair_chains_iptm": pair_chains_iptm.detach(), + } + + +def _inverse_softplus(value: float) -> float: + return value + math.log(-math.expm1(-value)) + + +def _convert_te_modules_to_fp8_inplace(module: nn.Module) -> None: + """Re-init each TE module via quantized_model_init so weights live as fp8. + + Must be called inside torch.no_grad(); covers nn.Linear, te.Linear, + te.LayerNormLinear, te.LayerNormMLP — the last two hold 99% of ESMC weight. + """ + if not TE_AVAILABLE: + raise RuntimeError("transformer_engine is not available; cannot use fp8.") + assert te is not None + from transformer_engine.pytorch import quantized_model_init + + def _walk(mod: nn.Module) -> None: + assert te is not None + for name, child in list(mod.named_children()): + replaced = False + if isinstance(child, nn.Linear): + in_f, out_f = child.in_features, child.out_features + has_bias = child.bias is not None + device = child.weight.device + dtype = child.weight.dtype + w = child.weight.data + b = child.bias.data if has_bias else None + setattr(mod, name, nn.Identity()) + del child + torch.cuda.empty_cache() + with quantized_model_init(enabled=True): + new_mod = te.Linear( + in_f, out_f, bias=has_bias, params_dtype=dtype + ).to(device) + 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) # ty:ignore[call-non-callable] + del w, b + replaced = True + 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 + device = child.weight.device + dtype = ( + child.weight.dtype + if not hasattr(child.weight, "_data") + else torch.bfloat16 + ) + state = {k: v.detach().clone() for k, v in child.state_dict().items()} + setattr(mod, name, nn.Identity()) + del child + torch.cuda.empty_cache() + with quantized_model_init(enabled=True): + 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 + ): + state = {k: v.detach().clone() for k, v in child.state_dict().items()} + hidden_size = child.in_features + out_features = child.out_features + has_bias = child.use_bias + device = next(child.parameters()).device + setattr(mod, name, nn.Identity()) + del child + torch.cuda.empty_cache() + with quantized_model_init(enabled=True): + new_mod = te.LayerNormLinear( + hidden_size, + out_features, + bias=has_bias, + params_dtype=torch.bfloat16, + ).to(device) + new_mod.load_state_dict(state, strict=False) + replaced = True + 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 # 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 + ) + 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( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + bias=has_bias, + activation="swiglu", + params_dtype=torch.bfloat16, + ).to(device) + new_mod.load_state_dict(state, strict=False) + replaced = True + + if replaced: + # Freeze via .eval()+.requires_grad_(False); per-param ops would unwrap Float8Tensor. + new_mod.eval().requires_grad_(False) + setattr(mod, name, new_mod) + torch.cuda.empty_cache() + else: + _walk(child) + + _walk(module) + torch.cuda.empty_cache() + + +@contextmanager +def _lm_precision_context(fp8: bool): + """bf16 autocast (+ optional TE fp8 autocast) around the LM forward. + + te.autocast keeps te.Linear outputs bf16 instead of the fp32 default + (~425 MB at L=1024 in the hidden-state cache). + """ + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + if fp8 and TE_AVAILABLE: + assert DelayedScaling is not None + assert Format is not None + assert te is not None + fp8_recipe = DelayedScaling( + fp8_format=Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + ) + with te.autocast(enabled=True, recipe=fp8_recipe): + yield + else: + yield + + +class EsmFold2Model(HubPreTrainedModel): + """ESMFold2 — all-atom structure prediction with an ESMC PLM backbone. + + This is the standard released ESMFold2 architecture (uses a linear- + recurrent trunk, internally referred to as "parcae"). + + Forward kwargs that callers commonly override: + + * ``num_loops`` (default ``config.num_loops``): trunk refinement + loops. + * ``num_diffusion_samples`` (default ``config.num_diffusion_samples``): + parallel structure samples; the confidence head re-runs once per + sample, so memory scales linearly. Pass ``1`` for cheap inference. + * ``num_sampling_steps`` (default ``config.structure_head.inference_num_steps``): + diffusion ODE solver steps. Lower for speed, higher for quality. + * ``noise_scale`` / ``step_scale`` / ``max_inference_sigma``: sampler + overrides, forwarded to ``DiffusionStructureHead.sample``. The two scales + default to ``config.structure_head``; the sigma cap truncates the schedule. + * ``msa_max_depth`` / ``msa_column_mask_rate``: inference-time MSA diversity, + defaulting to ``config.msa_encoder`` when ``None``. + + Unknown keywords raise ``TypeError``; only the inference-irrelevant keys the + featurizer emits (``_IGNORED_FEATURE_KEYS``) are accepted and dropped. + + Memory / perf knobs: + + * ``model.set_chunk_size(int|None)``: caps L² ops (triangle / OPM / + pair transition) at this token-axis chunk. Default 64 — fits + L≈2k on an 80 GB GPU. Pass ``None`` for faster inference at L<600. + * ``model.set_kernel_backend(None | "fused" | "cuequivariance")``: + select kernel backend (None = reference path). + """ + + config_class = EsmFold2Config + _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"] + + @classmethod + def _normalize_checkpoint_layout( + cls, raw: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Translate a bundled ESMC encoder out of the published tensor layout. + + Scoped to the ``esmc.`` subtree so the trunk's own keys can never be + caught by the encoder's key patterns. + """ + return published_to_native_subtree(raw, "esmc.") + + def __init__(self, config: EsmFold2Config) -> None: + super().__init__(config) + d_inputs = config.single_inputs_size + d_pair = config.pairwise_hidden_size + + self.inputs_embedder = InputsEmbedder(config) + self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False) + self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False) + self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding( + n_relative_residx_bins=config.n_relative_residx_bins, + n_relative_chain_bins=config.n_relative_chain_bins, + d_pair=d_pair, + ) + self.token_bonds = nn.Linear(1, d_pair, bias=False) + self.language_model = LanguageModelShim( + d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers + ) + # A bundled backbone is described by ``esmc_config`` and arrives in the + # same checkpoint, so it is built here and populated by the single + # ``from_pretrained`` pass. Otherwise ``load_esmc`` fetches it later. + self.esmc: nn.Module | None = ( + EsmcModel(config.esmc_config) if config.esmc_config is not None else None + ) + self._esmc_fp8: bool = False # set by load_esmc(fp8=True) + + self.folding_trunk = FoldingTrunk( + n_layers=config.folding_trunk_num_hidden_layers, + d_pair=d_pair, + expansion_ratio=4, + ) + if config.lm_encoder.enabled: + self.lm_encoder: FoldingTrunk | None = FoldingTrunk( + n_layers=config.lm_encoder.num_hidden_layers, + d_pair=d_pair, + expansion_ratio=4, + ) + else: + self.lm_encoder = None + + self.parcae_input_norm = nn.LayerNorm(d_pair) + self.parcae_log_a = nn.Parameter(torch.zeros(d_pair)) + parcae_decay_init = math.sqrt(1.0 / 5.0) + parcae_delta_init = -math.log(parcae_decay_init) + self.parcae_log_delta = nn.Parameter( + torch.full( + (d_pair,), _inverse_softplus(parcae_delta_init), dtype=torch.float32 + ) + ) + self.parcae_b_cont = nn.Parameter(torch.eye(d_pair)) + self.parcae_readout = nn.Linear(d_pair, d_pair, bias=False) + nn.init.eye_(self.parcae_readout.weight) + self.parcae_coda = FoldingTrunk( + n_layers=config.parcae_num_coda_layers, d_pair=d_pair, expansion_ratio=4 + ) + + # Heads -------------------------------------------------------------- + self.structure_head = DiffusionStructureHead(config) + self.distogram_head = nn.Linear( + d_pair, config.structure_head.distogram_bins, bias=True + ) + self.confidence_head = ConfidenceHead(config) + + msa_cfg = config.msa_encoder + self.msa_encoder = None + if msa_cfg.enabled: + self.msa_encoder = MSAEncoder( + d_msa=msa_cfg.hidden_size, + d_pair=d_pair, + d_inputs=d_inputs, + d_hidden=msa_cfg.outer_hidden_size, + n_layers=msa_cfg.num_hidden_layers, + n_heads_msa=msa_cfg.num_attention_heads, + msa_head_width=msa_cfg.head_width, + ) + + self.post_init() + + def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None: + """Fetch the ESMC LM from a separate repo and attach it. + + Only needed when the backbone is *not* bundled into this checkpoint; see + ``EsmFold2Config.esmc_config``. + """ + from esm.models.esmc import EsmcModel + + self.esmc = EsmcModel.from_pretrained(esmc_model_path) + self.set_esmc_precision(precision) + + def set_esmc_precision(self, precision: str = "bf16") -> None: + """Cast the attached ESMC LM and freeze it. + + ``precision``: ``"bf16"`` (default), ``"fp32"``, or ``"fp8"``. + ``"fp8"`` requires H100 + TransformerEngine ≥ 2.x and quantizes + every TE module's weights to fp8 storage. + """ + dtype_map = { + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp8": torch.bfloat16, # underlying weights stay bf16, TE re-quantizes to fp8 + } + if precision not in dtype_map: + raise ValueError( + f"precision must be one of {list(dtype_map)}, got {precision!r}" + ) + if self.esmc is None: + raise RuntimeError("no ESMC LM is attached; nothing to cast.") + + esmc = self.esmc.to(device=self.device, dtype=dtype_map[precision]).eval() + for p in esmc.parameters(): + p.requires_grad_(False) + + if precision == "fp8": + if not TE_AVAILABLE: + raise RuntimeError( + "transformer_engine is not available; cannot use fp8." + ) + with torch.no_grad(): + _convert_te_modules_to_fp8_inplace(esmc) + self._esmc_fp8 = True + else: + self._esmc_fp8 = False + + self.esmc = esmc + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path, + *, + load_esmc: bool = True, + esmc_precision: str = "bf16", + config: EsmFold2Config | None = None, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, + **kwargs, + ): + local_dir = resolve_model_dir(pretrained_model_name_or_path, **kwargs) + if config is None: + config = EsmFold2Config.from_pretrained(local_dir) + if cls is EsmFold2Model and config.type == "experimental": + from esm.models.esmfold2.experimental import EsmFold2ExperimentalModel + + return EsmFold2ExperimentalModel.from_pretrained( + local_dir, + load_esmc=load_esmc, + esmc_precision=esmc_precision, + config=config, + device=device, + dtype=dtype, + ) + model = cls._load_pretrained(local_dir, config, device=device, dtype=dtype) + # A bundled backbone came in with the trunk; only the separate-repo + # arrangement needs a second fetch. + if load_esmc and config.esmc_config is None: + model.load_esmc(model.config.esmc_id, precision=esmc_precision) + elif config.esmc_config is not None: + model.set_esmc_precision(esmc_precision) + return model + + def set_kernel_backend(self, backend: str | None) -> None: + """Select kernel backend. + + Args: + backend: ``None`` (reference path), ``"fused"`` (vendored Triton + kernels), or ``"cuequivariance"`` (cuequivariance kernels + where applicable; vanilla python fallback otherwise). + """ + self.folding_trunk.set_kernel_backend(backend) + if self.lm_encoder is not None: + self.lm_encoder.set_kernel_backend(backend) + self.parcae_coda.set_kernel_backend(backend) + self.confidence_head.set_kernel_backend(backend) + self.structure_head.set_kernel_backend(backend) + if self.msa_encoder is not None: + self.msa_encoder.set_kernel_backend(backend) + + def apply_torch_compile( + self, mode: str = "fixed_seqlen", dynamic: bool | None = None + ) -> None: + """Compile L²-heavy blocks. ``mode='fixed_seqlen'`` recompiles per L; ``'dynamic_seqlen'`` compiles once. + + Does NOT stack with our Triton kernels — call ``set_kernel_backend(None)`` + before compiling. + """ + import torch._dynamo.config as dynamo_config + + setattr(dynamo_config, "cache_size_limit", 512) + setattr(dynamo_config, "accumulated_cache_size_limit", 512) + # capture_scalar_outputs avoids graph breaks at .item() in atom-attention path. + dynamo_config.capture_scalar_outputs = True + + if dynamic is None: + dynamic = mode == "dynamic_seqlen" + kwargs: dict = {"dynamic": dynamic} + + from esm.models.esmfold2.layers import ( + DiffusionModule, + DiffusionTransformer, + PairUpdateBlock, + ) + + compile_targets = ( + PairUpdateBlock, + DiffusionTransformer, + DiffusionModule, + MSAEncoderBlock, + ) + + def _maybe_compile(module: nn.Module) -> None: + if isinstance(module, compile_targets): + module.forward = cast(Any, torch.compile(module.forward, **kwargs)) + + self.apply(_maybe_compile) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + if self.lm_encoder is not None: + self.lm_encoder.set_chunk_size(chunk_size) + self.parcae_coda.set_chunk_size(chunk_size) + self.confidence_head.set_chunk_size(chunk_size) + if self.msa_encoder is not None: + self.msa_encoder.set_chunk_size(chunk_size) + + def _compute_lm_hidden_states( + self, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + tok_mask: Tensor, + lm_mask_pct: float = 0.0, + ) -> Tensor: + assert self.esmc is not None + # fp8 TE kernels require prod(shape[:-1]) % 8 == 0. + pad_to = 8 if self._esmc_fp8 else None + with _lm_precision_context(self._esmc_fp8): + return compute_lm_hidden_states( + self.esmc, + input_ids, + asym_id, + residue_index, + mol_type, + tok_mask, + pad_to_multiple=pad_to, + lm_mask_pct=lm_mask_pct, + ) + + def _discretized_dynamics(self) -> tuple[Tensor, Tensor]: + delta = F.softplus(self.parcae_log_delta) + a = torch.exp(-delta * torch.exp(self.parcae_log_a)) + b = delta[:, None] * self.parcae_b_cont + return a, b + + def _init_pair_state(self, ref: Tensor) -> Tensor: + std = math.sqrt(2.0 / (5.0 * ref.shape[-1])) + state = torch.empty_like(ref, dtype=torch.float32) + nn.init.trunc_normal_(state, mean=0.0, std=std, a=-3 * std, b=3 * std) + return state.to(dtype=ref.dtype) + + def _run_one_loop( + self, + z: Tensor, + z_init: Tensor, + lm_z: Tensor | None, + _msa_inputs: dict | None, + pair_mask: Tensor, + a: Tensor, + b_mat: Tensor, + tok_mask: Tensor, + total_steps: int, + ) -> Tensor: + # Helper method (not inline) so per-iter locals free on return — + # otherwise leaks ~2 GB L²×c_z into distogram/sample scope. + # training=True forces dropout under eval(), matching the per-loop + # dropout strategy used at train time. + lm_cfg = self.config.lm_encoder + _per_loop_lm_dropout = ( + lm_z is not None + and getattr(lm_cfg, "per_loop_lm_dropout", False) + and getattr(lm_cfg, "lm_dropout", 0.0) > 0.0 + ) + _lm_dropout_p = getattr(lm_cfg, "lm_dropout", 0.0) + + for _ in range(total_steps): + if _per_loop_lm_dropout: + assert lm_z is not None # narrowed by _per_loop_lm_dropout + lm_z_i: Tensor | None = F.dropout(lm_z, p=_lm_dropout_p, training=True) + else: + lm_z_i = lm_z + + refined_lm_z: Tensor | None = None + if lm_z_i is not None and self.lm_encoder is not None: + refined_lm_z = self.lm_encoder( + lm_z_i.to(z_init.dtype), pair_attention_mask=pair_mask + ) + + z_inject_pair = z_init + if lm_z_i is not None and self.lm_encoder is None: + z_inject_pair = z_inject_pair + lm_z_i.to(z_inject_pair.dtype) + + if self.msa_encoder is not None and _msa_inputs is not None: + # Fresh row subsample each iteration (column mask was applied + # once in forward, before this loop). + msa_i, mask_i, hd_i, dv_i = maybe_subsample_msa( + _msa_inputs["msa"], + _msa_inputs["msa_attention_mask"], + _msa_inputs["has_deletion"], + _msa_inputs["deletion_value"], + max_depth=_msa_inputs["max_depth"], + enabled=_msa_inputs["subsample_enabled"], + ) + B_msa, M, L_msa = msa_i.shape + msa_oh = F.one_hot( + msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES + ).float() + msa_attn = ( + mask_i.permute(0, 2, 1).float() + if mask_i is not None + else tok_mask[:, :, None].expand(-1, -1, M).float() + ) + # Bias-free MSAEncoder.embed requires zeroed padding. + msa_oh = msa_oh * msa_attn.unsqueeze(-1) + hd = ( + hd_i.permute(0, 2, 1).float() + if hd_i is not None + else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + ) + dv = ( + dv_i.permute(0, 2, 1).float() + if dv_i is not None + else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + ) + msa_pair = self.msa_encoder( + x_pair=z_inject_pair, + x_inputs=_msa_inputs["x_inputs"], + msa_oh=msa_oh, + has_deletion=hd, + deletion_value=dv, + msa_attention_mask=msa_attn, + ).to(z_inject_pair.dtype) + z_inject_pair = ( + msa_pair + if self.config.msa_encoder.overwrite + else (z_inject_pair + msa_pair) + ) + + if refined_lm_z is not None: + z_inject_pair = z_inject_pair + refined_lm_z.to(z_inject_pair.dtype) + + injected_pair = self.parcae_input_norm(z_inject_pair) + z = a * z + F.linear(injected_pair.to(z.dtype), b_mat) + z = self.folding_trunk(z, pair_attention_mask=pair_mask) + + return z + + @torch.inference_mode() + def forward( + self, + token_index: Tensor, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + mol_type: Tensor, + res_type: Tensor, + token_bonds: Tensor, + token_attention_mask: Tensor, + ref_pos: Tensor, + ref_element: Tensor, + ref_charge: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + atom_attention_mask: Tensor, + atom_to_token: Tensor, + distogram_atom_idx: Tensor, + deletion_mean: Tensor | None = None, + msa: Tensor | None = None, + has_deletion: Tensor | None = None, + deletion_value: Tensor | None = None, + msa_attention_mask: Tensor | None = None, + input_ids: Tensor | None = None, + lm_hidden_states: Tensor | None = None, + num_loops: int | None = None, + num_diffusion_samples: int | None = None, + num_sampling_steps: int | None = None, + lm_mask_pct: float | None = None, + msa_max_depth: int | None = None, + msa_column_mask_rate: float | None = None, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: float | None = 256.0, + disto_cond: Tensor | None = None, + disto_cond_mask: Tensor | None = None, + **unused_features: Tensor, + ) -> dict[str, Tensor]: + unexpected = sorted(set(unused_features) - _IGNORED_FEATURE_KEYS) + if unexpected: + raise TypeError( + f"{type(self).__name__}.forward() got unexpected keyword " + f"argument(s) {unexpected}." + ) + + # Named rather than left to **kwargs so conditioning cannot be swallowed + # and the fold silently run unconditioned. No released checkpoint carries + # the disto_conditioning_proj weights that would consume these. + if disto_cond_mask is not None and disto_cond_mask.any(): + raise NotImplementedError( + "distogram conditioning is not implemented for ESMFold2; " + "fold without it." + ) + + tok_mask = token_attention_mask + atm_mask = atom_attention_mask + disto_idx = distogram_atom_idx + + n_loops: int = num_loops if num_loops is not None else self.config.num_loops + n_samples: int = ( + num_diffusion_samples + if num_diffusion_samples is not None + else self.config.num_diffusion_samples + ) + total_steps = max(1, n_loops + 1) + + if res_type.dim() == 2: + res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float() + res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float() + else: + res_type_oh = res_type.float() + + if msa is not None: + msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float() + if msa_attention_mask is not None: + mask_f = msa_attention_mask.float().unsqueeze(-1) + msa_oh_profile = msa_oh_profile * mask_f + valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1) + profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1) + else: + profile = msa_oh_profile.mean(dim=1) + else: + profile = res_type_oh + + if deletion_mean is None: + deletion_mean = torch.zeros( + res_type.shape[0], res_type.shape[1], device=res_type.device + ) + + ref_element_oh = F.one_hot( + ref_element.long(), num_classes=MAX_ATOMIC_NUMBER + ).float() + ref_atom_name_chars_oh = F.one_hot( + ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE + ).float() + # Bias-free downstream Linears require zeroed padding. + atm_mask_f = atm_mask.float() + ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1) + ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze( + -1 + ).unsqueeze(-1) + atom_to_token = atom_to_token * atm_mask.long() + + use_amp = ref_pos.device.type == "cuda" + with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16): + x_inputs = self.inputs_embedder( + aatype=res_type_oh, + profile=profile.float(), + deletion_mean=deletion_mean.float(), + ref_pos=ref_pos, + atom_attention_mask=atm_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element_oh, + ref_atom_name_chars=ref_atom_name_chars_oh, + atom_to_token=atom_to_token, + ) + + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2( + x_inputs + ).unsqueeze(1) + + relative_position_encoding = self.rel_pos( + residue_index=residue_index, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + token_index=token_index, + ) + token_bonds_encoding = self.token_bonds(token_bonds.float()) + z_init = z_init + relative_position_encoding + token_bonds_encoding + + if ( + lm_hidden_states is None + and input_ids is not None + and self.esmc is not None + ): + lm_hidden_states = self._compute_lm_hidden_states( + input_ids, + asym_id, + residue_index, + mol_type, + tok_mask, + lm_mask_pct=( + lm_mask_pct + if lm_mask_pct is not None + else self.config.lm_mask_pct + ), + ) + lm_z: Tensor | None = None + if lm_hidden_states is not None: + lm_z = self.language_model(lm_hidden_states.detach()) + del lm_hidden_states + + pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() + + z = self._init_pair_state(z_init) + + a, b = self._discretized_dynamics() + a = a.view(1, 1, 1, -1).to(device=z.device, dtype=z.dtype) + b_mat = b.to(device=z.device, dtype=z.dtype) + + # Inference-time MSA diversity: column mask is applied once here + # (shared across recycling loops); row subsampling is deferred to + # per-iter inside _run_one_loop (fresh subset per loop). + _msa_inputs: dict | None = None + if self.msa_encoder is not None and msa is not None: + # ``None`` means "use the checkpoint's value". + msa_cfg = self.config.msa_encoder + depth = msa_cfg.max_depth if msa_max_depth is None else msa_max_depth + rate = ( + msa_cfg.column_mask_rate + if msa_column_mask_rate is None + else msa_column_mask_rate + ) + msa_attention_mask = maybe_apply_msa_column_masking( + msa_attention_mask, rate=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=depth, + subsample_enabled=depth is not None, + ) + + # Method call (not inline loop) frees per-iter L²×c_z locals. + z = self._run_one_loop( + z=z, + z_init=z_init, + lm_z=lm_z, + _msa_inputs=_msa_inputs, + pair_mask=pair_mask, + a=a, + b_mat=b_mat, + tok_mask=tok_mask, + total_steps=total_steps, + ) + del z_init, lm_z, _msa_inputs, a, b_mat + + z = self.parcae_readout(z) + z = self.parcae_coda(z, pair_attention_mask=pair_mask) + + z = z.float() + distogram_logits = self.distogram_head(z + z.transpose(-2, -3)) + + structure_output = self.structure_head.sample( + z_trunk=z, + s_inputs=x_inputs, + s_trunk=None, + relative_position_encoding=relative_position_encoding, + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=atm_mask, + ref_element=ref_element_oh, + ref_atom_name_chars=ref_atom_name_chars_oh, + ref_space_uid=ref_space_uid, + tok_idx=atom_to_token, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=tok_mask, + num_diffusion_samples=n_samples, + num_sampling_steps=num_sampling_steps, + noise_scale=noise_scale, + step_scale=step_scale, + max_inference_sigma=max_inference_sigma, + return_atom_repr=False, + ) + + sample_coords = structure_output["sample_atom_coords"] + assert sample_coords is not None + 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(), + ) + output.update(confidence_output) + output["atom_pad_mask"] = ( + atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask + ) + output["residue_index"] = residue_index + output["entity_id"] = entity_id + return output + + @torch.no_grad() + def infer_protein(self, seq: str, **forward_kwargs) -> dict: + from esm.models.esmfold2.protein_utils import ( + OUTPUT_TO_PDB_FEATURE_KEYS, + prepare_protein_features, + ) + + features = prepare_protein_features(seq) + features = {k: v.to(self.device) for k, v in features.items()} + output = self(**features, **forward_kwargs) + for k in OUTPUT_TO_PDB_FEATURE_KEYS: + output[k] = features[k] + return output + + def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str: + return self.output_to_pdb(self.infer_protein(seq, **forward_kwargs)) + + @staticmethod + def output_to_pdb(output: dict) -> str: + from esm.models.esmfold2.protein_utils import output_to_pdb as _output_to_pdb + + return _output_to_pdb(output) + + +class MSAEncoderBlock(nn.Module): + """One MSA encoder block: OPM into pair, MSA pair-weighted averaging, triangle update.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_hidden: int, + n_heads_msa: int, + msa_head_width: int, + is_final_block: bool = False, + ) -> None: + super().__init__() + self.is_final_block = is_final_block + self.outer_product_mean = OuterProductMean(d_msa, d_hidden, d_pair) + if not is_final_block: + self.msa_pair_weighted_averaging = MSAPairWeightedAveraging( + d_msa, d_pair, n_heads_msa, msa_head_width + ) + self.msa_transition = PairTransition(d_msa, expansion_ratio=4) + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = PairTransition(d_pair, expansion_ratio=4) + + # Only the triangle updates take a backend: PairTransition chunks but has no + # kernel to choose. + def set_kernel_backend(self, backend: str | None) -> None: + self.tri_mul_out.set_kernel_backend(backend) + self.tri_mul_in.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.outer_product_mean.set_chunk_size(chunk_size) + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + if not self.is_final_block: + self.msa_transition.set_chunk_size(chunk_size) + self.pair_transition.set_chunk_size(chunk_size) + + def forward( + self, + m: Tensor, + pair: Tensor, + msa_attention_mask: Tensor, + pair_attention_mask: Tensor, + ) -> tuple[Tensor, Tensor]: + pair = pair + self.outer_product_mean(m, msa_attention_mask) + if not self.is_final_block: + m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask) + m = m + self.msa_transition(m) + pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask) + pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask) + pair = pair + self.pair_transition(pair) + return m, pair + + +class MSAEncoder(nn.Module): + """Stack of [`MSAEncoderBlock`] layers that conditions the pair on an MSA.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_inputs: int, + d_hidden: int = 32, + n_layers: int = 4, + n_heads_msa: int = 8, + msa_head_width: int = 16, + ) -> None: + super().__init__() + self.embed = nn.Linear(35, d_msa, bias=False) + self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False) + self.blocks = nn.ModuleList( + [ + MSAEncoderBlock( + d_msa=d_msa, + d_pair=d_pair, + d_hidden=d_hidden, + n_heads_msa=n_heads_msa, + msa_head_width=msa_head_width, + is_final_block=(i == n_layers - 1), + ) + for i in range(n_layers) + ] + ) + + def set_kernel_backend(self, backend: str | None) -> None: + for block in self.blocks: + cast(MSAEncoderBlock, block).set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(MSAEncoderBlock, block).set_chunk_size(chunk_size) + + def forward( + self, + x_pair: Tensor, + x_inputs: Tensor, + msa_oh: Tensor, + has_deletion: Tensor, + deletion_value: Tensor, + msa_attention_mask: Tensor, + ) -> Tensor: + # All inputs are pre-transposed to [B, L, M, ...] before calling. + m_feat = torch.cat( + [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1 + ) + m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2) + tok_mask = msa_attention_mask[:, :, 0].bool() + pair_attention_mask = tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1) + for block in self.blocks: + m, x_pair = block(m, x_pair, msa_attention_mask, pair_attention_mask) + return x_pair diff --git a/esm/models/esmfold2/prepare_input.py b/esm/models/esmfold2/prepare_input.py index 272ded59..84c48ff2 100644 --- a/esm/models/esmfold2/prepare_input.py +++ b/esm/models/esmfold2/prepare_input.py @@ -56,6 +56,7 @@ RNAInput, StructurePredictionInput, ) +from esm.utils.structure.input_builder import _entity_key # ============================================================================= # Lightweight data model @@ -672,35 +673,20 @@ def tokenize_ligand_smiles( # ============================================================================= -def _get_sequence_key(item) -> str: - """Get a hashable key for entity deduplication.""" - if isinstance(item, ProteinInput): - return f"PROTEIN:{item.sequence}" - elif isinstance(item, DNAInput): - return f"DNA:{item.sequence}" - elif isinstance(item, RNAInput): - return f"RNA:{item.sequence}" - elif isinstance(item, LigandInput): - if item.ccd: - return f"LIGAND_CCD:{','.join(item.ccd)}" - return f"LIGAND_SMILES:{item.smiles}" - raise ValueError(f"Unknown input type: {type(item)}") - - def build_chains_from_input( input: StructurePredictionInput, seed: int | None = None ) -> tuple[list[ChainInfo], list[TokenInfo], list[AtomInfo]]: """Build chains, tokens, and atoms from StructurePredictionInput. - Handles entity deduplication (identical sequences get same entity_id), - sym_id assignment, and delegates to type-specific tokenization functions. + Handles entity deduplication (chemically identical chains get the same + entity_id), sym_id assignment, and delegates to type-specific tokenization. """ chains: list[ChainInfo] = [] all_tokens: list[TokenInfo] = [] all_atoms: list[AtomInfo] = [] # Entity deduplication - sequence_to_entity: dict[str, int] = {} + sequence_to_entity: dict[tuple, int] = {} entity_sym_count: dict[int, int] = {} next_entity_id = 0 @@ -717,7 +703,7 @@ def build_chains_from_input( for item in input.sequences: # Entity deduplication - seq_key = _get_sequence_key(item) + seq_key = _entity_key(item) if seq_key in sequence_to_entity: entity_id = sequence_to_entity[seq_key] else: diff --git a/esm/models/esmfold2/prepare_input_test.py b/esm/models/esmfold2/prepare_input_test.py index 338a0625..b34a5fb0 100644 --- a/esm/models/esmfold2/prepare_input_test.py +++ b/esm/models/esmfold2/prepare_input_test.py @@ -4,10 +4,16 @@ from rdkit import Chem from esm.models.esmfold2.prepare_input import ( + _entity_key, build_chains_from_input, compute_token_bonds, ) -from esm.models.esmfold2.types import LigandInput, StructurePredictionInput +from esm.models.esmfold2.types import ( + LigandInput, + Modification, + ProteinInput, + StructurePredictionInput, +) @pytest.mark.parametrize( @@ -29,3 +35,71 @@ def test_smiles_ligand_bonds_match_molecular_graph(smiles: str): n_edges = int(token_bonds.sum().item()) // 2 # symmetric matrix assert n_edges == mol.GetNumBonds() assert n_edges < len(tokens) * (len(tokens) - 1) // 2 # not a clique + + +def test_entity_key_separates_chemically_distinct_chains(): + """Same sequence but different chemistry must not collapse to one entity. + + Entities that merge also become sym_id copies of each other, i.e. claim a + symmetry between two chains that are not copies. + """ + plain = ProteinInput(id="A", sequence="AAAGGG") + modified = ProteinInput( + id="B", sequence="AAAGGG", modifications=[Modification(position=1, ccd="MSE")] + ) + + assert _entity_key(plain) == _entity_key(ProteinInput(id="D", sequence="AAAGGG")) + assert _entity_key(plain) != _entity_key(modified) + + +def test_entity_key_ignores_smiles_when_ccd_is_present(): + """A redundant smiles must not split one ligand entity into two. + + Tokenization takes the ccd branch and warns that the smiles is unused, so a + key that reads both fields would call two identically tokenized chains + different entities. + """ + assert _entity_key(LigandInput(id="A", ccd=["HEM"])) == _entity_key( + LigandInput(id="B", ccd=["HEM"], smiles="CCO") + ) + assert _entity_key(LigandInput(id="A", ccd=["HEM"])) != _entity_key( + LigandInput(id="B", ccd=["NAG"]) + ) + + +def test_build_chains_assigns_entity_and_sym_ids_to_repeated_ligands(): + """Repeated copies are one entity numbered by sym_id, through the real builder. + + sym_id is what the relative position encoding turns into its cross-chain + feature, so entity dedup is only correct if the sym_id it drives is too. + """ + spi = StructurePredictionInput( + sequences=[ + LigandInput(id="A", smiles="CCO"), + LigandInput(id="B", smiles="c1ccccc1"), + LigandInput(id="C", smiles="CCO"), + ] + ) + chains, _, _ = build_chains_from_input(spi, seed=0) + + assert [c.entity_id for c in chains] == [0, 1, 0] + assert [c.sym_id for c in chains] == [0, 0, 1] + + +def test_build_chains_numbers_copies_across_both_id_forms(): + """One entry with several ids and a duplicate entry must keep counting. + + ``id=['A', 'B']`` expands to two chains inside one entry while chain C is a + separate entry for the same entity, so the three copies are sym_id 0, 1, 2. + """ + spi = StructurePredictionInput( + sequences=[ + LigandInput(id=["A", "B"], smiles="CCO"), + LigandInput(id="C", smiles="CCO"), + ] + ) + chains, _, _ = build_chains_from_input(spi, seed=0) + + assert [c.chain_id for c in chains] == ["A", "B", "C"] + assert [c.entity_id for c in chains] == [0, 0, 0] + assert [c.sym_id for c in chains] == [0, 1, 2] diff --git a/esm/models/esmfold2/processor.py b/esm/models/esmfold2/processor.py index 21f89d58..edb9e539 100644 --- a/esm/models/esmfold2/processor.py +++ b/esm/models/esmfold2/processor.py @@ -332,6 +332,7 @@ def fold( num_loops: int = 20, num_sampling_steps: int = 200, num_diffusion_samples: int = 1, + confidence_chunk_size: int | None = None, seed: int | None = None, noise_scale: float | None = None, step_scale: float | None = None, @@ -353,6 +354,10 @@ def fold( User-facing input specification. num_loops, num_sampling_steps, num_diffusion_samples : int Inference knobs forwarded to the model. + confidence_chunk_size : int, optional + Maximum number of diffusion samples processed together by the + confidence head. Lower values reduce peak memory without changing + outputs. seed : int, optional Seeds both input prep (SMILES conformer generation) and diffusion sampling. noise_scale, step_scale, max_inference_sigma, early_exit @@ -402,6 +407,7 @@ def fold( num_loops=num_loops, num_sampling_steps=num_sampling_steps, num_diffusion_samples=num_diffusion_samples, + confidence_chunk_size=confidence_chunk_size, early_exit=early_exit, msa_max_depth=msa_max_depth, msa_column_mask_rate=msa_column_mask_rate, diff --git a/esm/models/esmfold2/protein_utils.py b/esm/models/esmfold2/protein_utils.py new file mode 100644 index 00000000..6eac5bac --- /dev/null +++ b/esm/models/esmfold2/protein_utils.py @@ -0,0 +1,588 @@ +# coding=utf-8 +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Self-contained protein featurization for ESMFold2 inference.""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +from torch import Tensor + +MOL_TYPE_PROTEIN = 0 +PROTEIN_UNK_RES_TYPE = 22 +MSA_GAP_TOKEN_ID = 1 + +PROTEIN_RESIDUE_TO_RES_TYPE: dict[str, int] = { + "ALA": 2, + "ARG": 3, + "ASN": 4, + "ASP": 5, + "CYS": 6, + "GLN": 7, + "GLU": 8, + "GLY": 9, + "HIS": 10, + "ILE": 11, + "LEU": 12, + "LYS": 13, + "MET": 14, + "PHE": 15, + "PRO": 16, + "SER": 17, + "THR": 18, + "TRP": 19, + "TYR": 20, + "VAL": 21, +} + +PROTEIN_1TO3: dict[str, str] = { + "A": "ALA", + "R": "ARG", + "N": "ASN", + "D": "ASP", + "C": "CYS", + "Q": "GLN", + "E": "GLU", + "G": "GLY", + "H": "HIS", + "I": "ILE", + "L": "LEU", + "K": "LYS", + "M": "MET", + "F": "PHE", + "P": "PRO", + "S": "SER", + "T": "THR", + "W": "TRP", + "Y": "TYR", + "V": "VAL", + "X": "UNK", +} + +ESM_PROTEIN_VOCAB: dict[str, int] = { + "L": 4, + "A": 5, + "G": 6, + "V": 7, + "S": 8, + "E": 9, + "R": 10, + "T": 11, + "I": 12, + "D": 13, + "P": 14, + "K": 15, + "Q": 16, + "N": 17, + "F": 18, + "Y": 19, + "M": 20, + "H": 21, + "W": 22, + "C": 23, + "X": 3, +} + +# Heavy atoms per canonical residue, in training-time order. +PROTEIN_HEAVY_ATOMS: dict[str, list[str]] = { + "ALA": ["N", "CA", "C", "O", "CB"], + "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"], + "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"], + "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"], + "CYS": ["N", "CA", "C", "O", "CB", "SG"], + "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"], + "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"], + "GLY": ["N", "CA", "C", "O"], + "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"], + "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"], + "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"], + "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"], + "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"], + "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"], + "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"], + "SER": ["N", "CA", "C", "O", "CB", "OG"], + "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"], + "TRP": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "NE1", + "CE2", + "CE3", + "CZ2", + "CZ3", + "CH2", + ], + "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"], + "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"], + "UNK": ["N", "CA", "C", "O"], +} + +PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = { + "ALA": { + "N": (-0.01003183238208294, -1.2073018550872803, -1.0555061101913452), + "CA": (-0.04190138354897499, 0.17447763681411743, -0.5729365348815918), + "C": (1.2127548456192017, 0.4737588167190552, 0.19521640241146088), + "O": (1.9390329122543335, 1.4484562873840332, -0.13759790360927582), + "CB": (-1.276943325996399, 0.4288230538368225, 0.29937705397605896), + }, + "ARG": { + "N": (-2.0170421600341797, 0.6717798113822937, -1.1794233322143555), + "CA": (-2.0503084659576416, -0.5735036730766296, -0.4097220301628113), + "C": (-3.469440460205078, -1.0612813234329224, -0.2755832374095917), + "O": (-3.8218462467193604, -2.1369943618774414, -0.8294969797134399), + "CB": (-1.4193516969680786, -0.3735991418361664, 0.9852858781814575), + "CG": (0.11878877878189087, -0.3112654983997345, 0.963895857334137), + "CD": (0.6643245816230774, 1.0068185329437256, 0.3963329493999481), + "NE": (2.1090238094329834, 1.0977025032043457, 0.6120952367782593), + "CZ": (3.098905324935913, 0.3215920031070709, -0.09047172218561172), + "NH1": (4.461230278015137, 0.3844667971134186, 0.34141138195991516), + "NH2": (2.7856509685516357, -0.4166366159915924, -1.1148239374160767), + }, + "ASN": { + "N": (-0.7595629096031189, 0.7503494620323181, 1.1369825601577759), + "CA": (-0.76087886095047, 0.23876343667507172, -0.23573364317417145), + "C": (-1.9211044311523438, -0.6982439160346985, -0.42196929454803467), + "O": (-2.677666187286377, -0.5753439664840698, -1.4223182201385498), + "CB": (0.5504899024963379, -0.5078350305557251, -0.5390339493751526), + "CG": (1.7250099182128906, 0.4264017939567566, -0.5778228640556335), + "OD1": (1.9470350742340088, 1.1086392402648926, -1.613560438156128), + "ND2": (2.57365345954895, 0.5730618834495544, 0.5608599781990051), + }, + "ASP": { + "N": (-1.8452696800231934, -1.2169504165649414, 0.19437327980995178), + "CA": (-0.6379959583282471, -0.41974392533302307, 0.41681644320487976), + "C": (-0.9431572556495667, 1.0356197357177734, 0.18555717170238495), + "O": (-1.5183608531951904, 1.4045922756195068, -0.8739855885505676), + "CB": (0.48594576120376587, -0.8970447778701782, -0.5209363698959351), + "CG": (1.780342936515808, -0.19918935000896454, -0.2310730367898941), + "OD1": (2.5202910900115967, -0.6044584512710571, 0.7049641013145447), + "OD2": (2.1454880237579346, 0.9208861589431763, -0.9712985157966614), + }, + "CYS": { + "N": (0.0469963513314724, 1.190075159072876, -1.1607273817062378), + "CA": (0.11344368755817413, -0.09400428831577301, -0.45952197909355164), + "C": (-1.2652032375335693, -0.6832379698753357, -0.3594406247138977), + "O": (-1.4631439447402954, -1.8851220607757568, -0.6826791763305664), + "CB": (0.6919880509376526, 0.09034398198127747, 0.952482283115387), + "SG": (2.4619927406311035, 0.5235707759857178, 0.9020372629165649), + }, + "GLN": { + "N": (-2.370004653930664, -0.9637529850006104, -0.7942749261856079), + "CA": (-1.370002269744873, -0.6000258922576904, 0.2103111445903778), + "C": (-1.7545503377914429, 0.7091967463493347, 0.8433493971824646), + "O": (-1.8520662784576416, 0.7999289631843567, 2.0964975357055664), + "CB": (0.02040259726345539, -0.5004461407661438, -0.44764479994773865), + "CG": (1.1377512216567993, -0.28680720925331116, 0.582992434501648), + "CD": (2.4745187759399414, -0.24800164997577667, -0.09364881366491318), + "OE1": (3.1685523986816406, -1.2966246604919434, -0.1717153936624527), + "NE2": (2.947425603866577, 0.9601329565048218, -0.6888364553451538), + }, + "GLU": { + "N": (-1.5850872993469238, -1.337684154510498, 0.9490851163864136), + "CA": (-1.0560977458953857, 0.027459044009447098, 1.0306966304779053), + "C": (-1.7741456031799316, 0.9664392471313477, 0.09259600937366486), + "O": (-1.9012441635131836, 2.181349992752075, 0.402479350566864), + "CB": (0.4706551432609558, 0.048803869634866714, 0.8114414811134338), + "CG": (0.9133604764938354, -0.4219329059123993, -0.5830985307693481), + "CD": (2.398822069168091, -0.3097084164619446, -0.7210537791252136), + "OE1": (3.1389315128326416, -1.274524450302124, -0.39029765129089355), + "OE2": (2.9647817611694336, 0.8781346082687378, -1.1732689142227173), + }, + "GLY": { + "N": (-1.3942985534667969, -0.39875128865242004, -0.3370324671268463), + "CA": (-0.39974430203437805, 0.5488945245742798, 0.15242962539196014), + "C": (0.9440054893493652, -0.10314033925533295, 0.19859643280506134), + "O": (1.3352899551391602, -0.669218122959137, 1.2541258335113525), + }, + "HIS": { + "N": (-1.4532867670059204, -1.0689626932144165, 0.881072461605072), + "CA": (-1.3396095037460327, 0.24797579646110535, 0.24960045516490936), + "C": (-2.675257921218872, 0.6571555733680725, -0.30441102385520935), + "O": (-3.1311378479003906, 1.8079776763916016, -0.06785715371370316), + "CB": (-0.3041955828666687, 0.21721023321151733, -0.8885309100151062), + "CG": (1.0887513160705566, 0.028941065073013306, -0.36419469118118286), + "ND1": (1.840459942817688, 1.0411773920059204, 0.29804590344429016), + "CD2": (1.780855417251587, -1.1011489629745483, -0.3814258575439453), + "CE1": (2.9566943645477295, 0.4924798905849457, 0.6477115750312805), + "NE2": (3.0280203819274902, -0.8751969337463379, 0.26084381341934204), + }, + "ILE": { + "N": (-0.7167549729347229, -1.5426139831542969, -0.9983330368995667), + "CA": (-1.0636085271835327, -0.35169270634651184, -0.21393552422523499), + "C": (-1.3896740674972534, 0.8142145276069641, -1.1164065599441528), + "O": (-1.2377792596817017, 0.7302915453910828, -2.3656840324401855), + "CB": (0.061667006462812424, 0.01599610224366188, 0.8057394623756409), + "CG1": (1.502519965171814, -0.08899776637554169, 0.24154816567897797), + "CG2": (-0.053174979984760284, -0.8521055579185486, 2.0702083110809326), + "CD1": (1.7929610013961792, 0.899773120880127, -0.8863027691841125), + }, + "LEU": { + "N": (1.9657520055770874, -1.9763224124908447, -0.18391533195972443), + "CA": (1.3077669143676758, -0.6677430868148804, -0.19492436945438385), + "C": (1.9905058145523071, 0.24182087182998657, 0.7879968285560608), + "O": (2.06896710395813, -0.07880014181137085, 2.0048046112060547), + "CB": (-0.20306941866874695, -0.8093230128288269, 0.11243502795696259), + "CG": (-0.9916267395019531, 0.5234957337379456, 0.06723011285066605), + "CD1": (-2.4228057861328125, 0.29949337244033813, 0.573042094707489), + "CD2": (-1.0282856225967407, 1.1250264644622803, -1.346014380455017), + }, + "LYS": { + "N": (2.4221372604370117, -0.6473312377929688, 0.6370573043823242), + "CA": (2.0314927101135254, 0.2786507308483124, -0.4298512041568756), + "C": (2.7168593406677246, 1.595757246017456, -0.20924785733222961), + "O": (3.397681713104248, 2.116427421569824, -1.1332510709762573), + "CB": (0.5018402934074402, 0.4873858690261841, -0.49062973260879517), + "CG": (-0.25062066316604614, -0.7894009947776794, -0.9055535793304443), + "CD": (-1.769762635231018, -0.5552700161933899, -1.040329933166504), + "CE": (-2.576533555984497, -1.0221366882324219, 0.18493641912937164), + "NZ": (-2.269151210784912, -0.24293844401836395, 1.3849012851715088), + }, + "MET": { + "N": (1.8903918266296387, -1.5252995491027832, -0.42638593912124634), + "CA": (1.2630571126937866, -0.24417810142040253, -0.7626462578773499), + "C": (2.30391001701355, 0.8367712497711182, -0.7254616618156433), + "O": (2.465414524078369, 1.5928632020950317, -1.7207728624343872), + "CB": (0.10567972809076309, 0.10861825942993164, 0.19741646945476532), + "CG": (-1.0658042430877686, -0.8736631274223328, 0.08811883628368378), + "SD": (-2.4557132720947266, -0.3332225978374481, 1.1461700201034546), + "CE": (-3.265165090560913, 0.7033554911613464, -0.11588376015424728), + }, + "PHE": { + "N": (-2.8484435081481934, -1.525790810585022, 0.01789816841483116), + "CA": (-1.591969609260559, -0.8545162677764893, 0.35214468836784363), + "C": (-1.8900631666183472, 0.45833414793014526, 1.0232222080230713), + "O": (-1.3424992561340332, 0.74432373046875, 2.121629476547241), + "CB": (-0.760358452796936, -0.6342853307723999, -0.9257160425186157), + "CG": (0.604112982749939, -0.07200468331575394, -0.6148118376731873), + "CD1": (0.8468314409255981, 1.2480632066726685, -0.7146694660186768), + "CD2": (1.6827683448791504, -0.9758077263832092, -0.1423054188489914), + "CE1": (2.1801748275756836, 1.7875733375549316, -0.3744623064994812), + "CE2": (2.888307809829712, -0.48277512192726135, 0.16804970800876617), + "CZ": (3.149812936782837, 0.9656873941421509, 0.04440271109342575), + }, + "PRO": { + "N": (-0.836250364780426, -0.9899801015853882, 0.5561304688453674), + "CA": (0.32722190022468567, -0.6164458394050598, -0.25072571635246277), + "C": (1.6121541261672974, -1.1711241006851196, 0.31082412600517273), + "O": (1.6127740144729614, -2.2771971225738525, 0.9156193733215332), + "CB": (0.3248198926448822, 0.9028244018554688, -0.33368146419525146), + "CG": (-1.1425083875656128, 1.2730128765106201, -0.2590600252151489), + "CD": (-1.8495968580245972, 0.026575811207294464, 0.2681289613246918), + }, + "SER": { + "N": (0.674650251865387, 1.5018702745437622, -0.5367295145988464), + "CA": (0.00013792862591799349, 0.4966467022895813, 0.28510504961013794), + "C": (0.9941009879112244, -0.5374617576599121, 0.73505038022995), + "O": (1.0545241832733154, -0.8683545589447021, 1.9495396614074707), + "CB": (-1.1279288530349731, -0.1659376323223114, -0.5160963535308838), + "OG": (-1.8135979175567627, -1.085249662399292, 0.28947514295578003), + }, + "THR": { + "N": (-1.325830340385437, -1.3728225231170654, 0.6882233023643494), + "CA": (-0.5433306097984314, -0.16364754736423492, 0.41697052121162415), + "C": (-1.294381856918335, 0.7077372074127197, -0.5549946427345276), + "O": (-1.6939635276794434, 0.23654410243034363, -1.6540418863296509), + "CB": (0.853203296661377, -0.5363803505897522, -0.14109353721141815), + "OG1": (1.5220820903778076, -1.379003643989563, 0.7635167837142944), + "CG2": (1.7225933074951172, 0.7054727077484131, -0.3651331067085266), + }, + "TRP": { + "N": (3.686030864715576, 0.7599999904632568, 0.496155709028244), + "CA": (2.384092092514038, 0.09079249948263168, 0.5325262546539307), + "C": (2.1113572120666504, -0.6121063232421875, -0.7733646035194397), + "O": (1.796526312828064, -1.8323148488998413, -0.7775964140892029), + "CB": (1.281521201133728, 1.1139036417007446, 0.8559791445732117), + "CG": (-0.04292375594377518, 0.44645074009895325, 1.0942792892456055), + "CD1": (-0.42329534888267517, -0.15470874309539795, 2.2227554321289062), + "CD2": (-1.1023900508880615, 0.2158389836549759, 0.11529432237148285), + "NE1": (-1.7030320167541504, -0.7665823101997375, 2.0595016479492188), + "CE2": (-2.045644998550415, -0.4881173074245453, 0.710669219493866), + "CE3": (-1.2173502445220947, 0.6102271676063538, -1.300106406211853), + "CZ2": (-3.256009340286255, -0.9164394736289978, -0.00984987337142229), + "CZ3": (-2.315925121307373, 0.2306906282901764, -1.9776310920715332), + "CH2": (-3.3817875385284424, -0.5677337646484375, -1.3032053709030151), + }, + "TYR": { + "N": (-1.7900604009628296, -0.8409399390220642, 1.3180142641067505), + "CA": (-1.913882851600647, 0.23552845418453217, 0.330669641494751), + "C": (-3.347280740737915, 0.3588399887084961, -0.09830684959888458), + "O": (-3.967811346054077, -0.6449354290962219, -0.5423302054405212), + "CB": (-1.0093992948532104, 0.0004731413209810853, -0.8981552124023438), + "CG": (0.4520410895347595, 0.021162061020731926, -0.5305932760238647), + "CD1": (1.0992432832717896, 1.1877919435501099, -0.3579142987728119), + "CD2": (1.1803174018859863, -1.253401279449463, -0.31122180819511414), + "CE1": (2.5253450870513916, 1.1990256309509277, 0.029804613441228867), + "CE2": (2.471151113510132, -1.240687608718872, 0.043534230440855026), + "CZ": (3.180687665939331, 0.04672492295503616, 0.2214856892824173), + "OH": (4.523719787597656, 0.0671030730009079, 0.5877485871315002), + }, + "VAL": { + "N": (0.5987519025802612, -1.569443702697754, -0.7379124760627747), + "CA": (0.6014357209205627, -0.10503966361284256, -0.6336286664009094), + "C": (1.8391697406768799, 0.4067850410938263, 0.06351757049560547), + "O": (2.3952062129974365, -0.2666190266609192, 0.9731166958808899), + "CB": (-0.694736897945404, 0.4259096384048462, 0.03581475466489792), + "CG1": (-1.9276031255722046, 0.09515828639268875, -0.8172357082366943), + "CG2": (-0.8938426971435547, -0.08640842139720917, 1.472349762916565), + }, + "UNK": { + "N": (0.0, 0.0, 0.0), + "CA": (0.0, 0.0, 0.0), + "C": (0.0, 0.0, 0.0), + "O": (0.0, 0.0, 0.0), + }, +} + +# Protonated nitrogens at physiological pH (matches CHARGED_ATOMS in the +# opensource constants for the protein subset). +PROTEIN_CHARGED_ATOMS: dict[tuple[str, str], int] = { + ("LYS", "NZ"): 1, + ("ARG", "NH2"): 1, + ("HIS", "ND1"): 1, +} + +# Only the elements that appear in canonical protein heavy atoms. +_PROTEIN_ELEMENT_TO_ATOMIC_NUM: dict[str, int] = {"C": 6, "N": 7, "O": 8, "S": 16} + + +def _encode_atom_name(name: str) -> list[int]: + padded = name.ljust(4)[:4] + return [ord(c) - 32 if c != " " else 0 for c in padded] + + +def prepare_protein_features(sequence: str) -> dict[str, Tensor]: + """Featurize a single protein sequence for EsmFold2ExperimentalModel.forward. + + Returns the same keys with the same dtypes/shapes as + ``ESMFold2InputBuilder.prepare_input(StructurePredictionInput(...))`` + restricted to a single-chain protein with no MSA, modifications, + distogram conditioning, or covalent bonds. All tensors have a + leading batch dim of 1; the caller is responsible for moving them + to the model device. + """ + if not sequence: + raise ValueError("sequence must be non-empty") + + res_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence] + L = len(sequence) + + token_atom_starts: list[int] = [] + atom_records: list[tuple[int, str, str, int, tuple[float, float, float]]] = [] + res_type_vals: list[int] = [] + input_id_vals: list[int] = [] + distogram_rep_atom_idx: list[int] = [] + + atom_cursor = 0 + for t_idx, (letter, res_3) in enumerate(zip(sequence, res_3letter)): + atom_names = PROTEIN_HEAVY_ATOMS[res_3] + res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE) + input_id = ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"]) + + token_atom_starts.append(atom_cursor) + for name in atom_names: + charge = PROTEIN_CHARGED_ATOMS.get((res_3, name), 0) + element = name[0] # protein heavy atoms are all single-letter C/N/O/S + ref_pos = PROTEIN_REF_POS[res_3][name] + atom_records.append((t_idx, name, element, charge, ref_pos)) + atom_cursor += 1 + + rep_name = "CB" if "CB" in atom_names else "CA" + distogram_rep_atom_idx.append( + token_atom_starts[t_idx] + atom_names.index(rep_name) + ) + + res_type_vals.append(res_type) + input_id_vals.append(input_id) + + n_real_atoms = len(atom_records) + n_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32 + + ref_pos = torch.zeros(n_atoms, 3, dtype=torch.float32) + ref_element = torch.zeros(n_atoms, dtype=torch.int64) + ref_charge = torch.zeros(n_atoms, dtype=torch.int8) + ref_atom_name_chars = torch.zeros(n_atoms, 4, dtype=torch.int64) + ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64) + atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool) + atom_to_token = torch.zeros(n_atoms, dtype=torch.int64) + + for i, (t_idx, name, element, charge, pos) in enumerate(atom_records): + ref_pos[i] = torch.tensor(pos, dtype=torch.float32) + ref_element[i] = _PROTEIN_ELEMENT_TO_ATOMIC_NUM[element] + ref_charge[i] = charge + ref_atom_name_chars[i] = torch.tensor( + _encode_atom_name(name), dtype=torch.int64 + ) + ref_space_uid[i] = t_idx + atom_attention_mask[i] = True + atom_to_token[i] = t_idx + + token_index = torch.arange(L, dtype=torch.int64) + residue_index = torch.arange(L, dtype=torch.int64) + asym_id = torch.zeros(L, dtype=torch.int64) + sym_id = torch.zeros(L, dtype=torch.int64) + entity_id = torch.ones(L, dtype=torch.int64) + mol_type = torch.full((L,), MOL_TYPE_PROTEIN, dtype=torch.int64) + res_type = torch.tensor(res_type_vals, dtype=torch.int64) + input_ids = torch.tensor(input_id_vals, dtype=torch.int64) + token_bonds = torch.zeros(L, L, 1, dtype=torch.float32) + token_attention_mask = torch.ones(L, dtype=torch.bool) + distogram_atom_idx = torch.tensor(distogram_rep_atom_idx, dtype=torch.int64) + + # Single-sequence MSA: depth 1, row 0 is the sequence itself. + msa = res_type.unsqueeze(0) + msa_attention_mask = torch.ones(1, L, dtype=torch.bool) + has_deletion = torch.zeros(1, L, dtype=torch.bool) + deletion_value = torch.zeros(1, L, dtype=torch.float32) + deletion_mean = torch.zeros(L, dtype=torch.float32) + + features = { + "token_index": token_index, + "residue_index": residue_index, + "asym_id": asym_id, + "sym_id": sym_id, + "entity_id": entity_id, + "mol_type": mol_type, + "res_type": res_type, + "input_ids": input_ids, + "token_bonds": token_bonds, + "token_attention_mask": token_attention_mask, + "ref_pos": ref_pos, + "ref_element": ref_element, + "ref_charge": ref_charge, + "ref_atom_name_chars": ref_atom_name_chars, + "ref_space_uid": ref_space_uid, + "atom_attention_mask": atom_attention_mask, + "atom_to_token": atom_to_token, + "distogram_atom_idx": distogram_atom_idx, + "msa": msa, + "msa_attention_mask": msa_attention_mask, + "has_deletion": has_deletion, + "deletion_value": deletion_value, + "deletion_mean": deletion_mean, + } + return {k: v.unsqueeze(0) for k, v in features.items()} + + +# 0-32 res_type → 3-letter name (only protein indices 2-22 are populated). +_RES_TYPE_TO_3LETTER: dict[int, str] = { + rt: three for three, rt in PROTEIN_RESIDUE_TO_RES_TYPE.items() +} +_RES_TYPE_TO_3LETTER[PROTEIN_UNK_RES_TYPE] = "UNK" + +# Featurization keys that ``output_to_pdb`` reads off the forward output. +# ``infer_protein`` re-attaches them because ``forward`` does not echo them +# back; both ESMFold2 model classes share this list. +OUTPUT_TO_PDB_FEATURE_KEYS: tuple[str, ...] = ( + "res_type", + "atom_to_token", + "ref_atom_name_chars", + "atom_attention_mask", + "token_attention_mask", + "residue_index", +) + + +def output_to_pdb(output: dict) -> str: + """Convert an ESMFold2 protein forward output into a PDB string. + + Expects ``output`` to carry the featurization keys re-attached by + ``infer_protein`` (``res_type``, ``atom_to_token``, + ``ref_atom_name_chars``, ``atom_attention_mask``, + ``token_attention_mask``, ``residue_index``) alongside the predicted + ``sample_atom_coords`` and ``plddt``. Builds a 37-atom + ``OFProtein`` (per-atom pLDDT in the b-factor column) and renders it + with the OpenFold utilities shipped in ``transformers.models.esm``. + """ + from esm.utils import residue_constants as rc + from esm.utils.structure.protein_chain import ProteinChain + + coords = output["sample_atom_coords"] + if coords.dim() == 4: + coords = coords[:, 0] + coords = coords.detach().cpu().numpy()[0] + + plddt = output["plddt"].detach().cpu().numpy()[0] + atom_to_token = output["atom_to_token"].cpu().numpy() + ref_chars = output["ref_atom_name_chars"].cpu().numpy() + res_type = output["res_type"].cpu().numpy() + token_mask = output["token_attention_mask"].cpu().numpy().astype(bool) + atom_mask_in = output["atom_attention_mask"].cpu().numpy().astype(bool) + residue_index_arr = output["residue_index"].cpu().numpy() + + if atom_to_token.ndim == 2: + atom_to_token = atom_to_token[0] + ref_chars = ref_chars[0] + res_type = res_type[0] + token_mask = token_mask[0] + atom_mask_in = atom_mask_in[0] + residue_index_arr = residue_index_arr[0] + + valid_tok = np.where(token_mask)[0] + n_res = valid_tok.shape[0] + + residues = ["X"] * n_res + for new_i, t in enumerate(valid_tok): + rt = int(res_type[t]) + three = _RES_TYPE_TO_3LETTER.get(rt) + if three is not None and three != "UNK": + residues[new_i] = rc.restype_3to1.get(three, "X") + + atom_positions = np.zeros((n_res, 37, 3), dtype=np.float32) + atom_mask = np.zeros((n_res, 37), dtype=np.float32) + b_factors = np.zeros((n_res, 37), dtype=np.float32) + tok_to_new = {int(t): i for i, t in enumerate(valid_tok)} + + for a in range(atom_to_token.shape[0]): + if not atom_mask_in[a]: + continue + tok = int(atom_to_token[a]) + if tok not in tok_to_new: + continue + new_i = tok_to_new[tok] + name = "".join( + chr(int(c) + 32) if int(c) != 0 else " " for c in ref_chars[a] + ).strip() + idx37 = rc.atom_order.get(name) + if idx37 is None: + continue + atom_positions[new_i, idx37] = coords[a] + atom_mask[new_i, idx37] = 1.0 + b_factors[new_i, idx37] = float(plddt[tok]) + + # Per-residue confidence is the token pLDDT the per-atom column repeats. + residue_confidence = np.array( + [float(plddt[int(t)]) for t in valid_tok], dtype=np.float32 + ) + chain = ProteinChain( + id="pred", + sequence="".join(residues), + chain_id="A", + entity_id=None, + residue_index=residue_index_arr[valid_tok].astype(np.int32) + 1, + insertion_code=np.full(n_res, "", dtype=" None: + """Strip Transformer Engine's non-tensor ``_extra_state`` entries; TE returns + these as ``io.BytesIO`` objects the checkpoint writer cannot dtype-infer.""" + for key in list(state_dict): + if key.endswith("_extra_state"): + del state_dict[key] + + +def read_safetensors_dir(directory: str | os.PathLike) -> dict[str, torch.Tensor]: + """Read a (possibly sharded) safetensors checkpoint into a single dict.""" + directory = Path(directory) + index = directory / _SAFETENSORS_INDEX + if index.exists(): + with open(index) as f: + weight_map = json.load(f)["weight_map"] + state_dict: dict[str, torch.Tensor] = {} + for shard in sorted(set(weight_map.values())): + state_dict.update(load_file(str(directory / shard))) + return state_dict + single = directory / _SAFETENSORS_SINGLE + if single.exists(): + return load_file(str(single)) + raise FileNotFoundError( + f"No safetensors checkpoint found in {directory} " + f"(looked for {_SAFETENSORS_INDEX} and {_SAFETENSORS_SINGLE})." + ) + + +def resolve_model_dir( + pretrained_model_name_or_path: str | os.PathLike, + *, + revision: str | None = None, + cache_dir: str | os.PathLike | None = None, + token: str | None = None, + local_files_only: bool = False, + force_download: bool = False, +) -> str: + """Local dir -> return as-is; hub id -> ``snapshot_download`` it.""" + path = Path(pretrained_model_name_or_path) + if path.is_dir() and (path / CONFIG_NAME).exists(): + return str(path) + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=str(pretrained_model_name_or_path), + revision=revision, + cache_dir=None if cache_dir is None else str(cache_dir), + token=token, + allow_patterns=["*.json", "*.safetensors"], + local_files_only=local_files_only, + force_download=force_download, + ) + + +class HubPreTrainedModel(nn.Module): + """Base class handling weight initialisation and Hub loading.""" + + config_class: ClassVar[type] + + # Checkpoint keys with no counterpart in the module tree that are expected + # anyway, as regexes. Anything unexpected and unlisted is an error. + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [] + + # Parameters a checkpoint is expected not to carry, as regexes - a freshly + # added task head, say. Anything else the checkpoint misses is an error. + _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [] + + def __init__(self, config) -> None: + super().__init__() + self.config = config + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def _init_weights(self, module: nn.Module, device: torch.device | str = "cpu"): + """Initialise a freshly materialised module. Overridden per family.""" + + def post_init(self) -> None: + """Mirror ``PreTrainedModel.post_init``: run per-module initialisation. + + Families that do not override :meth:`_init_weights` get a no-op, which + is what HF did for these models too - their weights come entirely from + the checkpoint. + """ + self.apply(self._init_weights) + + def _materialize_uninitialized(self, device: torch.device | str = "cpu") -> None: + """Materialize and initialize parameters/buffers still on the meta + device after a partial checkpoint load (e.g. a freshly-added head or + non-persistent RoPE buffers).""" + for module in self.modules(): + direct = list(module._parameters.values()) + list(module._buffers.values()) + if any(t is not None and t.is_meta for t in direct): + module.to_empty(device=device, recurse=False) + self._init_weights(module, device=device) + + @classmethod + def _adapt_checkpoint_keys( + cls, raw: dict[str, torch.Tensor], model_keys: set[str] + ) -> dict[str, torch.Tensor]: + """Align published-checkpoint keys onto the target module tree. + + The default keeps every key untouched, so nothing is dropped silently. + """ + return dict(raw) + + @classmethod + def _normalize_checkpoint_layout( + cls, raw: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Translate a checkpoint out of the published tensor layout. + + Runs before any key accounting, so a family whose in-memory parameters + are packed differently from its published ones (ESMC fuses q/k/v and + gate/up for Transformer Engine) is still held to the same + nothing-dropped guarantee. The default is the identity. + """ + return raw + + @classmethod + def _load_pretrained( + cls, + local_dir: str | os.PathLike, + config, + *, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, + strict_keys: bool = True, + ) -> Self: + """Instantiate ``cls(config)`` and load the checkpoint in ``local_dir``. + + With ``strict_keys`` the load refuses to discard anything: checkpoint keys + that find no home raise, and so does any parameter the checkpoint failed + to cover. Families list the keys they expect to have no counterpart in + :attr:`_keys_to_ignore_on_load_unexpected`. + """ + with init_empty_weights(): + model = cls(config) + + raw = cls._normalize_checkpoint_layout(read_safetensors_dir(local_dir)) + adapted = cls._adapt_checkpoint_keys(raw, set(model.state_dict().keys())) + incompatible = model.load_state_dict(adapted, strict=False, assign=True) + + ignore = [re.compile(p) for p in cls._keys_to_ignore_on_load_unexpected] + allow_missing = [re.compile(p) for p in cls._keys_to_ignore_on_load_missing] + keep = [k for k in raw if not any(p.search(k) for p in ignore)] + unexpected = [ + k + for k in incompatible.unexpected_keys + if not any(p.search(k) for p in ignore) + ] + # Counted by how many checkpoint entries reached the model, not by name: + # key adaptation renames as well as drops, and comparing names counted + # every rename as a loss. Entries that landed are the adapted ones minus + # those the model rejected, so `unexpected` adds to the loss. + dropped = len(keep) - (len(adapted) - len(unexpected)) + # The failure that actually corrupts a model: a parameter the checkpoint + # never covered, left holding whatever `to_empty` found in memory. + missing = [ + k + for k in incompatible.missing_keys + if not any(p.search(k) for p in allow_missing) + ] + + if strict_keys and (unexpected or dropped > 0 or missing): + raise RuntimeError( + f"{cls.__name__}.from_pretrained refused to load the checkpoint " + f"in {local_dir} silently.\n" + f" unexpected keys ({len(unexpected)}): {unexpected[:20]}\n" + f" missing keys ({len(missing)}): {missing[:20]}\n" + f" checkpoint entries that reached nothing: {dropped}" + ) + + model._materialize_uninitialized(device="cpu") + model.to(device) + if dtype is not None: + model.to(dtype) + model.eval() + return model diff --git a/esm/pretrained.py b/esm/pretrained.py index 5c12bdf8..09389f06 100644 --- a/esm/pretrained.py +++ b/esm/pretrained.py @@ -1,16 +1,18 @@ import inspect +import warnings from typing import Callable import torch import torch.nn as nn from accelerate import init_empty_weights -from huggingface_hub import load_torch_model from esm.models.esm3 import ESM3 -from esm.models.esmc import ESMC +from esm.models.esmc.compatibility import ESMC +from esm.models.esmc.config import ESMC_6B_HF_REPO, ESMC_300M_HF_REPO, ESMC_600M_HF_REPO +from esm.models.esmc.model import EsmcForMaskedLM from esm.models.function_decoder import FunctionTokenDecoder from esm.models.vqvae import StructureTokenDecoder, StructureTokenEncoder -from esm.tokenization import get_esm3_model_tokenizers, get_esmc_model_tokenizers +from esm.tokenization import get_esm3_model_tokenizers from esm.utils.constants.esm3 import data_root from esm.utils.constants.models import ( ESM3_FUNCTION_DECODER_V0, @@ -63,48 +65,6 @@ def ESM3_function_decoder_v0(device: torch.device | str = "cpu"): return model -def ESMC_300M_202412(device: torch.device | str = "cpu", use_flash_attn: bool = True): - with init_empty_weights(): - model = ESMC( - d_model=960, - n_heads=15, - n_layers=30, - tokenizer=get_esmc_model_tokenizers(), - use_flash_attn=use_flash_attn, - ).eval() - load_torch_model(model, data_root("esmc-300")) - model = model.to(device) - return model - - -def ESMC_600M_202412(device: torch.device | str = "cpu", use_flash_attn: bool = True): - with init_empty_weights(): - model = ESMC( - d_model=1152, - n_heads=18, - n_layers=36, - tokenizer=get_esmc_model_tokenizers(), - use_flash_attn=use_flash_attn, - ).eval() - load_torch_model(model, data_root("esmc-600")) - model = model.to(device) - return model - - -def ESMC_6B_202412(device: torch.device | str = "cpu", use_flash_attn: bool = True): - with init_empty_weights(): - model = ESMC( - d_model=2560, - n_heads=40, - n_layers=80, - tokenizer=get_esmc_model_tokenizers(), - use_flash_attn=use_flash_attn, - ).eval() - load_torch_model(model, data_root("esmc-6b")) - model = model.to(device) - return model - - def ESM3_sm_open_v0(device: torch.device | str = "cpu"): with init_empty_weights(): model = ESM3( @@ -125,13 +85,65 @@ def ESM3_sm_open_v0(device: torch.device | str = "cpu"): return model +def _deprecated_esmc( + new_name: str, repo: str, device: torch.device | str, use_flash_attn: bool +) -> ESMC: + """Load an ESMC checkpoint behind one of the retired date-stamped names.""" + warnings.warn( + f"{new_name} is deprecated; load the model directly with " + f'EsmcForMaskedLM.from_pretrained("{repo}"). The weights now come from ' + f"the HuggingFace Hub rather than a local data root, so the first call " + f"downloads them.", + DeprecationWarning, + stacklevel=3, + ) + with warnings.catch_warnings(): + # The wrapper warns on its own; one deprecation per call is enough. + warnings.simplefilter("ignore", DeprecationWarning) + # dtype=None keeps fp32 on every device: the date-stamped factories + # loaded raw fp32 weights, and only ESMC.from_pretrained cast to bf16. + return ESMC( + model=EsmcForMaskedLM.from_pretrained( + repo, + device=device, + dtype=None, + attn_implementation="flash_attention_2" if use_flash_attn else "sdpa", + ) + ).eval() + + +def ESMC_300M_202412( + device: torch.device | str = "cpu", use_flash_attn: bool = True +) -> ESMC: + """Deprecated. Use ``EsmcForMaskedLM.from_pretrained(ESMC_300M_HF_REPO)``.""" + return _deprecated_esmc( + "ESMC_300M_202412", ESMC_300M_HF_REPO, device, use_flash_attn + ) + + +def ESMC_600M_202412( + device: torch.device | str = "cpu", use_flash_attn: bool = True +) -> ESMC: + """Deprecated. Use ``EsmcForMaskedLM.from_pretrained(ESMC_600M_HF_REPO)``.""" + return _deprecated_esmc( + "ESMC_600M_202412", ESMC_600M_HF_REPO, device, use_flash_attn + ) + + +def ESMC_6B_202412( + device: torch.device | str = "cpu", use_flash_attn: bool = True +) -> ESMC: + """Deprecated. Use ``EsmcForMaskedLM.from_pretrained(ESMC_6B_HF_REPO)``.""" + return _deprecated_esmc("ESMC_6B_202412", ESMC_6B_HF_REPO, device, use_flash_attn) + + LOCAL_MODEL_REGISTRY: dict[str, ModelBuilder] = { ESM3_OPEN_SMALL: ESM3_sm_open_v0, ESM3_STRUCTURE_ENCODER_V0: ESM3_structure_encoder_v0, ESM3_STRUCTURE_DECODER_V0: ESM3_structure_decoder_v0, ESM3_FUNCTION_DECODER_V0: ESM3_function_decoder_v0, - ESMC_600M: ESMC_600M_202412, ESMC_300M: ESMC_300M_202412, + ESMC_600M: ESMC_600M_202412, ESMC_6B: ESMC_6B_202412, } diff --git a/esm/tests/__init__.py b/esm/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/esm/tests/conftest.py b/esm/tests/conftest.py new file mode 100644 index 00000000..a241db2f --- /dev/null +++ b/esm/tests/conftest.py @@ -0,0 +1,235 @@ +"""Shared fixtures for the ESMC test suite. + +Tiny randomly-initialised models cover the structural contracts on CPU in +milliseconds. The published-weight fixtures download once per session and skip +when the Hub is unreachable, so an offline runner degrades to the structural +tests instead of failing. +""" + +import os +import random +from pathlib import Path + +import pytest +import torch + +REFERENCE_DIR = Path(__file__).parent / "reference" + +# Small enough to build instantly, wide enough that the head dim is sane. Every +# field the architecture depends on is present: `from_dict` defaults none of them. +TINY_ESMC = dict( + vocab_size=64, + hidden_size=32, + num_attention_heads=4, + num_hidden_layers=2, + pad_token_id=1, + mask_token_id=32, +) +TINY_SEQUENCES = ["MQIFVKTLTGKT", "MKV"] + +ESMC_300M_REPO = "biohub/ESMC-300M" +#: Subdirectory holding this model under ``$ESM_ESMC_WEIGHTS_ROOT``. +ESMC_300M_STAGED_DIR = "ESMC-300M" + + +# --------------------------------------------------------------------------- +# Sequences +# --------------------------------------------------------------------------- + +# Two real sequences, spliced into three length regimes. ESMC has no +# length-dependent branch other than the RoPE cache, so the regimes are picked +# where numerics and kernels change behaviour: +# +# short (10 aa, 12 tokens) Every softmax runs over a handful of keys. +# medium (129 aa, 131 tokens) Past ESMC-300M's head_dim (64), and 131 is prime, +# so it is not a multiple of any fused-attention tile (8/16/32/64/128) and +# every kernel has to handle a ragged tail. +# long (410 aa, 412 tokens) ~3x medium. Forces the RoPE cache to take its +# grow branch when a model sees medium first, and puts ~35x as many terms +# in each softmax denominator as short - which is where the CPU/GPU and +# fused/reference tolerances separate. +UBIQUITIN = ( + "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG" +) +LYSOZYME = ( + "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINSRWWCNDGRTPGSRNLC" + "NIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDVQAWIRGCRL" +) + +SHORT_SEQUENCE = UBIQUITIN[:10] +MEDIUM_SEQUENCE = LYSOZYME +# A tandem two-domain construct rather than a natural protein: length is the +# point, and repeating real domains keeps the tokens in-distribution so +# perplexity stays a meaningful number. +LONG_SEQUENCE = UBIQUITIN + LYSOZYME + UBIQUITIN + LYSOZYME + +# Ordered short -> long; tests parametrize over the keys. +SEQUENCES = {"short": SHORT_SEQUENCE, "medium": MEDIUM_SEQUENCE, "long": LONG_SEQUENCE} +LENGTH_NAMES = tuple(SEQUENCES) + +# A typo in either literal above would silently change every reference value. +assert [len(s) for s in SEQUENCES.values()] == [10, 129, 410] + +# Only the 20 canonical residues: X/B/U/Z/O and the gap, insertion and +# chain-break tokens are in the vocabulary but are not amino acids, and a random +# sequence containing them tests the vocabulary rather than the model. +CANONICAL_AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY" +RANDOM_SEQUENCE_SEED = 20241203 +# Spans the three regimes and the awkward edges: a single residue, an exact power +# of two and both its neighbours. +RANDOM_SEQUENCE_LENGTHS = (1, 2, 5, 31, 63, 64, 65, 127, 200, 410) + + +def random_amino_acid_sequences( + lengths: tuple[int, ...] | list[int] = RANDOM_SEQUENCE_LENGTHS, + seed: int = RANDOM_SEQUENCE_SEED, +) -> list[str]: + """Seeded random sequences over the canonical amino acids.""" + rng = random.Random(seed) + return ["".join(rng.choices(CANONICAL_AMINO_ACIDS, k=n)) for n in lengths] + + +@pytest.fixture(scope="session") +def random_sequences() -> list[str]: + return random_amino_acid_sequences() + + +# --------------------------------------------------------------------------- +# Published weights +# --------------------------------------------------------------------------- + + +def staged_model_dir(subdirectory: str) -> str | None: + """A local copy of a model, when ``ESM_ESMC_WEIGHTS_ROOT`` points at one. + + Read straight from the environment: this package is the bottom import layer. + """ + root = os.environ.get("ESM_ESMC_WEIGHTS_ROOT") + if not root: + return None + directory = Path(root) / subdirectory + return str(directory) if (directory / "config.json").exists() else None + + +def _skip_if_unreachable(repo: str) -> str: + """Resolve a Hub snapshot, or skip when the runner genuinely has no access. + + Only offline and transport failures are a skip. Anything else - a bad token, a + repo that has moved, a resolver bug - has to fail, or a regression there would + quietly turn the published-weight coverage green. + """ + from huggingface_hub.errors import LocalEntryNotFoundError, OfflineModeIsEnabled + from requests.exceptions import ConnectionError as RequestsConnectionError + from requests.exceptions import Timeout + + from esm.models.hub import resolve_model_dir + + try: + return resolve_model_dir(repo) + except ( + LocalEntryNotFoundError, + OfflineModeIsEnabled, + RequestsConnectionError, + Timeout, + ) as exc: + pytest.skip( + f"{repo} unreachable ({type(exc).__name__}); skipping" + ) # ty:ignore[too-many-positional-arguments] + + +#: Fixtures that resolve a real published checkpoint. Anything depending on one +#: needs network (or staged weights) and takes seconds rather than milliseconds. +WEIGHTS_FIXTURES = frozenset({"esmc_300m_dir"}) + + +def pytest_collection_modifyitems(items): + """Mark every published-weight test ``merge_only``. + + Derived from the fixture graph rather than written on each test, so a + weights-backed test cannot be added without the marker. + """ + for item in items: + if not WEIGHTS_FIXTURES.isdisjoint(getattr(item, "fixturenames", ())): + # `gpu` already routes to its own job; a second marker would only + # deselect it there. + if not item.get_closest_marker("gpu"): + item.add_marker(pytest.mark.merge_only) + + +@pytest.fixture(scope="session") +def tiny_esmc_config(): + from esm.models.esmc import EsmcConfig + + return EsmcConfig( + vocab_size=TINY_ESMC["vocab_size"], + hidden_size=TINY_ESMC["hidden_size"], + num_attention_heads=TINY_ESMC["num_attention_heads"], + num_hidden_layers=TINY_ESMC["num_hidden_layers"], + ) + + +@pytest.fixture +def tiny_esmc(tiny_esmc_config): + from esm.models.esmc import EsmcModel + + torch.manual_seed(0) + return EsmcModel(tiny_esmc_config).eval() + + +@pytest.fixture +def tiny_esmc_mlm(tiny_esmc_config): + from esm.models.esmc import EsmcForMaskedLM + + torch.manual_seed(0) + return EsmcForMaskedLM(tiny_esmc_config).eval() + + +@pytest.fixture(scope="session") +def esmc_tokenizer(): + from esm.models.esmc import EsmcTokenizer + + return EsmcTokenizer() + + +@pytest.fixture(scope="session") +def esmc_300m_dir(): + """A local copy of the weights when one is staged, else the published repo.""" + return staged_model_dir(ESMC_300M_STAGED_DIR) or _skip_if_unreachable( + ESMC_300M_REPO + ) + + +@pytest.fixture(scope="session") +def esmc_300m(esmc_300m_dir): + from esm.models.esmc import EsmcForMaskedLM + + return EsmcForMaskedLM.from_pretrained(esmc_300m_dir, device="cpu") + + +@pytest.fixture(scope="session") +def esmc_300m_cpu_bf16(esmc_300m_dir): + from esm.models.esmc import EsmcForMaskedLM + + return EsmcForMaskedLM.from_pretrained( + esmc_300m_dir, device="cpu", dtype=torch.bfloat16 + ) + + +@pytest.fixture(scope="session") +def esmc_300m_cpu_reference(esmc_300m, esmc_tokenizer): + """CPU / fp32 logits and hidden states for each length, computed once. + + The runtime counterpart of the checked-in reference file: that pins the + numbers across commits, this is the full tensor other configurations are + compared against rather than a stored summary. + """ + reference = {} + for name, sequence in SEQUENCES.items(): + enc = esmc_tokenizer(sequence, return_tensors="pt") + with torch.no_grad(): + out = esmc_300m(**enc) + reference[name] = { + "logits": out.logits.float(), + "last_hidden_state": out.last_hidden_state.float(), + } + return reference diff --git a/esm/tests/esmc_test.py b/esm/tests/esmc_test.py new file mode 100644 index 00000000..2f526ef4 --- /dev/null +++ b/esm/tests/esmc_test.py @@ -0,0 +1,1448 @@ +"""ESMC test suite. + +Structural tests run on CPU against tiny randomly-initialised models. Tests that +need published weights use ESMC-300M and skip when the Hub is unreachable. GPU +tests are marked ``gpu``. +""" + +import functools +import gzip +import json +import math +import pickle +import warnings +from dataclasses import replace +from pathlib import Path +from typing import NamedTuple + +import pytest +import torch +import torch.nn as nn +from safetensors.torch import save_file + +from esm.models.esmc import ( + ESMC, + EsmcConfig, + EsmcForMaskedLM, + EsmcForSequenceClassification, + EsmcMaskedLMOutput, + EsmcModel, + EsmcOutput, + EsmcTokenizer, +) +from esm.models.esmc import model as model_module +from esm.models.esmc.checkpoint_layout import native_to_published, published_to_native +from esm.models.esmc.kernels import FLASH_ATTN_INSTALLED, TE_INSTALLED +from esm.models.esmc.layers import ( + EsmcFlashMultiHeadAttention, + EsmcLayerNormLinear, + EsmcLayerNormMLP, + EsmcRotaryEmbedding, +) +from esm.models.hub import CONFIG_NAME, read_safetensors_dir +from esm.tests.conftest import ( + LENGTH_NAMES, + MEDIUM_SEQUENCE, + REFERENCE_DIR, + SEQUENCES, + SHORT_SEQUENCE, + TINY_ESMC, + TINY_SEQUENCES, + random_amino_acid_sequences, +) + +REFERENCE_FILE = REFERENCE_DIR / "esmc_300m_cpu_fp32.pkl.gz" + +# Each masked position costs one forward, so an uncapped pseudo-perplexity is +# O(L^2) tokens. 16 evenly spaced positions hold every length to a few seconds. +PPL_MAX_POSITIONS = 16 + +# Which layers ``build_esmc`` puts in the model. +FUSED_LAYERS = "fused" +REFERENCE_LAYERS = "reference" + + +def _tokens(tokenizer, sequences): + return tokenizer(sequences, return_tensors="pt", padding=True) + + +@functools.lru_cache(maxsize=1) +def reference_values() -> dict: + # A missing file is a broken checkout, not a reason to skip. + assert REFERENCE_FILE.exists(), ( + f"{REFERENCE_FILE} is missing. It is checked in; regenerate with " + f"`python -m esm.tests.regenerate_reference`." + ) + with gzip.open(REFERENCE_FILE, "rb") as f: + return pickle.load(f) + + +class PseudoPerplexity(NamedTuple): + perplexity: float + #: Fraction of masked residues the model's argmax recovers. + recovery: float + + +def masked_token_positions( + n_residues: int, limit: int = PPL_MAX_POSITIONS +) -> list[int]: + """Token indices to mask: every residue, or ``limit`` evenly spaced ones.""" + # Token 0 is and the last is , so residue i is token i + 1. + if n_residues <= limit: + return list(range(1, n_residues + 1)) + return [1 + (i * n_residues) // limit for i in range(limit)] + + +def pseudo_perplexity( + model, + tokenizer, + sequence: str, + *, + limit: int = PPL_MAX_POSITIONS, + batch_size: int = 8, +) -> PseudoPerplexity: + """Masked-LM pseudo-perplexity and top-1 recovery over the masked positions. + + Row ``i`` of a batch masks one residue, so one forward scores ``batch_size`` + positions. Rows are equal-length and never attend to each other, so batching + cannot change a row's result beyond reduction order. + """ + enc = tokenizer(sequence, return_tensors="pt") + ids = enc["input_ids"][0] + positions = masked_token_positions(len(sequence), limit) + + total_logprob = 0.0 + recovered = 0 + for start in range(0, len(positions), batch_size): + chunk = positions[start : start + batch_size] + batch = ids.unsqueeze(0).repeat(len(chunk), 1).to(model.device) + for row, position in enumerate(chunk): + batch[row, position] = model.config.mask_token_id + + with torch.no_grad(): + logits = model(input_ids=batch).logits + + for row, position in enumerate(chunk): + row_logits = logits[row, position].float().cpu() + target = int(ids[position]) + total_logprob += float(torch.log_softmax(row_logits, dim=-1)[target]) + recovered += int(int(row_logits.argmax()) == target) + + return PseudoPerplexity( + perplexity=math.exp(-total_logprob / len(positions)), + recovery=recovered / len(positions), + ) + + +def build_esmc( + directory, + *, + device: str, + dtype: torch.dtype | None = None, + attn: str = "sdpa", + layers: str = FUSED_LAYERS, + model_class=EsmcModel, +): + """Load ESMC with an explicit choice of fused vs pure-PyTorch layers. + + ``EsmcModel.__init__`` reads ``torch.get_default_device()`` to pick its + kernels, so building for the CPU and then moving is the only way to get the + reference layers onto CUDA. That also disables the flash-attention block, so + ``layers="reference"`` is only meaningful with ``attn="sdpa"`` / ``"eager"``. + """ + build_device = "cpu" if layers == REFERENCE_LAYERS else device + model = model_class.from_pretrained( + directory, device=build_device, attn_implementation=attn + ) + model = model.to(device) + if dtype is not None: + model = model.to(dtype) + return model.eval() + + +def module_output_shapes(model: nn.Module, **inputs) -> dict[str, tuple[int, ...]]: + """Shape of every module's (first) tensor output during one forward. + + Modules returning a dataclass are skipped; the caller asserts their fields. + """ + shapes: dict[str, tuple[int, ...]] = {} + handles = [] + + def record(name: str): + def hook(module, args, output): + tensor = output[0] if isinstance(output, tuple) else output + if isinstance(tensor, torch.Tensor): + shapes[name] = tuple(tensor.shape) + + return hook + + for name, module in model.named_modules(): + handles.append(module.register_forward_hook(record(name))) + try: + with torch.no_grad(): + model(**inputs) + finally: + for handle in handles: + handle.remove() + return shapes + + +def argmax_mismatches(logits: torch.Tensor, expected: list[int]) -> int: + """How many positions predict a different token than the reference.""" + predicted = logits.float().argmax(-1)[0].tolist() + assert len(predicted) == len(expected) + return sum(a != b for a, b in zip(predicted, expected)) + + +def drop_keys(directory: Path, predicate) -> list[str]: + """Rewrite a saved checkpoint without the keys matching ``predicate``.""" + from safetensors import safe_open + + with safe_open(str(directory / "model.safetensors"), framework="pt") as f: + tensors = {k: f.get_tensor(k) for k in f.keys()} + dropped = [k for k in tensors if predicate(k)] + assert dropped, f"predicate matched nothing among {sorted(tensors)[:5]}" + save_file( + {k: v.contiguous() for k, v in tensors.items() if k not in set(dropped)}, + str(directory / "model.safetensors"), + metadata={"format": "pt"}, + ) + return dropped + + +def safetensors_keys(directory) -> list[str]: + from safetensors import safe_open + + with safe_open(str(directory / "model.safetensors"), framework="pt") as f: + return list(f.keys()) + + +def read_config(directory: Path) -> dict: + return json.loads((directory / CONFIG_NAME).read_text()) + + +def write_config(directory: Path, raw: dict) -> None: + (directory / CONFIG_NAME).write_text(json.dumps(raw)) + + +# --------------------------------------------------------------------------- +# Construction and the forward shape contract +# --------------------------------------------------------------------------- + + +def test_construction(tiny_esmc, tiny_esmc_mlm, tiny_esmc_config): + config = tiny_esmc_config + assert len(tiny_esmc.transformer.blocks) == config.num_hidden_layers + assert tiny_esmc.embed.weight.shape == (config.vocab_size, config.hidden_size) + for block in tiny_esmc.transformer.blocks: + assert block.attn.n_heads == config.num_attention_heads + + assert isinstance(tiny_esmc_mlm.esmc, EsmcModel) + assert tiny_esmc_mlm.lm_head[-1].out_features == config.vocab_size + + +def test_forward_shape_contract(tiny_esmc, tiny_esmc_mlm, esmc_tokenizer): + enc = _tokens(esmc_tokenizer, TINY_SEQUENCES) + b, length = enc["input_ids"].shape + + with torch.no_grad(): + out = tiny_esmc(**enc, output_hidden_states=True, output_attentions=True) + mlm = tiny_esmc_mlm(**enc, output_hidden_states=True) + + assert isinstance(out, EsmcOutput) + assert out.last_hidden_state.shape == (b, length, TINY_ESMC["hidden_size"]) + # One entry per block input plus the final post-LayerNorm output. + assert out.hidden_states.shape == ( + TINY_ESMC["num_hidden_layers"] + 1, + b, + length, + TINY_ESMC["hidden_size"], + ) + assert len(out.attentions) == TINY_ESMC["num_hidden_layers"] + + assert isinstance(mlm, EsmcMaskedLMOutput) + assert mlm.logits.shape == (b, length, TINY_ESMC["vocab_size"]) + assert mlm.last_hidden_state.shape == (b, length, TINY_ESMC["hidden_size"]) + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_every_module_output_has_the_expected_shape( + tiny_esmc_mlm, esmc_tokenizer, tiny_esmc_config, name +): + """Pin the width of every intermediate, not just the model's output. + + A wrong QKV fan-out or FFN width still produces a correctly-shaped final + tensor. + """ + config = tiny_esmc_config + hidden = config.hidden_size + heads = config.num_attention_heads + enc = esmc_tokenizer(SEQUENCES[name], return_tensors="pt") + b, length = enc["input_ids"].shape + + expected: dict[str, tuple[int, ...]] = { + "esmc.embed": (b, length, hidden), + "esmc.transformer": (b, length, hidden), + "esmc.transformer.norm": (b, length, hidden), + # Linear -> GELU -> LayerNorm -> Linear, so only the last widens. + "lm_head": (b, length, config.vocab_size), + "lm_head.0": (b, length, hidden), + "lm_head.1": (b, length, hidden), + "lm_head.2": (b, length, hidden), + "lm_head.3": (b, length, config.vocab_size), + } + for layer in range(config.num_hidden_layers): + block = f"esmc.transformer.blocks.{layer}" + expected.update( + { + block: (b, length, hidden), + f"{block}.attn": (b, length, hidden), + # One fused projection for Q, K and V. + f"{block}.attn.layernorm_qkv": (b, length, 3 * hidden), + f"{block}.attn.q_ln": (b, length, hidden), + f"{block}.attn.k_ln": (b, length, hidden), + # RoPE runs on the unflattened heads. + f"{block}.attn.rotary": (b, length, heads, config.head_dim), + f"{block}.attn.out_proj": (b, length, hidden), + f"{block}.ffn": (b, length, hidden), + } + ) + + # Equality, not containment: a new sub-module has to be added here too. + assert module_output_shapes(tiny_esmc_mlm, **enc) == expected + + # EsmcLayerNormMLP has no child modules, so the SwiGLU widths are not + # observable through a hook. fc1 is doubled for the gate and value halves. + ffn = tiny_esmc_mlm.esmc.transformer.blocks[0].ffn + assert ffn.fc1_weight.shape == (2 * config.intermediate_size, hidden) + assert ffn.fc2_weight.shape == (hidden, config.intermediate_size) + + with torch.no_grad(): + out = tiny_esmc_mlm(**enc, output_attentions=True) + for attention in out.attentions: + assert attention.shape == (b, heads, length, length) + + +def test_gradients_reach_every_parameter(tiny_esmc_config, esmc_tokenizer): + """Fine-tuning is a supported use, so the backward pass has to work.""" + torch.manual_seed(0) + model = EsmcForMaskedLM(tiny_esmc_config).train() + batch = _tokens(esmc_tokenizer, TINY_SEQUENCES) + + model(**batch).logits.square().mean().backward() + + gradients = {name: p.grad for name, p in model.named_parameters()} + assert not [name for name, grad in gradients.items() if grad is None] + # A dead parameter has a gradient of exactly zero everywhere. + assert not [ + name + for name, grad in gradients.items() + if grad is not None and (not torch.isfinite(grad).all() or grad.eq(0).all()) + ] + + +def test_sequence_classification_head(tiny_esmc_config, esmc_tokenizer): + """The head reads ````, so it must emit one row per sequence.""" + torch.manual_seed(0) + config = replace(tiny_esmc_config, num_labels=3) + model = EsmcForSequenceClassification(config).eval() + batch = _tokens(esmc_tokenizer, TINY_SEQUENCES) + labels = torch.tensor([0, 2]) + + with torch.no_grad(): + out = model(**batch, labels=labels) + + assert out.logits.shape == (len(TINY_SEQUENCES), 3) + assert out.loss is not None and torch.isfinite(out.loss) + # A label-independent head would score both rows identically. + assert not torch.allclose(out.logits[0], out.logits[1]) + + +# --------------------------------------------------------------------------- +# Config round-trip +# --------------------------------------------------------------------------- + + +def test_config_round_trip(tmp_path, tiny_esmc_config): + """Every architectural field survives ``save_pretrained`` / ``from_pretrained``. + + Goes through the real ``save_pretrained``: ``to_dict`` deliberately drops the + run-scoped fields, so writing the dataclass verbatim would assert they + persist when in fact they must not. + """ + tiny_esmc_config.save_pretrained(tmp_path) + reloaded = EsmcConfig.from_pretrained(tmp_path) + + for key in tiny_esmc_config.to_dict(): + if key == "model_type": + continue + assert getattr(reloaded, key) == getattr(tiny_esmc_config, key), key + + +def test_attn_implementation_is_not_persisted(tmp_path, tiny_esmc_config): + """It describes the run, not the weights, so it must not survive a save.""" + replace(tiny_esmc_config, attn_implementation="eager").save_pretrained(tmp_path) + + written = read_config(tmp_path) + assert "attn_implementation" not in written + default = EsmcConfig(hidden_size=1, num_attention_heads=1).attn_implementation + assert EsmcConfig.from_pretrained(tmp_path).attn_implementation == default + + +def test_config_tolerates_unknown_fields(tmp_path): + """A newer published config.json must not break an older reader. + + ``EsmcConfig.from_pretrained`` reads known fields by name, so unknown ones + are ignored rather than retained as attributes. + """ + write_config(tmp_path, {**TINY_ESMC, "some_future_field": 123}) + cfg = EsmcConfig.from_pretrained(tmp_path) + assert cfg.hidden_size == TINY_ESMC["hidden_size"] + assert not hasattr(cfg, "some_future_field") + + +# --------------------------------------------------------------------------- +# Tokenizer +# --------------------------------------------------------------------------- + + +def test_tokenizer_round_trips_the_whole_alphabet(esmc_tokenizer): + canonical = "ACDEFGHIKLMNPQRSTVWY" + # X/B/U/Z/O and the gap and insertion tokens are real vocabulary entries. + extended = canonical + "XBUZO.-" + + ids = esmc_tokenizer(extended)["input_ids"] + assert ids == esmc_tokenizer(extended)["input_ids"] + assert esmc_tokenizer.unk_token_id not in ids + assert len(ids) == len(extended) + 2 + assert esmc_tokenizer.decode(torch.tensor(ids[1:-1])).replace(" ", "") == extended + + # Lowercase is not in the vocabulary and silently becomes . + assert set(esmc_tokenizer(canonical.lower())["input_ids"][1:-1]) == { + esmc_tokenizer.unk_token_id + } + + +def test_tokenizer_pads_to_longest(esmc_tokenizer): + enc = _tokens(esmc_tokenizer, TINY_SEQUENCES) + longest = max(len(s) for s in TINY_SEQUENCES) + 2 + + assert enc["input_ids"].shape == (len(TINY_SEQUENCES), longest) + for row, sequence in enumerate(TINY_SEQUENCES): + keep = len(sequence) + 2 + assert enc["attention_mask"][row].tolist() == [1] * keep + [0] * ( + longest - keep + ) + assert (enc["input_ids"][row, keep:] == esmc_tokenizer.pad_token_id).all() + + +def test_the_chain_break_token_survives_tokenisation(esmc_tokenizer): + """``|`` separates chains within one sequence, so it has to reach the model.""" + ids = esmc_tokenizer("AAAA|MKV", return_tensors="pt")["input_ids"][0] + break_id = esmc_tokenizer.convert_tokens_to_ids("|") + + assert break_id is not None + assert (ids == break_id).sum() == 1 + # Position 0 is , so the break lands after the first chain's 4 residues. + assert int((ids == break_id).nonzero()[0, 0]) == 5 + + +# --------------------------------------------------------------------------- +# Determinism and the RoPE cache, the only length-dependent state in the model +# --------------------------------------------------------------------------- + + +def test_eval_forward_is_deterministic(tiny_esmc, esmc_tokenizer): + enc = _tokens(esmc_tokenizer, TINY_SEQUENCES) + with torch.no_grad(): + a = tiny_esmc(**enc).last_hidden_state + b = tiny_esmc(**enc).last_hidden_state + assert torch.equal(a, b) + + +def _fill_rope_cache(rope: EsmcRotaryEmbedding, length: int): + q = torch.zeros(1, length, 1, rope.dim) + rope(q, q) + assert rope._cos_cached is not None and rope._sin_cached is not None + return rope._cos_cached.clone(), rope._sin_cached.clone() + + +def test_rope_cache_grows_to_the_longest_sequence_seen(tiny_esmc_config): + """Growing the table must extend it, not re-derive it on a new grid. + + ``_update_cos_sin_cache`` rebuilds only when asked for more positions, so the + rows it already holds have to stay correct for every shorter sequence. + """ + rope = EsmcRotaryEmbedding(tiny_esmc_config.head_dim) + assert rope._seq_len_cached == 0 + + short_cos, short_sin = _fill_rope_cache(rope, 12) + assert rope._seq_len_cached == 12 + + long_cos, long_sin = _fill_rope_cache(rope, 412) + assert rope._seq_len_cached == 412 + assert torch.equal(long_cos[:12], short_cos) + assert torch.equal(long_sin[:12], short_sin) + + # A shorter sequence reuses the long table rather than shrinking it. + _fill_rope_cache(rope, 12) + assert rope._seq_len_cached == 412 + + +@pytest.mark.parametrize( + "warmup,measured", [("short", "long"), ("long", "short"), ("long", "medium")] +) +def test_a_forward_does_not_depend_on_the_length_run_before_it( + tiny_esmc, tiny_esmc_config, esmc_tokenizer, warmup, measured +): + """Compared against a freshly-built model, not against itself, so a stale or + wrongly-extended RoPE cache cannot hide by being wrong both times. + """ + enc = esmc_tokenizer(SEQUENCES[measured], return_tensors="pt") + with torch.no_grad(): + tiny_esmc(**esmc_tokenizer(SEQUENCES[warmup], return_tensors="pt")) + warm = tiny_esmc(**enc).last_hidden_state + + # Same seed as the ``tiny_esmc`` fixture, so the weights are identical. + torch.manual_seed(0) + cold = EsmcModel(tiny_esmc_config).eval() + with torch.no_grad(): + expected = cold(**enc).last_hidden_state + assert torch.equal(warm, expected) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_cpu_dtypes(tiny_esmc_config, esmc_tokenizer, dtype, name): + torch.manual_seed(0) + model = EsmcModel(tiny_esmc_config).eval().to(dtype) + # Batched against a short sequence so every length also carries padding. + enc = _tokens(esmc_tokenizer, [SEQUENCES[name], SHORT_SEQUENCE]) + with torch.no_grad(): + out = model(**enc).last_hidden_state + assert out.dtype == dtype + assert torch.isfinite(out.float()).all() + + +# --------------------------------------------------------------------------- +# Published-weight loading and the key policy +# --------------------------------------------------------------------------- + + +def _checkpoint_state(directory) -> dict[str, torch.Tensor]: + """The published checkpoint, in the model's own (fused) tensor layout.""" + raw = published_to_native(read_safetensors_dir(directory)) + return {k: v for k, v in raw.items() if not k.endswith("_extra_state")} + + +@pytest.mark.parametrize("model_class", [EsmcForMaskedLM, EsmcModel]) +def test_every_state_dict_entry_comes_from_the_checkpoint(esmc_300m_dir, model_class): + """No parameter may keep the value ``init_empty_weights`` left behind. + + Bit-equality against the file on disk is the strongest form of this: a key the + checkpoint never covered holds whatever ``to_empty`` found in memory. + ``EsmcModel`` drops ``lm_head.*`` and strips the ``esmc.`` prefix, which is + what the key-set comparison encodes. + """ + model = model_class.from_pretrained(esmc_300m_dir, device="cpu") + checkpoint = _checkpoint_state(esmc_300m_dir) + if model_class is EsmcModel: + checkpoint = { + k.removeprefix("esmc."): v + for k, v in checkpoint.items() + if k.startswith("esmc.") + } + + state = model.state_dict() + assert not [k for k, t in state.items() if t.is_meta] + assert set(state) == set(checkpoint) + unpopulated = [k for k, t in state.items() if not torch.equal(t, checkpoint[k])] + assert not unpopulated, unpopulated + + +def test_rope_buffers_are_materialised_after_loading(esmc_300m): + """The one thing no checkpoint can cover. + + ``inv_freq`` is registered non-persistently, so it is absent from both the + checkpoint and ``state_dict``; it still has to come off the meta device. + """ + rotaries = [m for m in esmc_300m.modules() if isinstance(m, EsmcRotaryEmbedding)] + assert len(rotaries) == esmc_300m.config.num_hidden_layers + for rope in rotaries: + assert not rope.inv_freq.is_meta + assert torch.isfinite(rope.inv_freq).all() + + +def test_hidden_state_and_attention_shapes_on_real_weights(esmc_300m, esmc_tokenizer): + """The per-layer contract at the published width, not just the tiny one. + + Short sequence only: the attention tensor is + ``n_layers x batch x heads x L x L``. + """ + config = esmc_300m.config + enc = esmc_tokenizer(SHORT_SEQUENCE, return_tensors="pt") + length = enc["input_ids"].shape[1] + with torch.no_grad(): + out = esmc_300m(**enc, output_hidden_states=True, output_attentions=True) + + assert out.logits.shape == (1, length, config.vocab_size) + assert out.last_hidden_state.shape == (1, length, config.hidden_size) + assert out.last_hidden_state_prenorm.shape == (1, length, config.hidden_size) + assert out.hidden_states.shape == ( + config.num_hidden_layers + 1, + 1, + length, + config.hidden_size, + ) + assert len(out.attentions) == config.num_hidden_layers + for attention in out.attentions: + assert attention.shape == (1, config.num_attention_heads, length, length) + # Every row is a distribution over the keys it is allowed to see. + torch.testing.assert_close( + attention.sum(-1), torch.ones_like(attention.sum(-1)), atol=1e-5, rtol=0 + ) + + +# --------------------------------------------------------------------------- +# Config reading rules and the published tensor layout +# --------------------------------------------------------------------------- + + +def test_reads_pre_alignment_config(tmp_path): + """Checkpoints published before the field-name alignment must still resolve.""" + write_config( + tmp_path, + { + "vocab_size": 64, + "d_model": 64, + "n_heads": 4, + "n_layers": 3, + "pad_token_id": 1, + "mask_token_id": 32, + }, + ) + with pytest.warns(FutureWarning, match="pre-alignment"): + cfg = EsmcConfig.from_pretrained(tmp_path) + assert (cfg.hidden_size, cfg.num_attention_heads, cfg.num_hidden_layers) == ( + 64, + 4, + 3, + ) + + +@pytest.mark.parametrize( + "missing", ["vocab_size", "hidden_size", "num_attention_heads", "num_hidden_layers"] +) +def test_missing_architecture_field_raises(tmp_path, missing): + """Defaulting an architecture field would build the wrong model silently.""" + raw = {**TINY_ESMC} + del raw[missing] + write_config(tmp_path, raw) + with pytest.raises(KeyError, match=missing): + EsmcConfig.from_pretrained(tmp_path) + + +def test_rejects_inconsistent_intermediate_size(): + with pytest.raises(ValueError, match="intermediate_size"): + EsmcConfig(hidden_size=32, num_attention_heads=4, intermediate_size=999) + + +def test_honours_attn_implementation_from_config(tmp_path, tiny_esmc_config): + """It gates the flash-attention block, so a config that sets it must be read.""" + EsmcForMaskedLM(tiny_esmc_config).save_pretrained(tmp_path) + raw = read_config(tmp_path) + raw["attn_implementation"] = "eager" + write_config(tmp_path, raw) + + assert EsmcConfig.from_pretrained(tmp_path).attn_implementation == "eager" + loaded = EsmcForMaskedLM.from_pretrained(tmp_path, device="cpu") + explicit = EsmcForMaskedLM.from_pretrained( + tmp_path, device="cpu", attn_implementation="sdpa" + ) + assert loaded.config.attn_implementation == "eager" + # An explicit argument still wins. + assert explicit.config.attn_implementation == "sdpa" + + +def test_published_layout_round_trip_is_bit_identical(tiny_esmc_config): + """Fusing on load and splitting on save must compose to the identity. + + A fused weight is the row-wise concatenation of its parts, so any drift means + a checkpoint written by one implementation is read wrongly by the other. + """ + model = EsmcForMaskedLM(tiny_esmc_config) + native = model.state_dict() + published = native_to_published(native) + + # The published layout is one tensor per projection. + for suffix in ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj"): + assert any(suffix in k for k in published), suffix + assert not any("layernorm_qkv" in k or "fc1_weight" in k for k in published) + assert any(k.startswith("esmc.layers.0.") for k in published) + + back = published_to_native(published) + assert set(back) == set(native) + assert all(torch.equal(back[k], native[k]) for k in native) + + +def test_save_pretrained_round_trips_through_published_layout( + tmp_path, tiny_esmc_config, esmc_tokenizer +): + """A full save -> load cycle must not move the logits at all.""" + model = EsmcForMaskedLM(tiny_esmc_config).eval() + batch = _tokens(esmc_tokenizer, TINY_SEQUENCES) + with torch.no_grad(): + before = model(**batch).logits + + model.save_pretrained(tmp_path) + saved = set(safetensors_keys(tmp_path)) + assert any("self_attn.q_proj" in k for k in saved) + assert not any("layernorm_qkv" in k for k in saved) + + reloaded = EsmcForMaskedLM.from_pretrained(tmp_path, device="cpu") + with torch.no_grad(): + after = reloaded(**batch).logits + assert torch.equal(before, after) + + +def test_an_incomplete_checkpoint_is_refused(tmp_path, tiny_esmc_config): + """Strict loading is what turns an uncovered parameter into an error. + + Drop one block and the load must refuse rather than silently keep whatever + ``to_empty`` left in memory. + """ + EsmcForMaskedLM(tiny_esmc_config).save_pretrained(tmp_path) + drop_keys(tmp_path, lambda k: ".1." in k and ("layers" in k or "blocks" in k)) + + with pytest.raises(RuntimeError, match="missing keys"): + EsmcForMaskedLM.from_pretrained(tmp_path, device="cpu") + + +def test_a_fresh_classification_head_is_the_only_excused_gap( + tmp_path, tiny_esmc_config +): + """``_keys_to_ignore_on_load_missing`` must excuse the head and nothing else. + + Loading a backbone checkpoint into a classification model is the intended + fine-tuning entry point, so an absent ``classifier.*`` cannot raise. Both + sides of that hole are pinned: the head may be missing, a missing backbone + tensor must still refuse. + """ + EsmcForMaskedLM(tiny_esmc_config).save_pretrained(tmp_path) + + model = EsmcForSequenceClassification.from_pretrained(tmp_path, device="cpu") + assert not any(p.is_meta for p in model.classifier.parameters()) + assert all(torch.isfinite(p).all() for p in model.classifier.parameters()) + + drop_keys(tmp_path, lambda k: k.endswith("embed_tokens.weight")) + with pytest.raises(RuntimeError, match="missing keys"): + EsmcForSequenceClassification.from_pretrained(tmp_path, device="cpu") + + +# --------------------------------------------------------------------------- +# Numerical reference values on ESMC-300M (CPU, fp32) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_reference_logits(esmc_300m_cpu_reference, name): + case = reference_values()["cases"][name] + logits = esmc_300m_cpu_reference[name]["logits"] + hidden = esmc_300m_cpu_reference[name]["last_hidden_state"] + + assert list(logits.shape) == case["logits_shape"] + # Bit-identical across reruns on one host; the tolerance leaves room for a + # different CPU BLAS. + torch.testing.assert_close( + logits[0, 1, : len(case["first_residue_logits"])], + torch.tensor(case["first_residue_logits"]), + atol=1e-3, + rtol=1e-3, + ) + assert logits.argmax(-1)[0].tolist() == case["argmax"] + # Per-position L2 norms of the post-LayerNorm output: O(1) in magnitude, so + # one tolerance is meaningful at every dtype, and position-sensitive in a way + # the argmax over a 64-token vocabulary is not. + torch.testing.assert_close( + hidden[0].norm(dim=-1), + torch.tensor(case["last_hidden_state_norms"]), + atol=1e-3, + rtol=1e-4, + ) + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_reference_pseudo_perplexity(esmc_300m, esmc_tokenizer, name): + """The model's own loss: the cheapest number that moves when anything in the + stack changes numerically *and* has a scale worth reading. + """ + case = reference_values()["cases"][name] + measured = pseudo_perplexity(esmc_300m, esmc_tokenizer, SEQUENCES[name]) + assert measured.perplexity == pytest.approx(case["pseudo_perplexity"], rel=1e-3) + # Integer count over a fixed position set, so this is exact. + assert measured.recovery == case["top1_recovery"] + + if name == "long": + # The long fixture is a tandem repeat of two domains, so every masked + # residue has an exact copy ~205 positions away and full recovery is only + # reachable if attention spans the whole sequence. Asserted as a value, + # not against the stored reference: attention confined to a local window + # would still produce a stable reference. + assert measured.recovery == 1.0 + assert measured.perplexity < 1.1 + + +@pytest.mark.parametrize("name", ["medium", "long"]) +def test_real_sequences_score_better_than_random_ones(esmc_300m, esmc_tokenizer, name): + """The check a stored reference cannot make: the numbers have to be *good*. + + A model with its weights loaded into the wrong slots still produces stable + logits and a reproducible perplexity - it just cannot tell a real protein + from a uniformly random one. The short sequence is excluded: 10 residues + carry too little context for the gap to mean anything. + """ + sequence = SEQUENCES[name] + scrambled = random_amino_acid_sequences([len(sequence)], seed=7)[0] + real = pseudo_perplexity(esmc_300m, esmc_tokenizer, sequence, limit=8) + noise = pseudo_perplexity(esmc_300m, esmc_tokenizer, scrambled, limit=8) + + # Observed on ESMC-300M over 8 masked positions: real ppl 2.26 / recovery + # 0.75 vs random 32.6 / 0.00 (medium), 1.02 / 1.00 vs 24.0 / 0.00 (long). + # The thresholds sit an order of magnitude inside those margins. + assert real.perplexity < 0.5 * noise.perplexity + assert real.recovery >= 0.5 + assert noise.recovery <= 0.25 + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_cpu_bf16_reaches_the_fp32_reference(esmc_300m_cpu_bf16, esmc_tokenizer, name): + """bf16 cannot reproduce the reference logits, so pin what it must preserve: + the decision the model makes and its loss, not the logits themselves. + """ + case = reference_values()["cases"][name] + enc = esmc_tokenizer(SEQUENCES[name], return_tensors="pt") + with torch.no_grad(): + out = esmc_300m_cpu_bf16(**enc) + + # bf16 keeps 8 mantissa bits, so a near-tie between two residues can flip; + # exact equality is not available. Observed 0 / 2 / 0 mismatches. + assert argmax_mismatches(out.logits, case["argmax"]) <= 3 + # Observed max|Δ| on the hidden-state norms (~1.1-1.5): under 0.009. + torch.testing.assert_close( + out.last_hidden_state.float()[0].norm(dim=-1), + torch.tensor(case["last_hidden_state_norms"]), + atol=0.03, + rtol=0, + ) + measured = pseudo_perplexity(esmc_300m_cpu_bf16, esmc_tokenizer, SEQUENCES[name]) + # Observed relative shift vs the fp32 reference: under 6e-3. + assert measured.perplexity == pytest.approx(case["pseudo_perplexity"], rel=2e-2) + + +# --------------------------------------------------------------------------- +# Masking schemes: padded, unpadded, multi-chain, and -1 as padding +# --------------------------------------------------------------------------- + + +def test_pad_tokens_cannot_influence_real_positions(tiny_esmc, esmc_tokenizer): + """Overwriting the padding must not move a single valid output.""" + batch = _tokens(esmc_tokenizer, ["MQIFVKTLTGKT", "MKV"]) + keep = batch["attention_mask"].bool() + + with torch.no_grad(): + before = tiny_esmc(**batch).last_hidden_state + + scrambled = batch["input_ids"].clone() + scrambled[~keep] = 7 # any residue token, where padding used to be + with torch.no_grad(): + after = tiny_esmc( + input_ids=scrambled, attention_mask=batch["attention_mask"] + ).last_hidden_state + + for row in range(keep.shape[0]): + length = int(keep[row].sum()) + torch.testing.assert_close( + before[row, :length], after[row, :length], atol=0, rtol=0 + ) + + +def test_padding_attends_to_the_real_tokens(tiny_esmc, esmc_tokenizer): + """The one assertion that separates a key-axis mask from a query==key one. + + Padding hides a key from every query, so a pad *query* still attends to the + real tokens and its output moves when they change. + """ + batch = _tokens(esmc_tokenizer, ["MQIFVKTLTGKT", "MKV"]) + keep = batch["attention_mask"].bool() + padded_row = next( + row for row in range(keep.shape[0]) if int(keep[row].sum()) < keep.shape[1] + ) + tail = slice(int(keep[padded_row].sum()), None) + + changed = batch["input_ids"].clone() + changed[padded_row, 1] = 7 if changed[padded_row, 1] != 7 else 8 + + with torch.no_grad(): + before = tiny_esmc(**batch).last_hidden_state[padded_row, tail] + after = tiny_esmc( + input_ids=changed, attention_mask=batch["attention_mask"] + ).last_hidden_state[padded_row, tail] + + assert not torch.allclose(before, after, atol=1e-6) + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_padded_batch_matches_solo_runs_at_every_length( + tiny_esmc, esmc_tokenizer, name +): + """Padding a batch must not perturb the unpadded rows. + + The most direct check that attention masking is wired correctly, held at + every length regime: the padded tail here is hundreds of tokens long. + """ + sequences = [SEQUENCES[name], SHORT_SEQUENCE, MEDIUM_SEQUENCE[:7]] + batch = _tokens(esmc_tokenizer, sequences) + keep = batch["attention_mask"].bool() + + with torch.no_grad(): + padded = tiny_esmc(**batch).last_hidden_state + for row, sequence in enumerate(sequences): + alone = tiny_esmc( + **esmc_tokenizer(sequence, return_tensors="pt") + ).last_hidden_state[0] + length = int(keep[row].sum()) + # Not exact: the padded batch takes SDPA's masked path and the solo + # run its unmasked one, so the reduction orders differ. + torch.testing.assert_close( + padded[row, :length], alone[:length], atol=1e-5, rtol=1e-4 + ) + + +def test_padded_batch_matches_solo_runs_across_random_lengths( + tiny_esmc, esmc_tokenizer, random_sequences +): + """Same invariant over ten seeded lengths, including a single residue. + + A one-residue sequence in a 410-wide batch is the extreme case: 411 of its + 412 key positions are padding. + """ + batch = _tokens(esmc_tokenizer, random_sequences) + keep = batch["attention_mask"].bool() + + with torch.no_grad(): + padded = tiny_esmc(**batch).last_hidden_state + for row, sequence in enumerate(random_sequences): + alone = tiny_esmc( + **esmc_tokenizer(sequence, return_tensors="pt") + ).last_hidden_state[0] + length = int(keep[row].sum()) + torch.testing.assert_close( + padded[row, :length], alone[:length], atol=1e-5, rtol=1e-4 + ) + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_the_unmasked_fast_path_agrees_with_the_masked_one( + tiny_esmc, esmc_tokenizer, name +): + """Equal-length rows carry no padding, so ``forward`` drops the mask entirely + and the fused kernels can take over. That fast path has to agree with the + masked one, reached here by handing the same batch a single-chain + ``sequence_id`` - every position attends everywhere either way. + """ + sequence = SEQUENCES[name] + batch = esmc_tokenizer([sequence, sequence[::-1]], return_tensors="pt") + assert batch["attention_mask"].bool().all() + one_chain = torch.zeros_like(batch["input_ids"]) + + with torch.no_grad(): + unmasked = tiny_esmc(**batch).last_hidden_state + masked = tiny_esmc( + input_ids=batch["input_ids"], sequence_id=one_chain + ).last_hidden_state + + # Observed exactly equal, but which kernel SDPA picks is not contractual. + torch.testing.assert_close(unmasked, masked, atol=1e-5, rtol=1e-4) + + +@pytest.mark.parametrize("name", ["medium", "long"]) +def test_multi_chain_sequence_id_isolates_chains(tiny_esmc, esmc_tokenizer, name): + """Chains must not see each other, at lengths where each chain is real. + + Not compared against running each chain alone: RoPE positions are absolute, + so a chain's representation legitimately depends on where it sits in the + concatenated input. What must hold is that perturbing one chain leaves the + others exactly where they were. + """ + ids = esmc_tokenizer(SEQUENCES[name], return_tensors="pt")["input_ids"] + length = ids.shape[1] + bounds = [0, length // 3, 2 * length // 3, length] + sequence_id = torch.zeros(1, length, dtype=torch.long) + for chain, start in enumerate(bounds[:-1]): + sequence_id[:, start : bounds[chain + 1]] = chain + + middle = slice(bounds[1], bounds[2]) + changed = ids.clone() + changed[:, middle] = 7 # any residue token, uniformly + + with torch.no_grad(): + before = tiny_esmc(input_ids=ids, sequence_id=sequence_id).last_hidden_state + after = tiny_esmc(input_ids=changed, sequence_id=sequence_id).last_hidden_state + + outside = torch.ones(length, dtype=torch.bool) + outside[middle] = False + # Bit-identical: masked attention weights are exactly zero, so the middle + # chain contributes 0 * v to every other chain's context. + assert torch.equal(before[0, outside], after[0, outside]) + # And the perturbation has to do something inside its own chain, or this + # would pass with attention switched off entirely. + assert not torch.allclose(before[0, middle], after[0, middle]) + + +def test_chain_break_and_chain_ids_describe_the_same_split(tiny_esmc, esmc_tokenizer): + """Ties the two ways of expressing multiple chains together: the ``|`` token + marks the boundary, the ``sequence_id`` enforces it. + """ + ids = esmc_tokenizer("AAAA|MKV", return_tensors="pt")["input_ids"] + break_at = int( + (ids[0] == esmc_tokenizer.convert_tokens_to_ids("|")).nonzero()[0, 0] + ) + sequence_id = torch.zeros_like(ids) + sequence_id[:, break_at + 1 :] = 1 + + with torch.no_grad(): + isolated = tiny_esmc( + input_ids=ids, sequence_id=sequence_id + ).last_hidden_state.clone() + # Perturbing the second chain must not move the first. + perturbed = ids.clone() + perturbed[0, -2] = 7 if perturbed[0, -2] != 7 else 8 + after = tiny_esmc( + input_ids=perturbed, sequence_id=sequence_id + ).last_hidden_state + + torch.testing.assert_close( + isolated[:, : break_at + 1], after[:, : break_at + 1], atol=0, rtol=0 + ) + assert not torch.allclose(isolated[:, break_at + 1 :], after[:, break_at + 1 :]) + + +@pytest.mark.gpu +def test_multi_chain_under_flash_attention_2_raises(tiny_esmc_config, esmc_tokenizer): + """The varlen flash path cannot express a block-diagonal chain mask, and the + model documents a ValueError for it; a silent wrong answer would be worse. + """ + if not FLASH_ATTN_INSTALLED: + pytest.skip( + "flash-attn is not installed" + ) # ty:ignore[too-many-positional-arguments] + + with torch.device("cuda"): + model = EsmcForMaskedLM( + replace(tiny_esmc_config, attn_implementation="flash_attention_2") + ).eval() + assert model.esmc._use_flash_attn, "flash attention did not actually dispatch" + + ids = esmc_tokenizer("MQIFVKTLTGKT", return_tensors="pt")["input_ids"].cuda() + sequence_id = torch.zeros_like(ids) + sequence_id[:, ids.shape[1] // 2 :] = 1 + + with pytest.raises(ValueError, match="Multi-chain"): + model(input_ids=ids, sequence_id=sequence_id) + + +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_sequence_id_of_minus_one_masks_padding_like_attention_mask( + tiny_esmc, esmc_tokenizer, name +): + """``sequence_id == -1`` and ``attention_mask == 0`` hide the same keys. + + They do not build the same mask: ``sequence_id`` compares query against key, + so a *pad* query ends up attending only to other padding, while an + ``attention_mask`` hides padding from every query. Only the valid positions + can be required to agree - which is the half that feeds anything downstream. + """ + batch = _tokens(esmc_tokenizer, [SEQUENCES[name], SHORT_SEQUENCE]) + keep = batch["attention_mask"].bool() + sequence_id = torch.where(keep, 0, -1) + + with torch.no_grad(): + via_mask = tiny_esmc(**batch).last_hidden_state + via_ids = tiny_esmc( + input_ids=batch["input_ids"], sequence_id=sequence_id + ).last_hidden_state + + torch.testing.assert_close(via_mask[keep], via_ids[keep], atol=1e-5, rtol=1e-4) + + +def test_attention_mask_folds_into_a_multi_chain_sequence_id(tiny_esmc, esmc_tokenizer): + """Both arguments at once: the mask is folded in, not dropped. + + A caller with chains *and* ragged lengths supplies a ``sequence_id`` that + describes only the chains, leaving padding to ``attention_mask``. Folding it + in has to equal marking the padding ``-1`` by hand. + """ + batch = _tokens(esmc_tokenizer, [MEDIUM_SEQUENCE, SHORT_SEQUENCE]) + keep = batch["attention_mask"].bool() + assert not keep.all(), "fixture must actually pad for this to test anything" + + length = batch["input_ids"].shape[1] + chains = torch.zeros_like(batch["input_ids"]) + chains[:, length // 2 :] = 1 + folded_by_hand = torch.where(keep, chains, -1) + + with torch.no_grad(): + folded = tiny_esmc( + input_ids=batch["input_ids"], + attention_mask=batch["attention_mask"], + sequence_id=chains, + ).last_hidden_state + by_hand = tiny_esmc( + input_ids=batch["input_ids"], sequence_id=folded_by_hand + ).last_hidden_state + + torch.testing.assert_close(folded[keep], by_hand[keep], atol=0, rtol=0) + + # And the fold must not have cost the chain isolation it was folded into. + changed = batch["input_ids"].clone() + tail = length // 2 + changed[:, tail] = 7 if int(changed[0, tail]) != 7 else 8 + with torch.no_grad(): + after = tiny_esmc( + input_ids=changed, + attention_mask=batch["attention_mask"], + sequence_id=chains, + ).last_hidden_state + head = keep.clone() + head[:, tail:] = False + torch.testing.assert_close(folded[head], after[head], atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# Kernel selection follows the target device +# --------------------------------------------------------------------------- + + +def _pretend_installed(monkeypatch) -> None: + """Patch where the flags are read, not where they are defined.""" + monkeypatch.setattr(model_module, "TE_INSTALLED", True) + monkeypatch.setattr(model_module, "FLASH_ATTN_INSTALLED", True) + + +@pytest.mark.parametrize("attn", ["sdpa", "flash_attention_2"]) +def test_cpu_builds_never_get_cuda_only_kernels( + monkeypatch, tiny_esmc_config, esmc_tokenizer, attn +): + """Every direct constructor has to reach CPU-safe layers *and* return usable + numbers, with the CUDA-only kernels pretending to be importable. + """ + _pretend_installed(monkeypatch) + config = replace(tiny_esmc_config, attn_implementation=attn) + batch = _tokens(esmc_tokenizer, TINY_SEQUENCES) + + bare = EsmcModel(config).eval() + mlm = EsmcForMaskedLM(config).eval() + for encoder in (bare, mlm.esmc): + block = encoder.transformer.blocks[0] + assert not encoder._use_flash_attn + assert not isinstance(block.attn, EsmcFlashMultiHeadAttention) + assert isinstance( + block.attn.layernorm_qkv, EsmcLayerNormLinear + ) # ty:ignore[unresolved-attribute] + assert isinstance(block.ffn, EsmcLayerNormMLP) + + with torch.no_grad(): + assert torch.isfinite(bare(**batch).last_hidden_state).all() + assert torch.isfinite(mlm(**batch).logits).all() + + +def test_the_deprecated_constructor_runs_on_cpu(monkeypatch, esmc_tokenizer): + _pretend_installed(monkeypatch) + batch = _tokens(esmc_tokenizer, TINY_SEQUENCES) + + with pytest.warns(DeprecationWarning): + legacy = ESMC(d_model=32, n_heads=4, n_layers=2).eval() + with torch.no_grad(): + out = legacy(sequence_tokens=batch["input_ids"]) + assert torch.isfinite(out.sequence_logits).all() + + +def test_eager_and_sdpa_build_the_same_attention(tiny_esmc_config): + """``attn_implementation`` only selects the flash block, so "sdpa" and "eager" + produce an identical module tree and comparing their outputs proves nothing. + Pinned structurally instead, so a future divergence has to be deliberate. + """ + builds = [ + EsmcModel(replace(tiny_esmc_config, attn_implementation=attn)) + for attn in ("sdpa", "eager") + ] + first, second = (model.transformer.blocks[0].attn for model in builds) + assert type(first) is type(second) + assert not builds[0]._use_flash_attn and not builds[1]._use_flash_attn + + +# --------------------------------------------------------------------------- +# GPU +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_state_dict_round_trips_under_transformer_engine(tiny_esmc_config): + """TE adds ``_extra_state`` to ``state_dict`` and a hook strips it, so the + loader has to be told not to expect it back.""" + if not TE_INSTALLED: + pytest.skip( + "TransformerEngine unavailable" + ) # ty:ignore[too-many-positional-arguments] + with torch.device("cuda"): + model = EsmcForMaskedLM(tiny_esmc_config).eval() + + assert not [k for k in model.state_dict() if k.endswith("_extra_state")] + model.load_state_dict(model.state_dict()) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("name", LENGTH_NAMES) +def test_gpu_dtypes(tiny_esmc_config, esmc_tokenizer, dtype, name): + torch.manual_seed(0) + model = EsmcModel(tiny_esmc_config).eval().to("cuda", dtype) + # Batched against a short sequence so every length also carries padding. + enc = { + k: v.cuda() + for k, v in _tokens(esmc_tokenizer, [SEQUENCES[name], SHORT_SEQUENCE]).items() + } + with torch.no_grad(): + out = model(**enc).last_hidden_state + assert out.dtype == dtype + assert torch.isfinite(out.float()).all() + + +@pytest.mark.gpu +def test_fp8_backbone(esmc_300m_dir): + """ESMC is commonly run fp8 as a folding backbone.""" + if not TE_INSTALLED: + pytest.skip( + "TransformerEngine unavailable" + ) # ty:ignore[too-many-positional-arguments] + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip( + "fp8 needs compute capability >= 9.0 (Hopper)" + ) # ty:ignore[too-many-positional-arguments] + + from esm.models.esmfold2.model import _convert_te_modules_to_fp8_inplace + + model = EsmcModel.from_pretrained( + esmc_300m_dir, device="cuda", dtype=torch.bfloat16 + ) + with torch.no_grad(): + _convert_te_modules_to_fp8_inplace(model) + tokenizer = EsmcTokenizer() + enc = { + k: v.cuda() for k, v in tokenizer(SHORT_SEQUENCE, return_tensors="pt").items() + } + with torch.no_grad(): + out = model(**enc).last_hidden_state + assert torch.isfinite(out.float()).all() + + +@pytest.mark.gpu +def test_flash_attention_2_dispatches_and_agrees_with_sdpa( + esmc_300m_dir, esmc_tokenizer +): + """The only backend that is a distinct build, so the only one worth comparing. + + Dispatch is asserted first: ``from_pretrained`` silently falls back to sdpa + when flash-attn is missing, so without this the test could pass having + compared sdpa against itself. + """ + if not FLASH_ATTN_INSTALLED: + pytest.skip( + "flash-attn is not installed" + ) # ty:ignore[too-many-positional-arguments] + + reference = EsmcModel.from_pretrained( + esmc_300m_dir, device="cuda", dtype=torch.bfloat16, attn_implementation="sdpa" + ) + candidate = EsmcModel.from_pretrained( + esmc_300m_dir, + device="cuda", + dtype=torch.bfloat16, + attn_implementation="flash_attention_2", + ) + + assert candidate._use_flash_attn, "flash_attention_2 silently fell back" + assert isinstance(candidate.transformer.blocks[0].attn, EsmcFlashMultiHeadAttention) + assert not isinstance( + reference.transformer.blocks[0].attn, EsmcFlashMultiHeadAttention + ) + + for sequence in SEQUENCES.values(): + enc = { + k: v.cuda() + for k, v in esmc_tokenizer(sequence, return_tensors="pt").items() + } + with torch.no_grad(): + expected = reference(**enc).last_hidden_state.float() + actual = candidate(**enc).last_hidden_state.float() + # The flash path packs the batch and runs flash_attn_varlen where sdpa + # runs xformers, so the reduction orders differ and the gap compounds + # over 30 blocks. Held to a few bf16 ULP at activation scale. + torch.testing.assert_close( + actual, expected, atol=5 * torch.finfo(torch.bfloat16).eps, rtol=1e-2 + ) + + +@pytest.mark.gpu +def test_cpu_inference_works_on_a_cuda_machine(esmc_300m_dir, esmc_tokenizer): + """Transformer Engine is selected whenever it imports and CUDA exists, but its + modules reject CPU inputs, so a CPU-targeted load has to opt out of them. + """ + model = EsmcModel.from_pretrained(esmc_300m_dir, device="cpu") + enc = esmc_tokenizer(SHORT_SEQUENCE, return_tensors="pt") + with torch.no_grad(): + out = model(**enc).last_hidden_state + assert out.device.type == "cpu" + assert torch.isfinite(out).all() + + +@pytest.mark.gpu +def test_cpu_and_gpu_agree(esmc_300m_dir, esmc_tokenizer): + """Same weights, same input, both devices, fp32 on each side. + + Also a cross-implementation check: the CPU model uses the reference layers + while the CUDA model uses the fused Transformer Engine ones, which is where + almost all of the difference below comes from. + """ + on_cpu = EsmcModel.from_pretrained(esmc_300m_dir, device="cpu") + on_gpu = EsmcModel.from_pretrained(esmc_300m_dir, device="cuda") + + for sequence in SEQUENCES.values(): + enc = esmc_tokenizer(sequence, return_tensors="pt") + with torch.no_grad(): + expected = on_cpu(**enc).last_hidden_state.float() + actual = on_gpu(**{k: v.cuda() for k, v in enc.items()}) + actual = actual.last_hidden_state.float().cpu() + + # Observed max|Δ| at fp32, activations ~O(1): 2.4e-4 short to 1.3e-3 + # long. It grows with length because the difference is a LayerNorm + # reduction repeated per block over a lengthening residual stream. + torch.testing.assert_close(actual, expected, atol=3e-3, rtol=1e-4) + + +# Every build reachable on a GPU: dtype x attention backend x fused/reference +# layers. flash_attention_2 is only paired with fp16 / bf16 because the kernel +# rejects fp32, and only with the fused layers because a CPU-built model never +# gets the flash block at all (see ``build_esmc``). +GPU_BUILDS = [ + (torch.float32, "sdpa", FUSED_LAYERS), + (torch.float32, "sdpa", REFERENCE_LAYERS), + (torch.float32, "eager", FUSED_LAYERS), + (torch.bfloat16, "sdpa", FUSED_LAYERS), + (torch.bfloat16, "sdpa", REFERENCE_LAYERS), + (torch.bfloat16, "flash_attention_2", FUSED_LAYERS), + (torch.float16, "sdpa", FUSED_LAYERS), + (torch.float16, "flash_attention_2", FUSED_LAYERS), +] + +# What each dtype can be held to against the CPU / fp32 reference: max|Δ| on the +# per-position hidden-state norms (~1.1-1.5) and how many positions' argmax may +# differ. Worst observed across every build and length in GPU_BUILDS: +# max|Δ| norms argmax mismatches (of 12 / 131 / 412) +# fp32 4.5e-4 0 / 0 / 0 +# bf16 1.46e-2 1 / 3 / 3 +# fp16 1.62e-3 0 / 1 / 0 +# fp32 with the reference layers reproduces the CPU reference to 1.7e-6; the fp32 +# bound has to accommodate Transformer Engine. +REFERENCE_NORM_TOLERANCE = { + torch.float32: 1e-3, + torch.bfloat16: 0.03, + torch.float16: 5e-3, +} +REFERENCE_ARGMAX_MISMATCHES = {torch.float32: 0, torch.bfloat16: 5, torch.float16: 2} + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype,attn,layers", GPU_BUILDS) +def test_every_gpu_build_reaches_the_reference( + esmc_300m_dir, esmc_tokenizer, dtype, attn, layers +): + """Correct, not merely self-consistent: every GPU build vs the CPU reference. + + No GPU build can reach it bit-for-bit, so what is asserted is the decision + path (argmax over the vocabulary) plus a bound on the hidden-state norms. + Without Transformer Engine or flash-attn installed the builds that would use + them degenerate to the reference ones rather than failing, so this stays a + superset of whatever a given runner can reach. + """ + try: + model = build_esmc( + esmc_300m_dir, + device="cuda", + dtype=dtype, + attn=attn, + layers=layers, + model_class=EsmcForMaskedLM, + ) + # Only a genuinely absent package is a reason to skip. `from_pretrained` + # signals a strict-key failure with RuntimeError, so catching that here would + # turn a checkpoint regression into a green skip across every build. + except ImportError as exc: + pytest.skip( + f"build {dtype}/{attn}/{layers} unavailable: {exc}" + ) # ty:ignore[too-many-positional-arguments] + + for name, sequence in SEQUENCES.items(): + case = reference_values()["cases"][name] + enc = { + k: v.cuda() + for k, v in esmc_tokenizer(sequence, return_tensors="pt").items() + } + with torch.no_grad(): + out = model(**enc) + mismatches = argmax_mismatches(out.logits, case["argmax"]) + assert mismatches <= REFERENCE_ARGMAX_MISMATCHES[dtype], (name, mismatches) + torch.testing.assert_close( + out.last_hidden_state.float()[0].norm(dim=-1).cpu(), + torch.tensor(case["last_hidden_state_norms"]), + atol=REFERENCE_NORM_TOLERANCE[dtype], + rtol=0, + msg=lambda text: f"{name}: {text}", + ) + + +# --------------------------------------------------------------------------- +# The retired date-stamped factories still do real work +# --------------------------------------------------------------------------- + + +def test_date_stamped_factory_loads_real_weights(esmc_300m_dir, esmc_300m): + """``esm.pretrained.ESMC_300M_202412`` shipped publicly, so it has to keep + loading a usable model and agree with the model loaded directly. + """ + from esm import pretrained + + with pytest.warns(DeprecationWarning): + legacy = pretrained.ESMC_300M_202412(use_flash_attn=False) + + assert isinstance(legacy, ESMC) + # The date-stamped factories loaded raw fp32 weights on every device; only + # ESMC.from_pretrained ever cast to bf16. + assert next(legacy.parameters()).dtype is torch.float32 + assert not legacy.training + + tokens = legacy._tokenize(["MQIFVKTLTGKTITLEVEPS"]) + with torch.no_grad(): + legacy_out = legacy(sequence_tokens=tokens) + direct_out = esmc_300m(input_ids=tokens, return_dict=True) + + # Same weights, so the logits must be identical, not merely close. + assert torch.equal(legacy_out.sequence_logits, direct_out.logits) + # 300M has 30 blocks, and the old stack held one entry per block output. + assert legacy_out.hidden_states.shape[0] == 30 + assert legacy_out.embeddings.shape == direct_out.last_hidden_state.shape + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + via_registry = pretrained.load_local_model("esmc_300m", use_flash_attn=False) + assert isinstance(via_registry, ESMC) diff --git a/esm/tests/reference/esmc_300m_cpu_fp32.pkl.gz b/esm/tests/reference/esmc_300m_cpu_fp32.pkl.gz new file mode 100644 index 00000000..e2eaebc3 Binary files /dev/null and b/esm/tests/reference/esmc_300m_cpu_fp32.pkl.gz differ diff --git a/esm/tests/regenerate_reference.py b/esm/tests/regenerate_reference.py new file mode 100644 index 00000000..6feb81fe --- /dev/null +++ b/esm/tests/regenerate_reference.py @@ -0,0 +1,79 @@ +"""Regenerate the numerical reference values used by ``esmc_test.py``. + +Run from the repo root when a deliberate numerical change lands:: + + python -m esm.tests.regenerate_reference + +CPU + fp32, so the file reproduces anywhere; every other build configuration is +checked against it rather than against a reference of its own. +""" + +import gzip +import pickle + +import torch + +from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer +from esm.tests.conftest import ( + ESMC_300M_REPO, + ESMC_300M_STAGED_DIR, + REFERENCE_DIR, + SEQUENCES, + staged_model_dir, +) +from esm.tests.esmc_test import PPL_MAX_POSITIONS, REFERENCE_FILE, pseudo_perplexity + +# One residue's logit row plus the argmax path moves on any real numerical +# change, so there is no need to store the full logits tensor. +N_LOGITS = 16 + + +def main() -> None: + torch.manual_seed(0) + source = staged_model_dir(ESMC_300M_STAGED_DIR) or ESMC_300M_REPO + model = EsmcForMaskedLM.from_pretrained(source, device="cpu") + tokenizer = EsmcTokenizer() + + cases = {} + for name, sequence in SEQUENCES.items(): + enc = tokenizer(sequence, return_tensors="pt") + with torch.no_grad(): + out = model(**enc) + logits = out.logits.float() + measured = pseudo_perplexity(model, tokenizer, sequence) + cases[name] = { + "sequence": sequence, + "n_residues": len(sequence), + "logits_shape": list(logits.shape), + "first_residue_logits": logits[0, 1, :N_LOGITS].tolist(), + "argmax": logits.argmax(-1)[0].tolist(), + # Per-position L2 norms of the post-LayerNorm output: O(1) in + # magnitude, so one tolerance is meaningful at every dtype. + "last_hidden_state_norms": out.last_hidden_state.float()[0] + .norm(dim=-1) + .tolist(), + "pseudo_perplexity": measured.perplexity, + "top1_recovery": measured.recovery, + } + print( + f" {name:<7} L={len(sequence):<4} " + f"ppl={measured.perplexity:.6f} recovery={measured.recovery:.4f}" + ) + + payload = { + "repo": ESMC_300M_REPO, + "device": "cpu", + "dtype": "float32", + "n_logits": N_LOGITS, + "ppl_max_positions": PPL_MAX_POSITIONS, + "cases": cases, + } + + REFERENCE_DIR.mkdir(parents=True, exist_ok=True) + with gzip.open(REFERENCE_FILE, "wb", compresslevel=9) as f: + pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL) + print(f"wrote {REFERENCE_FILE}") + + +if __name__ == "__main__": + main() diff --git a/esm/utils/structure/input_builder.py b/esm/utils/structure/input_builder.py index ad5aef4e..a5716441 100644 --- a/esm/utils/structure/input_builder.py +++ b/esm/utils/structure/input_builder.py @@ -80,6 +80,44 @@ class StructurePredictionInput: covalent_bonds: list[CovalentBond] | None = None +def _entity_key(item: ProteinInput | RNAInput | DNAInput | LigandInput) -> tuple: + """Canonical identity of a chain's chemical entity, ignoring chain ids. + + Chains collapse to one entity only if they are chemically identical, so + modifications are part of the key: a modified chain and its unmodified twin + are separate entities, and merging them would make them ``sym_id`` copies of + each other. + + Entity identity is chemical, not per-instance. Two copies of one entity can + therefore still differ in token count -- a covalently bonded copy of a CCD + ligand drops its leaving atoms while a free copy keeps them -- so "same + entity" does not imply "same length". + """ + if isinstance(item, LigandInput): + # ccd wins over smiles at tokenization, so a redundant smiles alongside a + # ccd must not split one entity into two. + if item.ccd: + return ("NONPOLYMER", None, tuple(item.ccd)) + return ("NONPOLYMER", item.smiles, ()) + if isinstance(item, ProteinInput): + entity_type = "PROTEIN" + elif isinstance(item, RNAInput): + entity_type = "RNA" + elif isinstance(item, DNAInput): + entity_type = "DNA" + else: # pragma: no cover - the annotation excludes unsupported inputs + raise TypeError(f"Unsupported sequence input type: {type(item)}") + # `msa` is excluded on purpose: it is stored per chain name, so two copies of + # one entity may legitimately carry different MSAs. + return ( + entity_type, + item.sequence, + frozenset( + (mod.position, mod.ccd, mod.smiles) for mod in item.modifications or [] + ), + ) + + def serialize_structure_prediction_input(all_atom_input: StructurePredictionInput): def create_chain_data(seq_input, chain_type: str) -> dict[str, Any]: chain_data: dict[str, Any] = { diff --git a/esm/utils/structure/protein_chain.py b/esm/utils/structure/protein_chain.py index a5023a01..83b0435d 100644 --- a/esm/utils/structure/protein_chain.py +++ b/esm/utils/structure/protein_chain.py @@ -1048,6 +1048,19 @@ def from_pdb( atom_array = PDBFile.read(path).get_structure( model=1, extra_fields=["b_factor"] ) + return cls._from_atomarray( + atom_array, id=file_id, chain_id=chain_id, is_predicted=is_predicted + ) + + @classmethod + def _from_atomarray( + cls, + atom_array: bs.AtomArray, + id: str, + chain_id: str = "detect", + is_predicted: bool = False, + ) -> "ProteinChain": + """Shared body of :meth:`from_pdb` and :meth:`from_atomarray`.""" if chain_id == "detect": chain_id = atom_array.chain_id[0] atom_array = atom_array[ @@ -1076,9 +1089,6 @@ def from_pdb( confidence = np.ones([num_res], dtype=np.float32) for i, res in enumerate(bs.residue_iter(atom_array)): - chain = atom_array[atom_array.chain_id == chain_id] - assert isinstance(chain, bs.AtomArray) - res_index = res[0].res_id residue_index[i] = res_index insertion_code[i] = res[0].ins_code @@ -1101,7 +1111,7 @@ def from_pdb( assert all(sequence), "Some residue name was not specified correctly" return cls( - id=file_id, + id=id, sequence=sequence, chain_id=chain_id, entity_id=entity_id, @@ -1151,16 +1161,14 @@ def from_atomarray( cls, atom_array: bs.AtomArray, id: str | None = None, is_predicted: bool = False ) -> "ProteinChain": """A simple converter from bs.AtomArray -> ProteinChain. - Uses PDB file format as intermediate.""" - atom_array = atom_array.copy() - atom_array.box = None # remove surrounding box, from_pdb won't handle this - pdb_file = PDBFile() - pdb_file.set_structure(atom_array) - buf = io.StringIO() - pdb_file.write(buf) - buf.seek(0) - return cls.from_pdb(buf, id=id, is_predicted=is_predicted) + Must NOT round-trip through PDB text: with a blank chain id the fixed-width + columns push the thousands digit of ``res_id`` into the chain-id field, so + residues numbered 1000 and above are silently dropped. + """ + return cls._from_atomarray( + atom_array, id=id if id is not None else "null", is_predicted=is_predicted + ) def get_normalization_frame(self) -> Affine3D: """Given a set of coordinates, compute a single frame. diff --git a/pyproject.toml b/pyproject.toml index c4d32c19..84e2b1b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,8 +55,13 @@ addopts = """ --cov=esm --cov-report term-missing:skip-covered -n auto + -m 'not gpu' --ignore=tests/oss_pytests/test_oss_client.py """ +markers = [ + "gpu: requires a CUDA device", + "merge_only: needs a real published checkpoint (network or staged weights)", +] [tool.setuptools] package-dir = {"" = "."} @@ -67,7 +72,7 @@ where = ["."] include = ["esm*"] [tool.setuptools.package-data] -esm = ["data/*"] +esm = ["data/*", "tests/reference/*.pkl.gz"] [tool.pixi.workspace] channels = ["conda-forge"] @@ -160,13 +165,19 @@ unused-ignore-comment = "ignore" unused-type-ignore-comment = "ignore" [tool.ty.analysis] -# ty can't introspect optional/compiled deps absent from this lint env (flash_attn is -# GPU-only; zstd is a C-extension), so it falsely reports them unresolved. Treat as Any +# ty can't introspect optional compiled deps absent from this CPU lint environment, so +# it falsely reports them unresolved. Treat them as Any # rather than disabling unresolved-import, so genuinely-broken first-party imports still # surface. replace-imports-with-any = [ "flash_attn", "flash_attn.**", + "transformer_engine", + "transformer_engine.**", + "xformers", + "xformers.**", + "cuequivariance_torch", + "cuequivariance_torch.**", "zstd", ]