Uneven GPU utilization during training backpropagation

username: neonrights
cooked: <p>Hello PyTorch community,</p>
<p>I used pytorch’s DataParallel module in order to parallelize training of a large model.  I noticed the utilization of my primary GPU (coda:0) would be far larger than my other GPUs.  I checked previous topics and implemented recommended changes such as calculating loss in the forward pass.  However this problem persists.  I profiled my GPU usage and noticed it is evenly distributed until calling backward() on my loss.  Memory usage on the primary GPU then increases for several iterations until plateauing around the 10-20th iteration.  My training loop logic is pretty basic besides taking advantage of gradient accumulation so I am at a how to fix my problem.</p>
<pre><code class="lang-python"># model is a DataParallel module
for i, batch in enumerate(dataloader):
    if i % freq == 0:
        optimizer.zero_grad()

    loss = model(**batch) # even GPU utilization
    loss = loss.sum()
    loss.backward() # increase in primary GPU utilization

    if (i+1) % freq == 0:
        optimizer.step()
</code></pre>
<p>GPU usage before calling loss.backward()</p>
<pre><code class="lang-auto">+-------------------------------+----------------------+----------------------+
|   2  GeForce GTX 108...  Off  | 00000000:08:00.0 Off |                  N/A |
| 25%   26C    P8    15W / 250W |   3599MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GeForce GTX 108...  Off  | 00000000:09:00.0 Off |                  N/A |
| 25%   35C    P8    16W / 250W |   3633MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   4  GeForce GTX 108...  Off  | 00000000:83:00.0 Off |                  N/A |
| 25%   35C    P8    15W / 250W |   3633MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   5  GeForce GTX 108...  Off  | 00000000:84:00.0 Off |                  N/A |
| 25%   33C    P8    15W / 250W |   3633MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
</code></pre>
<p>GPU usage after calling loss.backward()</p>
<pre><code class="lang-auto">+-------------------------------+----------------------+----------------------+
|   2  GeForce GTX 108...  Off  | 00000000:08:00.0 Off |                  N/A |
| 25%   28C    P2    72W / 250W |   7101MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GeForce GTX 108...  Off  | 00000000:09:00.0 Off |                  N/A |
| 25%   37C    P2    72W / 250W |   4335MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   4  GeForce GTX 108...  Off  | 00000000:83:00.0 Off |                  N/A |
| 25%   36C    P2    66W / 250W |   4335MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   5  GeForce GTX 108...  Off  | 00000000:84:00.0 Off |                  N/A |
| 25%   35C    P2    70W / 250W |   4335MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
</code></pre>
<p>GPU usage a couple iterations later</p>
<pre><code class="lang-auto">+-------------------------------+----------------------+----------------------+
|   2  GeForce GTX 108...  Off  | 00000000:08:00.0 Off |                  N/A |
| 25%   32C    P2    73W / 250W |   9957MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GeForce GTX 108...  Off  | 00000000:09:00.0 Off |                  N/A |
| 25%   42C    P2    74W / 250W |   4335MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   4  GeForce GTX 108...  Off  | 00000000:83:00.0 Off |                  N/A |
| 25%   41C    P2    73W / 250W |   4335MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   5  GeForce GTX 108...  Off  | 00000000:84:00.0 Off |                  N/A |
| 25%   39C    P2    70W / 250W |   4335MiB / 11178MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
</code></pre>

username: neonrights
cooked: <p>Upon further profiling, it seems almost all the remaining space on primary GPU is being consumed by gradient accumulation.  When calling step() on the optimizer, the remaining space is used resulting in an out of memory error.  Is this intended behavior to store gradients and optimizers on GPU 0?  If so this seems highly inconvenient…</p>

username: rasbt
cooked: <p>This is normal and expected as you load the data onto the primary GPU in DataParallel and also the gradient accumulation happens on the primary GPU as you mentioned.</p>

username: neonrights
cooked: <p>Then is there a way to store accumulated gradients and/or the optimizer on a separate GPU or perform the forward pass without using the primary GPU?  As it stands, 3/4 of my GPUs are not even half utilized.</p>

