Bug/issue 5 rotary emb mismatch

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<prefix>.+\.esm\.)encoder\.layer\.(?P<layer>\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
+    ``<prefix>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