How to run with multiple GPU when having 2 optimizers?

username: tcsn_wty
cooked: <p>I think I totally understood the tutorial on using the following line of code to do data parallel.</p>
<pre><code class="lang-auto">model = nn.DataParallel(model) 
</code></pre>
<p>I tested it on my device with 2 GPUs and it worked. However, for a complicated model I recently wrote, it doesn’t truly train on both GPUs when I printed the ‘Inside’ ‘Outside’ debugging message.</p>
<p>I managed to figure out that it was because of the optimizers. To my purpose, I need to have two optimizers each contains some shared part and some unique part.</p>
<pre><code class="lang-auto">model = model.module
optimizer_1 = optim.Adam([{'params': model.shared.parameters()},
                              {'params': model.unique_a.parameters()}],
                             lr=5e-5
                             )
optimizer_2 = optim.Adam([{'params': model.encoder.parameters()},
                              {'params': model.unique_b.parameters()}],
                             lr=5e-5
                             )
</code></pre>
<p>I deleted anything related to optimizers and loss criterions and it worked just fine on 2 GPUs according to the printed debugging message. But whenever the optimizers and loss criterions are added, it can only be trained on one single GPU.</p>
<p>I wonder if there’s any suggestions on this to make it work?</p>
<p>Thanks!</p>

username: ptrblck
cooked: <p>Could you post an executable code snippet, as the optimizers shouldn’t change the behavior of <code>nn.DataParallel</code> in any way, please?</p>

username: tcsn_wty
cooked: <p>I’ll post lines that are correlated to this question below.</p>
<pre><code class="lang-auto">device = "cuda"

model = nn.DataParallel(model)  # For multi-gpu
model = model.module  # For multi-gpu
model.to(device)

optimizer_1 = optim.Adam([{'params': model.shared.parameters()},
                              {'params': model.unique_a.parameters()}],
                             lr=5e-5
                             )
optimizer_2 = optim.Adam([{'params': model.encoder.parameters()},
                              {'params': model.unique_b.parameters()}],
                             lr=5e-5
                             )
criterion = nn.L1Loss()
epochs = 30
for epoch in range(epochs):
        output = model(...)
        loss = criterion(output, target)
        loss.backward()
        optimizer_1.step()
        print("A Outside: output size", output.shape)
</code></pre>
<p>Within the model I have a print statement to print “Inside”, and I wish to see “inside” printed twice and “outside” printed once as I have 2 GPUs.<br>
Please note: when I comment out anything related to optimizers and loss, I got the output I wanted. However, when optimizers and loss are included, it runs but I can easily see the output to show that it was only utilizing one GPU.</p>
<p>Please let me know if you need me to post any other piece of code.</p>
<p>Thanks!</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="tcsn_wty" data-post="3" data-topic="117693">
<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/tcsn_wty/48/11956_2.png" class="avatar"> tcsn_wty:</div>
<blockquote>
<p>Please let me know if you need me to post any other piece of code.</p>
</blockquote>
</aside>
<p>Yes, an executable code snippet would be needed to debug this issue further. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: tcsn_wty
cooked: <aside class="quote no-group" data-username="tcsn_wty" data-post="3" data-topic="117693">
<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/tcsn_wty/48/11956_2.png" class="avatar"> tcsn_wty:</div>
<blockquote>
<p><code>criterion = nn.L1Loss()</code></p>
</blockquote>
</aside>
<p>Here it is.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable

# Parameters and DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100

class Model(nn.Module):
    # Our model
    
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc_a = nn.Linear(input_size, output_size)  # Unique A
        self.fc_b = nn.Linear(input_size, output_size)  # Unique B
        self.fc_out = nn.Linear(output_size, output_size)  # Shared layer
    
    def forward(self, input, select='A'):
        if select is 'A':
            input = self.fc_a(input)
        else:
            input = self.fc_b(input)
        output = self.fc_out(input)
        print("\tIn Model: output size", output.size())
        return output

if torch.cuda.device_count() &gt; 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # dim = 0 [30, xxx] -&gt; [10, ...], [10, ...], [10, ...] on 3 GPUs

