Repository: pytorch
Issue: Update _symbolic_trace.py for proxy_placeholder function
Labels: triaged,module: fx

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
A bug in proxy_placeholder function cause custom nn.Module with concrete_args not working for torch 1.9.0.
This is because torch force to create proxy for each concrete_args. For specialized function, is it correct to create proxy?

This bug would cause the huggingface/transformers is not able to support torch.fx with torch 1.9.0. https://github.com/huggingface/transformers/blob/3ff2cde5ca4a2d3c622b827d9edf7e3d0b7f4fb7/src/transformers/utils/fx.py#L366
Then, we want to propose a new feature to solve this bug and implement it.

## To Reproduce

Steps to reproduce the behavior:

1. install torch 1.9.0
1. install huggingface/transformers by `pip install transformers`
1. force huggingface/transformers to support torch.fx of 1.9.0 by changing https://github.com/huggingface/transformers/blob/ab7551cd7ff84cb5b7328bc37a06e06fa19f02bb/src/transformers/file_utils.py#L278 as TORCH_FX_REQUIRED_VERSION = version.parse("1.9")
1. run the follow code sample.

```
import torch
from transformers import BertModel, BertConfig
from transformers.utils.fx import symbolic_trace as transformers_symbolic_trace

config = BertConfig(
    hidden_size=1024,
    num_hidden_layers=24,
    num_attention_heads=16,
    intermediate_size=4096
)

module = BertModel(config).cuda()
symbolic_traced_module = transformers_symbolic_trace(
    module,
    batch_size=1,
    sequence_length=224,
)
```

error massage
```
UserWarning: Was not able to add assertion to guarantee correct inputs to specialized function. It is up to the user to make sure that your inputs match the inputs you specialized the function with.
  torch.warnings.warn(
Traceback (most recent call last):
  File "script.py", line 13, in <module>
    symbolic_traced_module = transformers_symbolic_trace(
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/transformers/utils/fx.py", line 378, in symbolic_trace
    traced = torch.fx.GraphModule(model, traced_graph)
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/torch/fx/graph_module.py", line 235, in __init__
    self.graph = graph
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1178, in __setattr__
    object.__setattr__(self, name, value)
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/torch/fx/graph_module.py", line 260, in graph
    self.recompile()
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/torch/fx/graph_module.py", line 472, in recompile
    cls.forward = _forward_from_src(self._code, python_code.globals)
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/torch/fx/graph_module.py", line 44, in _forward_from_src
    exec_with_source(src, globals_copy)
  File "/home/t-chengluo/.local/lib/python3.8/site-packages/torch/fx/graph_module.py", line 28, in exec_with_source
    exec(compile(src, key, 'exec'), globals)
  File "<eval_with_key_0>", line 2
SyntaxError: non-default argument follows default argument

```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment
Collecting environment information...
PyTorch version: 1.9.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  2 2021, 10:49:15)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-80-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.2.67
GPU models and configuration:
GPU 0: Tesla V100-PCIE-32GB
GPU 1: Tesla V100-PCIE-32GB
GPU 2: Tesla V100-PCIE-32GB
GPU 3: Tesla V100-PCIE-32GB

Nvidia driver version: 460.27.04
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.1
[pip3] torch==1.9.0
[pip3] torchvision==0.10.0
[conda] Could not collect

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

## Additional context

<!-- Add any other context about the problem here. -->
We want to propose a solution for this issue, that is simply moving one line of https://github.com/pytorch/pytorch/blob/3df4870343d212134a92934207bccee91fa6518f/torch/fx/_symbolic_trace.py#L434 as

```
        def proxy_placeholder(name: str):
            if concrete_args is not None and name in concrete_args :
                cnt = 0

                def replace_ph(x):
                    nonlocal cnt
                    cnt += 1
                    if x == PH:
                        out = self.create_proxy('placeholder', f'{name}_{str(cnt)}', (), {})
                        return out
                    # Union[int, bool] == bool in Python <= 3.6
                    if type(x) == bool or type(x) in base_types and type(x) != torch.Tensor:
                        out = self.create_proxy('placeholder', f'{name}_{str(cnt)}', (), {})
                        torch._assert(out == x, f"{name} has been specialized to have value {x}")
                    else:
                        torch.warnings.warn(
                            "Was not able to add assertion to guarantee correct inputs to "
                            "specialized function. It is up to the user to make sure that your inputs match the "
                            "inputs you specialized the function with."
                        )

                    return x

                return pytree.tree_map(replace_ph, concrete_args[name])

```
our solution logic is create_proxy is only required when the user passes in `fx.PH` for values.

cc @ezyang @SherlockNoMad @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @mruberry