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

class Model(nn.Module): def init(self): super().init() # Conv feature extraction self.conv = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.pool = nn.AdaptiveAvgPool2d(4) # Embedding + permute path self.embedding = nn.Embedding(1000, 128) self.emb_proj = nn.Linear(128, 1024) # Multiple F.linear with different weights (batch linear fusion target) self.w1 = nn.Parameter(torch.randn(256, 1024)) self.b1 = nn.Parameter(torch.randn(256)) self.w2 = nn.Parameter(torch.randn(256, 1024)) self.b2 = nn.Parameter(torch.randn(256)) self.w3 = nn.Parameter(torch.randn(256, 1024)) self.b3 = nn.Parameter(torch.randn(256)) self.w4 = nn.Parameter(torch.randn(256, 1024)) self.b4 = nn.Parameter(torch.randn(256)) self.out_proj = nn.Linear(256, 10)

def forward(self, conv_input, embedding_input):
    # Conv path
    cx = self.pool(F.relu(self.conv(conv_input)))   # [B,64,4,4]
    cx = cx.flatten(1)                               # [B,1024]

    # Embedding path
    ex = self.embedding(embedding_input)             # [B,16,128]
    ex = ex.mean(dim=1)                              # [B,128]
    ex = self.emb_proj(ex)                           # [B,1024]

    # Fuse both paths
    x = cx + ex                                      # [B,1024]

    # Multiple F.linear with same input → PreGradBatchLinearFusion target
    h1 = F.linear(x, self.w1, self.b1)
    h1 = F.relu(h1)
    h2 = F.linear(x, self.w2, self.b2)
    h2 = F.relu(h2)
    h3 = F.linear(x, self.w3, self.b3)
    h3 = F.relu(h3)
    h4 = F.linear(x, self.w4, self.b4)
    h4 = F.relu(h4)

    combined = h1 + h2 + h3 + h4
    return self.out_proj(combined)

device = "cuda" torch.manual_seed(42) model = Model().to(device).eval() conv_input = torch.randn(4, 3, 32, 32, device=device) embedding_input = torch.randint(0, 1000, (4, 16), device=device)

Eager: deterministic

with torch.no_grad(): ref1 = model(conv_input, embedding_input) ref2 = model(conv_input, embedding_input) print(f"Eager deterministic: {(ref1 - ref2).abs().max().item():.6e}")

Compiled: different result

torch._dynamo.reset() compiled = torch.compile(model, backend="inductor") with torch.no_grad(): out = compiled(conv_input, embedding_input)

diff = (ref1.float() - out.float()).abs() print(f"max_diff={diff.max().item():.6e}") print(f"Output range: [{ref1.min().item():.2f}, {ref1.max().item():.2f}]") print(f"Relative error: {diff.max().item() / ref1.abs().max().item():.6e}")

TL;DR

The most likely fix is to disable the PreGradBatchLinearFusion optimization in the inductor backend or use a different backend that does not perform this optimization.

Guidance

• Identify the specific optimization causing the issue: PreGradBatchLinearFusion batches independent linear operations, changing the accumulation order and resulting in different numerical values.
• Consider disabling this optimization or using a different backend to avoid the batched GEMM operation.
• Verify the fix by comparing the output of the compiled model with the eager mode output to ensure the numerical values match within the expected tolerance.
• If disabling the optimization is not feasible, explore other workarounds, such as using a different numerical precision or scaling the weights to reduce the impact of the accumulation order change.

Example

No code example is provided as the issue is related to the optimization pass in the inductor backend, and the fix would involve modifying the backend or disabling the optimization.

Notes

The provided information suggests that the issue is specific to the inductor backend and the PreGradBatchLinearFusion optimization. Disabling this optimization or using a different backend may resolve the issue, but it may also impact performance.

Recommendation

Apply a workaround, such as disabling the PreGradBatchLinearFusion optimization, as the issue is specific to this optimization and the inductor backend. This workaround may have performance implications, but it should resolve the numerical mismatch issue.