EsmForProteinFolding.from_pretrained("esmfold_v1", device_map="auto") has error when using multiple gpus
Author: zhengdd0422
Pull request: no
Labels: name: bug

### System Info

`transformers` version: 4.45.2
- Platform: Linux-5.4.0-150-generic-x86_64-with-glibc2.27
- Python version: 3.9.0
- Huggingface_hub version: 0.24.7
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- Accelerate config:    not found
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): 2.9.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA GeForce RTX 2080 Ti

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

I want to use multiple gpus to fold single amino acid sequence by using EsmForProteinFolding.from_pretrained("/Softwares/transformers_pretrained/esmfold_v1", device_map="auto"). But error occurs: 

 _File "/bin/TTSEs_predictor/transformers_esmfold_predict3d_4eachsequence.py", line 88, in <module>
    output = model(tokenized_input)
  File "/Softwares/anaconda3/envs/py39_cuda11.3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Softwares/anaconda3/envs/py39_cuda11.3/lib/python3.9/site-packages/accelerate/hooks.py", line 170, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/Softwares/anaconda3/envs/py39_cuda11.3/lib/python3.9/site-packages/transformers/models/esm/modeling_esmfold.py", line 2135, in forward
    esm_s = self.compute_language_model_representations(esmaa)
  File "/Softwares/anaconda3/envs/py39_cuda11.3/lib/python3.9/site-packages/transformers/models/esm/modeling_esmfold.py", line 2236, in compute_language_model_representations
    esm_s = torch.stack(esm_hidden_states, dim=2)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_cat)._

This is how i run my code: **CUDA_VISIBLE_DEVICES=0,1,2,3 python transformers_esmfold_predict3d_4eachsequence.py**

transformers_esmfold_predict3d_4eachsequence.py looks like this : 
_import os
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
import torch
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--effector_type", type=str, default="T1SE", choices=['T1SE', 'T2SE', 'T3SE', 'T4SE', 'T6SE', 'Non-effector'])
parser.add_argument("--data_path", default='/Work/TTSEs/', type=str)
args = parser.parse_args()

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

# mkdir output dir
output_dir = args.data_path + "data_preprocess/esm_esmfold/" + args.effector_type + "/esmfold_out/"
if not os.path.exists(output_dir):
	os.makedirs(output_dir)
