Parallelize computation between GPUs or utilize TPUs?

[Tony363]

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,

[surisdi]

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.

[Tony363]

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

[surisdi]

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