Transformers are untraceable with FX after 4.38
Author: Li357
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.48.0
- Platform: Linux-6.8.0-1020-gcp-x86_64-with-glibc2.39
- Python version: 3.12.3
- Huggingface_hub version: 0.27.1
- Safetensors version: 0.4.5
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.6.0a0+df5bbc09d1.nv24.12 (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?: Yes
- GPU type: NVIDIA H100 80GB HBM3

### Who can help?

@michaelbenayoun since I've seen them on some FX fixes

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

Use `transformers.utils.fx.symbolic_trace` on any AutoModelForCausalLM:

```
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.utils.fx import symbolic_trace

model = AutoModelForCausalLM.from_pretrained('deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B')
symbolic_trace(model)
```

You'll get an error about unpacking a Proxy object:

```
Traceback (most recent call last):
  File "...", line 5, in <module>
    symbolic_trace(model)
  File "/usr/local/lib/python3.12/dist-packages/transformers/utils/fx.py", line 1506, in symbolic_trace
    traced_graph = tracer.trace(model, concrete_args=concrete_args)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/utils/fx.py", line 1329, in trace
    self.graph = super().trace(root, concrete_args=concrete_args)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 823, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/models/qwen2/modeling_qwen2.py", line -1, in forward
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 801, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/utils/fx.py", line 1193, in call_module
    return super().call_module(m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 519, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 794, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/models/qwen2/modeling_qwen2.py", line 574, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 801, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/utils/fx.py", line 1193, in call_module
    return super().call_module(m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 519, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 794, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/models/qwen2/modeling_qwen2.py", line 259, in forward
    hidden_states, self_attn_weights = self.self_attn(
                                       ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 801, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/utils/fx.py", line 1193, in call_module
    return super().call_module(m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 519, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 794, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/models/qwen2/modeling_qwen2.py", line 159, in forward
    hidden_shape = (*input_shape, -1, self.head_dim)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/proxy.py", line 461, in __iter__
    return self.tracer.iter(self)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/utils/fx.py", line 901, in iter
    return super().iter(prxy)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/proxy.py", line 330, in iter
    raise TraceError('Proxy object cannot be iterated. This can be '
torch.fx.proxy.TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors
```
Blame on that line shows that it was due to this refactor commit on a ton of models: https://github.com/huggingface/transformers/commit/2c47618c1a282f925446506d53108dc6e82d9ef0 that introduced `*splat` in a bunch of places. This breaks a lot of tracing, seems related to https://github.com/huggingface/transformers/issues/35622



### Expected behavior

We should be supporting FX tracing. I came up with a shoddy fix since it seems like `(*a_tuple, ...)` is used in a lot of places. This corresponds to a `LIST_EXTEND` operation. We can implement `HFTracer.iter` to explicitly index the tuple instead of iterating over it (lifted pretty much from `torch.fx.Proxy.__iter__` implementation`:

```py
def iter(self, proxy):
    import dis
    frame = inspect.currentframe()
    assert frame is not None
    calling_frame = frame.f_back.f_back # we need to go back two frames, because it goes through Proxy.__iter__ first
    assert calling_frame is not None
    inst_list = list(dis.get_instructions(calling_frame.f_code))
    if sys.version_info >= (3, 11):
        from bisect import bisect_left
        inst_idx = bisect_left(inst_list, calling_frame.f_lasti, key=lambda x: x.offset)
    else:
        inst_idx = calling_frame.f_lasti // 2
    inst = inst_list[inst_idx]
    if inst.opname == 'LIST_EXTEND':
        return (proxy[i] for i in range(len(proxy)))

    return super().iter(proxy)
```

This seems to work for the few models I've spot checked. I'm not familiar with Python bytecode, so I'm not sure if `LIST_EXTEND` has other usecases.

Comments:
cc @arthurzucker @cyrilvallez because this change appeared in #35235 

I think torch fx is kinda deprecated no? 

I guess a lot of it has been replaced by `torch.compile` - should we be explicit that we're deprecating support for it?

We can, torch has being explicit about it as well AFAIK!

I totally missed this, but makes sense. I'll go ahead and close this out.

Where does it say that `torch.fx` is deprecated? Nothing is indicated here https://pytorch.org/docs/2.6/fx.html#

@ArthurZucker @Rocketknight1 Also interested in this. It seems like `transformers.utils.fx` has been broken for the past few releases for some models which are supported (Llama being one of them).

I want to replace torch.add with torch.mul in huggingface model like tutorial ：https://huggingface.co/docs/optimum/torch_fx/usage_guides/optimization  how can i do the transform with torch.compile?