torch._dynamo.export does not work with bert model
Labels: triaged, ezyang's list, oncall: pt2, module: export

### 🐛 Describe the bug

I'm interested in using aot inductor. However, I'm getting an error while running the following script. Note torch.compile is able to generate one graph for this model.

```py
import torch
import transformers
import logging
torch._logging.set_logs(dynamo=logging.INFO,inductor=logging.INFO)


device = torch.device("cuda:0")
model = transformers.BertForMaskedLM.from_pretrained("bert-base-uncased")

model.eval()
model.cuda()
model.half()

bs = 1
ins = {'input_ids': torch.randint(0, 10, size=(bs, 512)).to(device), 'attention_mask': torch.ones(bs, 512, dtype=torch.int64).to(device)}

with torch.no_grad():
    module, tmp = torch._dynamo.export(model,  **ins)
```

Error message:
```
[2023-07-06 00:16:23,006] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']
- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
[2023-07-06 00:16:29,274] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing forward /home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py:1326
[2023-07-06 00:16:33,244] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo done tracing forward (RETURN_VALUE)
[2023-07-06 00:16:33,275] torch._dynamo.output_graph: [INFO] Step 2: calling compiler function dynamo_normalization_capturing_compiler
[2023-07-06 00:16:33,275] torch._dynamo.output_graph: [INFO] Step 2: done compiler function dynamo_normalization_capturing_compiler
[2023-07-06 00:16:33,951] torch._dynamo.eval_frame: [INFO] Summary of dimension constraints:
The following dimensions have been specialized and CANNOT be dynamic.
```
```
def specializations(input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None):
    # input_ids:
    assert input_ids.size()[0] == 1
    assert input_ids.size()[1] == 512

    # attention_mask:
    assert attention_mask.size()[0] == 1
    assert attention_mask.size()[1] == 512
```
```
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/ubuntu/src/AITemplate/fx2ait/bert_inductor.py:19 in <module>                               │
│                                                                                                  │
│   16 ins = {'input_ids': torch.randint(0, 10, size=(bs, 512)).to(device), 'attention_mask': t    │
│   17                                                                                             │
│   18 with torch.no_grad():                                                                       │
│ ❱ 19 │   module, tmp = torch._dynamo.export(model,  **ins)                                       │
│   20                                                                                             │
│                                                                                                  │
│ /home/ubuntu/src/pytorch/torch/_dynamo/eval_frame.py:1007 in export                              │
│                                                                                                  │
│   1004 │                                                                                         │
│   1005 │   assert graph_captured_result is not None                                              │
│   1006 │   flat_both = list(graph_captured_result) + flat_args                                   │
│ ❱ 1007 │   matched_output_elements_positions = produce_matching(flat_both, flat_results_traced)  │
│   1008 │                                                                                         │
│   1009 │   if aten_graph:                                                                        │
│   1010 │   │   # Running graph with interpreter is needed for propagating the stack_trace        │
│                                                                                                  │
│ /home/ubuntu/src/pytorch/torch/_dynamo/eval_frame.py:888 in produce_matching                     │
│                                                                                                  │
│    885 │   │   │   │   │   │   "Dynamo input/output is not consistent with traced input/output"  │
│    886 │   │   │   │   │   )                                                                     │
│    887 │   │   │   else:                                                                         │
│ ❱  888 │   │   │   │   assert (                                                                  │
│    889 │   │   │   │   │   id(arg) in dict_of_source_args                                        │
│    890 │   │   │   │   ), "Dynamo input and output is a strict subset of traced input/output"    │
│    891 │   │   │   │   matched_elements_positions.append(dict_of_source_args[id(arg)])           │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
AssertionError: Dynamo input and output is a strict subset of traced input/output
[2023-07-06 00:16:34,090] torch._dynamo.utils: [INFO] TorchDynamo compilation metrics:
Function                          Runtimes (s)
------------------------------  --------------
_compile                                4.129
OutputGraph.call_user_compiler          0.0003
```

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0a0+git4baac20
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-6ubuntu2) 7.5.0
Clang version: Could not collect
CMake version: version 3.25.2
Libc version: glibc-2.31

Python version: 3.9.13 (main, Aug 25 2022, 23:26:10)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-1038-aws-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305