Weird CUDA sync behavior involving MSELoss

iacolippo: <p>I am trying to perform some asynchronous processing where part of an algorithm is executed on an accelerator and then some processing is done on the GPU asynchronously.</p>
<p>I have a weird behavior involving <code>F.mse_loss</code>: I have commented out each line and it seems to cause a synchronization (defeating my goal of running async processing).</p>
<p>This is solved simply by switching the line</p>
<pre><code class="lang-auto">loss = F.mse_loss(y_hat, target)
</code></pre>
<p>with</p>
<pre><code class="lang-auto">loss = ((y_hat - target) ** 2).mean()
</code></pre>
<p><code>y_hat</code> and <code>target</code> are CUDA Tensors of size <code>(batch_size, 1)</code></p>
<p>I only call</p>
<pre><code class="lang-auto">loss.backward()
optimizer.step()
optimizer.zero_grad()
</code></pre>
<p>afterward and the sync issue goes away if I write the MSE explicitly.</p>
<p>I have tried to reproduce in an isolated way the problem, but running timings on <code>F.mse_loss</code> alone does not show this problem.</p>
<p>Is there any reason why a synchronization might be happening?</p>

ptrblck: <p>How are you detecting the synchronization?<br>
Are you seeing some syncs in <code>nvprof</code> or <code>nsight</code>?</p>

iacolippo: <p>I have added a dummy operation that artificially increases the running time when not done asynchronously.</p>
<pre><code class="lang-auto">for i in range(3):
            torch.mm(input, input.t())
</code></pre>
<p>input array is shape <code>(3000,  50000)</code>. Unfortunately the code contains some proprietary stuff so I can’t share, I’ll see if I can make a version I can post here. And I’ll run <code>nvprof</code></p>

iacolippo: <p>I have issues with <code>nvprof</code> and I’m not sysadmin so it’ll take some time to fix that, but I managed to reproduce the issue on a dual GPU machine (2 x V100)</p>
<aside class="onebox githubgist">
  <header class="source">
      <a href="https://gist.github.com/iacolippo/d7e7c26e980cbd4a9b86c24db6ca922e" target="_blank" rel="nofollow noopener">gist.github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://gist.github.com/iacolippo/d7e7c26e980cbd4a9b86c24db6ca922e" target="_blank" rel="nofollow noopener">https://gist.github.com/iacolippo/d7e7c26e980cbd4a9b86c24db6ca922e</a></h4>
<h5>async_processing_2gpus.ipynb</h5>
<pre><code class="Jupyter Notebook">{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",</code></pre>
This file has been truncated. <a href="https://gist.github.com/iacolippo/d7e7c26e980cbd4a9b86c24db6ca922e" target="_blank" rel="nofollow noopener">show original</a>

<p>
</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>hope this helps in figuring this out!</p>

iacolippo: <p>I have tried to reproduce in an even simpler setting, but apparently this happens only when using two GPUs. Doing all the operations on a single GPU gets the same timing for both <code>F.mse_loss</code> and writing it explicitly in terms of difference, power and mean.</p>