Repository: pytorch
Issue: fx.Tracer with param_shapes_constant=True not working for RobertaForMaskedLM
Labels: triaged,module: fx

### 🐛 Describe the bug

param_shapes_constant=True can't bypass issue mentioned by @jamesr66a 
https://github.com/pytorch/pytorch/issues/61733

```python
from transformers import RobertaForMaskedLM
from transformers import RobertaConfig
import torch.fx as fx
import inspect

config = RobertaConfig(
 vocab_size=52_000,
 max_position_embeddings=514,
 num_attention_heads=12,
 num_hidden_layers=12,
 type_vocab_size=1,
)

model = RobertaForMaskedLM(config)
input_names = ["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask",'token_type_ids']
sig = inspect.signature(model.forward)
concrete_args = {p.name: None for p in sig.parameters.values() if p.name not in input_names}

tracer = fx.Tracer(param_shapes_constant=True)
graph = tracer.trace(model, concrete_args)

```

Error Message:
```
---------------------------------------------------------------------------
TraceError                                Traceback (most recent call last)
Input In [20], in <cell line: 2>()
      1 tracer = fx.Tracer(param_shapes_constant=True)
----> 2 graph = tracer.trace(model, concrete_args)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py:615, in Tracer.trace(self, root, concrete_args)
    613         for module in self._autowrap_search:
    614             _autowrap_check(patcher, module.__dict__, self._autowrap_function_ids)
--> 615         self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
    616                          type_expr=fn.__annotations__.get('return', None))
    618 self.submodule_paths = None
    620 return self.graph

File ~/anaconda3/lib/python3.9/site-packages/transformers/models/roberta/modeling_roberta.py:1098, in RobertaForMaskedLM.forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, output_attentions, output_hidden_states, return_dict)
   1088 r"""
   1089 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
   1090     Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
   (...)
   1094     Used to hide legacy arguments that have been deprecated.
   1095 """
   1096 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-> 1098 outputs = self.roberta(
   1099     input_ids,
   1100     attention_mask=attention_mask,
   1101     token_type_ids=token_type_ids,
   1102     position_ids=position_ids,
   1103     head_mask=head_mask,
   1104     inputs_embeds=inputs_embeds,
   1105     encoder_hidden_states=encoder_hidden_states,
   1106     encoder_attention_mask=encoder_attention_mask,
   1107     output_attentions=output_attentions,
   1108     output_hidden_states=output_hidden_states,
   1109     return_dict=return_dict,
   1110 )
   1111 sequence_output = outputs[0]
   1112 prediction_scores = self.lm_head(sequence_output)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py:604, in Tracer.trace.<locals>.module_call_wrapper(mod, *args, **kwargs)
    600     return _orig_module_call(mod, *args, **kwargs)
    602 _autowrap_check(patcher, getattr(getattr(mod, "forward", mod), "__globals__", {}),
    603                 self._autowrap_function_ids)
--> 604 return self.call_module(mod, forward, args, kwargs)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py:422, in Tracer.call_module(self, m, forward, args, kwargs)
    420 module_qualified_name = self.path_of_module(m)
    421 if not self.is_leaf_module(m, module_qualified_name):
--> 422     return forward(*args, **kwargs)
    423 return self.create_proxy('call_module', module_qualified_name, args, kwargs)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py:600, in Tracer.trace.<locals>.module_call_wrapper.<locals>.forward(*args, **kwargs)
    599 def forward(*args, **kwargs):
--> 600     return _orig_module_call(mod, *args, **kwargs)

File ~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py:1102, in Module._call_impl(self, *input, **kwargs)
   1098 # If we don't have any hooks, we want to skip the rest of the logic in
   1099 # this function, and just call forward.
   1100 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102     return forward_call(*input, **kwargs)
   1103 # Do not call functions when jit is used
   1104 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/lib/python3.9/site-packages/transformers/models/roberta/modeling_roberta.py:824, in RobertaModel.forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
    820         token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
    822 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
    823 # ourselves in which case we just need to make it broadcastable to all heads.
--> 824 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
    826 # If a 2D or 3D attention mask is provided for the cross-attention
    827 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
    828 if self.config.is_decoder and encoder_hidden_states is not None:

File ~/anaconda3/lib/python3.9/site-packages/transformers/modeling_utils.py:559, in ModuleUtilsMixin.get_extended_attention_mask(self, attention_mask, input_shape, device)
    543 """
    544 Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
    545 
   (...)
    555     `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
    556 """
    557 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
    558 # ourselves in which case we just need to make it broadcastable to all heads.
--> 559 if attention_mask.dim() == 3:
    560     extended_attention_mask = attention_mask[:, None, :, :]
    561 elif attention_mask.dim() == 2:
    562     # Provided a padding mask of dimensions [batch_size, seq_length]
    563     # - if the model is a decoder, apply a causal mask in addition to the padding mask
    564     # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/proxy.py:251, in Proxy.__bool__(self)
    250 def __bool__(self) -> bool:
--> 251     return self.tracer.to_bool(self)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/proxy.py:152, in TracerBase.to_bool(self, obj)
    145 @compatibility(is_backward_compatible=True)
    146 def to_bool(self, obj: 'Proxy') -> bool:
    147     """Called when a proxy object is being converted to a boolean, such as
    148     when used in control flow.  Normally we don't know what to do because
    149     we don't know the value of the proxy, but a custom tracer can attach more
    150     information to the graph node using create_node and can choose to return a value.
    151     """
--> 152     raise TraceError('symbolically traced variables cannot be used as inputs to control flow')

TraceError: symbolically traced variables cannot be used as inputs to control flow
```

### Versions

PyTorch version: 1.10.2
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.9.12 (main, Apr  5 2022, 06:56:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.102.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA T500
Nvidia driver version: 472.91
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.5
[pip3] numpydoc==1.2
[pip3] torch==1.10.2
[pip3] torchvision==0.11.3
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py39h7f8727e_0
[conda] mkl_fft                   1.3.1            py39hd3c417c_0
[conda] mkl_random                1.2.2            py39h51133e4_0
[conda] numpy                     1.21.5           py39he7a7128_1
[conda] numpy-base                1.21.5           py39hf524024_1
[conda] numpydoc                  1.2                pyhd3eb1b0_0
[conda] pytorch                   1.10.2          cpu_py39hfa7516b_0
[conda] torchvision               0.11.3               py39_cu113    pytorch

cc @ezyang @SherlockNoMad