Backward hangs on Torch for CPU

NikolayL: <p>Very tiny model with two convolutions hangs on <code>loss.backward()</code> call.</p>
<pre><code class="lang-auto">UNet(
  (down_path): ModuleList(
    (0): UNetConvBlock(
      (block): Conv2d(3, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    )
  )
  (last): Conv2d(2, 13, kernel_size=(1, 1), stride=(1, 1))
)
</code></pre>
<p>Snippet from stacktrace in GDB:</p>
<pre><code class="lang-auto">#0  0x00007fca0afed99f in __GI___poll (fds=0x7fc9ebbd0040, nfds=1, timeout=5000) at ../sysdeps/unix/sysv/linux/poll.c:29
</code></pre>
<p>It’s reproduced on torch with CPU only, GPU one works fine.</p>
<p>The batch size is also important. No hang with batch size = 1, but freezes with batch size = 4.<br>
Profiling of the script for different batches shows different implementation called:<br>
No hang with <code>_slow_conv2d_forward</code> and <code>aten::_slow_conv2d_backward</code>.<br>
Hang happens with <code>aten::mkldnn_convolution</code>, <code>aten::convolution_backward</code>.</p>
<p>Snapshot of reproducing script:</p>
<pre><code class="lang-auto">def main():
    train_set = MockDataset()
    train_loader = DataLoader(train_set, batch_size=4, num_workers=1, drop_last=True)
    model = UNet(n_classes=13)
    print(model)
    device = 'cpu'
    model.to(device)
    optimizer = SGD(model.parameters(), lr=1e-3)
    for epoch in range(1):
        model.train()
        for step, batch_data in enumerate(train_loader):
            inputs = batch_data[0].to(device)
            labels = batch_data[1].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            print(outputs.shape, labels.shape)
            loss = nn.CrossEntropyLoss()(outputs, labels)
            print(f"\nBefore HANG {loss}\n")
            loss.backward()
            print("\nAFTER HANG\n")
</code></pre>
<p>CPU: Intel(R) Core™ i9-10980XE CPU @ 3.00GHz (can’t reproduce on Intel(R) Core™ i9-10920X CPU @ 3.50GHz)<br>
Ubuntu 20.04<br>
Python 3.8.10<br>
torch==1.13.1+cpu<br>
torchvision==0.14.1+cpu</p>
<p><a href="https://gist.github.com/ljaljushkin/c6daf65ea0c6e119b9336c5edba09a0c" class="inline-onebox" rel="noopener nofollow ugc">minimal reproducer for hang on backward · GitHub</a></p>

ptrblck: <p>Thanks for reporting this issue! Could you create another issue on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a> so that MKL devs could take a look at it and try to reproduce it?<br>
I’ve checked a few nodes I can lease but couldn’t find any with an <code>i9-10980XE</code> CPU and thus cannot reproduce the hang.</p>

NikolayL: <p>Thank you for trying!<br>
I’ve created the issue on GitHub: <a href="https://github.com/pytorch/pytorch/issues/91547" class="inline-onebox" rel="noopener nofollow ugc">The backward call hangs on Torch for CPU · Issue #91547 · pytorch/pytorch · GitHub</a><br>
Didn’t mention that the issue is reproduced in the docker when the python script is launched in this way:</p>
<pre><code class="lang-bash">sh -c "(script.sh &gt; log.txt 2&gt;&amp;1; wait)" &gt;&amp;- 2&gt;&amp;- &amp;
</code></pre>
<p>Content of the script.sh:</p>
<pre><code class="lang-bash">#!/bin/bash -xe
. venv/bin/activate
python3 main.py
</code></pre>
<p>Added it to description.</p>

NikolayL: <p>The reason of the hang is described here: <a href="https://github.com/pytorch/pytorch/issues/91547#issuecomment-1471361061" class="inline-onebox" rel="noopener nofollow ugc">The backward call hangs on Torch for CPU · Issue #91547 · pytorch/pytorch · GitHub</a><br>
<code>OMP_THREAD_LIMIT</code> should set to a number equal to or larger than <code>OMP_NUM_THREADS</code> to avoid the hanging issue.</p>