How to run ESMFold Inference on Multiple GPUs

[maxall41]

Hello,

Im trying to run ESMFold on a cluster of 8 Nvidia A100s but I'm running into a lot of issues with trying to get the model to run on multiple GPUs. I'm basing my code off of the fold.py script. I've tried wrapping model.esm in a nn.DataParallel but this causes internal issues inside of esm as it seems that it was not designed with this usecase in mind. So how can I run ESMFold inference on multiple GPUs?

[maxall41]

For anyone looking back on this is what I ended up doing was splitting the FASTA file I wanted to run into 4 chunks and then just modifying the fold.py script to add a new parameter called --device and updating the usages of model.cuda to model.to(args.device). I then just created 4 terminal sessions and ran each chunk on a different GPU.