Multi-GPU backward error

username: doliolarzz
cooked: <p>Hello,<br>
A model has time-series structure.<br>
So, I’ve split each time-step into each gpu, due to insufficiency of gpu memory.</p>
<p>But I got an error when calling loss.backward():</p>
<pre><code class="lang-auto">File "/home/xxx/.local/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function CatBackward returned an invalid gradient at index 1 - expected device cuda:2 but got cuda:0
</code></pre>
<p>Please find the minimum code below:</p>
<pre><code class="lang-auto">    self.config['DEVICE_ALL'] = [0, 2, 3]
    self.backbone = UNet(n_channels=config['IN_LEN']*3 - 2)
    self.geo = nn.Parameter(data=torch.randn(1, 4, self.h, self.w), requires_grad=True)
    self.outConv = nn.Conv2d(68, 1, kernel_size=1)
    self.output_time_length = 3

    def forward(self, input):

        outputs = []
        cur_input = input
        cur_state = self.get_state(...)

        for i in range(self.output_time_length):
            
            dev = self.config['DEVICE_ALL'][i]
            cur_input = cur_input.cuda(dev)
            cur_state = cur_state.cuda(dev)
            x = torch.cat([cur_input, cur_state], 1).cuda(dev)

            self.backbone = self.backbone.cuda(dev)
            x = self.backbone(x).cuda(dev)

            geo = self.geo.expand(x.shape[0], -1, -1, -1).cuda(dev)
            x = torch.cat([x, geo], 1).cuda(dev)

            self.outConv = self.outConv.cuda(dev)
            x = self.outConv(x).cuda(dev)

            outputs.append(x)
            cur_input = torch.cat([cur_input[:, 1:], x], 1)
            cur_state = self.get_next_state(...)

        return torch.cat(outputs, 1).cuda(dev)
</code></pre>
<p>Thanks in advance!</p>

username: albanD
cooked: <p>Hi,</p>
<p>Could you try enabling anomaly detection mode to see which of the cat is problematic please?<br>
When you have it, make sure that all the inputs are properly on the same GPU.</p>

username: doliolarzz
cooked: <p>Hello, thank you for you help.</p>
<blockquote>
<p>make sure that all the inputs are properly on the same GPU.</p>
</blockquote>
<p>The inputs and the labels(for calculating losses) are in device[cuda:3].</p>
<hr>
<p>I’ve enabled the anomaly detection by adding this line.</p>
<pre><code class="lang-auto">torch.autograd.set_detect_anomaly(True)
</code></pre>
<p>The warning was appeared as below:</p>
<pre><code class="lang-auto">Warning: Traceback of forward call that caused the error:
  File "train.py", line 77, in &lt;module&gt;
    main()
  File "train.py", line 70, in main
    trainer.train()
  File "../../utils/trainer.py", line 186, in train
    self.train_iteration()
  File "../../utils/trainer.py", line 162, in train_iteration
    output = self.model(train_data)
  File "/home/xxx/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "../../models/model.py", line 64, in forward
    return torch.cat(outputs, 1).cuda(dev)
 (print_stack at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:57)
</code></pre>
<p>Thank you.</p>

username: albanD
cooked: <p>Awesome.<br>
Can you print the size and device of each element in outputs juste before that line?</p>

username: doliolarzz
cooked: <p>The outputs seem to be like this:</p>
<pre><code class="lang-auto">cuda:0 torch.Size([2, 1, 840, 640])
cuda:2 torch.Size([2, 1, 840, 640])
cuda:3 torch.Size([2, 1, 840, 640])
</code></pre>
<p>These size are in format of (batch_size, channel, Height, Width).</p>
<p>The labels data for back propagation are:<br>
<code>cuda:3, torch.Size([2, 3, 1, 840, 640])</code></p>
<p>The ‘MSELoss’ function is also sent to ‘cuda:3’ too.</p>

