(LLama-2) (4.36.0) TensorParallelPreTrainedModel does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention
Author: VasilGeorgiev39
Pull request: no

### System Info

- `transformers` version: 4.36.0
- Platform: Linux-5.4.0-150-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.1 (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?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker 
@fxmarty 

### Information

- [ ] The official example scripts
- [X] 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
import transformers
import tensor_parallel as tp

tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf")
model = transformers.AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf")

modelp = tp.tensor_parallel(model)
```

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/root/code/cot-unfaithfulness/test-hf.py in line 8
      [5](file:///root/code/cot-unfaithfulness/test-hf.py?line=4) tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf")
      [6](file:///root/code/cot-unfaithfulness/test-hf.py?line=5) model = transformers.AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf")
----> [8](file:///root/code/cot-unfaithfulness/test-hf.py?line=7) modelp = tp.tensor_parallel(model)

File /opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py:61, in tensor_parallel(module, device_ids, tensor_parallel_config, distributed, sharded, sharded_param_names, **kwargs)
     [59](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=58) else:
     [60](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=59)     if isinstance(module, PreTrainedModel):
---> [61](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=60)         return TensorParallelPreTrainedModel(
     [62](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=61)             module,
     [63](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=62)             device_ids=device_ids,
     [64](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=63)             tensor_parallel_config=tensor_parallel_config,
     [65](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=64)             distributed=distributed,
     [66](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=65)             sharded=sharded,
     [67](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=66)             sharded_param_names=sharded_param_names,
     [68](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=67)             **kwargs,
     [69](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=68)         )
     [70](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=69)     else:
     [71](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=70)         return TensorParallel(
     [72](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=71)             module,
     [73](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=72)             device_ids=device_ids,
   (...)
     [78](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=77)             **kwargs,
     [79](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/factory.py?line=78)         )

File /opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py:47, in TensorParallelPreTrainedModel.__init__(self, module, device_ids, output_device, output_device_index, tensor_parallel_config, **kwargs)
     [38](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=37) def __init__(
     [39](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=38)     self,
     [40](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=39)     module: PreTrainedModel,
   (...)
     [45](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=44)     **kwargs,
     [46](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=45) ):
---> [47](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=46)     super().__init__(module.config)  # Temporary empty config. Gets replaced in from_pretrained
     [49](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=48)     if hasattr(module, "_hf_hook"):
     [50](file:///opt/conda/lib/python3.10/site-packages/tensor_parallel/pretrained_model.py?line=49)         from accelerate.hooks import remove_hook_from_module

File /opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py:1190, in PreTrainedModel.__init__(self, config, *inputs, **kwargs)
   [1184](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1183)     raise ValueError(
   [1185](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1184)         f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class "
   [1186](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1185)         "`PretrainedConfig`. To create a model from a pretrained model use "
   [1187](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1186)         f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`"
   [1188](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1187)     )
   [1189](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1188) # Save config and origin of the pretrained weights if given in model
-> [1190](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1189) config = self._autoset_attn_implementation(
   [1191](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1190)     config, torch_dtype=torch.get_default_dtype(), check_device_map=False
   [1192](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1191) )
   [1193](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1192) self.config = config
   [1195](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1194) self.name_or_path = config.name_or_path

File /opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py:1311, in PreTrainedModel._autoset_attn_implementation(cls, config, use_flash_attention_2, torch_dtype, device_map, check_device_map)
   [1302](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1301)     cls._check_and_enable_flash_attn_2(
   [1303](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1302)         config,
   [1304](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1303)         torch_dtype=torch_dtype,
   (...)
   [1307](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1306)         check_device_map=check_device_map,
   [1308](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1307)     )
   [1309](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1308) elif requested_attn_implementation in [None, "sdpa"]:
   [1310](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1309)     # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif.
-> [1311](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1310)     config = cls._check_and_enable_sdpa(
   [1312](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1311)         config, hard_check_only=False if requested_attn_implementation is None else True
   [1313](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1312)     )
   [1314](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1313) else:
   [1315](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1314)     config._attn_implementation = "eager"