username: rasbt
cooked: <p>What you can do is to designate a different GPU to do this. However, if the accumulated gradients fit on one GPU, it’s not worth thinking about hacks to do that on separate GPUs imho, because doing it on one device should be the fasted way to do it. Actually, I think they are accumulated to calculate the loss, and the gradients, I’ve sketched it out once …</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/e/e153c9fa07a2688c77fc1d7a572f26d12828528d.png" data-download-href="https://discuss.pytorch.org/uploads/default/e153c9fa07a2688c77fc1d7a572f26d12828528d" title="dataparallel"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/e/e153c9fa07a2688c77fc1d7a572f26d12828528d_2_275x500.png" alt="dataparallel" data-base62-sha1="w9kWLHLiy0iUDAkiTTrmktaGqLj" width="275" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/e/e153c9fa07a2688c77fc1d7a572f26d12828528d_2_275x500.png, https://discuss.pytorch.org/uploads/default/optimized/2X/e/e153c9fa07a2688c77fc1d7a572f26d12828528d_2_412x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/e/e153c9fa07a2688c77fc1d7a572f26d12828528d_2_550x1000.png 2x" data-dominant-color="352842"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">dataparallel</span><span class="informations">1275×2312 660 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The reason why it is computed on the default GPU is that you typically have something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">start_time = time.time()
for epoch in range(num_epochs):
    
    model.train()
    for batch_idx, (features, targets) in enumerate(train_loader):
        
        features = features.to(device)
        targets = targets.to(device)
            
        ### FORWARD AND BACK PROP
        logits, probas = model(features)
        cost = F.cross_entropy(logits, targets)
        optimizer.zero_grad()
        
        cost.backward()
        
        ### UPDATE MODEL PARAMETERS
        optimizer.step()
</code></pre>
<p>Note that <code>targets</code> (e.g., the class labels) is on the default GPU. Hence, if you want to do what you suggest you’d need to rewrite the code to split the class labels over the GPUs as well and then compute the loss on these separately and then combine the loss to compute the gradients. It’s not worth the hassle, imho, if the <code>targets</code> fit on the default GPU. Splitting it as well will only make it slower as the loss computation is super fast anyways.</p>

username: neonrights
cooked: <p>Oh sorry, let me clarify.  My problem is that I cannot fit a single mini-batch, gradients, and optimizer on a single GPU (in this case the primary GPU).  I originally thought, wrongly, the extra memory usage on the primary GPU was some sort of bug.</p>
<p>I now want to try to get around this by having a separate GPU dedicated to storing the gradients and optimizer which is not used in the forward or backward pass.  So if I have 5 GPUs, I would use 1-4 to calculate the loss and gradients, then send the resulting gradients to the 5th GPU in order to update the parameters.</p>
<p>Thanks for the help so far though!</p>

