Sampling with multiple GPUs? about jukebox HOT 3 CLOSED

We don't currently support it. One way to do so would be to keep half the layers of your model on one GPU, half on another, and spin up the process with two visible GPU's. For eg like here: https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html

You might have to insert .to("cuda:0/1") calls at the correct places for it to work

from jukebox.

@tmlbl Did you ever figure out a workable solution?

from jukebox.

@prafullasd - we need this to have a chance of running 2 cards - to get 48gb of VRAM.
https://www.nvidia.com/en-us/geforce/graphics-cards/30-series/rtx-3090/

related - #142

Could you please reopen this ticket? as an investigation.
Need to sift through the forks to see if anyone has made any progress.

can use this shell script to find-forks.sh
https://gist.githubusercontent.com/joeytwiddle/abee0610244d283f1ebe2f76bf608086/raw/094ef14c0b83125149698b234e95cb525751f25e/github_get_all_forks.sh

chmod +x find-forks.sh 
/find-forks.sh -do

seems like they exist

there is actual an official branch / this should be mentioned on readme.
https://github.com/openai/jukebox/tree/multi_gpu_sampling

so did you guys get anywhere with this? I don't have multiple gpus to support

c6617dd

@heewooj - do you know if multiple gpus can access more VRAM? if I buy 2 x 3090 nvidia cards / 24gb - will it give 48gb and allow fine tune of 5 billion parameters?

from jukebox.