fx.symbolic_trace not working for Roberta
Author: WeiHao97
Pull request: no
Labels: name: bug

### System Info

```shell
transformers version: 4.18.0

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
```


### Who can help?

@LysandreJi

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python
from transformers import RobertaForMaskedLM
from transformers import RobertaConfig
from transformers.utils import 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}
gm = fx.symbolic_trace(model,concrete_args)
```

### Expected behavior

```shell
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [11], in <cell line: 1>()
----> 1 gm = fx.symbolic_trace(model,concrete_args)

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:594, in symbolic_trace(model, input_names)
    592 # Tracing.
    593 tracer = HFTracer()
--> 594 traced_graph = tracer.trace(model, concrete_args=concrete_args)
    595 traced = torch.fx.GraphModule(model, traced_graph)
    597 return traced

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:470, in HFTracer.trace(self, root, concrete_args, method_names)
    467 sig = inspect.signature(root.forward)
    468 input_names = sig.parameters.keys() - concrete_args.keys()
--> 470 self.record(root, input_names, method_names=method_names)
    472 # TODO: adapt the way leaf function are wrapped with the "autowrap function" feature from Tracer.
    473 autowrap_functions = [patched for (_, _, patched) in self._leaf_functions_register.values()]

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:426, in HFTracer.record(self, model, input_names, method_names)
    423 cache_names, original_methods = self._monkey_patch_tensor_methods_for_model_recording(model, method_names)
    424 self.original_methods = original_methods
--> 426 model(**inputs)
    428 _reset_tensor_methods(original_methods)
    430 self.recorded_methods = {
    431     method_name: cache_name for method_name, cache_name in cache_names.items() if hasattr(model, cache_name)
    432 }

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: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/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:851, 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)
    842 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
    844 embedding_output = self.embeddings(
    845     input_ids=input_ids,
    846     position_ids=position_ids,
   (...)
    849     past_key_values_length=past_key_values_length,
    850 )