device = "cuda"
model = Model(input_size, output_size)
model = nn.DataParallel(model)  # For multi-gpu
model = model.module  # For multi-gpu
model.to(device)

optimizer_1 = optim.Adam([{'params': model.fc_a.parameters()},
                          {'params': model.fc_out.parameters()}],
                         lr=5e-5
                         )
optimizer_2 = optim.Adam([{'params': model.fc_b.parameters()},
                          {'params': model.fc_out.parameters()}],
                         lr=5e-5
                         )
criterion = nn.L1Loss()

batch_size = 16

print(model)

epochs = 10
for epoch in range(epochs):
    data_a, target_a = torch.randn(batch_size, data_size, input_size), torch.randn(batch_size, data_size, output_size)
    data_b, target_b = torch.randn(batch_size, data_size, input_size), torch.randn(batch_size, data_size, output_size)
    input_a, target_a = Variable(data_a.float()).to(device), Variable(target_a.float()).to(device)
    input_b, target_b = Variable(data_b.float()).to(device), Variable(target_b.float()).to(device)
    
    output_a = model(input_a, select='A')
    loss1 = criterion(output_a, target_a)
    loss1.backward()
    optimizer_1.step()
    print("Outside: output size", output_a.size())
    
    output_b = model(input_b, select='B')
    loss2 = criterion(output_b, target_b)
    loss2.backward()
    optimizer_2.step()
    print("Outside: output size", output_b.size())
</code></pre>
<p>Sorry for lacking of comments, please let me know if any explanation is needed.</p>
<p>By running this piece of code, I got the following kind of output:</p>
<pre><code class="lang-auto">        In Model: output size torch.Size([16, 100, 2])
Outside: output size torch.Size([16, 100, 2])
        In Model: output size torch.Size([16, 100, 2])
Outside: output size torch.Size([16, 100, 2])
        In Model: output size torch.Size([16, 100, 2])
Outside: output size torch.Size([16, 100, 2])
</code></pre>
<p>But I was expecting the following:</p>
<pre><code class="lang-auto">        In Model: output size torch.Size([8, 100, 2])
        In Model: output size torch.Size([8, 100, 2])
Outside: output size torch.Size([16, 100, 2])
</code></pre>

username: tcsn_wty
cooked: <p>Man just now I commented out everything related to optimizers and loss criterions but it suddenly still got the same output. (I swear in my real large project code it appears normal output when disable optimizers and loss.)<br>
But in general, the current output doesn’t look like 2 GPUs are utilized. Any suggestions?</p>
<p>FYI, I do have 2 GPUs available.</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 455.23.05    Driver Version: 455.23.05    CUDA Version: 11.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce GTX 108...  On   | 00000000:05:00.0 Off |                  N/A |
| 27%   26C    P8     9W / 250W |      1MiB / 11178MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  GeForce GTX 108...  On   | 00000000:06:00.0 Off |                  N/A |
| 27%   25C    P8     8W / 250W |      1MiB / 11178MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
</code></pre>

username: ptrblck
cooked: <p>The issue with your code is that you are removing the <code>nn.DataParallel</code> wrapper by calling:</p>
<pre><code class="lang-python">model = model.module  # For multi-gpu
</code></pre>
<p>so remove this line of code.<br>
After removing it, you would also have to change how the optimizers are created, as you need to either access the internal layers via <code>model.module.fc_x.parameters()</code> or create the optimizers before wrapping the model into <code>nn.DataParallel</code>.</p>

username: tcsn_wty
cooked: <p>Thanks a lot! That solves the whole problem.</p>
<p>A quick follow-up question: why this line of code does the removing on <code>nn.DataParallel</code> wrapper? I wasn’t expecting that to happen.</p>

username: ptrblck
cooked: <p><code>nn.DataParallel</code> assigns the passed model into the <code>.module</code> attribute and uses it internally to push to the different devices etc.<br>
Accessing this attribute via <code>model.module</code> gives you the initial model back and is used to e.g. store the <code>state_dict</code> without the <code>nn.DataParallel</code> wrapper.<br>
However, if you access this attribute and assign it back to <code>model</code>, you are basically just removing the <code>nn.DataParallel</code> usage.</p>