username: rasbt
cooked: <p>Hm, in that case, instead of wasting one GPU, it is maybe cheaper to just compute the loss on the CPU. I.e. you could try</p>
<pre><code>    features = features.to(device)
    targets = targets
    logits = model(features)
    cost = F.cross_entropy(logits.to(torch.device('cpu'), targets)
</code></pre>
<p>where “model” is still a model via dataparallel.</p>
<p>Generally though, I would recommend scaling the batchsize by 1/num_cuda_devices if you use dataparallel and want results that are somewhat comparable to the 1-GPU version.</p>

username: neonrights
cooked: <p>I am able to compute the loss and gradients on the GPU.  I am simply unable to store the resultant gradients in a single GPU while also using the same GPU to compute future losses and gradients.  So I simply want to gather and store the gradients elsewhere, like another GPU or CPU memory.</p>
<p>I tried computing loss on CPU, however the gradients still end up on the primary GPU.</p>

username: rasbt
cooked: <p>Hm, that’s tricky then.<br>
What you could try is to set <code>output_device=torch.device('cpu')</code> in DataParallel, but I am not sure if this will work and how fast this will be. I’d probably rather go with a smaller batch size instead of doing that.</p>
<p>Another thing that I don’t know if it would work is e.g., using device ids for cuda 0-2 via <code>device_ids</code> and then setting <code>output_device=torch.device('cuda:3')</code>. Let me know if that works, I’d be curious.</p>

username: neonrights
cooked: <p>Unfortunately, setting output to CPU results in <code>ValueError: Expected a cuda device, but got: cpu</code>.  Setting output to GPU 3 seems to only collect the loss to that GPU.  Gradients seem to still be gathered in GPU 0.</p>
<p>Reducing batch size won’t work since I am already processing a single sample on each GPU due to the large size of my model.</p>

username: rasbt
cooked: <p>oh that’s too bad. Unfortunately, I don’t have any better idea how to tackle this case.</p>
<blockquote>
<p>Setting output to GPU 3 seems to only collect the loss to that GPU. Gradients seem to still be gathered in GPU 0.</p>
</blockquote>
<p>I think this would be a useful thing to get to work properly such that the gradients are also gathered on the output device. I am not sure if this is a bug or a limitation, but I think this could be a useful feature/fix in PyTorch. Maybe that could be posted on the issue tracker for further discussion/investigation.</p>

username: neonrights
cooked: <p>Opened an issue on <a href="https://github.com/pytorch/pytorch/issues/16700" rel="nofollow noopener">github</a>.</p>

username: rasbt
cooked: <p>I actually just found that someone implemented a solution to this problem. While this is exactly what you were looking for, i haven’t tried it myself yet and am not sure how well it works – in theory though, it addresses the problem of unbalanced loads in DataParallel. They describe it in the context of Semantic Segmentation, but I assume it should work in more general and encompass other objective functions:</p>
<p><a href="https://hangzhang.org/PyTorch-Encoding/parallel.html" class="onebox" target="_blank" rel="nofollow noopener">https://hangzhang.org/PyTorch-Encoding/parallel.html</a></p>

username: colllin
cooked: <p>I got lucky today and found a way to get near-equal utilization across GPUs.</p>
<p>The two main changes were:</p>
<ol>
<li>
<p>Make sure my outer-most module is wrapped by DataParallel.  Before I had an inner module (like a layer) wrapped only.</p>
<p>From:</p>
<pre><code> class BundledNet(nn.Module):
     def __init__(self, net, preproc=None):
         self.net = net
         if torch.cuda.device_count() &gt; 1:
             self.net = torch.nn.DataParallel(self.net)
         self.preproc = preproc

     def forward(self, x):
         if self.preproc != None:
             x = self.preproc(x)
         x = self.net(x)
         return x

 net = BundledNet(ResNet(), preproc=resnet_norm)
</code></pre>
<p>To:</p>
<pre><code> class BundledNet(nn.Module):
     def __init__(self, net, preproc=None):
         self.net = net
         self.preproc = preproc

     def forward(self, x):
         if self.preproc != None:
             x = self.preproc(x)
         x = self.net(x)
         return x

 net = BundledNet(ResNet(), preproc=resnet_norm)
 if torch.cuda.device_count() &gt; 1:
     net = torch.nn.DataParallel(net)
</code></pre>
<p>It appears that the preprocessing and post-processing in the forward pass of <code>BundledNet</code> was enough to cause an imbalance across the GPUs.</p>
</li>
<li>
<p>Compute my loss function inside a DataParallel module.</p>
<p>From:</p>
<pre><code> loss = torch.nn.CrossEntropyLoss()
</code></pre>
<p>To:</p>
<pre><code> loss = torch.nn.CrossEntropyLoss()
 if torch.cuda.device_count() &gt; 1:
     loss = CriterionParallel(loss)
</code></pre>
<p>Given:</p>
<pre><code> class ModularizedFunction(torch.nn.Module):
     """
     A Module which calls the specified function in place of the forward pass.
     Useful when your existing loss is functional and you need it to be a Module.
     """
     def __init__(self, forward_op):
         super().__init__()
         self.forward_op = forward_op
     
     def forward(self, *args, **kwargs):
         return self.forward_op(*args, **kwargs)

 class CriterionParallel(torch.nn.Module):
     def __init__(self, criterion):
         super().__init__()
         if not isinstance(criterion, torch.nn.Module):
             criterion = ModularizedFunction(criterion)
         self.criterion = torch.nn.DataParallel(criterion)
 
     def forward(self, *args, **kwargs):
         """
         Note the .mean() here, which is required since DataParallel 
         gathers any scalar outputs of forward() into a vector with 
         one item per GPU (See DataParallel docs).
         """
         return self.criterion(*args, **kwargs).mean()
</code></pre>
</li>
</ol>
<p>The general rule seems to be that you should try not to do any cuda tensor computation whatsoever outside the <code>forward</code> pass of a DataParallel module.  Outside of a DataParallel module, I am only passing tensors around and moving them on and off the GPU.  See this before/after image of <code>nvidia-smi</code> to see the results:</p>
<p>Before:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/b/bdb4a0bca632be29f70d28a4953fe89f242f86d7.png" data-download-href="https://discuss.pytorch.org/uploads/default/bdb4a0bca632be29f70d28a4953fe89f242f86d7" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/2X/b/bdb4a0bca632be29f70d28a4953fe89f242f86d7.png" alt="image" data-base62-sha1="r4deAIQMKQGOeKuKymkzojMPSej" width="237" height="500" data-dominant-color="3D3D3D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">394×830 24 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>After:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/f/fcd825880497d339cbceb7da5a3d85d109eb7fac.png" data-download-href="https://discuss.pytorch.org/uploads/default/fcd825880497d339cbceb7da5a3d85d109eb7fac" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/f/fcd825880497d339cbceb7da5a3d85d109eb7fac_2_435x500.png" alt="image" data-base62-sha1="A4LpHqgMmkNP98wdjA7yQgQ3oWo" width="435" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/f/fcd825880497d339cbceb7da5a3d85d109eb7fac_2_435x500.png, https://discuss.pytorch.org/uploads/default/optimized/2X/f/fcd825880497d339cbceb7da5a3d85d109eb7fac_2_652x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/f/fcd825880497d339cbceb7da5a3d85d109eb7fac_2_870x1000.png 2x" data-dominant-color="343434"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1172×1346 98.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: Idan_Azuri
cooked: <aside class="quote no-group" data-username="colllin" data-post="14" data-topic="36117">
<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/colllin/48/2918_2.png" class="avatar"> colllin:</div>
<blockquote>
<p>torch.nn.CrossEntropyLoss()</p>
</blockquote>
</aside>
<p>Thanks! The trick of computing the loss function inside DataParallel actually works!</p>