search

openai / jukebox

Code for the paper "Jukebox: A Generative Model for Music"
https://openai.com/blog/jukebox/
Other
7.81k stars 1.4k forks source link

Is it possible to make JukeBox use more of my GPUs? #97

Open ratsaby opened 4 years ago

ratsaby commented 4 years ago

I installed JukeBox and it is running straight from the Jupiter notebook.

I would like to know if it is possible to make the code use more than one GPU?

I am unfamiliar with PyTorch, but I could see in this regards, the following which seems to cause

JukeBox to run on a single GPU (whose ID is 0)

================================== def setup_dist_from_mpi( master_addr="127.0.0.1", backend="nccl", port=29500, n_attempts=5, verbose=False ): if dist.is_available(): return _setup_dist_from_mpi(master_addr, backend, port, n_attempts, verbose) else: use_cuda = torch.cuda.is_available() print(f'Using cuda {use_cuda}')

    mpi_rank = 0
    local_rank = 0

    device = torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
    torch.cuda.set_device(local_rank)

    return mpi_rank, local_rank, device

=====================================

Can you advise how can I make it use all my 4 GPUs ?

Thanks J

prafullasd commented 4 years ago

The usual way of running the sample.py script with multiple GPU's would be to launch with mpiexec -n {ngpus} python sample.py ... (also make sure to change logdir here to {hps.name}_{dist.get_rank()}/level_{level} so the ranks dont overwrite each other)

For local jupyter notebooks, you can try running 4 separate notebooks, and for each notebook launch it with CUDA_VISIBLE_DEVICES={gpu_id} jupyter notebook so that it uses that specific GPU.