Using register_buffer with DataParallel and cuda

deJQK: <p>I am using DataParallel to allow the program run on several GPUs, and use register_buffer to define some variable for my modules, like this:</p>
<pre><code class="lang-auto">class A(nn.Module):
    def __init__(self):
        ...
        self.register_buffer('foo', torch.empty_like(self.weight, dtype=torch.long))

    def forward(self):
        ...
        print(self.foo)
        self.foo = torch.tensor(2018)
        print(self.foo)
</code></pre>
<p>For the parallel, I use</p>
<pre><code class="lang-auto">device = torch.device('cuda:0')
model = ...
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])
model = model.to(device)
</code></pre>
<p>I will get the following output and clearly the user-defined buffer “foo” is not assigned correctly:</p>
<blockquote>
<p>a:  tensor([140174715505560], device=‘cuda:3’)<br>
b:  tensor(2018)<br>
a:  tensor([140174715505560], device=‘cuda:2’)<br>
b:  tensor(2018)<br>
tensor([140174715505560], device=‘cuda:0’)<br>
b:  tensor(2018)<br>
a:  a:  tensor([140174715505560], device=‘cuda:1’)<br>
b:  tensor(2018)<br>
a:  tensor([140174715505528], device=‘cuda:3’)<br>
b:  tensor(2018)<br>
tensor([140174715505528], device=‘cuda:2’)<br>
b:  tensor(2018)<br>
tensor([140174715505528], device=‘cuda:0’)<br>
b:  tensor(2018)</p>
</blockquote>
<p>I am using nvidia 2080 and cuda 9.2. There is some other bugs reported, but does not influence the program running. The bugs are like this:</p>
<blockquote>
<p>THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/THCGeneral.cpp line=663 error=11 : invalid argument<br>
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/THCGeneral.cpp line=663 error=11 : invalid argument<br>
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/THCGeneral.cpp line=663 error=11 : invalid argument</p>
</blockquote>
<p>Also, I did not use the following part (by comment these two lines)</p>
<pre><code class="lang-auto">import torch.backends.cudnn
cudnn.benchmark = True
</code></pre>
<p>because otherwise there will be some other bugs and the program cannot run.</p>

ptrblck: <p>Could you try to assign the device to your new tensor using the device of another parameter?<br>
E.g. if your <code>A</code> module has some layers or other parameters stored, try the following:</p>
<pre><code class="lang-python">print(self.foo)
self.foo = torch.tensor(2018, device=self.another_parameter.device)
print(self.foo)
</code></pre>

deJQK: <p>Thanks for your reply. I tried now, but the result is the same.</p>
<p>Here is the modified part:</p>
<pre><code class="lang-auto">self.foo = torch.tensor(2018, device=self.weight.device)
</code></pre>
<p>Because A is a layer, I use self.weight.</p>
<p>It seems to me that inside the method of the module class, the buffer variable can be assigned correctly, but after that, it is reinitialized, or maybe not modified at all in the method of the class.</p>
<p>I also print the buffer variable outside the code file which contains the definition of class is as following:</p>
<pre><code class="lang-auto">for m in model.modules():
    if isinstance(m, A):
        print('c: ', m.foo.item)
</code></pre>
<p>and will output the random large number, instead of 2018.</p>
<p>With single GPU, I do not need to use the buffer, and there is no such problem.</p>

ptrblck: <p>Thanks for the explanation,<br>
In that case, could you try to use an in-place method like <code>self.foo.fill_(2018.)</code>?</p>

deJQK: <p>Thanks for your help. It doesn’t work either.<img src="https://discuss.pytorch.org/images/emoji/apple/frowning_face.png?v=6" title=":frowning_face:" class="emoji" alt=":frowning_face:"></p>

