Questions about loss and backward process in Dataparallel

Jianxin_Wang: <p>Hi,</p>
<p>I am trying to understand how data parallel distributes loss and gradient back to each device for a backward pass after the forward pass.</p>
<p>I did a little experiment.</p>
<pre><code class="lang-auto">class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        modules = [
            torch.nn.Linear(10, 3),
            torch.nn.Linear(3, 4),
            torch.nn.Linear(4, 5),
        ]
        self.net = torch.nn.ModuleList(modules)

    def forward(self, inputs):
        for i, n in enumerate(self.net):
            inputs = n(inputs)

        return inputs

def main():
    X = np.random.uniform(-1, 1, (20, 10)).astype(np.float32)
    y = np.random.randint(0, 5, (20,))
    print(X.shape)
    print(y.shape)

    model = Net()
    loss = torch.nn.CrossEntropyLoss()
    #print('Model:', type(model))
    #print('Loss:', type(loss))

    X = torch.from_numpy(X)
    y = torch.from_numpy(y)
    print('X', X.size(), 'y', y.size())

    if torch.cuda.is_available():
        model = torch.nn.DataParallel(model)
        print('Model:', type(model))
        print('Devices:', model.device_ids)

        model = model.cuda()

        loss = loss.cuda()
        X = X.cuda()
        y = y.cuda()

    else:
        print('No devices available')

    X = Variable(X)
    y = Variable(y)

    outputs = model(X)
    l = loss(outputs, y)
    print('Loss:', l)

if __name__ == '__main__':
    main()
</code></pre>
<p>I got following results,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/5/05ce2402799aa11c27abac73cb94344b006c0193.png" data-download-href="https://discuss.pytorch.org/uploads/default/05ce2402799aa11c27abac73cb94344b006c0193" title="Screen Shot 2022-07-11 at 10.38.52 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/5/05ce2402799aa11c27abac73cb94344b006c0193_2_690x138.png" alt="Screen Shot 2022-07-11 at 10.38.52 PM" data-base62-sha1="Pm2AeMu7yN9Px8a9gDH3RpCmC7" width="690" height="138" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/5/05ce2402799aa11c27abac73cb94344b006c0193_2_690x138.png, https://discuss.pytorch.org/uploads/default/original/3X/0/5/05ce2402799aa11c27abac73cb94344b006c0193.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/5/05ce2402799aa11c27abac73cb94344b006c0193.png 2x" data-dominant-color="EDEDED"><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">Screen Shot 2022-07-11 at 10.38.52 PM</span><span class="informations">944×190 57.8 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>In my understanding, the loss should be in the format of [loss 1, loss 2, loss3, loss4] from the data parallel working diagram (first picture in the second row).<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/3/f35503fcf153b96dc59210cf738fbe1933b11111.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/f35503fcf153b96dc59210cf738fbe1933b11111" title="Screen Shot 2022-07-11 at 10.43.45 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f35503fcf153b96dc59210cf738fbe1933b11111_2_690x220.jpeg" alt="Screen Shot 2022-07-11 at 10.43.45 PM" data-base62-sha1="yICayQwTDefdS75DQCBzUtxwuvD" width="690" height="220" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f35503fcf153b96dc59210cf738fbe1933b11111_2_690x220.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f35503fcf153b96dc59210cf738fbe1933b11111_2_1035x330.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/3/f35503fcf153b96dc59210cf738fbe1933b11111_2_1380x440.jpeg 2x" data-dominant-color="D9E2E9"><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">Screen Shot 2022-07-11 at 10.43.45 PM</span><span class="informations">1920×613 104 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><br>
Why it is just a single number from the experiment code?</p>

ptrblck: <p>The outputs will be gathered on the default device and thus the loss will also be calculated on the default device. Afterwards the gradients are scattered to all GPUs as shown in the diagram.</p>

Jianxin_Wang: <p>Thanks for the response! I wonder how the scattering process decomposes the loss on the default device to all GPU?</p>

ptrblck: <p>I don’t fully understand the question as the gradients are scattered as shown in your figure.</p>

Jianxin_Wang: <p>Sorry for the confusion. What I mean was how the loss in the default device turns to [loss1, loss2, loss3, loss4]. How the process divides the loss into 4 losses? (In my understanding,  it can’t just evenly divide the loss in the default device because the losses for the batch in 4 GPUs are different)</p>

ptrblck: <p><code>nn.DataParallel</code> will create chunks of the input initially which refers to the <code>[i1, i2, i3, i4]</code> tensors. Note that the input tensor on the default device will not be a <code>list</code> of 4 tensors in this case, but will be a single one. The figure uses these <code>list</code> notation to explain that this single tensor will be split.<br>
The same applies for the output tensor, which will be a single tensor on the default device containing the splits from all GPUs which are concatenated in the batch dimension (<code>dim0</code>).</p>