File /opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py:1464, in PreTrainedModel._check_and_enable_sdpa(cls, config, hard_check_only)
   [1462](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1461) if hard_check_only:
   [1463](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1462)     if not cls._supports_sdpa:
-> [1464](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1463)         raise ValueError(
   [1465](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1464)             f"{cls.__name__} does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet. Please open an issue on GitHub to "
   [1466](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1465)             "request support for this architecture: https://github.com/huggingface/transformers/issues/new"
   [1467](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1466)         )
   [1468](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1467)     if not is_torch_sdpa_available():
   [1469](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1468)         raise ImportError(
   [1470](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1469)             "PyTorch SDPA requirements in Transformers are not met. Please install torch>=2.1.1."
   [1471](file:///opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py?line=1470)         )

ValueError: TensorParallelPreTrainedModel does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet. Please open an issue on GitHub to request support for this architecture: https://github.com/huggingface/transformers/issues/new
```

### Expected behavior

Wrap the model using the tensor_parallel library https://github.com/BlackSamorez/tensor_parallel

This succeeds with transformers==4.35.2

The exception seems to be raised from this line added in 4.36 https://github.com/huggingface/transformers/commit/80377eb018c077dba434bc8e7912bcaed3a64d09#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaR1435

Comments:
Thank you @VasilGeorgiev39, a quick fix is to load your transformers model with

```python
import transformers
import tensor_parallel as tp

tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf")
model = transformers.AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf", attn_implementation="eager")

modelp = tp.tensor_parallel(model)
```

It looks like `tp.tensor_parallel` is overloading Transformers classes and does not have `_supports_sdpa = True`, which results in this error. I'll see what we can do.

For reference: https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-and-memory-efficient-attention-through-pytorchs-scaleddotproductattention

I am having the same issue with `Llama-2-70b-chat-hf model`. Tried to add `attn_implementation="eager`, but got the following new "no attribute 'to_legacy_cache'" errors:

```
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl                                                                                                                                 [1/1003]
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensor_parallel/pretrained_model.py", line 76, in forward
    return self.wrapped_model(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensor_parallel/tensor_parallel.py", line 159, in forward
    return parallel_apply(self.module_shards, inputs, kwargs_tup, self.devices)[self.output_device_index]
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/parallel_apply.py", line 110, in parallel_apply
    output.reraise()
  File "/usr/local/lib/python3.8/dist-packages/torch/_utils.py", line 694, in reraise
    raise exception
AttributeError: Caught AttributeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/parallel_apply.py", line 85, in _worker
    output = module(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/llama/modeling_llama.py", line 1174, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/llama/modeling_llama.py", line 1086, in forward
    next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
AttributeError: 'tuple' object has no attribute 'to_legacy_cache'
```


cc @gante as well! 

@taozhang9527 since `tensor_parallel` overloads the models, it is possible their classes are not compatible with the models that got the new cache format in `transformers` v4.36. I'd open an issue in `tensor_parallel` :)

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 everyone, has the problem been solved? I encountered the same problem as @taozhang9527 did. I am using the newest Llama-3.1-8B model, so downgrading the `transformers` library would cause imcompatibilities. Is there a workaround to allow me to use `tensor_parallel==2.0.0` with `transformers==4.43.3` ? 

If sdpa is the use for now I would recommend:
- to use flash attention
- to use eager mode
Also not sure if this is a transformers issue or a tensor parallel issue?

@ArthurZucker Thank you for your reply. I also opened an issue in `tensor_parallel` codebase.

I tried and found that the workarounds for now could be commetting out [these two lines](https://github.com/huggingface/transformers/blob/2e3f8f74747deeeead6cf1f0c12cf01bd7169b82/src/transformers/models/llama/modeling_llama.py#L1021C4-L1022C5) when using tensor_parallel. I am not very sure about it, but it stops throwing errors after this change is applied.

that is deprecated thankfully! https://github.com/huggingface/transformers/blob/241a9a037af10cdf6564dc8897a8f39c63237d36/src/transformers/models/llama/modeling_llama.py#L947-L955 will be removed

Yes, cool!