From 2cd0f3366c39033fb115d2b0d7373033d519623a Mon Sep 17 00:00:00 2001 From: Jarrid Rector-Brooks Date: Wed, 13 May 2026 15:12:20 +0000 Subject: [PATCH 1/5] Add code to check if we're on the newer ESM version and remap the erroring rotary embedding keys --- runner/inference.py | 67 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 66 insertions(+), 1 deletion(-) diff --git a/runner/inference.py b/runner/inference.py index 5c4b606..0ea603f 100644 --- a/runner/inference.py +++ b/runner/inference.py @@ -16,6 +16,7 @@ import copy import logging import os +import re import traceback from collections import OrderedDict from collections.abc import Mapping @@ -51,6 +52,69 @@ logger = logging.getLogger(__name__) +# 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.+\.)encoder\.layer\.(?P\d+)" + r"\.attention\.self\.rotary_embeddings\.inv_freq$" +) + + +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] + assert other_inv_freq.shape == first_inv_freq.shape and torch.equal( + other_inv_freq, first_inv_freq + ), ( + 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 seq_tnsr_to_str(pred_tnsr: torch.Tensor) -> str: """Converts a residue index tensor to a one-letter amino acid sequence string. @@ -189,8 +253,9 @@ def load_checkpoint(self) -> None: } current = self.model.state_dict() - filtered = OrderedDict() + checkpoint["model"] = _remap_esm_rotary_inv_freq(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 From 0e4cc5e970c2833862f27788f30584187d71dd0f Mon Sep 17 00:00:00 2001 From: Jarrid Rector-Brooks Date: Wed, 13 May 2026 11:47:16 -0400 Subject: [PATCH 2/5] Minor --- runner/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runner/inference.py b/runner/inference.py index 0ea603f..c67d2cf 100644 --- a/runner/inference.py +++ b/runner/inference.py @@ -171,7 +171,7 @@ 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" + #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}") From 2ba9924ac4b465f82a3f6c0b9470f49b6dc5e7fc Mon Sep 17 00:00:00 2001 From: Jarrid Rector-Brooks Date: Wed, 13 May 2026 12:28:31 -0400 Subject: [PATCH 3/5] Add things for backwards compat --- runner/inference.py | 49 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/runner/inference.py b/runner/inference.py index c67d2cf..211d3d6 100644 --- a/runner/inference.py +++ b/runner/inference.py @@ -18,6 +18,8 @@ import os import re import traceback +import warnings + from collections import OrderedDict from collections.abc import Mapping from contextlib import nullcontext @@ -51,6 +53,13 @@ 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". @@ -59,6 +68,12 @@ 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"^.+\.embeddings\.position_embeddings\.weight$" +) + def _remap_esm_rotary_inv_freq( checkpoint_state: Mapping[str, torch.Tensor], @@ -115,6 +130,37 @@ def _remap_esm_rotary_inv_freq( 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 + remapped[key] = tensor.detach().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. @@ -254,6 +300,9 @@ def load_checkpoint(self) -> None: current = self.model.state_dict() 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(): From c9b55dd0fea6a3ec2d001f0647b925558b4ed01f Mon Sep 17 00:00:00 2001 From: Jarrid Rector-Brooks Date: Wed, 13 May 2026 14:02:31 -0400 Subject: [PATCH 4/5] Some small fixes for cleanliness --- runner/inference.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/runner/inference.py b/runner/inference.py index 211d3d6..9fb19e8 100644 --- a/runner/inference.py +++ b/runner/inference.py @@ -71,14 +71,14 @@ # 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"^.+\.embeddings\.position_embeddings\.weight$" + 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]": +) -> OrderedDict[str, torch.Tensor]: """Promotes per-layer ESM rotary inv_freq buffers to the shared key. Newer transformers releases (the EsmModel rewrite introduced around @@ -115,13 +115,14 @@ def _remap_esm_rotary_inv_freq( first_inv_freq = checkpoint_state[first_layer_key] for _, other_key in layer_keys[1:]: other_inv_freq = checkpoint_state[other_key] - assert other_inv_freq.shape == first_inv_freq.shape and torch.equal( + if other_inv_freq.shape != first_inv_freq.shape or not torch.equal( other_inv_freq, first_inv_freq - ), ( - f"Rotary inv_freq mismatch under prefix '{prefix}': " - f"'{other_key}' differs from '{first_layer_key}'; cannot safely " - "promote to a single shared buffer." - ) + ): + 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: @@ -133,10 +134,10 @@ def _remap_esm_rotary_inv_freq( def _fill_unused_esm_position_embeddings( checkpoint_state: Mapping[str, torch.Tensor], current_state: Mapping[str, torch.Tensor], -) -> "OrderedDict[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 , + 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 @@ -156,7 +157,17 @@ def _fill_unused_esm_position_embeddings( continue if key in remapped: continue - remapped[key] = tensor.detach().clone() + # 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 @@ -217,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}") From 4dbd128c380caaf6ac2be7b9790dbcc593b557da Mon Sep 17 00:00:00 2001 From: Jarrid Rector-Brooks Date: Wed, 13 May 2026 17:31:33 -0400 Subject: [PATCH 5/5] Minor --- runner/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runner/inference.py b/runner/inference.py index 9fb19e8..8607e77 100644 --- a/runner/inference.py +++ b/runner/inference.py @@ -64,7 +64,7 @@ # 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.+\.)encoder\.layer\.(?P\d+)" + r"^(?P.+\.esm\.)encoder\.layer\.(?P\d+)" r"\.attention\.self\.rotary_embeddings\.inv_freq$" )