EsmForProteinFolding.from_pretrained("esmfold_v1", device_map="auto") has error when using multiple gpus

[zhengdd0422]

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?:
• Using GPU in script?:
• 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 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 wrappercat).

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 = tokenizedinput.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.

[zhengdd0422]

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?