transformers

c8df3fc9 - Port ESMC and ESMFold2 to Transformers (#46419)

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>