CPU parallelization across batch has random faulty behavior on backward

ysig: <p>Hi,</p>
<p>I am trying to implement a CPU loss that applies parallelization across batch.<br>
When running the backward function I receive a very strange error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "test.py", line 8, in &lt;module&gt;
    torch.sum(loss(a, b)).backward()
  File "/home/.local/lib/python3.7/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/.local/lib/python3.7/site-packages/torch/autograd/__init__.py", line 149, in backward
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
RuntimeError: res[i].defined()INTERNAL ASSERT FAILED at "../torch/csrc/autograd/functions/tensor.cpp":97, please report a bug to PyTorch.
</code></pre>
<p>I’ve managed to reduce my complicated function to something simple in order to provide you with a reproducible example.</p>
<p>Here are the instructions on how to set it up (I use linux/ubuntu 18.04 and pytorch 1.5.0)<br>
To use this create a dir named ‘loss’ and add the following files below it:</p>
<p>setup.py</p>
<pre><code class="lang-auto">from setuptools import setup, Extension
from torch.utils import cpp_extension
from __init__ import __version__

cpp_module = cpp_extension.CppExtension(
    'loss',
    sources=['loss.cpp'],
    extra_compile_args=['-fopenmp'],
    extra_link_args=['-lgomp']
)

setup(name='loss',
      ext_modules=[cpp_module],
      cmdclass={'build_ext': cpp_extension.BuildExtension},
      )
</code></pre>
<p>test.py</p>
<pre><code class="lang-auto">import torch
torch.set_num_threads(3)

from loss import loss

a = torch.nn.Parameter(torch.ones((3, 1000)))
b = torch.zeros((3, 1000))

torch.sum(loss(a, b)).backward()
</code></pre>
<p>loss.cpp</p>
<pre><code class="lang-auto">#include &lt;torch/extension.h&gt;
#include &lt;iostream&gt;
#include &lt;ATen/Parallel.h&gt;

torch::Tensor loss(torch::Tensor z, torch::Tensor x)
{
    torch::Tensor z_out = at::empty({z.size(0), z.size(1)}, z.options());
    int64_t batch_size = z.size(0);
    at::parallel_for(0, batch_size, 0, [&amp;](int64_t start, int64_t end) {
        for (int64_t b = start; b &lt; end; b++)
        {
            z_out[b] = z[b] - x[b];
        }
    });
    return z_out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("loss", &amp;loss, "loss");
}
</code></pre>
<p>test.sh</p>
<pre><code class="lang-auto">rm -rf build dist loss.egg-info __pycache__
pip uninstall loss -y
python setup.py install --user
python test.py
</code></pre>
<p>To run it move inside the ‘loss’ directory and run: <code>sh test.sh</code><br>
Note: the error doesn’t always appear. You may need to run the script multiple times to see it. Let me know if this is a version problem.</p>

ptrblck: <p>Could you check, if your code is still failing in the latest stable release (<code>1.10.1</code>) or the nightly? If so, would you mind creating an issue on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a> so that we could track and fix it, please?</p>

ysig: <p>Hi. Yes I still have it on version 1.10.1. I’ve raised an <a href="https://github.com/pytorch/pytorch/issues/70241" rel="noopener nofollow ugc">issue</a>. Hope it’s not something trivial <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>Thank you! Cross-posting: based on the response it seems to be a known limitation disallowing tensor operations inside <code>at::parallel_for</code> as described <a href="https://github.com/pytorch/pytorch/issues/56794">here</a> and <code>data_ptr</code>s must be used.</p>

ysig: <p>Can you give an example?</p>
<p>How would you write to this so it is correct?</p>
<pre><code class="lang-auto">    at::parallel_for(0, batch_size, 0, [&amp;](int64_t start, int64_t end) {
        for (int64_t b = start; b &lt; end; b++)
        {
            z_out[b] = z[b] - x[b];
        }
    });
</code></pre>