Spooky action at a distance between model compilations

username: avitase
cooked: <p>We have a strange problem with a (JIT) compiled model. We have two <code>pytest</code> unit tests of our model. In both tests the model is initialized with the same parameters, compiled via <code>torch.script.jit(model)</code> and fed with toy data, but only the second test uses the output to calculate a loss and eventually calls <code>backward</code> on it. If both tests are run, the second test throws:</p>
<pre><code class="lang-auto">&gt;       Variable._execution_engine.run_backward(
            tensors, grad_tensors_, retain_graph, create_graph, inputs,
            allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
E       RuntimeError: The following operation failed in the TorchScript interpreter.
E       Traceback of TorchScript (most recent call last):
E         File "&lt;string&gt;", line 162, in &lt;backward op&gt;
E               def abs(self):
E                   def backward(grad_output):
E                       return grad_output * self.sign()
E                                            ~~~~~~~~~ &lt;--- HERE
E
E                   return torch.abs(self), backward
E       RuntimeError: Unlike NumPy, torch.sign is not intended to support complex numbers. Please use torch.sgn instead.

../miniconda3/envs/nops/lib/python3.8/site-packages/torch/autograd/__init__.py:147: RuntimeError
</code></pre>
<p>which seems like a legit error (although we don’t understand where exactly this happens since all leaf parameters are real-valued scalars) but, more importantly, everything works fine if we either</p>
<ul>
<li>skip the first test,</li>
<li>or skip one of the (JIT) compilation steps,</li>
<li>or use different parameters to initialize the model in the first or second test.</li>
</ul>
<p>There is no shared state between both test cases and we don’t know what could possible causes this problem. Our model code is a bit convoluted but the test code is rather short:</p>
<pre><code class="lang-auto">def test_a():
    x = ...
    model_args = ...

    model = fno.FNO(model_args)
    model = torch.jit.script(model)

    out = model(x)


def test_b():
    x = torch.rand(...)
    y = torch.rand(...)
    model_args = ...

    model = Model(model_args)
    model = torch.jit.script(model)

    model.train()
    out = model(x)
    mse = F.mse_loss(out, y, reduction="mean")
    mse.backward()
</code></pre>
<p>We are tempted to assume that this is a <code>pytorch</code> bug but we don’t know how to proceed from here on. Do you have any ideas how to nail down the exact problem?</p>
<p><strong>EDIT:</strong> we managed to narrow down the error:</p>
<pre><code class="lang-auto">def test_c():
    x = torch.rand(...)
    y = torch.rand(...)
    model_args = ...

    model = fno.FNO(model_args)
    model = torch.jit.script(model)  # (1)
    model(x)                         # (2)

    model.train()
    out = model(x)
    mse = F.mse_loss(out, y, reduction="mean")
    mse.backward()
</code></pre>
<p><code>test_c()</code> now triggers the error on its own. Removing either line <code>(1)</code> or <code>(2)</code> suppresses the error…</p>

username: tom
cooked: <p>What happens is that in the second run you get the JIT autodiffed abs and that has been broken when complex support was added to the derivative of backward. Sadly, it seems that while the autograd that you get on the first run had been updated 10 months ago, the JIT has not been getting as much love from the complex number transition.</p>
<p>The fix is trivial in line with the exception: Change sign to sgn here: <a href="https://github.com/pytorch/pytorch/blob/03a79f43e33c1bef65fc8912c27160d01e0e15d5/torch/csrc/jit/runtime/symbolic_script.cpp#L972" class="inline-onebox" rel="noopener nofollow ugc">pytorch/symbolic_script.cpp at 03a79f43e33c1bef65fc8912c27160d01e0e15d5 · pytorch/pytorch · GitHub</a><br>
However, I have not managed to find a reproducing case, or maybe the nightly versions have a fix where the JIT avoids complex numbers altogether, I would not know.</p>
<p>Best regards</p>
<p>Thomas</p>

username: avitase
cooked: <p>Thanks a lot for your help! I am afraid, I am still a bit lost with your explanation about “autodiffed abs”. Could you explain this in a bit more detail? In particular, why does this error only occurs after the second compilation / after the second forward evaluation?</p>

username: tom
cooked: <p>The short answer is that this is because the JIT collects information on (by default only) the first run and after this does its optimizations. If you want plenty of detail, I can offer a <a href="https://lernapparat.de/jit-runtime-overview/" rel="noopener nofollow ugc">blog post on what happens when you call a JIT function</a> and <a href="https://lernapparat.de/jit-optimization-intro/" rel="noopener nofollow ugc">one on optimization in particular</a>.</p>
<p>Best regards</p>
<p>Thomas</p>

username: avitase
cooked: <p>Ok. Thanks for providing the links and your swift help:)</p>
<p><strong>EDIT:</strong> using the nightly build I get the same error</p>

username: tom
cooked: <aside class="quote no-group" data-username="avitase" data-post="5" data-topic="126709">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/94ad74/48.png" class="avatar"> avitase:</div>
<blockquote>
<p>using the nightly build I get the same error</p>
</blockquote>
</aside>
<p>Thank you for checking!<br>
So if you can get the example down to where the model is a short sequential model of standard PyTorch or somesuch, either of us could do a PR to fix it.</p>

username: avitase
cooked: <p>We are currently quite busy but I will see what we can do. Besides, shouldn’t fixing <code>torch.sign</code> to <code>torch.sgn</code> in <code>symbolic_script.cpp</code> a good idea anyhow?</p>