llama3 not fx traced
Author: myungjin
Pull request: no
Labels: name: bug

name: torch fx

### System Info

- `transformers` version: 4.45.1
- Platform: Linux-5.4.247-162.350.amzn2.x86_64-x86_64-with-glibc2.26
- Python version: 3.10.12
- Huggingface_hub version: 0.24.0
- Safetensors version: 0.4.3
- Accelerate version: 0.32.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.4.0+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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: Tesla V100-SXM2-16GB

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from transformers.utils.fx import symbolic_trace
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
gm = symbolic_trace(model, input_names=["input_ids", "attention_mask", "past_key_values"])
```

```bash
We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class (https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/utils/fx.py", line 1503, in symbolic_trace
    traced_graph = tracer.trace(model, concrete_args=concrete_args)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/utils/fx.py", line 1326, in trace
    self.graph = super().trace(root, concrete_args=concrete_args)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 822, in trace
    (self.create_arg(fn(*args)),),
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1189, in forward
    outputs = self.model(
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 800, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/utils/fx.py", line 1190, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 518, in call_module
    ret_val = forward(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 793, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1000, in forward
    layer_outputs = decoder_layer(
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 800, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/utils/fx.py", line 1190, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 518, in call_module
    ret_val = forward(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 793, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 729, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 800, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/utils/fx.py", line 1190, in call_module
    return super().call_module(m, forward, args, kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 518, in call_module
    ret_val = forward(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 793, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 646, in forward
    if query_states.device.type == "cuda" and causal_mask is not None:
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/transformers/utils/fx.py", line 669, in __bool__
    return super().__bool__()
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/proxy.py", line 447, in __bool__
    return self.tracer.to_bool(self)
  File "~/miniconda3/envs/pytorch-dev/lib/python3.10/site-packages/torch/fx/proxy.py", line 307, 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

it should run without error. `symbolic_trace()` with `["input_ids", "attention_mask"]` runs fine. However, when ["input_ids", "attention_mask", "past_key_values"] is fed as  input_names, the error occurs. If using past_key_values is incorrect, it should be warned and aborted before trying to trace the model.

While a fix on a related error (https://github.com/huggingface/transformers/issues/29923) is included in the released version, it seems there is still some bug.

Comments:
Hey! Thanks for reporting! cc @michaelbenayoun I can reproduce, but no idea how to avoid this. It's also not new, so I think past key value path was not tested!

@ArthurZucker I’ve experienced the same issue. It seems that although the fix on https://github.com/huggingface/transformers/issues/29923 resolves the problem related to cache, in https://github.com/huggingface/transformers/blob/70b07d97cf2c5f61fff55700b65528a1b6845cd2/src/transformers/utils/fx.py#L1052-L1068, we’re still creating dummy_inputs for past_key_values using the old tuple method. I replaced this part with:
```
elif "past_key_values" in input_name:
    if model.config.model_type not in _FX_SUPPORTED_MODELS_WITH_KV_CACHE:
        raise NotImplementedError(
            f"Symbolic tracing with past_key_values input is not supported yet for the model {model.config.model_type}. Please open an issue or a PR in the Transformers repository if you'd like to see this support added."
        )
    inputs_dict[input_name] = DynamicCache()
```
It seems we can now trace Llama. However, I’m not sure if this is the correct way to fix the issue.

I think there are fixes for direct tracing in `optimum`, but happy to have something more general here!

@ArthurZucker Thanks for looking into this.
Can you provide a link or pointer to the fixes you refer to?
I went through optimum [document](https://huggingface.co/docs/optimum/en/torch_fx/overview).
Examples in the doc still rely on `symbolic_trace`, so I don't know how `optimum` can help addressing this issue.

The linked PR should help!

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.