pdb_file = os.path.join(output_dir, 1.pdb")
if (os.path.exists(pdb_file) and os.path.getsize(pdb_file) > 0):
	print(f"{pdb_file} is existed, jump\n")
else:
	sequences = "NEGEVVLTPEQIQTLRGYASRGDTYGGWRYLANLGDRYADNAAAIVGKDTNLNGLNLWMKKGVENLWDDTVGKKTRLEKFDRVALQHFSQYVDLINKNNGRLPNTSEIERSYYKAVTYHGVSSSAAIDLVINRSLPDMADGYWALGLGIEAERIHNEQAVNNPNGSERDNRKQLISALDKGFDGSFKEKHFTFLQSVMMDLTKLGVEYTIDGWQKIGGWGNGIINDLYKSVVKREWTGIFEIVNNNIKQGNEAFKNEINSLVHDMKAAGKEFGDDLNTQWNNLTQAAEIIYNDIVDNTSQGIEKGVKAIKELSEKMKNAASDLADGSAEKAKQVVEDLAQAAKEAYENAKSTAEKAAQAAREFFKGLPSFKDLAEKFRDLFPNPEGWIDDGHQCFAPWVKETKKRNGKYHVYDPLALDLDGDGIETVATKGFSGSLFDHNRDGIRTATGWVAADDGLLVRDLNGNGIIDNGAELFGDNTKLADGSFAKHGYAALAELDSNGDNIINAADAAFQTLRVWQDLNQDGISQANELRTLEELGIQSLDLAYKDVNKNLGNGNTLAQQGSYTKTDGTTAKMGDLLLAADNLHSRFKDKVELTAEQAKAANLAGIGRLRDLREAAALSGDLANMLKAYSAAETKEAQLALLDNLIHKWAETDSNWGKKSPMRLSTDWTQTANEGIALTPSQVAQLKKNALVSLSDKAKAAIDAARDRIAVLDAYTGQDSSTLYYMSEEDALNIVKVTNDTYDHLAKNIYQNLLFQTRLQPYLNQISFKMENDTFTLDFSGLVQAFNHVKETNPQKAFVDLAEMLAYGELRSWYEGRRLMADYVEEAKKAGKFEDYQKVLGQETVALLAKTSGTQADDILQNVGFGHNKNVSLYGNDGNDTLIGGAGNDYLEGGSGSDTYVFGKGFGQDTVYNYDYATGRKDIIRFTDGITADMLTFTREGNHLLIKAKDDSGQVTVQSYFQNDGSGAYRIDEIHFDNGKVLDVATVKELVQQSTDGSDRLYAYQSGNTLNGGLGDDYLYGADGDDLLNGDAGNDSIYSGNGNDTLNGGEGNDALYGYNGNDALNGGEGNDHLNGEDGNDTLIGGAGNDYLEGGSGSDTYVFGKGFGQDTVYNYDYATGRKDIIRFTDGITADMLTFTREGNHLLIKAKDGSGQVTVQSYFQNDGSGAYRIDEIHFDNGKVLDVATVKELVQQSTDGSDRLYAYQSGNTLNGGLGDDYLYGADGDDLLNGDAGNDSIYSGNGNDTLDGGEGNDALYGYNGNDALNGGEGNDHLNGEDGNDTLIGGAGNDYLEGGSGSDTYVFGKGFGQDTVYNYDYATGRKDIIRFTDGITADMLTFTREGNHLLIKAKDDSGQVTVQSYFQNDGSGAYRIDEIHFDNGKVLDVATVKELVQQSTDGSDRLYAYQSGSTLNGGLGDDYLYGADGDDLLNGDAGNDSIYSGNGNDTLDGGEGNDALYGYNGNDALNGGEGNDHLNGEDGNDTLIGGAGNDYLEGGSGSDTYVFGKGFGQDTVYNYDYATGRKDIIRFTDGITADMLTFTREGNHLLIKAKDGSGQVTVQSYFQNDGSGAYRIDEIHFDNGKVLDVATVKKLVQQSTDGSDRLYAYQSGNTLNGGLGDDYLYGADGDDLLNGDAGNDSIYSGNGNDTLNGGEGNDALYGYNGNDVLNGGEGNDHLNGEDGNDTLIGGAGNDYLEGGSGSDTYVFGKGFGQDTVYNYHVDKNSDTMHFKGFKAADVHFIRSGSDLVLSASEQDNVRISGFFYGENHRVDTFVFDDAAISNPDFAKYINAGNNLVQSMSVFGSNTAATGGNVDANTQSVQQPLLVTPSA"	
	model = EsmForProteinFolding.from_pretrained("/Softwares/transformers_pretrained/esmfold_v1", device_map="auto")
	tokenized_input = tokenizer([sequences], return_tensors="pt", add_special_tokens=False)['input_ids']
	tokenized_input = tokenized_input.cuda()
	with torch.no_grad():
		output = model(tokenized_input)
	
	torch.cuda.empty_cache()
	pdb = convert_outputs_to_pdb(output)
	with open(pdb_file, "w") as f:
		f.write("".join(pdb))

	print("All structures have been predicted and saved.")_
I also tried"_device = next(model.parameters()).device
tokenized_input = tokenized_input.to(device)_"， it did not work. This error kept occured. 

### Expected behavior

Expected behavior is that the model should train in a multi-GPU setting without throwing any errors. 

Comments:
My single GPU has 16GB of memory. When I fold a single sequence of length more than 1000, I got a "RuntimeError: Out of CUDA memory" error. And when I used "model = EsmForProteinFolding.from_pretrained("/Softwares/transformers_pretrained/esmfold_v1", device_map="auto")", I got  a "RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_cat)." error.  Does anyone know how to use multiple GPUs to predict structures with transformers' esmfold_v1? Can you tell me,please?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.