torch._subclasses.fake_tensor.DataDependentOutputException: aten._local_scalar_dense.default with `_prepare_4d_attention_mask_for_sdpa(
Author: amodab01
Pull request: no


> Hello @fxmarty 
> 
> When I try using torch.compile by using `_attn_implementation="sdpa"` in `BertConfig`, I get the error coming from `_prepare_4d_attention_mask_for_sdpa()` whichis because of the data dependent flow. 
> 
> Specifically, 
> 
> 
> 
> ```
> 
>  File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/transformers/models/bert/modeling_bert.py", line 1108, in forward
> 
>     extended_attention_mask = _prepare_4d_attention_mask_for_sdpa(
> 
>                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/transformers/modeling_attn_mask_utils.py", line 448, in _prepare_4d_attention_mask_for_sdpa
> 
>     if not is_tracing and torch.all(mask == 1):
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/functional_tensor.py", line 411, in __torch_dispatch__
> 
>     outs_unwrapped = func._op_dk(
> 
>                      ^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/utils/_stats.py", line 20, in wrapper
> 
>     return fn(*args, **kwargs)
> 
>            ^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 896, in __torch_dispatch__
> 
>     return self.dispatch(func, types, args, kwargs)
> 
>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1241, in dispatch
> 
>     return self._cached_dispatch_impl(func, types, args, kwargs)
> 
>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 974, in _cached_dispatch_impl
> 
>     output = self._dispatch_impl(func, types, args, kwargs)
> 
>              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1431, in _dispatch_impl
> 
>     op_impl_out = op_impl(self, func, *args, **kwargs)
> 
>                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/fake_impls.py", line 150, in dispatch_to_op_implementations_dict
> 
>     return op_implementations_dict[func](fake_mode, func, *args, **kwargs)
> 
>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> 
>   File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/_subclasses/fake_impls.py", line 284, in local_scalar_dense
> 
>     raise DataDependentOutputException(func)
> 
> torch._subclasses.fake_tensor.DataDependentOutputException: aten._local_scalar_dense.default
> 
> ```
> 
> Is this related to https://github.com/pytorch/pytorch/pull/120400, and do you anticipate there's any solution to this? Ofcourse turning SDPA off works
> 
>  

 _Originally posted by @amodab01 in [221aaec](https://github.com/huggingface/transformers/commit/221aaec6ecf7558e4956dadd662d7d3adb22e420#r152370315)_

Comments:
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.