Error SIGABRT when running esmfold_v1 on TPU 
Author: conchaeloko
Pull request: no

### System Info

tpu-vm-pt-2.0 (torch ; torchvision ; torch-xla : 2.0)
accelerator : v2-8 
transformers: 4.27.4

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Create a TPUvm and connect to it : 
```
gcloud compute tpus tpu-vm create ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --accelerator-type=v2-8 --version=tpu-vm-pt-2.0
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project=${PROJECT_ID}  

```
Load the model on the TPU : 
```
import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()

from transformers import AutoTokenizer, EsmForProteinFolding
from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
from transformers.models.esm.openfold_utils.feats import atom14_to_atom37

tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")

import gc
gc.collect()

model = model.half()
model = model.to(device)
model.trunk.set_chunk_size(64)
```
Try running the first part of the esmfold_v1 script : 

```
def esmfold_prediction(tokenized_sequences , path_out) :
    """
    The function takes as an input :
    - 'tokenized_sequences', the output of the function tokenize_fasta
    - 'path_out', the path of the directory when the pdb files are to be written
    The function generates the pdb files in the path_out
        
    """
    for protein in tokenized_sequences :
        pdb_files = []
        with torch.no_grad():
            prot_to_pred = protein[1].to(device)
            output = model(prot_to_pred)
```

Error : 
```
src/tcmalloc.cc:332] Attempt to free invalid pointer 0x7ffc17c89fc0
https://symbolize.stripped_domain/r/?trace=7f99dc4ff00b,7f99dc4ff08f,fffffffffb6affff,e900000002bffe88&map=
*** SIGABRT received by PID 132988 (TID 132988) on cpu 36 from PID 132988; stack trace: ***
PC: @     0x7f99dc4ff00b  (unknown)  raise
    @     0x7f988e574a1a       1152  (unknown)
    @     0x7f99dc4ff090  (unknown)  (unknown)
    @ 0xfffffffffb6b0000  (unknown)  (unknown)
    @ 0xe900000002bffe89  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7f99dc4ff00b,7f988e574a19,7f99dc4ff08f,fffffffffb6affff,e900000002bffe88&map=ceee8fa20ddf9c34af43f587221e91de:7f988164c000-7f988e78b840
E0410 13:36:27.439930  132988 coredump_hook.cc:414] RAW: Remote crash data gathering hook invoked.
E0410 13:36:27.439944  132988 coredump_hook.cc:453] RAW: Skipping coredump since rlimit was 0 at process start.
E0410 13:36:27.439952  132988 client.cc:278] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0410 13:36:27.439956  132988 coredump_hook.cc:512] RAW: Sending fingerprint to remote end.
E0410 13:36:27.439962  132988 coredump_socket.cc:120] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0410 13:36:27.439970  132988 coredump_hook.cc:518] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0410 13:36:27.439974  132988 coredump_hook.cc:580] RAW: Dumping core locally.
E0410 13:36:27.833706  132988 process_state.cc:784] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
```
Disabling tcmalloc  that way : 
`export LD_PRELOAD=""`

Running the script again returns the error : 
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 14, in esmfold_prediction
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/robbyconchaeloko/.local/lib/python3.8/site-packages/transformers/models/esm/modeling_esmfold.py", line 2154, in forward
    structure: dict = self.trunk(s_s_0, s_z_0, aa, position_ids, attention_mask, no_recycles=num_recycles)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/robbyconchaeloko/.local/lib/python3.8/site-packages/transformers/models/esm/modeling_esmfold.py", line 1965, in forward
    structure = self.structure_module(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/robbyconchaeloko/.local/lib/python3.8/site-packages/transformers/models/esm/modeling_esmfold.py", line 1782, in forward
    rigids = rigids.compose_q_update_vec(self.bb_update(s))
  File "/home/robbyconchaeloko/.local/lib/python3.8/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py", line 917, in compose_q_update_vec
    new_rots = self._rots.compose_q_update_vec(q_vec)
  File "/home/robbyconchaeloko/.local/lib/python3.8/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py", line 518, in compose_q_update_vec
    return Rotation(
  File "/home/robbyconchaeloko/.local/lib/python3.8/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py", line 289, in __init__
    quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)
RuntimeError: Error while lowering: [] aten::div, xla_shape=f32[1,214,2560]{2,1,0}
Error: /pytorch/xla/torch_xla/csrc/convert_ops.cpp:86 : Unsupported XLA type 10
Frames:
```


### Expected behavior

Generates the input for the next function : 
```
def convert_outputs_to_pdb(outputs):
    final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
    outputs = {k: v.to("cpu").numpy() for k, v in outputs.items()}
    final_atom_positions = final_atom_positions.cpu().numpy()
    final_atom_mask = outputs["atom37_atom_exists"]
    pdbs = []
    for i in range(outputs["aatype"].shape[0]):
        aa = outputs["aatype"][i]
        pred_pos = final_atom_positions[i]
        mask = final_atom_mask[i]
        resid = outputs["residue_index"][i] + 1
        pred = OFProtein(
            aatype=aa,
            atom_positions=pred_pos,
            atom_mask=mask,
            residue_index=resid,
            b_factors=outputs["plddt"][i],
            chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None,
        )
        pdbs.append(to_pdb(pred))
    return pdbs
```

And writes the pdb file

Comments:
cc @Rocketknight1 

I'm not an expert on torch XLA, but I think the problem here is that TPUs do not support `float16`, only `bfloat16`. `model.half()` converts the model parameters to `float16`, and the error you're seeing is caused by TPUs not having a division operation that can work on `float16` inputs.

You could try removing the `model.half()` line, and/or using some of the PyTorch environment variables for downcasting to BF16 on TPU instead, such as `XLA_USE_BF16`. Please see the docs [here](https://pytorch.org/xla/release/2.0/index.html#xla-tensors-and-bfloat16).

Thank you for your suggestion @Rocketknight1. I got to step further by adding : 
```
model = model.half()
model = model.to(dtype=torch.bfloat16)
model = model.to(device)
```
However I run into some memory issues. Two options for me : put my hand on a more powerful accelerator or try model parallelism. Trying my luck, how much have you guys played with model parallelism on TPUs ? 

Thanks again for the help

That's interesting - I admit I haven't tried it on PyTorch + TPU! However, in our testing, we were able to get ESMFold to run on a GPU with 16-24GB of memory. This meant we were able to generate protein folds fine using our [ESMFold Colab notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/protein_folding.ipynb), even with the free GPU, and if you want to do longer or batch predictions then the premium GPUs should be more than enough. Have you tried running on Colab already?

@Rocketknight1, I tried but got issues with GPUs with vram<14GB (free colab and free Kaggle notebooks). I think I'll do it on a GPU V100 32GB.

Thanks again for the help