Control flow issue with symbolic_trace when using inputs_embeds in MistralForCausalLM
Author: Hongjie1Chu
Pull request: no

### System Info

- `transformers` version: 4.41.2
- Platform: Linux-5.15.0-88-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.23.0
- Safetensors version: 0.4.3
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+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 GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

 @ArthurZucker  @younesbelkada @Narsil

### Information

- [X] The official example scripts
- [ ] 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

```
# from model_mistral import MistralForCausalLM
from transformers import MistralForCausalLM,MistralConfig
import torch
from transformers.utils.fx import symbolic_trace


cfg = MistralConfig.from_pretrained('mistralai/Mistral-7B-v0.1')
cfg.num_hidden_layers = 2
model = MistralForCausalLM(cfg)
print(model)
batch_size = 1
sequence_length = 10
hidden_dim = cfg.hidden_size

# create dummy_input_ids
dummy_input_ids = {
    'input_ids': torch.rand(batch_size, sequence_length)
}
traced = symbolic_trace(model,input_names=list(dummy_input_ids.keys()))

# create dummy_inputs_embeds
dummy_input_embeds = {
    'inputs_embeds': torch.rand(batch_size, sequence_length, hidden_dim)
}


traced = symbolic_trace(model,input_names=list(dummy_input_embeds.keys()))
```

when i use dummy_input_ids to trace ,no error,but when i use dummy_inputs_embeds to trace ,error:
```
Traceback (most recent call last):
  File "/root/chj/geesibling-pytorch_5.23/examples/pytorch/test.py", line 26, in <module>
    traced = symbolic_trace(model,input_names=list(dummy_input_embeds.keys()))
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 1483, in symbolic_trace
    traced_graph = tracer.trace(model, concrete_args=concrete_args)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 1306, in trace
    self.graph = super().trace(root, concrete_args=concrete_args)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 817, in trace
    (self.create_arg(fn(*args)),),
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 1139, in forward
    outputs = self.model(
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 795, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 1170, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 479, in call_module
    ret_val = forward(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 788, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 1024, in forward
    layer_outputs = decoder_layer(
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 795, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 1170, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 479, in call_module
    ret_val = forward(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 788, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 738, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 795, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 1170, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 479, in call_module
    ret_val = forward(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 788, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 650, in forward
    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 795, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 1170, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 479, in call_module
    ret_val = forward(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 788, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py", line 119, in forward
    if seq_len > self.max_seq_len_cached:
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/transformers/utils/fx.py", line 668, in __bool__
    return super().__bool__()
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/proxy.py", line 437, in __bool__
    return self.tracer.to_bool(self)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/fx/proxy.py", line 300, in to_bool
    raise TraceError('symbolically traced variables cannot be used as inputs to control flow')
torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow
```



### Expected behavior

I would like to know why using input_ids allows for normal tracing, and there are no errors when executing if seq_len > self.max_seq_len_cached;, but switching to inputs_embeds causes an issue at this point?


"If I must trace using `inputs_embeds`, how should I proceed? I am working on a pipeline job and need to construct parts of the model as follows:

model1:
```python
MistralForCausalLM(
  (model): MistralModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0): MistralDecoderLayer(
        (self_attn): MistralSdpaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): MistralRotaryEmbedding()
        ),
        (mlp): MistralMLP(
          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLU()
        ),
        (input_layernorm): MistralRMSNorm()
        (post_attention_layernorm): MistralRMSNorm()
      )
    )
  )
)
```

model2:
```python
MistralForCausalLM(
  (model): MistralModel(
    (layers): ModuleList(
      (0): MistralDecoderLayer(
        (self_attn): MistralSdpaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): MistralRotaryEmbedding()
        ),
        (mlp): MistralMLP(
          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLU()
        ),
        (input_layernorm): MistralRMSNorm()
        (post_attention_layernorm): MistralRMSNorm()
      )
    ),
    (norm): MistralRMSNorm()
  ),
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
```

So, model2 can only be traced using `inputs_embeds`. How should I modify it?"

Comments:
cc @fxmarty and @michaelbenayoun 

cc @fxmarty and @michaelbenayoun

can u help me ? @fxmarty and @michaelbenayoun


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.

Hi @Hongjie1Chu, this should be fixed on the main branch with https://github.com/huggingface/transformers/pull/31574.