diff --git a/runner/inference.py b/runner/inference.py index 5c4b606..8607e77 100644 --- a/runner/inference.py +++ b/runner/inference.py @@ -16,7 +16,10 @@ import copy import logging import os +import re import traceback +import warnings + from collections import OrderedDict from collections.abc import Mapping from contextlib import nullcontext @@ -50,6 +53,124 @@ logger = logging.getLogger(__name__) +warnings.filterwarnings( + "ignore", + message="Using a non-tuple sequence for multidimensional indexing.*", + category=UserWarning, + module="openfold.model.triangular_multiplicative_update", +) + + +# Matches per-layer ESM/DPLM rotary inv_freq keys from old checkpoints, e.g. +# "lm_module.lm_model.model.esm.encoder.layer.23.attention.self.rotary_embeddings.inv_freq". +_ESM_PER_LAYER_INV_FREQ_RE = re.compile( + r"^(?P.+\.esm\.)encoder\.layer\.(?P\d+)" + r"\.attention\.self\.rotary_embeddings\.inv_freq$" +) + +# Matches ESM embedding position_embeddings keys, e.g. +# "lm_module.lm_model.model.esm.embeddings.position_embeddings.weight". +_ESM_POSITION_EMBEDDINGS_RE = re.compile( + r"^.+\.esm\.embeddings\.position_embeddings\.weight$" +) + + +def _remap_esm_rotary_inv_freq( + checkpoint_state: Mapping[str, torch.Tensor], + current_state: Mapping[str, torch.Tensor], +) -> OrderedDict[str, torch.Tensor]: + """Promotes per-layer ESM rotary inv_freq buffers to the shared key. + + Newer transformers releases (the EsmModel rewrite introduced around + transformers 5.x) store a single shared rotary inv_freq buffer at + ``rotary_embeddings.inv_freq`` instead of one per attention + layer. Old DISCO checkpoints carry the per-layer copies, which are + identical across layers, so the first one can safely be promoted. + + The remap is only applied when the current model actually expects the + new shared key, which makes the fix a no-op on older transformers + versions where the per-layer buffers still exist. + """ + remapped: OrderedDict[str, torch.Tensor] = OrderedDict(checkpoint_state) + + grouped: dict[str, list[tuple[int, str]]] = {} + for key in checkpoint_state: + match = _ESM_PER_LAYER_INV_FREQ_RE.match(key) + if match: + grouped.setdefault(match.group("prefix"), []).append( + (int(match.group("layer")), key) + ) + + for prefix, layer_keys in grouped.items(): + new_key = f"{prefix}rotary_embeddings.inv_freq" + if new_key not in current_state: + # Current model still uses the per-layer layout; leave untouched. + continue + if new_key in remapped: + continue + + layer_keys.sort(key=lambda lk: lk[0]) + _, first_layer_key = layer_keys[0] + + first_inv_freq = checkpoint_state[first_layer_key] + for _, other_key in layer_keys[1:]: + other_inv_freq = checkpoint_state[other_key] + if other_inv_freq.shape != first_inv_freq.shape or not torch.equal( + other_inv_freq, first_inv_freq + ): + raise RuntimeError( + f"Rotary inv_freq mismatch under prefix '{prefix}': " + f"'{other_key}' differs from '{first_layer_key}'; cannot safely " + "promote to a single shared buffer." + ) + + remapped[new_key] = first_inv_freq + for _, key in layer_keys: + remapped.pop(key, None) + + return remapped + + +def _fill_unused_esm_position_embeddings( + checkpoint_state: Mapping[str, torch.Tensor], + current_state: Mapping[str, torch.Tensor], +) -> OrderedDict[str, torch.Tensor]: + """Fills missing ESM ``position_embeddings.weight`` keys from current state. + + On older transformers versions like 4.50.0, + ``EsmEmbeddings.position_embeddings`` is unconditionally + instantiated as an ``nn.Embedding``, despite the weight being unused when + the model is configured with ``position_embedding_type="rotary"``. The + original DISCO checkpoint was saved against an even older transformers + release where this layer was conditional on the absolute-position branch, + so the key is absent from the checkpoint and ``strict=True`` loads fail + with a missing-key error. + + The remap is only applied when the checkpoint is missing the key but the + current model has one, so it is a no-op for checkpoints that do carry an + absolute-position embedding. + """ + remapped: OrderedDict[str, torch.Tensor] = OrderedDict(checkpoint_state) + + for key, tensor in current_state.items(): + if not _ESM_POSITION_EMBEDDINGS_RE.match(key): + continue + if key in remapped: + continue + # The fill assumes position_embedding_type="rotary", in which case the + # weight is unused at inference. Warn so that a future operator running + # this against an absolute-position ESM model notices the silent fill. + logger.warning( + "Filling missing ESM position embedding '%s' from the freshly " + "initialized model weights. This is only safe when the ESM module " + "is configured with position_embedding_type='rotary'; an absolute " + "configuration would receive random init values.", + key, + ) + remapped[key] = tensor.clone() + + return remapped + def seq_tnsr_to_str(pred_tnsr: torch.Tensor) -> str: """Converts a residue index tensor to a one-letter amino acid sequence string. @@ -107,7 +228,6 @@ def init_env(self) -> None: self.fabric.launch() self.device = self.fabric.device torch.cuda.set_device(self.device) - os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0,8.9" if self.configs.use_deepspeed_evo_attention: env = os.getenv("CUTLASS_PATH", None) self.print(f"env: {env}") @@ -189,8 +309,12 @@ def load_checkpoint(self) -> None: } current = self.model.state_dict() - filtered = OrderedDict() + checkpoint["model"] = _remap_esm_rotary_inv_freq(checkpoint["model"], current) + checkpoint["model"] = _fill_unused_esm_position_embeddings( + checkpoint["model"], current + ) + filtered = OrderedDict() for k, v in checkpoint["model"].items(): if k in current and v.shape == current[k].shape: filtered[k] = v # → OK: same name & same shape