Torch.compile _softmax in MultiheadAttention does not return same value as eager

Eldalier: <p><strong>Hello,</strong></p>
<p>I’m trying to implement torch.compile in my program, but I’m encountering reproducibility issues with this script using PyTorch 2.7.1:</p>
<pre data-code-wrap="python"><code class="lang-python">import os
from functools import partial

import torch
from torch.nn import MultiheadAttention
from torch.testing import assert_close

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"  # for deterministic cdist

torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms(True, warn_only=False)

torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)
#torch.backends.cuda.enable_math_sdp(False)

strict_assert_close = partial(assert_close, atol=0,rtol=0)
device  = "cuda"

def fn():
    torch._dynamo.reset()

    self_attn_builder = partial(MultiheadAttention,embed_dim=4,num_heads=1,dropout=0.1,bias=True,batch_first=True)

    torch.manual_seed(0)
    tf_raw = self_attn_builder().to(device)

    torch.manual_seed(0)
    tf_compiled = torch.compile(self_attn_builder().to(device),fullgraph=True)

    torch.manual_seed(0)
    tranfomer_input = torch.randn(torch.Size([128, 50, 4])).to(device)

    torch.manual_seed(0)
    output_1 = tf_raw(tranfomer_input,tranfomer_input,tranfomer_input,attn_mask=None,key_padding_mask=None)

    torch.manual_seed(0)
    output_compiled = tf_compiled(tranfomer_input,tranfomer_input,tranfomer_input,attn_mask=None,key_padding_mask=None)

    try:
        strict_assert_close(output_compiled, output_1)
        return True
    except AssertionError:
        return False
from torch._inductor.compiler_bisector import CompilerBisector

CompilerBisector.do_bisect(fn)
</code></pre>
<p>I can see that the output of <code>aten._softmax</code> differs between compiled and eager modes, which breaks my reproducibility tests.</p>
<p>Is this expected behavior, or is it a bug? Are there any options to ensure identical results between torch.compile and eager execution?</p>
<p>Thanks for your help!</p>

ptrblck: <p>How large is the absolute and relative error?</p>
<aside class="quote no-group" data-username="Eldalier" data-post="1" data-topic="221058">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/e/e47774/48.png" class="avatar"> Eldalier:</div>
<blockquote>
<p>Are there any options to ensure identical results between torch.compile and eager execution?</p>
</blockquote>
</aside>
<p>No, there is no guarantee different algorithms will return bitwise identical outputs.</p>

Eldalier: <p>Yes, but the error is not so small.</p>
<pre><code class="lang-auto">AssertionError: Tensor-likes are not equal!

Mismatched elements: 25600 / 25600 (100.0%)
Greatest absolute difference: 0.4337801933288574 at index (3, 1, 0)
Greatest relative difference: 2057.0537109375 at index (97, 14, 0)

The failure occurred for item [0]
</code></pre>
<p>In training, it completely changes my training results.</p>

ptrblck: <p>I can reproduce this larger error, but get the expected numerical mismatches after disabling <code>dropout</code>. I don’t know if random operations are guaranteed to return the same outputs as I would guess the <code>mask</code> sampling differs.</p>

Eldalier: <p>Ah, thank you, I will think that torch.manual_seed(0) will make dropout reproducible with torch compile, and CompilerBisector.do_bisect led me to the wrong path.</p>