Port ESMC and ESMFold2 to Transformers (#46419)

* Port ESMC + ESMFold2 model code from fork onto v5 main (baseline, unadapted)

Moves the purely-additive model code from the Biohub fork
(github.com/Biohub/transformers @ f9a5a374b, based on v4.57.6) onto a
branch off current main (v5.10.0.dev0). This is the verbatim fork code as
a starting point; v5 convention adaptation (attention interface, modular,
__all__, nested-config round-trip) is follow-up work per the port plan.

Contents:
- src/transformers/models/esmc/    (6 files: config, sae config, modeling,
  sae modeling, tokenizer) — imports and exports cleanly under v5.
- src/transformers/models/esmfold2/ (24 files incl. deferred kernels/,
  distributed/, experimental) — config + modeling import; ESMFold2Model is
  defined but not yet exported (no __all__ — known adaptation item).

Auto-registration adapted to the v5 layout (NOT a verbatim copy of the
fork's 4 hook diffs):
- models/__init__.py: from .esmc/.esmfold2 import *
- auto/auto_mappings.py: CONFIG_MAPPING_NAMES + SPECIAL_MODEL_TYPE_TO_MODULE_NAME
  (these moved out of configuration_auto.py in v5; MODEL_NAMES_MAPPING was
  dropped in v5 so the fork's hunk for it has no equivalent).
- auto/modeling_auto.py: base + masked-LM + seq-cls + token-cls maps.
- auto/tokenization_auto.py: flat ("esmc", "ESMCTokenizer") form (v5 dropped
  the (slow, fast) tuple format).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Strip vendored Triton kernels + set_kernel_backend selector from esmfold2

Transformers loads fused kernels from the Hub via the `kernels` library
(@use_kernel_forward_from_hub / register_kernel_mapping), not by vendoring
Triton in a model dir. Remove the fork's bespoke acceleration stack and
leave a single pure-PyTorch path; fused ops can return later via a Hub
kernels repo as an opt-in follow-up.

Removed:
- src/transformers/models/esmfold2/kernels/ (8 Triton files, ~3.3k LOC).
- The shared `set_kernel_backend` selector across the module tree, which
  drove BOTH backends: "fused" (vendored Triton) and "cuequivariance"
  (external lib). Both gone — the cueq import block and BACKEND_*/
  _VALID_BACKENDS/_fused_active/_cueq_active helpers with them.
- Per-module fused/cueq branches and now-dead helpers (_can_use_*,
  _fused_trimul_with_residual, split_kernel_weights, _kernel_flow_direction,
  Transition._swiglu_pre_w3/_addmm_residual, DropoutResidual fused impl).
- The vestigial no-op set_kernel_backend hooks in distributed/utils.py and
  modeling_esmfold2_experimental.py.

Kept intact: the independent `set_chunk_size` memory knob, and the optional
flash-attn / transformer_engine guards.

Verified: both modeling modules import; CPU smoke test runs AttentionPairBias,
TriangleMultiplicativeUpdate (both orientations), Transition (chunked ==
unchunked), DropoutResidual, and FoldingTrunk end-to-end on the pure-PyTorch
path. No new ruff errors introduced (9 pre-existing fork lint items remain for
the later `make style` pass).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove vendored distributed/ (2D context-parallel) stack from esmfold2

The distributed/ package is a NVIDIA/MIT-licensed 2D context-parallel
implementation of the folding trunk (DTensor + DeviceMesh + NCCL) for
multi-GPU 6B inference. It is dropped from the port because:

- It is not imported by any model code (core, config, __init__, or the
  experimental file) — fully inert in the package.
- It is broken on import: all 7 files import from
  `projects.huggingface.transformers.models.esmfold2...` (the fork's
  internal monorepo path), so `import transformers.models.esmfold2.distributed`
  raises ModuleNotFoundError. It never worked in the standalone layout.
- It is NVIDIA/MIT-licensed, unlike the Apache/Biohub model code.
- Transformers expresses parallelism declaratively via `base_model_tp_plan`
  / `tp_plan="auto"`, not a vendored per-model DTensor/NCCL stack.

Nothing unique is lost: the math it shards already exists as the
pure-PyTorch reference in modeling_esmfold2_common.py. If multi-GPU
inference is needed later, author a tp_plan on ESMFold2Model fresh.

Verified: nothing references distributed/; `import transformers` and the
esmfold2 modeling module still import cleanly.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Refactor ESMC attention onto the v5 ALL_ATTENTION_FUNCTIONS interface

Replace ESMC's bespoke attention dispatch with the standard v5 interface,
mirroring models/esm (a bidirectional encoder). Behaviour is preserved
bit-exactly at all real-token positions.

Before: a hand-rolled `_scaled_dot_product_attention` choosing xformers ->
flash-attn-2 -> SDPA, a `_FlashMultiHeadAttention` subclass with manual
unpad_input/pad_input in ESMCModel.forward, a `_TritonRotaryEmbedding`, and
a `seq_id`-threaded chain mask.

After:
- Module-level `eager_attention_forward`; `MultiHeadAttention` dispatches via
  `ALL_ATTENTION_FUNCTIONS.get_interface(config._attn_implementation, ...)`
  with q/k/v shaped (B, H, S, Dh) and `scaling=head_dim**-0.5` (RoPE is
  rotation-invariant to scaling, so it stays in the module). output_attentions
  forces the eager interface so probabilities remain observable.
- ESMCModel.forward builds the 4D mask once: `create_bidirectional_mask` for
  the padding case (works for eager/sdpa/flash), and a block-diagonal additive
  bias for multi-chain `sequence_id` (eager/sdpa; flash multi-chain still
  raises). Removes the unpad/pad varlen path.
- Drops the xformers / flash-attn / triton-rotary import machinery and their
  warnings. transformer_engine LN/MLP fusion is unchanged. The
  `_supports_sdpa/_supports_flash_attn/_supports_attention_backend` flags
  (already declared) are now actually honoured.

flash-attn-2 is still supported — now via the standard attn_implementation
backend rather than bespoke dispatch.

Verified: loading identical weights into the refactored model reproduces the
pre-refactor outputs to 0.000e+00 at every non-padding position (plain,
padding-mask, and multi-chain cases); the only differences are at padding
positions, which are masked out downstream. eager and sdpa agree bit-exactly
on valid-token logits; MaskedLM/SequenceClassification/TokenClassification
and output_attentions all work.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove transformer_engine dependency from ESMC (pure-PyTorch only)

No other Transformers model depends on transformer_engine; drop it from
ESMC. TE provided fused fp32-reduction LayerNorm+Linear / LayerNorm+MLP
kernels, but the model already shipped pure-PyTorch fallbacks whose
parameter names (`layer_norm_weight`, `fc1_weight`, `fc2_weight`, ...) match
the published-checkpoint layout. Make those the only path.

- Drop the `transformer_engine` import + `_te_available` guard + the
  "TE not installed" warning.
- `_swiglu_ln_ffn` / `_make_attn_layernorm_qkv` / `_make_attn_out_proj` now
  unconditionally return the pure-PyTorch modules (`_PyTorchLayerNormMLP`,
  `_PyTorchLayerNormLinear`, `nn.Linear`).
- Remove dead `_SwiGLU` class (the MLP fallback inlines silu(x1)*x2).

If exact TE numerics (fp32-reduction LayerNorm) are ever required, that
belongs in a Hub kernel via the `kernels` library, not a hard dependency.

Verified: strict state_dict load from the pre-change baseline succeeds
(parameter names unchanged -> published checkpoints still load), and
last_hidden_state is bit-identical (0.0) at all valid positions for the
plain, padding-mask, and multi-chain cases. Locally TE was never installed,
so this is the exact path that already ran.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Modernize ESMC rotary to the standard (cos, sin) + apply_rotary_pos_emb convention

Replace the flash-attn-style cache-based rotary with the standard
Transformers convention used by esm/llama, as the last code-shaping step
before authoring modular_esmc.py.

Before: a stateful `RotaryEmbedding` (per-attention-module) caching cos/sin,
`forward(q, k)` returning rotated tensors, a custom `_apply` override to keep
`inv_freq` fp32 across device casts, plus `_rotate_half` / `_apply_rotary_emb_torch`.

After:
- `rotate_half` + `apply_rotary_pos_emb` (identical to esm/llama).
- `ESMCRotaryEmbedding(config)` -> `(cos, sin)`, computed once in
  `ESMCModel.forward` and threaded down (position_embeddings) through the
  stack/block to attention, mirroring esm. `inv_freq` is fp32 and
  non-persistent (matches the old behaviour: no rotary tensors in the
  checkpoint), and cos/sin are built in fp32 then cast.
- Add `config.rope_theta` (default 10000.0, the previously-hardcoded base).
- `_init_weights` recomputes `inv_freq` for `ESMCRotaryEmbedding` (meta-init safe).

Verified: strict state_dict load from the saved baseline succeeds (rotary
buffers are non-persistent, so keys are unchanged -> published checkpoints
still load), and last_hidden_state is bit-identical (0.0) at all valid
positions for plain, padding-mask, and multi-chain. The fp32 matmul-based
freqs equal the old `outer(t, inv_freq)`; same RoPE math, idiomatic shape.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add modular_esmc.py; generate modeling_esmc.py from it

ESMC now follows the modular convention. modular_esmc.py is the source of
truth; modeling_esmc.py is generated by utils/modular_model_converter.py and
carries the auto-generated header.

Reuse from esm (the natural parent — also a bidirectional protein encoder):
`eager_attention_forward`, `rotate_half`, and `apply_rotary_pos_emb` are now
imported from ..esm.modeling_esm and inlined into the generated file with
`# Copied from` headers (so they stay in sync). `rotate_half` is pulled in
transitively as a dependency of `apply_rotary_pos_emb`, matching the qwen3
pattern.

Everything else stays ESMC-specific and is defined in the modular file: the
SAE-integrated ESMCModel + ForMaskedLM/SequenceClassification/
TokenClassification, the fused-LN MultiHeadAttention, SwiGLU FFN,
TransformerStack, ESMCRotaryEmbedding, and the SAE-carrying output
dataclasses. As expected for this architecture the dedup is modest; the win
is convention compliance + auto-sync of the shared functions.

The modular file was ruff-fixed/formatted (Optional[X] -> X | None, import
order) before regeneration, so both files are now ruff-clean.

Verified: `check_modular_conversion.py` passes (files in sync); `transformers`
imports; and loading identical weights reproduces the pre-conversion
last_hidden_state bit-for-bit (0.0) at all valid positions for plain,
padding-mask, and multi-chain inputs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix ESMC checkpoint loading + bidirectional attention (verified vs fork on ESMC-6B)

Two bugs surfaced only when loading the real biohub/ESMC-6B checkpoint and
running an unpadded sequence (prior parity tests all used padded inputs):

1. _init_weights re-initialized loaded nn.Linear weights. It used
   `module.weight.data.normal_()`, which writes through `.data` and bypasses
   the `_is_hf_initialized` flag that transformers sets on loaded params. So
   `from_pretrained` clobbered out_proj / lm_head with random init (silently —
   no missing-keys warning), while custom/LayerNorm/Embedding modules survived
   only because _init_weights had no branch for them. Switch to the
   flag-respecting `transformers.initialization` helpers (init.normal_,
   init.zeros_, init.copy_), matching esm/base. (This bug is latent in the
   v4 fork too; it only bites under the v5 init-after-load flow.)

2. Attention defaulted to causal on unpadded inputs. ESMC is a bidirectional
   encoder, but MultiHeadAttention never set `is_causal`, so the sdpa/flash
   interface fell back to `getattr(module, "is_causal", True)` and applied
   causal masking whenever `attention_mask is None`. Set `self.is_causal =
   False`. (Introduced by the attention-interface refactor; missed because
   every earlier parity test passed a padded 4D mask.)

Source of truth is modular_esmc.py; modeling_esmc.py regenerated.

Verified: loading biohub/ESMC-6B (and ESMC-300M) into the refactored model
and into the original fork code yields BIT-IDENTICAL logits (max|Δ| = 0.0)
and identical argmax predictions on an 80-residue sequence. State-dict loads
with no remapping (240 TE `_extra_state` keys ignored as before).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix ESMCTokenizer docstring example + add tokenizer tests

The ESMCTokenizer behaves correctly under v5 (verified against the published
biohub/ESMC-6B tokenizer.json), but its docstring doctest example was wrong:
it listed 21 ids for a 20-residue sequence with two residues mis-ordered/
dropped. Correct it to the actual output (22 ids, <cls> ... <eos>).

Add tests/models/esmc/test_tokenization_esmc.py (fast-only tokenizer, so it
mirrors models/esm's plain-TestCase style rather than the slow vocab-file
setup): documented example, character-level tokenize, <cls>/<eos> wrapping,
special-token ids (incl. bos==cls alias), batch padding + attention_mask,
chain-break token, mask token, unknown-residue -> <unk>, decode round-trip,
save/load round-trip, and a @slow integration test asserting AutoTokenizer
resolves to ESMCTokenizer and the hub tokenizer matches the code-built one.

Verified: 10 passed, 1 slow-skipped (the slow test passes with RUN_SLOW=1).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop vestigial vocab_file entry from ESMCTokenizer VOCAB_FILES_NAMES

ESMCTokenizer is fast-only (built from a tokenizer object); it never reads or
writes a slow `vocab.txt`. Declare only `tokenizer_file: tokenizer.json`.
save_pretrained already emits just tokenizer.json + tokenizer_config.json;
tokenizer tests still pass (10 passed, 1 slow-skipped).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add ESMC model tests + make _init_weights cover all modules

Add tests/models/esmc/test_modeling_esmc.py (ModelTesterMixin +
PipelineTesterMixin), mirroring models/esm: ESMCModelTester +
create_and_check for ESMCModel / ForMaskedLM / ForSequenceClassification /
ForTokenClassification, plus a @slow masked-LM integration test on
biohub/ESMC-300M. ESMCModel has no pooler, so the base-model check asserts
only last_hidden_state.

Fix surfaced by test_can_init_all_missing_weights: _init_weights only handled
nn.Linear, so the fused-LN modules, nn.LayerNorm and nn.Embedding were left
uninitialized on a from-scratch / meta-device init. Handle the custom
_PyTorchLayerNormLinear / _PyTorchLayerNormMLP explicitly (their `weight` is a
Linear weight, not a norm — the base initializer matches norms by class-name
substring and would wrongly set it to ones), the rotary buffer explicitly,
and delegate nn.Linear / nn.Embedding / nn.LayerNorm to super()._init_weights.

Skip test_retain_grad_hidden_states_attentions: ESMC returns `hidden_states`
as one stacked tensor (consumed by the SAE feature), not the live per-layer
tensors, so grad cannot flow back to the returned copy — intentional.

Verified: full model test file = 92 passed, 102 skipped, 0 failed; the @slow
integration test passes; and from_pretrained still loads biohub/ESMC weights
bit-exactly (the flag-respecting init helpers don't clobber loaded params).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add ESMC docs + resolve auto_docstring checkpoint for the classification heads

Add docs/source/en/model_doc/esmc.md (short: overview, links to the
biohub/ESMC-{300M,600M,6B} checkpoints and the evolutionaryscale/esm repo, a
usage snippet, and [[autodoc]] for ESMCConfig, ESMCSAEConfig, ESMCTokenizer,
ESMCModel, ESMCFor{MaskedLM,SequenceClassification,TokenClassification},
ESMCSAEModel) and register it in _toctree.yml (alphabetically after ESM).

ESMCForSequenceClassification/ForTokenClassification have no Examples block in
their forward docstrings, so auto_docstring could not find a checkpoint and
errored on import. Add the standard checkpoint sentence with a Hub link
([Biohub/ESMC-600M-2024-12](...)) to the ESMCConfig docstring; auto_docstring
falls back to the config's checkpoint, resolving it for all model classes at
once (no per-class decorator needed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove the SAE (esmc_sae) from ESMC — deferred to a follow-up PR

The sparse-autoencoder is architecturally distinct from ESMC (a linear
encoder + top-k sparsity + linear decoder operating on frozen hidden states,
not a transformer) and has a non-standard per-layer from_pretrained/
save_pretrained path. Keep it out of the core ESMC PR; it will land as its
own follow-up (recoverable from this commit's parent).

Removed:
- modeling_esmc_sae.py and configuration_esmc_sae.py.
- esmc_sae auto-registration (CONFIG_MAPPING_NAMES, MODEL_MAPPING_NAMES, and
  the esmc_sae -> esmc SPECIAL_MODEL_TYPE_TO_MODULE_NAME mapping).
- The SAE integration woven through ESMCModel and the heads: add_sae_models,
  _get_sae_outputs / _get_sae_layer_num_requested / _validate_sae_inputs /
  _SAE_KEY_RE / _sae_models, the compute_sae/normalize_sae params, and the
  sae_outputs field on all four output dataclasses. The forward simplifies
  accordingly (layers_to_collect now only serves output_hidden_states; the
  SAE-only bool_mask is gone).
- ESMCSAEConfig/ESMCSAEModel autodoc from the model doc.

Source of truth is modular_esmc.py; modeling_esmc.py regenerated.

Verified: ESMCSAE* no longer importable and esmc_sae deregistered; ruff clean;
model + tokenizer tests pass (102 passed, 103 skipped); and fork-vs-new logits
on biohub/ESMC-300M remain bit-identical (max|Δ| = 0.0) — the removal is
output-neutral for the core model.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Normalize ESMC docstrings to satisfy check_docstrings

`make fix-repo`/check_docstrings normalizations on the ESMC public objects:
- ESMCTokenizer: default values use single backticks (`"<unk>"`) per the
  canonical "*optional*, defaults to ..." format.
- ESMCModel / ESMCForMaskedLM / ESMCForSequenceClassification: reorder the
  forward docstring argument entries to match the signature order
  (output_attentions before labels) and add the missing blank line before the
  Examples block.

Docstring-only (no logic change); modular and generated modeling stay in sync.
ESMC now passes check_repo, check_copies, check_inits, check_dummies,
check_modular_conversion, and check_docstrings (remaining repo-wide failures
are pre-existing Qwen output-doc errors and Phase-B esmfold2, neither ESMC).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Get ESMFold2 importing: add __all__ + load ESMC via the Auto registry

- Add `__all__ = ["ESMFold2Model"]` to modeling_esmfold2.py so the model is
  exported (`from transformers import ESMFold2Model` now works; previously the
  class was defined but the module had no `__all__`).
- Replace the hard cross-model import `from ..esmc.modeling_esmc import
  ESMCModel` in `load_esmc` (both modeling_esmfold2.py and the experimental
  file) with `AutoModel.from_pretrained(...)`. ESMC is a shared, frozen 6B
  backbone loaded separately from its own repo (`config.esmc_id`); resolving it
  through the Auto registry (model_type "esmc" -> ESMCModel) keeps esmc and
  esmfold2 as separate model directories without a runtime cross-dir import.

Verified: ESMFold2Model + ESMFold2Config export; core + experimental modules
import; no `..esmc` imports remain in the esmfold2 dir; AutoModel.from_config
resolves ESMCConfig -> ESMCModel. (Also removed the empty kernels/ and
distributed/ dirs left behind by their earlier git rm.)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop the experimental ESMFold2 variant + esmfold2_v2 mapping (release-only)

Scope the core ESMFold2 PR to the release model. The experimental
(legacy/dev binder-design) variant is deferred like the ESMC SAE.

- Delete modeling_esmfold2_experimental.py and remove it from the package
  __init__ and from the ESMFold2Model.from_pretrained `config.type ==
  "experimental"` routing (from_pretrained now always builds ESMFold2Model).
- ESMFold2Config: keep the `type` field for checkpoint compat but accept only
  "release"; fix the docstring example to use ESMFold2Model. Update the two
  stale `ESMFold2ExperimentalModel` references in comments/docstrings.
- Remove the `esmfold2_v2 -> esmfold2` SPECIAL_MODEL_TYPE_TO_MODULE_NAME entry
  (no v2 variant ships).

Verified: ESMFold2Model + ESMFold2Config import; ESMFold2ExperimentalModel
gone; no experimental/v2 references remain.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove TransformerEngine fp8 path from ESMFold2

ESMFold2's only TransformerEngine dependency was the optional fp8
quantization of the (now TE-free) ESMC backbone. Drop the import guard,
the dead _convert_te_modules_to_fp8_inplace walker, the "fp8" precision
option, and the fp8 padding/autocast plumbing. _lm_precision_context is
now a plain bf16 autocast; ESMC loads at bf16 (default) or fp32.

No other Transformers model depends on transformer_engine.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Route ESMFold2 plain self-attention through the v5 attention interface

SWA3DRoPEAttention's plain softmax(QKᵀ)V core now dispatches through
ALL_ATTENTION_FUNCTIONS / a local eager_attention_forward, keyed on
config._attn_implementation, with the sliding window expressed as an
additive attention mask. The custom flash-attention path (native
bidirectional window_size + varlen for packed inputs) is kept as an
opt-in backend, now gated on _attn_implementation == "flash_attention_2"
instead of auto-selecting whenever flash-attn is importable — so the
default is sdpa (matching the fork's SDPA fallback bit-for-bit) and
flash is opt-in, per v5 conventions.

ESMFold2Model declares _supports_sdpa / _supports_flash_attn /
_supports_attention_backend and, after construction, attaches its shared
config to every SWA3DRoPEAttention (the atom encoders/decoders build them
from explicit dims), so dispatch stays live under set_attn_implementation.
is_causal=False guards against the interface defaulting to causal when no
mask is passed. Pair-bias (AttentionPairBias) and triangular math are left
untouched.

Validated on CPU vs the pre-refactor forward (random weights): sdpa
max|Δ|=0.0 (bit-exact), eager max|Δ|=1.3e-3 (bf16 softmax precision).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Convert ESMFold2Config to the v5 @strict / PreTrainedConfig style

The config was one of ~4 holdouts still using the old `def __init__(**kwargs)`
+ hand-rolled `to_dict` style; 460/463 configs use `@strict` dataclass
`PreTrainedConfig`. Rewrite it to match, mirroring `models/esm` (ESMFold v1):

- Every sub-config (MSAEncoder/Parcae/LMEncoder/AtomAttention/FoldingTrunk/
  InputsEmbedder/DiffusionModule/DiffusionStructureHead/ConfidenceHead) is now a
  `@strict PreTrainedConfig` with typed fields + defaults, instead of a plain
  `@dataclass`.
- Nesting is declared via `sub_configs = {...}` on each parent + a `__post_init__`
  that turns `dict -> SubConfig(**dict)`. This deletes the brittle hand-rolled
  `ESMFold2Config.to_dict()` (base class now serializes via `sub_configs`) and
  gives recursive `_attn_implementation` propagation to sub-configs for free.
- Top config switches to `@auto_docstring(checkpoint="biohub/ESMFold2")`, which
  resolves the outstanding `check_docstrings` failure for ESMFold2Config.
- `__all__` is reduced to `["ESMFold2Config"]` (sub-configs are implementation
  detail, matching ESM which exports only `EsmConfig`).

check_config_attributes: top config keeps a precise 4-item allow-list (type +
three training/experimental recipe knobs not read by the core inference path);
the 9 sub-configs are allow-listed wholesale because their fields are threaded
into submodules as explicit dims (e.g. `d_atom=cfg.inputs.atom_encoder.d_atom`),
which the checker's `config.<attr>` heuristic cannot trace.

Verified: default + nested-dict construction, save/load round-trip (to_dict
identical, sub-config types preserved), recursive attn-impl propagation, type
validation, unknown-kwarg tolerance; the tiny ESMFold2Model still builds and the
SWA attention equivalence is unchanged. ruff + check_docstrings +
check_config_attributes all clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add ESMFold2 tests + sub-config model_types

ESMFold2 is an all-atom structure predictor whose forward takes ~18 structural
feature tensors and returns a plain dict (not a ModelOutput), so it doesn't fit
ModelTesterMixin. Following the sanctioned pattern for such models, register the
test file in check_repo's TEST_FILES_WITH_NO_COMMON_TESTS and provide focused
coverage:

- ESMFold2ConfigTest: ConfigTester common tests (incl. composite sub-config
  save/load), `type` validation, nested round-trip, attn-impl propagation.
- ESMFold2ModelTest (CPU): full pure-PyTorch forward via infer_protein with no
  ESMC backbone (LM conditioning skipped) under both sdpa and eager; SWA config
  dispatch; weight-level save/load fidelity + a usable reloaded model.
- ESMFold2IntegrationTest: @slow real-weight fold on biohub/ESMFold2 (GPU-gated).

To make ConfigTester's composite test pass, each sub-config now declares a
unique `model_type` (e.g. "esmfold2_inputs_embedder") — the CLIP pattern — so
that `SubConfig.from_pretrained(<composite dir>)` extracts the matching nested
dict (configuration_utils keys this off model_type). ESM dodges this test
because its sub-config is None by default; ESMFold2's are always present.

check_repo: `modeling_esmfold2_common` (shared building blocks, no public model)
is added to get_model_modules' _ignore_modules. The remaining check_repo item
for esmfold2 is the model-doc page, which is a separate pending task.

The tiny test config encodes two real sizing constraints discovered via the
forward smoke: 3D RoPE needs 3*n_spatial + n_uid <= head_dim//2, and
inputs.d_inputs == 67 + d_token//2 == diffusion_module.c_s_inputs.

7 non-slow tests pass; ruff + check_config_attributes + check_docstrings clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add ESMFold2 model doc page

Short overview page (mirrors esmc.md): describes the all-atom structure
predictor and its separately-loaded ESMC backbone, a usage snippet
(infer_protein_as_pdb), and autodoc for ESMFold2Config + ESMFold2Model. Registered
in _toctree.yml after ESMC. Resolves the final check_repo "objects documented"
item for esmfold2.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Ruff-format the ESMFold2 fork files (deferred style sweep)

Runs the previously-deferred ruff sweep over the esmfold2 files now that the
substantive work (TE removal, attention interface, config) is done: drops the
`# coding=utf-8` lines (UP009), rewrites the `_msa_kwargs = dict(...)` as a
literal (C408), sorts imports, and applies `ruff format`. No logic changes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Convert ESMCConfig to the v5 @strict / @auto_docstring style

ESMCConfig was the last non-@strict direct config (old def __init__ style),
flagged by the TRF010 model-structure rule (make typing). Convert it to the
dataclass `@strict` form with class-level typed fields, mirroring ESM, and
switch the docstring to `@auto_docstring(checkpoint="biohub/ESMC-600M")` so the
standard/base fields are auto-documented and the classification heads keep a
checkpoint to fall back on. Verified: defaults, override, save/load round-trip,
real-config load (biohub/ESMC-300M), and the 92 esmc model tests still pass;
make typing + check_docstrings + check_config_attributes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix ESMFold2 integration test: ubiquitin + correct 0-1 pLDDT/pTM scale

The @slow integration test used a poorly-folding sequence and asserted pLDDT on
a 0-100 scale, but this model emits pLDDT/pTM on a 0-1 scale
(_categorical_mean(..., start=0, end=1)) — so the old `mean > 50` assertion
could never pass. Switch to ubiquitin (PDB 1UBQ, a textbook well-folding 76-mer),
draw 8 diffusion samples and assert on the best one (these folders rank N samples
and return the best): best pLDDT > 0.7 and best pTM > 0.6.

Validated end-to-end on real weights (biohub/ESMFold2 + the 6B ESMC backbone,
CPU/fp32 since this box's GPU can't JIT-compile): weights load with 0 missing/
unexpected keys, and ubiquitin folds to best pLDDT ~0.80 / pTM ~0.74, GB1 to
~0.85 / ~0.78 — confirming the port produces correct, confident structures.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Make ESMFold2 dtype-honest: drop in-model autocast, support from_pretrained(dtype=)

ESMFold2 was the only ported model relying on in-model `torch.amp.autocast`
(`use_amp = device.type=="cuda"`) to (a) run matmuls in bf16 and (b) keep
norms/softmax in fp32. That made it GPU-only for bf16 and silently broke CPU and
`from_pretrained(dtype=bfloat16)`. Replace it with the standard Transformers
idiom — the model runs in its loaded dtype, with numerically-sensitive ops
explicitly pinned to fp32:

- Add an fp32-pinned `LayerNorm(nn.LayerNorm)` (computes in fp32, returns the
  input dtype, keeps fp32 weights under `from_pretrained(dtype=bf16)`) and use it
  for every `nn.LayerNorm`; the affine-free `RMSNorm`/`F.rms_norm` stay in the
  activation dtype (they ran bf16 under autocast too). Pin the remaining
  softmaxes with `dtype=torch.float32`.
- Remove the three in-model autocast regions (trunk, confidence trunk, ESMC
  `_lm_precision_context`) and convert the Kabsch/SVD `enabled=False` islands to
  plain `.float()`.
- The model interleaves fp32 islands (geometry, coords, one-hot/continuous input
  features, the fp32-accumulated confidence pair, the fp32 triangular contract)
  with bf16 compute, which autocast used to bridge; add explicit `.to(dtype)`
  casts at each island→matmul boundary (atom encoder, inputs embedder, rel_pos,
  token_bonds, distogram head, MSA embed, diffusion s/z/coords conditioning,
  confidence pair). Also align the ESMC hidden states to the LM projection dtype,
  which additionally lets a bf16 backbone feed an fp32 trunk (no more
  `esmc_precision="fp32"` requirement on CPU).

Validated vs the autocast baseline (ubiquitin, seed 0): GPU bf16 0.798/0.737,
GPU fp32 0.800/0.740, CPU fp32 0.801/0.740 — all match 0.80/0.74. CPU bf16 runs
correctly too (its matmul accumulates in fp32, ~4e-3 rel-err), just slower.
make typing / check_docstrings / check_config_attributes / ruff clean; 7 tests +
the @slow GPU integration test (now bf16) pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Apply make fix-repo sweep + bf16 ESMFold2 usage example

`make fix-repo` adds the `add_dates` "contributed on …" stamp to the esmc/esmfold2
doc pages and ruff-reformats a few esmc files (line-length only, no logic change).
Also update the esmfold2.md usage example to the recommended `dtype=torch.bfloat16`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Keep ESMFold2 sub-configs internal (drop their model_types)

The sub-config model_types added earlier (for the ConfigTester composite test)
made `make fix-repo` auto-register all 9 in CONFIG_MAPPING_NAMES, which then
failed check_repo (a registered config must be importable from `transformers`,
but these are intentionally not exported — `__all__ = ["ESMFold2Config"]`, like
ESM exporting only EsmConfig). Rather than expose 9 internal architecture configs
publicly (with the docs/docstring burden that implies), drop the sub-config
model_types so they stay internal, and skip the composite "load each sub-config
standalone from the parent dir" ConfigTester check, which doesn't apply to
internal sub-configs. Serialization still round-trips via `sub_configs` +
`__post_init__` (no model_type needed); config tests pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Add kernel, update docs

* Some cleanup

* More cleanup

* dtype cleanups to make outputs match the original fork

* More chasing dtypes

* More bf16 to match the original + bump speed

* Remove more float upcasts for speeds + closer fork match

* Remove more float upcasts for speeds + closer fork match

* More trunk norm matching

* Big general cleanup, no more _common.py, lots of repo standardization

* Remove a lot of redundant dtype casts now that we no longer need them

* Rename a lot of config attributes to the standard ones

* Modular cleanup, use standard output types

* More modular cleanup

* More modular cleanup

* More modular cleanup

* Even more modular cleanup

* No more modular ESMFold2

* Dead code cleanup

* No more apply_torch_compile

* Simplify the obvious einsums but leave the very messy ones

* Config cleanup, stop passing naked kwargs around

* Small docs cleanups

* Get rid of CPU-only path that we don't need

* Test cleanup

* Doc dates

* Comments cleanup

* Re-add the kernel after rebase

* Fix dates

* Push ESMC fixes

* Fix the tokenizer to match Llama

* Make the token classifier generic, attention cleanup, big modular reductions for ESMC

* Import the sequence classifier

* More review fixes

* make fix-repo

* tokenizer fixup

* Remove FA2 for esmfold2, lift some constants into the config

* Big refactor to address reviewer comments

* Bundle args together, drop some dead args

* date fixup for CI

* date fixup for CI

* Move more stuff to config, merge more swiglus

* Cleaning up some constants

* Lots of review fixes

* More review cleanup

* Fix chain-pair ipTM to use max-of-row-mean in ESMFold2 head (#47221)

pair_chains_iptm was a flat double-mean of tm_expected over all residue
pairs (i in c1, j in c2), whereas the global iptm (and AF3) use
max_i(mean_j TM(i,j)). Because mean <= max, the matrix read lower than
iptm and, for a heterodimer, iptm no longer equalled the max off-diagonal
entry.

Compute pair_chains_iptm[c1, c2] as max over rows i in chain c2 of the
mean over columns j in chain c1, matching the confidence head's
chain-pair convention and restoring iptm == max off-diagonal.

Co-authored-by: Fausto Milletari <fmilletari@tenant-ac-biohub-login-fmilletari-b42c3be-0.tenant-ac-biohub-login-all-users.tenant-ac-biohub.svc.cluster.local>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Matt <rocketknight1@gmail.com>

* Small dtype bugfix

* Doc dates

* Cleanup, add the missing ESMFold2 conversion script, add some missing auto_docstring, fixed copyright headers

* More review fixes

* More review fixes

* More review fixes

* Committing review fixes so far

* i declare an end to the norm dance

* Move the trunk into forward(), collapse the step-invariants machinery

EsmFold2Model.forward was a six-line delegation to the diffusion module while the
folding trunk lived in the generation mixin, so modeling_esmfold2.py defined modules
it never composed. forward() now runs the trunk end to end (featurize -> inputs
embedder -> pair init -> recycling -> distogram) and returns EsmFold2TrunkOutput;
generation_esmfold2.py keeps only the diffusion sampling loop, the confidence-head
call and the infer_protein entry points. main_input_name is now genuinely forward's
first argument, and the circular-import dance is down to one local import.

Step invariants:
  - one flat EsmFold2DiffusionStepInvariants instead of two nested dataclasses
  - prepare_step_invariants is the single expansion boundary; num_diffusion_samples
    drops out of 7 signatures and all three "am I already expanded?" guards go away
  - the expansion sits after the unexpanded work, so the atom featurization and the
    pair-bias projection still run at the unexpanded batch (the bias projection now
    pushes num_heads channels instead of pairwise_hidden_size)
  - the single-inputs projection is hoisted out of the per-step loop
  - n_tokens comes from the token axis instead of atom_to_token.max().item(), which
    removes a device sync and fixes a crash on batches where every sequence is
    shorter than the padded axis

ESMC: EsmcLayerNorm restores its input dtype. Unfusing the reference's
LayerNormLinear into a standalone nn.LayerNorm meant autocast handed back fp32,
promoting the residual stream while the rotary cos/sin stayed bf16 -- invisible to
eager RoPE, which promotes, but it applied bf16-precision cos/sin to fp32 queries
and hard-failed the fused rotary kernel. fp32 output is unchanged (the cast is a
no-op there).

Also: drop the dead sigma_data override, remove 10 casts that were provably no-ops
in both fp32 and bf16, and add trust_remote_code to the trimul kernel mapping
(a personal namespace has no organization overview, so the publisher trust check
cannot resolve it; drop this once the kernel moves to kernels-community).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* Let the 3D rotary and the triangular contraction follow the activation dtype

Two hardcoded-precision spots, both flagged in review, both now matching what the
reference actually does.

EsmFold2RotaryEmbedding3D forced cos/sin to bf16, so an fp32 model silently got
bf16-precision rotary and every later op that touched cos/sin was promoted back up.
It now takes the activation dtype from the caller, the same contract the
sequence-RoPE modules use. bf16 output is bit-identical; only fp32 changes, and it
changes by gaining the precision it asked for.

The triangular contraction ran in fp32 under a bf16 load, because `visibility` is an
fp32 pair mask and the masking multiply promoted `routed`. The reference looks like
it upcasts here too -- it calls `.float()` before the einsum -- but einsum is
autocast-eligible and its whole-fold bf16 autocast overrides that, so a bf16
reference model contracts in bf16 and an fp32 one in fp32. Casting the mask to the
activation dtype reproduces both. fp32 is unchanged; bf16 changes, and moves toward
the reference: vs the v4 fork on ubiquitin, distogram 1.65e-2 -> 1.47e-2 relative,
plddt 8.65e-2 -> 6.78e-2, ptm 3.61e-2 -> 3.23e-2. It is also ~1.40x faster on the
trunk's dominant op (L=384: 14.1 -> 10.1 ms) at ~5% less peak memory, and it makes
the following downcast into proj_emit a genuine no-op, so that one is gone.

Refresh the bf16 integration expectations for the contraction change. Fold quality
is unaffected (ptm 0.7433 vs 0.742, best pLDDT 0.8034); the fp32 expectations still
pass unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* More dtype chasing

* Fold more precomputation into the invariants step

* Cut down on arg passing, read config as much as possible

* Remove cross-model openfold imports

* Lots of folding into trunk_kwargs

* Big comment cleanup

* Broadcast cleanup to save memory

* Attention refactor, cleanup, remove _swa_window_mask_function

* More attention refactor, enable the FA2 path

* Last few fixes, update convert script

* Ship conversion mapping changes, drop a lot of weight renaming. Do everything at checkpoint creation now!

* Use the staging repo for now

* Add ESMC conversion script

* Big variable renaming pass

* The rest of the big renaming pass

* Go back to modular for esmfold2

* Modular refactor, generation stops accessing model attributes

* Get rid of EsmFold2AdaLNModulationLinear

* More confusing names gone

* Use config validation between parents and subconfigs

* Rename the GenerationMixin so it stops triggering generation checks, and shuffle configs a bit

* Regenerate modular ESMC after rebase onto main

Picks up upstream llama changes: nn.Buffer instead of register_buffer,
deprecate_kwarg on the rotary device arg, and the
use_kernel_func_from_hub -> use_kernel_forward_from_hub rename.

* Rename layers->blocks, decorate with no_grad() to reduce the need for detach()

* Batch of smaller review fixups

* Updated kernel decorator name

* more block->layer

* Make the attention layers a bit more similar (still can't merge them tho)

* dates

* Small fixes, comments, fewer magic numbers

* Precision correctness fix

* make fix-repo

* Update kernel path

* mlinter cleanup

* Subconfig passing cleanup

* Review fixes

* Review fixes

* Review fixes, split lists into _0 and _1

* More review fixes, reducing the big pile of args

* More review fixes, more big pile of args cleanup

* fix-dates

* Review fixes

* Review fixes

* Bugfixes for keep-in-fp32

* Drop comment

* mlinter fixes

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Fausto Milletari <fausto.milletari@gmail.com>
Co-authored-by: Fausto Milletari <fmilletari@tenant-ac-biohub-login-fmilletari-b42c3be-0.tenant-ac-biohub-login-all-users.tenant-ac-biohub.svc.cluster.local>