torch.nn.Transformer cannot be compiled

[yxd886]

Issue description

I intended to use torch.nn.Transformer to build my model, and use torch.compile to compile it. But I found that the model is still eager mode even if I used torch.compile

Code example

import torch
device = "cpu"
model = torch.nn.Transformer(nhead=2, num_encoder_layers=2, d_model=6, dim_feedforward=4).to(device)
print(model)
src = torch.rand(2, 4, 6).to(device)
tgt = torch.rand(4, 4, 6).to(device)
compiled_model = torch.compile(model)
out = compiled_model(src, tgt)
print(out)

• PyTorch or Caffe2: Pytorch
• PyTorch version: 2.0.0
• Python version: 3.8

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @ezyang @msaroufim @wconstab @ngimel @bdhirsh @anijain2305

[yanboliang]

It's compiled successfully with your example, how can you say it's still eager mode?

[yxd886]

It's compiled successfully with your example, how can you say it's still eager mode?

when I set torch._inductor.config.debug = True, It does not show any IR created.

[ezyang]

cc @jansel

[ngimel]

TORCH_LOGS=+inductor python transformer.py shows expected debugging info

[Hurray0]

Hello, I have the same problem.

import torch
import torch._inductor.config
print("torch-version: ", torch.__version__)
torch._inductor.config.debug = True
device = "cpu"
model = torch.nn.Transformer(nhead=2, num_encoder_layers=2, d_model=6, dim_feedforward=4).to(device)
#print(model)
src = torch.rand(2, 4, 6).to(device)
tgt = torch.rand(4, 4, 6).to(device)
compiled_model = torch.compile(model)
out = compiled_model(src, tgt)
print(out)

Compare with Transformer, the other model is fine:

code:

import torch
import torch._inductor.config
print("torch-version: ", torch.__version__)
torch._inductor.config.debug = True
device = "cpu"

def func(x):
    return x.sin() + x.cos()

model = func
compiled_model = torch.compile(model)
data = torch.rand(20)
out = compiled_model(data)
print(out)

[Hurray0]

cc @ngimel

[yiliu30]

Hello, I have the same problem.

import torch
import torch._inductor.config
print("torch-version: ", torch.__version__)
torch._inductor.config.debug = True
device = "cpu"
model = torch.nn.Transformer(nhead=2, num_encoder_layers=2, d_model=6, dim_feedforward=4).to(device)
#print(model)
src = torch.rand(2, 4, 6).to(device)
tgt = torch.rand(4, 4, 6).to(device)
compiled_model = torch.compile(model)
out = compiled_model(src, tgt)
print(out)

Compare with Transformer, the other model is fine:

code:

import torch
import torch._inductor.config
print("torch-version: ", torch.__version__)
torch._inductor.config.debug = True
device = "cpu"

def func(x):
    return x.sin() + x.cos()

model = func
compiled_model = torch.compile(model)
data = torch.rand(20)
out = compiled_model(data)
print(out)

I tried the nightly version torch-version: 2.1.0.dev20230526+cpu, it works for me.