Single-machine, single-GPU: distributed best practices

ArchieGertsman: <p>I’ve been reading the docs about PyTorch’s distributed features such as <code>torch.distributed</code> and DDP models. It seems like these features are geared towards multi-node and/or multi-GPU settings, so I’m wondering if there is a set of best practices for the simplest case: multiprocessing using one CPU and one GPU.</p>
<p>Should the distributed package be used at all? Are there overheads that one should worry about? The reason that this package appeals to me is that I need to do a reduce operation at the end of each RL episode, which is easily handled by <code>dist.all_reduce</code>. And DDP also makes model updates super clean. I know that I can purely use <code>torch.multiprocessing</code>, but I guess the code would be messier.</p>
<p>Could asynchronous, DDP model training with a single GPU be sped up by having each subprocess use a different CUDA stream? My thinking is that it will allow multiple subprocesses run CUDA kernels in parallel, rather than queuing them up.</p>

ptrblck: <p>I don’t think DDP supports single GPU use cases as it would assume multiple NCCL ProcessGroups could communicate on the same device (I don’t know if this is supported in GLOO or another backend), so using different streams might be the right approach for you.</p>
<aside class="quote no-group" data-username="ArchieGertsman" data-post="1" data-topic="169243">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/archiegertsman/48/51153_2.png" class="avatar"> ArchieGertsman:</div>
<blockquote>
<p>Could asynchronous, DDP model training with a single GPU be sped up by having each subprocess use a different CUDA stream?</p>
</blockquote>
</aside>
<p>Each process would already be a separate Python process and would not share any CUDA streams.</p>
<aside class="quote no-group" data-username="ArchieGertsman" data-post="1" data-topic="169243">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/archiegertsman/48/51153_2.png" class="avatar"> ArchieGertsman:</div>
<blockquote>
<p>My thinking is that it will allow multiple subprocesses run CUDA kernels in parallel, rather than queuing them up.</p>
</blockquote>
</aside>
<p>Parallel execution requires free compute resources. If one kernel uses all compute resources you won’t be able to execute anything in parallel.</p>

ArchieGertsman: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="169243">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>I don’t think DDP supports single GPU use cases as it would assume multiple NCCL ProcessGroups could communicate on the same device (I don’t know if this is supported in GLOO or another backend), so using different streams might be the right approach for you.</p>
</blockquote>
</aside>
<p>DDP is working with the GLOO backend and single device for me:</p>
<pre><code class="lang-auto">''' in main process '''
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'

device = torch.device('cuda:0')

# model object sent to each subprocess
model = ...
model.to(device)

''' in each subprocess '''
dist.init_process_group('gloo', rank=rank, world_size=num_envs)

# use the same device in each instance of DDP
model = DDP(model, device_ids=[device])
</code></pre>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="169243">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Each process would already be a separate Python process and would not share any CUDA streams.</p>
</blockquote>
</aside>
<p>I see! Didn’t realize that different processes use different streams.</p>

copythatpasta: <p><a class="mention" href="/u/archiegertsman">@ArchieGertsman</a>  can you elaborate on how you achieved this? What do you mean by your comments “in each subprocess” ? Is there a way to achieve this by running a single script that process batches on the same GPU with multiple processes?</p>

ArchieGertsman: <p>The code block under the first comment runs in the main process, while the code block under the second comment runs in each subprocess that spawns from the main process. Here is some more complete pseudo-code where each process handles its own batch of data (I just threw this together without running):</p>
<pre><code class="lang-auto">import os
import multiprocessing as mp
import torch
from torch import nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


def f(rank, world_size, device, model):
    dist.init_process_group('gloo', rank=rank, world_size=world_size)

    # use the same device in each instance of DDP
    model = DDP(model, device_ids=[device])

    # whatever
    model(torch.rand(16, 4)).sum().backward()

    # sum gradients
    for param in model.parameters():
        dist.all_reduce(param.grad)


if __name__ == "__main__":
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'

    # single GPU will be shared among all subprocesses
    device = torch.device('cuda:0')

    model = nn.Linear(4, 4)
    model.to(device)

    # start subprocesses
    num_procs = 8
    procs = []
    for rank in range(num_procs):
        proc = mp.Process(target=f, args=(rank, num_procs, device, model))
        procs += [proc]
        proc.start()

    # terminate subprocesses
    for proc in procs:
        proc.join()
</code></pre>