FX tracer doen't work when requesting non-default input argument
Author: pgimenes
Pull request: no

### System Info

- `transformers` version: 4.43.0.dev0
- Platform: Linux-5.15.0-113-generic-x86_64-with-glibc2.35
- Python version: 3.11.4
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.2
- Accelerate version: 0.29.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.0.dev20240714+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: no
- GPU type: NVIDIA RTX A6000

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
import torch
from transformers.models.opt.configuration_opt import OPTConfig
from transformers.models.opt.modeling_opt import OPTModel

from transformers.utils.fx import symbolic_trace

config = OPTConfig()
config.num_hidden_layers = 1

model = OPTModel(config)

# Check the forward pass works
out = model(
    inputs_embeds=torch.randn(1, 1, 768),
)

traced = symbolic_trace(model, input_names=["inputs_embeds"])

```

### Expected behavior

This is the generated output:

```
Traceback (most recent call last):
  File "/home/pedrogimenes/hf_issue.py", line 17, in <module>
    traced = symbolic_trace(model, input_names=["inputs_embeds"])
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/utils/fx.py", line 1483, in symbolic_trace
    traced_graph = tracer.trace(model, concrete_args=concrete_args)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/utils/fx.py", line 1306, in trace
    self.graph = super().trace(root, concrete_args=concrete_args)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 822, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/models/opt/modeling_opt.py", line 969, in forward
    decoder_outputs = self.decoder(
                      ^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 800, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/utils/fx.py", line 1170, in call_module
    return super().call_module(m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 518, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 793, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1716, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1727, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/models/opt/modeling_opt.py", line 829, in forward
    causal_attention_mask = _prepare_4d_causal_attention_mask(
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/modeling_attn_mask_utils.py", line 320, in _prepare_4d_causal_attention_mask
    if attention_mask is not None and len(attention_mask.shape) == 2:
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/transformers/utils/fx.py", line 663, in __len__
    return super().__len__()
           ^^^^^^^^^^^^^^^^^
  File "/home/pedrogimenes/anaconda3/envs/mase/lib/python3.11/site-packages/torch/fx/proxy.py", line 456, in __len__
    raise RuntimeError("'len' is not supported in symbolic tracing by default. If you want "
RuntimeError: 'len' is not supported in symbolic tracing by default. If you want this call to be recorded, please call torch.fx.wrap('len') at module scope
```

When I pass the `input_names=["inputs_embeds"]` argument, I'd expect the behaviour to be the same as setting input_names=None, with the only difference being that the output graph should have an `input_embeds` placeholder node, instead of `input_ids`. However with inputs_embeds, it seems the metadata for some Proxy objects is not set, so their boolean attribute cannot be evaluated from the `transformers/models/opt/modeling_opt.py` forward pass code. For some reason, the proxy metadata is correctly generated with the input_ids argument, but not inputs_embeds.

Comments:
cc @fxmarty 

Hi @pgimenes, I can reproduce this issue on transformers==24.1.2, but I can not on transformers installed from main (ac946aac257cadfa8264fa4a284cd0ea1061c5b5). Are you sure you installed transformers from main as mentioned in the bug repro?

Maybe https://github.com/huggingface/transformers/pull/31574 (that is not released) helped there.



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.