Dist.all_gather() and gradient preservation in multi-GPU training

cbaziotis: <p>I want to use the <a href="https://paperswithcode.com/method/nt-xent" rel="noopener nofollow ugc">NT-Xent</a> loss from the <a href="https://arxiv.org/abs/2002.05709v3" rel="noopener nofollow ugc">SimCLR</a> paper and I am unsure about what is the correct implementation in a multi-GPU setting, specifically how to properly use <code>dist.all_gather()</code>.</p>
<p>Each batch is divided into smaller parts and distributed across the different GPUs, and each GPU contains only a certain partition of the full batch. After each GPU computes the embedding for each input, I want to recombine them so I can properly compute the contrastive loss. I know that I have to use <code>dist.all_gather()</code> to achieve that and that this function does not maintain the <code>grad_fn</code> property in the combined data.</p>
<p>I have found two approaches that address this issue.</p>
<h3>
<a name="1-extend-torchautogradfunctionhttpspytorchorgdocsstableautogradhtmlfunction-such-as-1" class="anchor" href="#1-extend-torchautogradfunctionhttpspytorchorgdocsstableautogradhtmlfunction-such-as-1"></a><span class="hashtag">#1</span> Extend  <a href="https://pytorch.org/docs/stable/autograd.html#function" rel="noopener nofollow ugc"><code>torch.autograd.Function</code></a>, such as:</h3>
<ol>
<li>In <a href="http://github.com/Spijkervet/SimCLR:" rel="noopener nofollow ugc">github.com/Spijkervet/SimCLR:</a> this is the implementation of the <a href="https://github.com/Spijkervet/SimCLR/blob/master/simclr/modules/nt_xent.py" rel="noopener nofollow ugc">loss</a> and this is its <a href="https://github.com/Spijkervet/SimCLR/blob/master/simclr/modules/gather.py" rel="noopener nofollow ugc">autograd.Function</a>.</li>
<li>In <a href="http://github.com/PyTorchLightning:" rel="noopener nofollow ugc">github.com/PyTorchLightning:</a> this is the implementation of the <a href="https://github.com/PyTorchLightning/lightning-bolts/blob/master/pl_bolts/models/self_supervised/simclr/simclr_module.py#L224" rel="noopener nofollow ugc">loss</a> and this is its <a href="https://github.com/PyTorchLightning/lightning-bolts/blob/5577453a6d7072724d9ae24184daf8f45d4baff7/pl_bolts/models/self_supervised/simclr/simclr_module.py#L20" rel="noopener nofollow ugc">autograd.Function</a>.</li>
</ol>
<pre><code class="lang-python">https://github.com/Spijkervet/SimCLR/blob/master/simclr/modules/nt_xent.py
def nt_xent_loss(self, ...):
        ...

        if self.world_size &gt; 1:
            z = torch.cat(GatherLayer.apply(z), dim=0)

        sim = self.similarity_f(z.unsqueeze(1), z.unsqueeze(0)) / self.temperature
        ...
</code></pre>
<pre><code class="lang-auto"># https://github.com/Spijkervet/SimCLR/blob/master/simclr/modules/gather.py
class GatherLayer(torch.autograd.Function):
    """Gather tensors from all process, supporting backward propagation."""

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        output = [torch.zeros_like(input) for _ in range(dist.get_world_size())]
        dist.all_gather(output, input)
        return tuple(output)

    @staticmethod
    def backward(ctx, *grads):
        (input,) = ctx.saved_tensors
        grad_out = torch.zeros_like(input)
        grad_out[:] = grads[dist.get_rank()]
        return grad_out
</code></pre>
<h3>
<a name="2-overwrite-the-gathered-list-using-the-partition-from-the-local-replica-2" class="anchor" href="#2-overwrite-the-gathered-list-using-the-partition-from-the-local-replica-2"></a><span class="hashtag">#2</span> Overwrite the gathered list using the partition from the local replica</h3>
<p>I’ve seen this approach described by <strong><a href="https://github.com/JohnGiorgi" rel="noopener nofollow ugc">JohnGiorgi</a></strong> in this <a href="https://github.com/KevinMusgrave/pytorch-metric-learning/issues/10#issuecomment-593170720" rel="noopener nofollow ugc">comment</a>. The idea is to use the local tensor that has the proper <code>grad_fn</code> to overwrite the corresponding partition in the list returned by <code>dist.all_gather()</code>.</p>
<p>(I am copying his pseudo-code here to make the question self-contained)</p>
<pre><code class="lang-python"># From: https://github.com/KevinMusgrave/pytorch-metric-learning/issues/10
import torch
import torch.distributed as dist

# Dummy code representing the forward pass for some batch of text on one replica.
embeddings = model(batch)

# Gather the embeddings from every replica.
embeddings_list = [torch.ones_like(embeddings) for _ in range(dist.get_world_size())]
dist.all_gather(embeddings_list, embeddings)

# Gathered tensors have no gradient, so we overwrite the gathered tensor for the current replica with the embeddings produced on this replica, which do have gradients.
embeddings_list[dist.get_rank()] = embeddings

# Finally, concatenate the list of embeddings before computing a loss.
embeddings = torch.cat(embeddings_list)

# I didn't demonstrate how to generate the labels, this will be task-dependent.
loss = some_contrastive_loss(embeddings, labels)
</code></pre>
<p>Also, this <a href="https://amsword.medium.com/gradient-backpropagation-with-torch-distributed-all-gather-9f3941a381f8" rel="noopener nofollow ugc">blog post</a> claims that when using this approach, we should re-scale the gradients by the number of GPUS (i.e., word_size) to obtain the correct gradients.</p>
<h3>
<a name="question-3" class="anchor" href="#question-3"></a>Question</h3>
<p>Which one should I use? Is it just a matter of style or are there meaningful differences between these approaches?</p>

zihui_wu: <p>Hi, did you find the difference? I am also troubled by this question</p>

MicPie: <p>I’m also currently investigating a multi-GPU contrastive setup and I’m currently debugging the difference between a 1 GPU full bs and x GPU bs/x setup to find the difference between the loss implementations.</p>
<p>Therefore, I would be very curious on how others approached this and with what you ended up in the end. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

111319: <p>For the three approaches that you provided, in my opinion, the <a href="http://github.com/Spijkervet/SimCLR:" rel="noopener nofollow ugc">SimCLR loss implemented by Spijkervet</a> is the same as the overwriting method of the gathered list.  The loss implemented by PyTorchLightning is different, which all_reduce the gradient in backward function, thus the gradient will be “dist.word_size()” times than the former implementation. I do not know which one is right.</p>