cvlab-columbia/viper

Parallelize computation between GPUs or utilize TPUs?

Tony363 opened this issue · 3 comments

Hello,

I am wondering if its possible to parallelize the GLIP and BLIP2 computation over multiple GPUs or to utilize TPUs for inference? I currently have access to 3 NVIDIA RTX A4500 each with 20GiBs of Vram memory. I suspect I should enough Vram memory to do inference with viperGPT but when I run the note book, I get out of memory errors for the GPU.

Alternatively, I am contemplating to perform inference using Kaggles TPUs with a Vram of 128 GiBs for 20 hours a week. However, I would need to utilize pytorch xla to make the torch tensors compatible with TPUs. Where in the code might I make the change? Should I edit the source code to perform inference?

Tony,

In that case, the file you will have to modify is vision_models.py. Every class in that file defines a module. As input, they receive the GPU ID (gpu_number), which assumes every model only uses a single GPU. In your implementation, you can either ignore that parameter, or you can modify the code in this line to provide more than one GPU per module.

Just to confirm, 3 NVIDIA RTX A4500 Gpus each with 20 vram of memory should be able to run viperGPT right?

Yes, with that you should be able to run ViperGPT. Make sure all GPUs are visible/available when running the code.