[Longformer] Issues with "is_index_masked" when using single encoder layer
Author: NVukobrat
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.19.0
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.9.7
- Huggingface_hub version: 0.6.0
- PyTorch version (GPU?): 1.11.0 (False)
- 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?

@ydshieh

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Steps to reproduce the behavior:

1. Just run the code attached bellow

### Sample script:
```python
import torch

from transformers import LongformerModel

model = LongformerModel.from_pretrained(
    "allenai/longformer-base-4096", torchscript=True
)

submodel = model.encoder.layer[0]

input_shape = (1, 512, 768)
activations = torch.rand(input_shape)
attention_mask = torch.ones((1, 512), dtype=torch.long)

results = submodel(activations, attention_mask=attention_mask)
```

### Traceback
```
Traceback (most recent call last):
  File "/Users/nvukobrat/miniconda3/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/Users/nvukobrat/miniconda3/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/Users/nvukobrat/.vscode/extensions/ms-python.python-2022.6.2/pythonFiles/lib/python/debugpy/__main__.py", line 45, in <module>
    cli.main()
  File "/Users/nvukobrat/.vscode/extensions/ms-python.python-2022.6.2/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 444, in main
    run()
  File "/Users/nvukobrat/.vscode/extensions/ms-python.python-2022.6.2/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 285, in run_file
    runpy.run_path(target_as_str, run_name=compat.force_str("__main__"))
  File "/Users/nvukobrat/miniconda3/lib/python3.9/runpy.py", line 268, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/Users/nvukobrat/miniconda3/lib/python3.9/runpy.py", line 97, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/Users/nvukobrat/miniconda3/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/Users/nvukobrat/Desktop/Python/pytorch_longformer_huggingface_bug.py", line 15, in <module>
    results = submodel(activations, attention_mask=attention_mask)
  File "/Users/nvukobrat/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/nvukobrat/miniconda3/lib/python3.9/site-packages/transformers/models/longformer/modeling_longformer.py", line 1209, in forward
    self_attn_outputs = self.attention(
  File "/Users/nvukobrat/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/nvukobrat/miniconda3/lib/python3.9/site-packages/transformers/models/longformer/modeling_longformer.py", line 1145, in forward
    self_outputs = self.self(
  File "/Users/nvukobrat/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/nvukobrat/miniconda3/lib/python3.9/site-packages/transformers/models/longformer/modeling_longformer.py", line 642, in forward
    attn_probs = torch.masked_fill(attn_probs, is_index_masked[:, :, None, None], 0.0)
TypeError: 'NoneType' object is not subscriptable
```

### Proposed solution
The issue here is that the `is_index_masked` isn't populated when a single encoder layer is extracted (problem occurs in the Longformer self-attention layer). The proposed solution could be to check and populate `is_index_masked` dynamically. 

File: `transformers/models/longformer/modeling_longformer.py`
Class : `LongformerSelfAttention`
Function: `forward`
Code:
```python
# rest of the code...

if layer_head_mask is not None:
    assert layer_head_mask.size() == (
        self.num_heads,
    ), f"Head mask for a single layer should be of size {(self.num_heads,)}, but is {layer_head_mask.size()}"
    attn_probs = layer_head_mask.view(1, 1, -1, 1) * attn_probs

# softmax sometimes inserts NaN if all positions are masked, replace them with 0

# Proposed fix
if is_index_masked is None:
    is_index_masked = attention_mask < 0

attn_probs = torch.masked_fill(attn_probs, is_index_masked[:, :, None, None], 0.0)
attn_probs = attn_probs.type_as(attn_scores)

# rest of the code...
```

### Expected behavior

```shell
I would expect to be able to get single encoder layer outputs once I run the provided script. 

Let me know if this fix is valid. If yes, I can open a pull request if needed.
```


Comments:
Hey @ydshieh, do you have any news/comments on this issue?

Hi @NVukobrat 

After looking more closely, here is my thought:

What you suggests could be achieved by adding 
```
 is_index_masked = attention_mask < 0 
 is_index_global_attn = attention_mask > 0 
 is_global_attn = is_index_global_attn.flatten().any().item() 
```
to `LongformerSelfAttention`. However, the input `attention_mask` is not as simple as `1 or 0` anymore, and the way you prepare it (`attention_mask = torch.ones`) as input to `LongformerLayer` is incorrectly. See the details below.

I have to discuss with the team members about the design, but so far my personal understanding is that we encourage the users to interact with the models at the `Model` level (for example `LongformerModel`) instead of the intermediate `layer`s - to avoid these kinds of incorrect inputs.

(Of course, the code is in open source, and the users could customize it if there is a real necessity - but they should be careful and responsible for the inputs) 

Again, **let me have a discussion with the team members** first and come back to this thread.

### More details here:

- The (base model) `LongformerModel` receives `attention_mask` which is what we are familiar with:
  - It is usually prepared by a tokenizer
  - If not provided, we use `attention_mask = torch.ones(...)`

- However, `attention_mask` will be processed in `LongformerModel`
  -  dealing with global attention:
https://github.com/huggingface/transformers/blob/b9bb417324c0d9013c505dc39c016ab9ca0e23c8/src/transformers/models/longformer/modeling_longformer.py#L1692-L1694
  - padding to window size:
https://github.com/huggingface/transformers/blob/b9bb417324c0d9013c505dc39c016ab9ca0e23c8/src/transformers/models/longformer/modeling_longformer.py#L1696-L1703
  - change to additive attention mask (changing shape and using `-10000.0`):
https://github.com/huggingface/transformers/blob/b9bb417324c0d9013c505dc39c016ab9ca0e23c8/src/transformers/models/longformer/modeling_longformer.py#L1705-L1709

- In `LongformerEncoder`, the following are computed from the processed `attention_mask`
https://github.com/huggingface/transformers/blob/b9bb417324c0d9013c505dc39c016ab9ca0e23c8/src/transformers/models/longformer/modeling_longformer.py#L1263-L1265

- The inputs to `LongformerLayer`, `LongformerAttention` and `LongformerSelfAttention` are the processed version shown above

Hey @NVukobrat,

Couldn't you just prepare the following args:

`is_index_masked=None`
`is_index_global_attn=None`
`is_global_attn=None`

before you pass your inputs to the LongformerSelfAttentionLayer. 

Note that `LongformerSelfAttentionLayer` is a non-public class which is subject to breaking changes,
so we don't recommend directly importing it. If you do however, I think it also shouldn't be too difficult to create the necessary inputs before calling it no? 

Hey @ydshieh @patrickvonplaten, thanks a lot for providing the details! Very informative and helpful for our use case!

Selecting and setting mentioned attributes (`attention_mask`, `layer_head_mask`, `is_index_masked`, `is_index_global_attn`, and `is_global_attn`) before passing activations to the `LongformerSelfAttentionLayer` works for us. 

Thanks once again for your help!