Nondeterminstic behaviour of sigmoid and tanh in JIT script

luke1110: <p>I recently discovered some nondeterministic behaviour when using a jit script which is supposed to make the usage of  <code>torch.sigmoid</code> more memory efficient, as described in <a href="https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" class="inline-onebox" rel="noopener nofollow ugc">Performance Tuning Guide — PyTorch Tutorials 2.2.0+cu121 documentation</a> .</p>
<p>When executed on CPU both <code>torch.tanh</code> and <code>torch.sigmoid</code> lead to different results for the same input when comparing the first and second invocation. After that all invocations lead to the same result. The mean average error between first and second result is around <code>1e-11</code>, but when used in a simple convolutional network with only 4 layers, this quickly becomes <code>1e-8</code>.</p>
<p><strong>This does not happen on GPU.</strong></p>
<p>The fact that it stays the same result after the second invocation makes me think, it’s some sort of caching problem, but I wonder, why the first result wouldn’t be cached.</p>
<p>Has someone ever experienced something similar? Is there a way to fix this? Would it be affected by setting <code>torch.use_deterministic_algorithms</code>? I haven’t tried that yet…</p>
<p>Here is the jit script I invoke:</p>
<pre><code class="lang-auto">@torch.jit.script
def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels, debug_name: str="None"):
    n_channels_int = n_channels[0]
    in_act = input_a + input_b
    t_act = torch.tanh(in_act[:, :n_channels_int, :])
    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
    acts = t_act * s_act
</code></pre>
<p>And this is the MAE between first and second invocation in the same input (<code>in_act</code>) in the same layer (<code>0</code>) in network <code>WN</code> (forward and backward have no meaning here, except for that backward holds the second invocation):<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7bf4e3def33edcd2d4a2ad8e4b9891955021a3c6.png" data-download-href="https://discuss.pytorch.org/uploads/default/7bf4e3def33edcd2d4a2ad8e4b9891955021a3c6" title="Screenshot 2024-01-19 at 12.15.46"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bf4e3def33edcd2d4a2ad8e4b9891955021a3c6_2_690x112.png" alt="Screenshot 2024-01-19 at 12.15.46" data-base62-sha1="hGzp1agXIgwAun1OU25dnO7xr82" width="690" height="112" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bf4e3def33edcd2d4a2ad8e4b9891955021a3c6_2_690x112.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bf4e3def33edcd2d4a2ad8e4b9891955021a3c6_2_1035x168.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bf4e3def33edcd2d4a2ad8e4b9891955021a3c6_2_1380x224.png 2x" data-dominant-color="0A0A08"><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">Screenshot 2024-01-19 at 12.15.46</span><span class="informations">1824×298 62.4 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>

ptrblck: <p><code>torch.jit.script</code> is able to optimize the computation graph and add e.g. kernel fusions, which might not output bitwise-identical results compared to the eager execution. I don’t know what exactly is used on the CPU, but note that <code>TorchScript</code> is in maintenance mode and the current recommendation is to use <code>torch.compile</code>. Also note that <code>torch.compile</code> might still yield the same mismatches due to optimizations.</p>