--> 851 encoder_outputs = self.encoder(
    852     embedding_output,
    853     attention_mask=extended_attention_mask,
    854     head_mask=head_mask,
    855     encoder_hidden_states=encoder_hidden_states,
    856     encoder_attention_mask=encoder_extended_attention_mask,
    857     past_key_values=past_key_values,
    858     use_cache=use_cache,
    859     output_attentions=output_attentions,
    860     output_hidden_states=output_hidden_states,
    861     return_dict=return_dict,
    862 )
    863 sequence_output = encoder_outputs[0]
    864 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

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:492, in RobertaEncoder.forward(self, hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
    479 def forward(
    480     self,
    481     hidden_states: torch.Tensor,
   (...)
    490     return_dict: Optional[bool] = True,
    491 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
--> 492     all_hidden_states = () if output_hidden_states else None
    493     all_self_attentions = () if output_attentions else None
    494     all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:326, in HFTracer._wrap_method_for_model_recording.<locals>.wrapped(*args, **kwargs)
    324     setattr(model, cache_name, [])
    325 cache = getattr(model, cache_name)
--> 326 res = method(*args, **kwargs)
    327 cache.append(res)
    328 return res

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:326, in HFTracer._wrap_method_for_model_recording.<locals>.wrapped(*args, **kwargs)
    324     setattr(model, cache_name, [])
    325 cache = getattr(model, cache_name)
--> 326 res = method(*args, **kwargs)
    327 cache.append(res)
    328 return res

RuntimeError: Boolean value of Tensor with more than one value is ambiguous
```


Comments:
After specifying: 
```python
input_names = ['input_ids',
 'attention_mask',
 'token_type_ids',
 'position_ids',
 'encoder_hidden_states',
 'encoder_attention_mask',
 'labels'] 
```
I saw another bug in:
```python
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [2], in <cell line: 22>()
     14 model = RobertaForMaskedLM(config)
     15 input_names = ['input_ids',
     16  'attention_mask',
     17  'token_type_ids',
   (...)
     20  'encoder_attention_mask',
     21  'labels']
---> 22 gm = symbolic_trace(model,input_names)

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:595, in symbolic_trace(model, input_names)
    593 tracer = HFTracer()
    594 traced_graph = tracer.trace(model, concrete_args=concrete_args)
--> 595 traced = torch.fx.GraphModule(model, traced_graph)
    597 return traced

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph_module.py:312, in GraphModule.__init__(self, root, graph, class_name)
    309 else:
    310     raise RuntimeError('Unsupported type ' + str(root) + ' passed for root!')
--> 312 self.graph = graph
    314 # Store the Tracer class responsible for creating a Graph separately as part of the
    315 # GraphModule state, except when the Tracer is defined in a local namespace.
    316 # Locally defined Tracers are not pickleable. This is needed because torch.package will
    317 # serialize a GraphModule without retaining the Graph, and needs to use the correct Tracer
    318 # to re-create the Graph during deserialization.
    319 self._tracer_cls = None

File ~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py:1225, in Module.__setattr__(self, name, value)
   1223     buffers[name] = value
   1224 else:
-> 1225     object.__setattr__(self, name, value)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph_module.py:346, in GraphModule.graph(self, g)
    344 self._graph = g
    345 g.owning_module = self
--> 346 self.recompile()

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph_module.py:562, in GraphModule.recompile(self)
    560     self._in_spec = self._graph._pytree_info.in_spec
    561     self._out_spec = self._graph._pytree_info.out_spec
--> 562 python_code = self._graph.python_code(root_module='self')
    563 self._code = python_code.src
    565 cls = type(self)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:834, in Graph.python_code(self, root_module)
    831             node._repr_fn = orig_repr_fns[node]
    833 with override_node_repr(self):
--> 834     return self._python_code(root_module, namespace)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:988, in Graph._python_code(self, root_module, namespace)
    983     raise NotImplementedError(f'node: {node.op} {node.target}')
    985 for node in self.nodes:
    986     # NOTE: emit_node does not emit a string with newline. It depends
    987     # on delete_unused_values to append one
--> 988     emit_node(node)
    989     delete_unused_values(node)
    991 if len(body) == 0:
    992     # If the Graph has no non-placeholder nodes, no lines for the body
    993     # have been emitted. To continue to have valid Python code, emit a
    994     # single pass statement

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:928, in Graph._python_code.<locals>.emit_node(node)
    927 def emit_node(node : Node):
--> 928     maybe_type_annotation = '' if node.type is None else f' : {type_repr(node.type)}'
    929     if node.op == 'placeholder':
    930         assert isinstance(node.target, str)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:885, in Graph._python_code.<locals>.type_repr(o)
    882     origin_typename = add_global(_type_repr(origin_type), origin_type)
    884     # Assign global names for each of the inner type variables.
--> 885     args = [type_repr(arg) for arg in o.__args__]
    887     return f'{origin_typename}[{",".join(args)}]'
    889 # Common case: this is a regular module name like 'foo.bar.baz'

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:885, in <listcomp>(.0)
    882     origin_typename = add_global(_type_repr(origin_type), origin_type)
    884     # Assign global names for each of the inner type variables.
--> 885     args = [type_repr(arg) for arg in o.__args__]
    887     return f'{origin_typename}[{",".join(args)}]'
    889 # Common case: this is a regular module name like 'foo.bar.baz'

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:885, in Graph._python_code.<locals>.type_repr(o)
    882     origin_typename = add_global(_type_repr(origin_type), origin_type)
    884     # Assign global names for each of the inner type variables.
--> 885     args = [type_repr(arg) for arg in o.__args__]
    887     return f'{origin_typename}[{",".join(args)}]'
    889 # Common case: this is a regular module name like 'foo.bar.baz'

File ~/anaconda3/lib/python3.9/typing.py:711, in _BaseGenericAlias.__getattr__(self, attr)
    709 if '__origin__' in self.__dict__ and not _is_dunder(attr):
    710     return getattr(self.__origin__, attr)
--> 711 raise AttributeError(attr)

AttributeError: __args__
```

@michaelbenayoun Do you mind take a look? Thank you very much!

Hi @WeiHao97 ,

To trace the model, you need to specify the input names and not the concrete args, the input_names represent the inputs your traced model will take.
If you want to trace Roberta, you most likely need to only provide "input_ids" and "attention_mask":

```python
model = RobertaForMaskedLM(config)
input_names = ["input_ids", "attention_mask"]
gm = fx.symbolic_trace(model, input_names)
```

> Hi @WeiHao97 ,
> 
> To trace the model, you need to specify the input names and not the concrete args, the input_names represent the inputs your traced model will take. If you want to trace Roberta, you most likely need to only provide "input_ids" and "attention_mask":
> 
> ```python
> model = RobertaForMaskedLM(config)
> input_names = ["input_ids", "attention_mask"]
> gm = fx.symbolic_trace(model, input_names)
> ```

After specifying input_names = ['input_ids',
 'attention_mask',
 'token_type_ids',
 'position_ids',
 'encoder_hidden_states',
 'encoder_attention_mask',
 'labels'] or input_names = ["input_ids", "attention_mask"], I got the second error.

Could you share the error message please?
Also what are you trying to do? It seems to be an odd set of inputs you want to use.

That's weird that you get an error with this:
```python
model = RobertaForMaskedLM(config)
input_names = ["input_ids", "attention_mask"]
gm = fx.symbolic_trace(model, input_names)
```
I am able to get this to work on my end.


> Could you share the error message please? Also what are you trying to do? It seems to be an odd set of inputs you want to use.
> 
> That's weird that you get an error with this:
> 
> ```python
> model = RobertaForMaskedLM(config)
> input_names = ["input_ids", "attention_mask"]
> gm = fx.symbolic_trace(model, input_names)
> ```
> 
> I am able to get this to work on my end.

I ran with transformers version '4.18.0' :
```python
from transformers import RobertaForMaskedLM
from transformers import RobertaConfig
from transformers.utils import fx

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"]
gm = fx.symbolic_trace(model, input_names)
```
and got:

```python
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [1], in <cell line: 15>()
     13 model = RobertaForMaskedLM(config)
     14 input_names = ["input_ids", "attention_mask"]
---> 15 gm = fx.symbolic_trace(model, input_names)

File ~/anaconda3/lib/python3.9/site-packages/transformers/utils/fx.py:595, in symbolic_trace(model, input_names)
    593 tracer = HFTracer()
    594 traced_graph = tracer.trace(model, concrete_args=concrete_args)
--> 595 traced = torch.fx.GraphModule(model, traced_graph)
    597 return traced

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph_module.py:312, in GraphModule.__init__(self, root, graph, class_name)
    309 else:
    310     raise RuntimeError('Unsupported type ' + str(root) + ' passed for root!')
--> 312 self.graph = graph
    314 # Store the Tracer class responsible for creating a Graph separately as part of the
    315 # GraphModule state, except when the Tracer is defined in a local namespace.
    316 # Locally defined Tracers are not pickleable. This is needed because torch.package will
    317 # serialize a GraphModule without retaining the Graph, and needs to use the correct Tracer
    318 # to re-create the Graph during deserialization.
    319 self._tracer_cls = None

File ~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py:1225, in Module.__setattr__(self, name, value)
   1223     buffers[name] = value
   1224 else:
-> 1225     object.__setattr__(self, name, value)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph_module.py:346, in GraphModule.graph(self, g)
    344 self._graph = g
    345 g.owning_module = self
--> 346 self.recompile()

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph_module.py:562, in GraphModule.recompile(self)
    560     self._in_spec = self._graph._pytree_info.in_spec
    561     self._out_spec = self._graph._pytree_info.out_spec
--> 562 python_code = self._graph.python_code(root_module='self')
    563 self._code = python_code.src
    565 cls = type(self)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:834, in Graph.python_code(self, root_module)
    831             node._repr_fn = orig_repr_fns[node]
    833 with override_node_repr(self):
--> 834     return self._python_code(root_module, namespace)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:988, in Graph._python_code(self, root_module, namespace)
    983     raise NotImplementedError(f'node: {node.op} {node.target}')
    985 for node in self.nodes:
    986     # NOTE: emit_node does not emit a string with newline. It depends
    987     # on delete_unused_values to append one
--> 988     emit_node(node)
    989     delete_unused_values(node)
    991 if len(body) == 0:
    992     # If the Graph has no non-placeholder nodes, no lines for the body
    993     # have been emitted. To continue to have valid Python code, emit a
    994     # single pass statement

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:928, in Graph._python_code.<locals>.emit_node(node)
    927 def emit_node(node : Node):
--> 928     maybe_type_annotation = '' if node.type is None else f' : {type_repr(node.type)}'
    929     if node.op == 'placeholder':
    930         assert isinstance(node.target, str)

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:885, in Graph._python_code.<locals>.type_repr(o)
    882     origin_typename = add_global(_type_repr(origin_type), origin_type)
    884     # Assign global names for each of the inner type variables.
--> 885     args = [type_repr(arg) for arg in o.__args__]
    887     return f'{origin_typename}[{",".join(args)}]'
    889 # Common case: this is a regular module name like 'foo.bar.baz'

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:885, in <listcomp>(.0)
    882     origin_typename = add_global(_type_repr(origin_type), origin_type)
    884     # Assign global names for each of the inner type variables.
--> 885     args = [type_repr(arg) for arg in o.__args__]
    887     return f'{origin_typename}[{",".join(args)}]'
    889 # Common case: this is a regular module name like 'foo.bar.baz'

File ~/anaconda3/lib/python3.9/site-packages/torch/fx/graph.py:885, in Graph._python_code.<locals>.type_repr(o)
    882     origin_typename = add_global(_type_repr(origin_type), origin_type)
    884     # Assign global names for each of the inner type variables.
--> 885     args = [type_repr(arg) for arg in o.__args__]
    887     return f'{origin_typename}[{",".join(args)}]'
    889 # Common case: this is a regular module name like 'foo.bar.baz'

File ~/anaconda3/lib/python3.9/typing.py:711, in _BaseGenericAlias.__getattr__(self, attr)
    709 if '__origin__' in self.__dict__ and not _is_dunder(attr):
    710     return getattr(self.__origin__, attr)
--> 711 raise AttributeError(attr)

AttributeError: __args__
```

Could you try with transformers sources?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.