username: albanD
cooked: <p>You should make sure that all the inputs are on the same device before giving them to cat. I am surprised cat did not raise an error during the forward…</p>

username: doliolarzz
cooked: <p>Sorry, I’ve tried to make all of inputs to be on the same device.<br>
by changing</p>
<pre><code class="lang-auto">return torch.cat(outputs, 1).cuda(dev)
</code></pre>
<p>to</p>
<pre><code class="lang-auto">return torch.cat([o.cuda(dev) for o in outputs], 1).cuda(dev)
</code></pre>
<p>All of tensors will be on cuda:3 before cat.<br>
But it raise a new error. <strong>Is it necessary to compute loss separately in each output? (Because it’s in separate device)</strong></p>
<pre><code class="lang-auto">RuntimeError: expected device cuda:3 but got device cuda:2
</code></pre>
<p>and below is the new anomaly detect warning:</p>
<pre><code class="lang-auto">Warning: Traceback of forward call that caused the error:
  File "train.py", line 77, in &lt;module&gt;
    main()
  File "train.py", line 70, in main
    trainer.train()
  File "../../utils/trainer.py", line 186, in train
    self.train_iteration()
  File "../../utils/trainer.py", line 162, in train_iteration
    output = self.model(train_data)
  File "/home/xxx/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "../../models/mrunet/model.py", line 58, in forward
    x = self.outConv(x).cuda(dev)
  File "/home/xxx/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/xxx/.local/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 345, in forward
    return self.conv2d_forward(input, self.weight)
  File "/home/xxx/.local/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 342, in conv2d_forward
    self.padding, self.dilation, self.groups)
 (print_stack at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:57)
</code></pre>

username: albanD
cooked: <p>Well when you concatenate the Tensor, you will end up with a single Tensor. That single Tensor can be only on a single device, So you already compute the loss on a single device no?</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> do we automatically send Tensors accross devices in torch.cat? That sounds wrong no?</p>

username: ptrblck
cooked: <p>I think this is the case.<br>
<a href="https://github.com/pytorch/pytorch/blob/480d1849b04548148bc7c7cf9c883408e8753104/aten/src/ATen/native/cuda/Shape.cu#L260">Here</a> the output tensor will be initialized with the device (options) of the first tensor.<br>
The device check is <a href="https://github.com/pytorch/pytorch/blob/480d1849b04548148bc7c7cf9c883408e8753104/aten/src/ATen/native/cuda/Shape.cu#L365-L368">here</a> and dispatches to <code>parallel_cat</code>.<br>
If different devices are found (or any other condition returns <code>false</code>), the <code>copy_</code> should push the tensors to the first device in <a href="https://github.com/pytorch/pytorch/blob/480d1849b04548148bc7c7cf9c883408e8753104/aten/src/ATen/native/cuda/Shape.cu#L387-L396">these lines of code</a>.</p>
<p>Could this explain the issue using multiple GPUs?</p>

username: doliolarzz
cooked: <aside class="quote no-group" data-username="albanD" data-post="8" data-topic="73427">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>So you already compute the loss on a single device no?</p>
</blockquote>
</aside>
<p>Yes, the loss is computed in a single device. But I really don’t understand why it still raise that error.</p>
<p>I’m very newbie in PyTorch. Could you please guide me what I’ve done wrong. Thank you.</p>

username: albanD
cooked: <p>Well the issue is that the backward formula has not been updated to reflect this… You can see <a href="https://github.com/pytorch/pytorch/blob/51d969e86ac0fad226f4e70889df0d3c6114ae4e/tools/autograd/templates/Functions.cpp#L523">here</a> that it assumes that all inputs are on the same device.</p>

username: doliolarzz
cooked: <p>Hello,<br>
I see. I think I’ll try to calculate loss and backprops in separate devices.<br>
Do you have any other suggestion or workaround?</p>
<p>Thank you. I’m really appreciated your help : )</p>