The latest transformer.utils.fx does not working on llama. Only far older version(4.45.1) works.
Author: Sangyu-Han
Pull request: no
Labels: name: bug

### System Info


- `transformers` version: 4.52.3
- Platform: Linux-6.8.0-59-generic-x86_64-with-glibc2.39
- Python version: 3.12.2
- Huggingface_hub version: 0.31.4
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.5.1+cu124 (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 GeForce RTX 3090


### 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

from transformers import AutoModelForCausalLM
from transformers.utils.fx import symbolic_trace
import torch
import sys, types, torch.fx._symbolic_trace as st
from types import CodeType
from transformers import AutoTokenizer
import sys, types
import torch.fx._symbolic_trace as sym
import transformers.utils.fx as tfx



    
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
model.eval()

prompt = """\
Context: Mount Everest attracts many climbers, including highly experienced mountaineers. There are two main climbing routes, one approaching the summit from the southeast in Nepal (known as the standard route) and the other from the north in Tibet. While not posing substantial technical climbing challenges on the standard route, Everest presents dangers such as altitude sickness, weather, and wind, as well as hazards from avalanches and the Khumbu Icefall. As of November 2022, 310 people have died on Everest. Over 200 bodies remain on the mountain and have not been removed due to the dangerous conditions. The first recorded efforts to reach Everest's summit were made by British mountaineers. As Nepal did not allow foreigners to enter the country at the time, the British made several attempts on the north ridge route from the Tibetan side. After the first reconnaissance expedition by the British in 1921 reached 7,000 m (22,970 ft) on the North Col, the 1922 expedition pushed the north ridge route up to 8,320 m (27,300 ft), marking the first time a human had climbed above 8,000 m (26,247 ft). The 1924 expedition resulted in one of the greatest mysteries on Everest to this day: George Mallory and Andrew Irvine made a final summit attempt on 8 June but never returned, sparking debate as to whether they were the first to reach the top. Tenzing Norgay and Edmund Hillary made the first documented ascent of Everest in 1953, using the southeast ridge route. Norgay had reached 8,595 m (28,199 ft) the previous year as a member of the 1952 Swiss expedition. The Chinese mountaineering team of Wang Fuzhou, Gonpo, and Qu Yinhua made the first reported ascent of the peak from the north ridge on 25 May 1960. \
Question: How high did they climb in 1922? According to the text, the 1922 expedition reached 8,"""


input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(model.device)

traced = tfx.symbolic_trace(model) <-- ValueError

print(traced.graph)      

The error message is :

File "/home/sangyu/Desktop/Master/delphi-main/trace_test.py", line 28, in <module>
traced = tfx.symbolic_trace(model)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/sangyu/anaconda3/envs/py312/lib/python3.12/site-packages/transformers/utils/fx.py", line 1503, in symbolic_trace
traced_graph = tracer.trace(model, concrete_args=concrete_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/sangyu/anaconda3/envs/py312/lib/python3.12/site-packages/transformers/utils/fx.py", line 1326, in trace
self.graph = super().trace(root, concrete_args=concrete_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/sangyu/anaconda3/envs/py312/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 776, in trace
fn, args = self.create_args_for_root(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/sangyu/anaconda3/envs/py312/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 663, in create_args_for_root
root_fn = _patch_function(root_fn, len(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/sangyu/anaconda3/envs/py312/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 189, in _patch_function
new_code = CodeType(*co_args)  # type: ignore[arg-type]
^^^^^^^^^^^^^^^^^^
ValueError: code: co_varnames is too small

### Expected behavior

It should not return error.

Even though this is the simplest implemention using transformers.utils.fx.symbolic_trace, it returns error.

However, when I run the same code with python==3.10 and transformers==4.45.1 It runs successfully.



Comments:
Hi @Sangyu-Han, we're not actively maintaining `utils.fx` anymore because PyTorch is also not really maintaining `torch.fx`, and are switching to `torch.export()` instead

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.