pytorch - 💡(How to fix) Fix `torch.compile` produces different output for Transformer using `torch.einsum` projections compared to eager mode [1 participants]

import os
os.environ["TRITON_BACKENDS_IN_TREE"] = "1"

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, batch_size=4, seq_len=16):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # Batch-specific projection weights (triggers batched einsum pattern)
        self.q_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model) * 0.02)
        self.k_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model) * 0.02)
        self.v_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model) * 0.02)
        self.out_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model) * 0.02)

        # Batch-specific biases
        self.q_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model) * 0.02)
        self.k_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model) * 0.02)
        self.v_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model) * 0.02)
        self.out_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model) * 0.02)

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape

        # einsum projections + element-wise bias add (target pattern)
        q = torch.einsum("b i d, b d h -> b i h", x, self.q_proj) + self.q_bias[:batch_size, :seq_len, :]
        k = torch.einsum("b i d, b d h -> b i h", x, self.k_proj) + self.k_bias[:batch_size, :seq_len, :]
        v = torch.einsum("b i d, b d h -> b i h", x, self.v_proj) + self.v_bias[:batch_size, :seq_len, :]

        # Reshape to multi-head
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention
        scale = self.head_dim ** 0.5
        attn = torch.matmul(q, k.transpose(-2, -1)) / scale
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        # Merge heads and output projection
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        out = torch.einsum("b i d, b d h -> b i h", out, self.out_proj) + self.out_bias[:batch_size, :seq_len, :]
        return out


class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, num_heads=8, batch_size=4, seq_len=16):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, batch_size, seq_len)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model),
        )

    def forward(self, x):
        # Pre-norm residual attention
        x_norm = self.norm1(x)
        attn_out = self.attn(x_norm)
        x = x + attn_out
        # Pre-norm residual FFN
        x = x + self.ffn(self.norm2(x))
        return x


device = "cuda"
torch.manual_seed(42)
model = TransformerBlock(d_model=512, num_heads=8, batch_size=4, seq_len=16).to(device).eval()
x = torch.randn(4, 16, 512, device=device)

# Eager: deterministic
with torch.no_grad():
    ref = model(x)
    ref2 = model(x)
print(f"Eager deterministic: {(ref - ref2).abs().max().item():.6e}")

# Compiled
torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    comp = compiled(x)

diff = (ref.float() - comp.float()).abs()
print(f"Max diff: {diff.max().item():.6e}")
print(f"Mean diff: {diff.mean().item():.6e}")
print(f"Match (atol=1e-5): {torch.allclose(ref, comp, atol=1e-5, rtol=1e-4)}")

---

PyTorch version: 2.12.0.dev20260327+cu126
Python: 3.10.12
OS: Ubuntu 22.04.5 LTS (WSL2)
GPU: NVIDIA GeForce RTX 3060 Laptop GPU
CUDA: 12.6