deJQK: <p>Actually the print out is a little messy, and it seems that there is some conflicts among the four GPUs to print out values:</p>
<blockquote>
<p>a:  a: a:   tensor([140339592887176], device=‘cuda:7’)<br>
tensor([140339592887176], device=‘cuda:6’)tensor([140339592887176], device=‘cuda:5’)a:</p>
<p>b:  b: tensor([140339592887176], device=‘cuda:4’)b:  tensor([2018], device=‘cuda:7’)</p>
<p>tensor([2018], device=‘cuda:5’)tensor([2018], device=‘cuda:6’)</p>
<p>b:  tensor([2018], device=‘cuda:4’)<br>
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/THCGeneral.cpp line=663 error=11 : invalid argument<br>
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/THCGeneral.cpp line=663 error=11 : invalid argument<br>
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/THCGeneral.cpp line=663 error=11 : invalid argument<br>
a:  tensor([140339592887176], device=‘cuda:4’)a:<br>
a: a:   b:  tensor([140339592887176], device=‘cuda:6’)<br>
tensor([2018], device=‘cuda:4’)<br>
b:  tensor([2018], device=‘cuda:6’)<br>
tensor([140339592887176], device=‘cuda:5’)tensor([140339592887176], device=‘cuda:7’)</p>
<p>b: b:   tensor([2018], device=‘cuda:5’)tensor([2018], device=‘cuda:7’)</p>
<p>a:  tensor([140339592887176], device=‘cuda:4’)<br>
a:  b:  tensor([140339592887176], device=‘cuda:6’)</p>
</blockquote>
<p>I can only understand this by that different layers are assigned to different GPUs, or maybe the model is assigned multiple times to different GPUs. I do not quite understand parallel computing.</p>

ptrblck: <p>The output look fine. Since CUDA calls are asynchronously, you won’t necessarily see a specific order of the print statements.<br>
However, if looks like <code>self.foo</code> is now on different devices.<br>
Are you getting any other error message now?</p>

deJQK: <p>Thanks. Actually no, the only error messages are the three ones posted above, and the program can run until finished.</p>
<p>I was trying to understand how to use register_buffer by checking the code of batch_norm, but it use some function torch.batch_norm, which seems not a python function, but implemented in C/C++.</p>

deJQK: <p>I tried a very tiny toy example. The output seems to be correct. There is no such problem in the tiny example. Since there are around 5 files even in this tiny example, I cannot copy all of them here.</p>
<p>Thanks very much. I will try to debug a little.</p>

deJQK: <p>Thank you very much. The problem is solved. The most important thing is that I can only use self.foo.fill_(some_tensor.item()), and avoid ANY using of self.foo = some_tensor. Even if I use self.foo.fill_(value) after self.foo = some_tensor, the result is incorrect. I don’t know why. Anyway, thank you very much.</p>

ptrblck: <p>Good to hear it’s working now!<br>
I still think it should be working if you specify the device in the tensor creation.</p>

deJQK: <p>Thanks. The following does not work:</p>
<pre><code class="lang-auto">self.foo = self.weight.new(2018).long()
</code></pre>

deJQK: <p>Hi, I feel sorry to bother you again, but after a long time simulation, the results seem to be incorrect. Do you have any idea of the reason? I think there should be no other bugs as the program is simple and the result is correct when I use single GPU. Thanks again.</p>
<p>With single GPU, I comment the following line</p>
<pre><code class="lang-auto">model=nn.DataParallel(model, device_ids=....)
</code></pre>

SimonW: <p>The way DataParallel works is that at each iteration, it creates a replica of the input model on each device. So each device has it’s own replica of the model, a different python object than the input model. So your original approach of directly assigning doesn’t work, as it is assigning as an attribute on a replica, which is freed after the iteration.</p>
<p>The weights &amp; buffers of these replicas are linked to the original weights &amp; buffers in input model by autograd graph, so backward still works. Moreover, weights &amp; buffers of the replica on device[0] share storage with those of the input model! So the inplace version “works” in the sense that changes done on the first replica (on device[0]) will propagate through. However, changes from other replicas are lost.</p>

deJQK: <p>Thank you very much for the explanation. I partly understand what you said now. So if I want to, say check some self defined status, like “foo” of the weights of layers in the model, using multiple GPUs (as I need to do experiments on very large dataset), is there any way to achieve what I want? Thank you again.</p>

deJQK: <p>It seems that if I only specify one GPU, even if I use <code>nn.DataParallel</code>, for example, as the following</p>
<pre><code class="lang-auto">model = nn.DataParallel(model, device_ids=[0])
</code></pre>
<p>the result will be correct. But if I use more than one, say</p>
<pre><code class="lang-auto">model = nn.DataParallel(model, device_ids=[0, 1])
</code></pre>
<p>the result will make no sense.</p>

Cosimo_Rulli: <p>Hi,<br>
did you manage to solve this problem? I am facing something simillar.</p>

deJQK: <p>I think DistributedDataParallel does not have this problem. I have to say I totally forgot how did I solve this. Sorry for this. I should write down the final answer back then.</p>