Averaging Gradients in DistributedDataParallel

ankahira: <p>I am a bit confused about averaging gradients in distributed data-parallel. It seems there are two examples from the PyTorch documentation that are different. In one example, you create the model and just pass it to the GPU available then create a separate function to average gradients.</p>
<pre><code class="lang-auto">def average_gradients(model):
    """ Gradient averaging. """
    size = float(dist.get_world_size())
    for param in model.parameters():
        dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM)
        param.grad.data /= size
</code></pre>
<p>and its executed as follows.</p>
<pre><code class="lang-auto">        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        average_gradients(model)
        optimizer.step()

</code></pre>
<p>The other approach I have seen doesn’t create a separate function and just calls DPP.</p>
<pre><code class="lang-auto">    model = ToyModel().cuda(device_ids[0])
    ddp_model = DDP(model, device_ids)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_ids[0])
    loss_fn(outputs, labels).backward()
    optimizer.step()
</code></pre>
<p>I would like to know whats the difference between the two approaches and which one should one use for distributed training in a HPC cluster. I specifically want to use two nodes, each with 4 GPUs.</p>

mrshenli: <p>Hey <a class="mention" href="/u/ankahira">@ankahira</a>, usually, there are 4 steps in distributed data parallel training:</p>
<ol>
<li>local forward to compute loss</li>
<li>local backward to compute local gradients</li>
<li>allreduce (communication) to compute global gradients. This would be allreduce with SUM + divide by world size to calculate average</li>
<li>optimizer step to use global gradients to update parameters</li>
</ol>
<p>Both examples you mentioned above conduct the same four steps and are mathematically equivalent. The difference is that DDP would allow step 2 (backward computation) and 3 (allreduce communication) to overlap and therefore DDP is expected to be faster than the <code>average_gradients</code> approach.</p>
<p>More specifically, in the first example with <code>average_gradients</code>, there is a hard barrier between backward and allreduce, i.e., no comm can start before computation finishes. In the second example, DDP organizes gradients into buckets, and will launch comm as soon as a bucket of gradients are ready, so that computation and communication can run in parallel. <a href="https://pytorch.org/docs/master/notes/ddp.html" rel="nofollow noopener">This</a> would help explain that.</p>
<p>I would recommend DDP. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Vanjoy: <p>Why would we need to divide the gradients by word size in step 3?</p>

mrshenli: <p>Because in many cases, the loss function computes per-sample loss (e.g., default <code>MSELoss</code>) instead of aggregated loss (e.g., default <code>sum</code>). So, DDP calculates the average and tries to keep the gradients scale consistent with local training.</p>