Port ESMC and ESMFold2 to Transformers
Author: Rocketknight1
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46419)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46419)
<!-- ci-dashboard-badge:end -->

This is just about ready at this point! I've deferred the experimental ESMFold2 models and the SAEs to a follow-up PR so we can launch this as soon as possible. It's got a lot of very verbose Claude comments and docstrings too, but those are forgivable and we can always clean them up later :sweat_smile: 

The main thing for reviewers to know is that there are two model classes here, ESMC (a masked protein LM) and ESMFold2 (a protein folding model that uses ESMC as a backbone model). ESMC is mostly just a 'normal' masked LM, but I wasn't able to find exact modular matches for a lot of methods. Still, it should be easy enough to review. ESMFold2 is much more challenging and is not like any other model in the codebase, even the original ESMFold. I have some tests confirming that we get the same output as BioHub's original code, so I think reviewers should probably just trust that for correctness and focus on the API, although you can try going through all of the diffusion code if you want!

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46419). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46419&sha=464744

@vasqu writing this in one big comment rather than breaking it up to every individual location, because there was a lot of repeated stuff! I just did a big pass for ESMFold2, now that ESMC is mostly okay.

I dropped the flashattention path for now, inlined a bunch of methods, broke RoPE out into a single class, tried to merge as many SwiGLUs as I could (not all were possible, some have little embedded norms). A lot of stuff got renamed and moved to configs instead of just being constants, and I tried to pull out private methods from the big long forward methods.

There's a few tricky bits still:

1) DiffusionModule still has massive signatures with 20 args. Want me to bundle that into one big object we just pass around and access attributes on?
2) Renamed a lot of things and broke up a lot of `nn.Sequential()` and single-letter vars. This required lots of weight renaming and `attribute_map`. We'll need to regenerate the checkpoint anyway (the existing checkpoints don't contain the ESMC backbone weights), so we could rename everything in the checkpoints and then delete the rename lines before merging. Marked those with TODO.
3) Some bits I left untouched. In particular, some single-letter names were kept because they're kind of standard, like splitting a quaternion into `r, i, j, k`, and some `q` and `v` keys. I can rename those if you want! Other stuff like `plddt` comes from AlphaFold and [people doing protein folding know what it means](https://www.ebi.ac.uk/training/online/courses/alphafold/inputs-and-outputs/evaluating-alphafolds-predicted-structures-using-confidence-scores/plddt-understanding-local-confidence/) even if it's quite weird for most ML engineers!
4) Some other bits I left include an `F.rms_norm()` call with no affine scaling and so no weights. You suggested nanochat there but I'm not sure why we need a module at all! It works fine as one functional line.

Overall I tried to make it more readable, but it's a very complex model, so there's only so much I can do! :sweat_smile: 

@Rocketknight1 Thanks a lot! Yes the model is super complex so I don't expect perfect stuff. I think the esmc model should be fairly good now. I will try to take a look at esmfold2 in the coming days.

Re 
1. Yea that might be more reasonable, having too long signatures is really hard to read imo - similar to how we pass configs around and only access the portions we need. Just a suggestion and I'm not 100% sure it's a good one, would not put it high prio
2. Gotcha, sounds good. Yea we have a few models that have a long list of renames as well dw.
3. Yea, hard to catch the abbreviations sometimes. Sometimes it's me quickly glancing, sometimes my missing background for biology etc. It's not too bad if there really is meaning; I just would avoid too many where possible
4. This is mainly for inheritance with modular + potentially applying kernels on them. Won't be possible if we rely on an F. portion of torch

On 4 I'm unsure - we're not doing modular for ESMFold2 because the code is so weird compared to other models. I'd leave that one for a future refactor, if there are kernels that turn out to really justify it, or if we get a lot of similar folding models in the codebase.

1 is done, though, and I managed to drop some dead args in the process!

@cyrilvallez review mostly addressed by now! A lot of the comments became outdated as a result, so I'll just address it here in one big block rather than individual comments that might have disappeared:

- config flattened, no more subconfigs
- Weight renames removed, we do that all in the conversion script now. I rebuilt `Rocketknight1/ESMFold2-merged-temp` with the new conversion script and we can make it the official checkpoint after merging this PR
- Inlined loads of wrappers or modules that were only used once
- Sliding-window mask is built once per forward instead of per layer
- Routed things through the attention interface where possible
- Less device juggling, tensors are built in the right place with fewer `.to()` calls
- Dropped lots of globals like `_DEFAULT_CHUNK_SIZE` and `_CONFIDENCE_EPS`, mostly by inlining values
- Lots and lots of renames, especially spelling out abbreviations like `det` and `denom`
- No more `copied from`
- Rearranging done, doc linter stuff fixed
- Using `no_grad()` to avoid grad computation instead of computing with grad and `detach()`
- Also merged a patch from the BioHub team to the ipTM calculations, this was a bug in their fork too
 - Fixed the casting of the backbone (the casting was right but the comment was very stale and wrong)
 
There's some stuff I didn't change, though:
- Kept the custom `_swa_window_mask_function`, since their behaviour differs a bit from ours depending on padding (they count the number of unpadded tokens, our functions just use the index)
- I mostly kept the einsums, because the non-einsum versions used a lot more memory. Einsums for complex contractions really are just better, unless they map to something torch's compiler recognizes and can optimize (like self-attention). I can edit these for readability but the matmul/bmm versions are imo less readable, as well as using more FLOPs/memory

And the biggest thing I'm wary of changing, the dtype/casting behaviour: The way they ran the original model was to have a frozen ESMC backbone (they used a separate ESMC checkpoint with fp32 norm weights) cast to `bf16` with `to()`. This casts **all** weights to bf16, including norms. Then they ran it in an `autocast` context, which will *upcast* for ops like norms. A bf16 norm layer in an autocast context will upcast both weights and computations to fp32! Therefore, the only way to actually match their output is to downcast weights to `bf16` and then upcast them at runtime. The only alternative would be to move this to the conversion script, downconvert weights to bf16 then upconvert to fp32, then save as fp32. Or we could just add `autocast()` back in where they used it in the fork and let it handle things so I stop doing this manually, which gets rid of a lot of the manual weight casting. WDYT?

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46419)
**Latest run:** [29344705619:1](https://github.com/huggingface/transformers/actions/runs/29344705619)
**Result:** `failure` | **Jobs:** 15 | **Tests:** 169,136 | **Failures:** 1 | **Duration:** 15h 43m

<!-- ci-dashboard-recap:end -->

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, esmc, esmfold2

Some more fixes, caught some missing docstrings, and realized I hadn't committed the ESMFold2 converter script, just the ESMC one. So that's in now!