Distributed InfoNCE Loss (CLIP)

Kerem_Can_Turgutlu: <p>I am trying to implement InfoNCE Loss from <a href="https://github.com/openai/CLIP" rel="noopener nofollow ugc">CLIP</a> in a distributed way. InfoNCE is a loss function which is used for contrastive learning and it favors large batch size during calculation. In CLIP, a batch is composed of image-text pairs, there is an image encoder and a text encoder. These encoders are then used to extract image and text embeddings, further these embeddings are used to calculate the pairwise cosine similarity and finally the loss.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/f/6f636f0ac022e07488aa53ac2db4eef6ea96c846.png" data-download-href="https://discuss.pytorch.org/uploads/default/6f636f0ac022e07488aa53ac2db4eef6ea96c846" title="Screen Shot 2021-02-20 at 8.00.03 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6f636f0ac022e07488aa53ac2db4eef6ea96c846_2_590x500.png" alt="Screen Shot 2021-02-20 at 8.00.03 PM" data-base62-sha1="fTo2aMAzo5rI2Uacz9LMs8liGN0" width="590" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/f/6f636f0ac022e07488aa53ac2db4eef6ea96c846_2_590x500.png, https://discuss.pytorch.org/uploads/default/original/3X/6/f/6f636f0ac022e07488aa53ac2db4eef6ea96c846.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/f/6f636f0ac022e07488aa53ac2db4eef6ea96c846.png 2x" data-dominant-color="F0F3F6"><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 2021-02-20 at 8.00.03 PM</span><span class="informations">864×732 40.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 the original CLIP paper 256 GPUs with batch size of 128 per GPU is used. Also, authors mentioned that they calculate the pairwise cosine similarity using all the batches across all GPUs, you can refer to this <a href="https://github.com/openai/CLIP/issues/29" rel="noopener nofollow ugc">issue</a>. If we use the vanilla DistributedDataParallel then the cosine similarity and the loss would be calculated only the pairs within a single GPU and later their gradients would be averaged and synced. What we really want is to gather all embeddings across all GPUs and calculate cosine similarity for a given batch across all the other batches like it’s mentioned in the issue.</p>
<p>For this purpose, I started writing some dummy code to test this idea on CPU before spending GPU compute. I created 2 gists, <a href="https://gist.github.com/KeremTurgutlu/4d8e1e2812a7c16fa329bbd732e764d3" rel="noopener nofollow ugc">1 script</a> as a test case which does loss calculation and backward without distributed context and <a href="https://gist.github.com/KeremTurgutlu/4a6f7078dc62f292c85b9903197c75f7" rel="noopener nofollow ugc">1 other script</a> which tries to implement this idea in distributed context. Although, it looks like <code>image_encoder.weight.grad</code> correctly matches the non-distributed version <code>text_encoder.weight.grad</code> is None.</p>
<p>Regarding above, I have a few questions:</p>
<ol>
<li>Are there already any existing code or best practices for implementing distributed contrastive losses which requires communicating model outputs across devices before loss calculation?</li>
<li>If <strong>1)</strong> is a no, then how can I fix the issue I see in my gists? Is it a good start for implementing this idea?</li>
<li>Currently I use <code>dist.all_gather</code> since this is a dummy example, but in a real training setting would using <code>dist.gather</code> allow to handle possible OOM issues better?</li>
</ol>
<p>Thanks a lot in advance!</p>

Kerem_Can_Turgutlu: <p><strong>Update</strong>: I updated code <a href="https://gist.github.com/KeremTurgutlu/4a6f7078dc62f292c85b9903197c75f7" rel="noopener nofollow ugc">here</a> to use <a href="https://github.com/open-mmlab/OpenSelfSup/blob/696d04950e55d504cf33bc83cfadbb4ece10fbae/openselfsup/models/utils/gather_layer.py" rel="noopener nofollow ugc">GatherLayer</a>. Now <code>text_encoder.weight.grad</code> exists but values are not same as the non-distributed version.</p>