DataParallel and nested functions in nn.Module

helange23: <p>I am encountering a very strange bug. The code that produces the bug is fairly long so below is a sketch of what is going.</p>
<p>I have an nn.Module that describes a feedforward NN. This module has another nn.Module as a submodule that describes a custom layer. The layer module has a forward() function, some nn.Linear layers and several functions.</p>
<pre><code class="lang-auto">class feedforward(nn.Module):
     def __init__(self, num_kernels, levels):
           super(feedforward, self).__init__()
           self.layer = custom_layer()

     def forward(self, inpt):
           return self.layer(inpt)


class custom_layer(nn.Module):
     def __init__(self, num_kernels, levels):
           super(custom_layer, self).__init__()
           self.l1 = nn.Linear(...)

     def func_that_uses_l1(self, other):
           print('Devices #2',  other.device, self.l1.weight.device)
           return self.l1(other)

     def somefunc(self, inpt):
           device = self.l1.weight.device
           other = torch.arange(10, device = device)
           print('Devices #1',  other.device, self.l1.weight.device)
           something = self.func_that_uses_l1(other)
           return function of inpt and something

     def forward(self, inpt):
           return self.somefunc(inpt)

net = nn.DataParallel(feedforward()).to('cuda:0'), device_ids=[0,1,2,3])
outpt = net(inpt)
</code></pre>
<p>The strange thing that happens is that in <em>somefunc</em> the devices are equal and correctly distributed over the GPUs, i.e. self.l1.weight is cuda:0, cuda:1, cuda:2 and so on but in <em>func_that_uses_l1</em> self.l1 is always on device cuda:0, i.e. the output is the following:</p>
<pre><code class="lang-auto">Devices#1 cuda:1 cuda:1
Devices#2 cuda:1 cuda:0
Devices#1 cuda:2 cuda:2
Devices#2 cuda:2 cuda:0
Devices#1 cuda:3 cuda:3
Devices#2 cuda:3 cuda:0
</code></pre>
<p>This obviously raises a “RuntimeError: arguments are located on different GPUs”</p>
<p>How can something like this happen? Am I not allowed to nest functions when I am using nn.DataParallel? It is so odd because I make sure that other is on the correct device (and it works properly) but then when <em>func_that_uses_l1</em> is called, self.l1 is magically back on cuda:0.</p>
<p>Note: The code is a sketch and not a minimal working example.</p>
<p>Update: The error vanishes when I get rid of the nesting, i.e. when I change <em>custom_layer</em> to this:</p>
<pre><code class="lang-auto">class custom_layer(nn.Module):
     def __init__(self, num_kernels, levels):
           super(custom_layer, self).__init__()
           self.l1 = nn.Linear(...)


     def somefunc(self, inpt):

          def func_that_uses_l1(other):
                print('Devices #2',  other.device, self.l1.weight.device)
                return self.l1(other)

           device = self.l1.weight.device
           other = torch.arange(10, device = device)
           print('Devices #1',  other.device, self.l1.weight.device)
           something = func_that_uses_l1(other)
           return function of inpt and something

     def forward(self, inpt):
           return self.somefunc(inpt)
</code></pre>
<p>But this is problematic for my code because I would like to reuse <em>func_that_uses_l1</em>. Is this behavior intended?</p>

ptrblck: <p>The issue is most likely raised, as <code>nn.DataParallel</code> uses the <code>forward</code> method to send each chunk of the data to the appropriate device.<br>
I assume that you are calling <code>func_that_uses_l1</code> directly, which won’t work.<br>
You could try to pass a flag to <code>forward</code> to call the corresponding method.<br>
Alternatively you would have to use scatter and gather the data manually in <code>fun_that_uses_l1</code>.</p>

helange23: <p>I am not calling <em>func_that_uses_l1</em> directly. It is being called by forward method. Specifically, the forward function calls <em>somefunc</em> which calls <em>func_that_uses_l1</em>. If I remove the nesting, the problem vanishes.</p>
<p>I am 95% certain this is a bug.</p>

ptrblck: <p>Could you post a reproducible code snippet, which shows this bug, so that we could debug and fix it?</p>