Model parallelism in Multi-GPUs: forward/backward graph

liangbright: <p>In Model parallelism, A DNN is divided into sub-modules and each module is handled by a GPU.</p>
<p>The forward graph, I assume, spans multiple GPUs.</p>
<p>Will the backward graph along with any internal data also span multiple GPUs?</p>

ptrblck: <p>Yes, this should be the case.<br>
You can check if by using model sharding and then print the device of the gradients for each submodule.<br>
They should be located on the corresponding device id.</p>

rasbt: <p>Genuine question, is that “model sharding” or model parallelism implemented, yet? So far, I’ve only been doing the multi-GPU stuff with data parallelism (e.g., like <a href="https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html" rel="nofollow noopener">https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html</a>).</p>

Dhorka: <p>Yes, you can do “model parallelism” for instance with 2 gpus like this:</p>
<pre><code class="lang-auto">class Network(nn.Module):
    def __init__(self, split_gpus):
        self.module1 = (some layers)
        self.module2 = (some layers)

        self.split_gpus = split_gpus
        if self.split_gpus:
            self.module1.device("cuda:0")
            self.module2.device("cuda:1")

    def forward(self, x):
        x = self.module1(x)
        if self.split_gpus:
            x = x.device("cuda:1") 
        return self.module2(x)
</code></pre>

rasbt: <p>Oh I see, thanks!</p>
<p>Regarding the speed-up. It is probably less than for data parallelism, right? Because for module 2 to compute the forward pass, you need to know the results of module 1. (And vice versa for backprop.)</p>
<p>So the advantage would be though that it would allow for fitting basically large models into memory rather than training speed?</p>

ptrblck: <p>Yes exactly. At least that’s what I’ve used it for.<br>
The transfer between both GPUs is done via P2P so that no host communication is needed as far as I know.</p>

Decker: <p>I implemented something like <a class="mention" href="/u/dhorka">@Dhorka</a>’s example above, and the (simplified) following works in the forward pass, but…</p>
<pre><code class="lang-auto">class Network(nn.Module): #simplified
    def __init__(self, split_gpus):
        self.layerseq1 = nn.Sequential(OrderedDict(some layers))
        self.layerseq2 = nn.Sequential(OrderedDict(some layers))
        self.split_gpus = split_gpus
        if self.split_gpus:
            self.layerseq1.cuda(device=0)
            self.layerseq2.cuda(device=1)

    def forward(self, x):
        x = self.layerseq1(x)
        if self.split_gpus:
            x = x.cuda(device=1) 
        return self.layerseq2(x)
</code></pre>
<p>…it runs into problems when it tries to run the nn.MSELoss, claiming:</p>
<pre><code class="lang-auto">Assertion THCTensor_(checkGPU) [...] failed. Some of the weight/gradient/input tensors 
are on different GPUs. Please move them to a single one.
 at /opt/software/pytorch/aten/src/THCUNN/generic/MSECriterion.cu:13
</code></pre>
<p>While running an epoch such as:</p>
<pre><code class="lang-auto">criterion = nn.MSELoss()
for batch_idx, (data, target) in enumerate(some_dataset()):
     data= Variable(data).cuda()
     output = Network(data)
     loss = criterion(output, data)   # &lt;-- fails at runtime
     optimizer.zero_grad()
     ...
</code></pre>
<p>I thought that couldn’t possibly mean that all trained layers need to be on the same GPU, so I tried moving both arguments to the same gpu before loss evaluation, (below), but that seems to have no effect</p>
<pre><code class="lang-auto">output.cuda(device=1)
data.cuda(device=1)
</code></pre>
<p>I’m still getting up to speed on pytorch, so any guidance would be appreciated…</p>

halahup: <p>What is wrong with wrapping the model into <code>DataParallel</code>, as <a class="mention" href="/u/rasbt">@rasbt</a> mentioned above, besides the memory motivation?</p>

rasbt: <p>With DataParallel, you accumulate the gradients from all GPUs on one of the main GPUs temporarily. Sometimes it can be an issue if all your GPUs use up a lot of memory (then this will basically result in an out-of-memory error). So, currently it is not possible to use a GPU for accumulating the gradients that is not also already used during training and is already holding a model in memory. In practice, I haven’t encountered an issue with that yet but it could be nicer (although maybe a bit slower) to use a dedicated GPU for the gradient accumulation such that memory spikes can be avoided</p>

halahup: <p>I haven’t had this issue either yet, but I see where you are coming from. It would be great to have an option to have a smaller GPU for just gradient accumulation as a central hub.</p>

rasbt: <p>I just realized that my answer was in a different context … regarding your question, I thought you were referring to <a href="https://discuss.pytorch.org/t/uneven-gpu-utilization-during-training-backpropagation/36117/13" class="inline-onebox">Uneven GPU utilization during training backpropagation</a> So my answer may be a bit out of context given the problem discussed in this thread <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=6" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>

Decker: <p><a class="mention" href="/u/halahup">@halahup</a> - It’s 100% the memory motivation. Loading this full model for training exceeds the amount of vram in a V100, so data parallelism won’t help.</p>

halahup: <p>Ok, gotcha, then yes!</p>

Decker: <p>Could the issue be related to this?  On pytorch’s <a href="https://pytorch.org/docs/stable/notes/cuda.html" rel="nofollow noopener">CUDA Semantics</a> page, it says: “Unless you enable peer-to-peer memory access, any attempts to launch ops on tensors spread across different devices will raise an error.”</p>
<p>Does running MSELoss count as “launching ops”? I’d test it by enabling it, but there doesn’t seem to be any documentation on how to enable GPU peer-to-peer memory access on pytorch either here or on Google, apart from one <a href="https://discuss.pytorch.org/t/how-to-enable-gpu-kernel-peer-to-peer-access/3324">post</a> where <a class="mention" href="/u/smth">@smth</a> indicates it’s on by default…</p>

shinx: <p>Seems like your labels and outputs are on different GPUs. Move them to the same GPU and it should work. Something like:</p>
<pre><code class="lang-auto">output = model(data)
target = target.to(output.device)
loss = F.nll_loss(output, target)
loss.backward()
</code></pre>