Add Multi Resolution Analysis (MRA) (New PR)
# Add Multi Resolution Analysis (MRA) for Approximate Self-Attention

This PR adds the MRA model to the repository. 

Paper: [https://arxiv.org/pdf/2207.10284.pdf](https://arxiv.org/pdf/2207.10284.pdf)
Code: [https://github.com/mlpen/mra-attention](https://github.com/mlpen/mra-attention)

To-do:

- [x] Improve loading cuda kernels
- [x] Improve formatting and documentation
- [x] Upload checkpoints
Copied all files over from #20573 <|||||>_The documentation is not available anymore as the PR was closed or merged._<|||||>Could you fix the failing tests?<|||||>Hello @sgugger, I've made sure all checks pass and fixed conflicts. <|||||>Hello @amyeroberts, I've addressed your comments and made some code changes. Please take a look at the updated files. <|||||>Hi @amyeroberts, I've addressed the suggestions from the code review. Please take a look at the updated code. <|||||>Thanks for catching these errors @amyeroberts! I've applied both changes. <|||||>@novice03 

It seems the CI get

```bash
(line 403)  ValueError: sequence length must be divisible by the block_size.
```
when `load_cuda_kernels` loads successfully.

It's likely due to `seq_length=8` from `MraModelTester`, but I am not able to set the correct combination of `seq_length`, `block_size`, `num_blocks` to make it works.

Note, our daily CI (with torch 2.0.1 + CUDA 11.8) fails to load `custom CUDA kernels` and the execution goes to 
```python
    if cuda_kernel is None:
        return torch.zeros_like(query).requires_grad_()
```
in `mra2_attention` and tests pass.

However, in our CI with torch 1.13 (and with CUDA 11.6.2), kernel is loaded, but the tests fail.

It would be great if you can help us to find the correct settings where the CI will pass when kernel is loaded.

Thanks in advance 🤗 .<|||||>You can run

```python
python3 -m pytest -v tests/models/mra/test_modeling_mra.py::MraModelTest::test_for_masked_lm
```

The full error log is (if custom cuda kernal is loaded successfully)

```bash
self = <tests.models.mra.test_modeling_mra.MraModelTest testMethod=test_for_masked_lm>

    def test_for_masked_lm(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
>       self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)

tests/models/mra/test_modeling_mra.py:322: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/models/mra/test_modeling_mra.py:210: in create_and_check_for_masked_lm
    result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/mra/modeling_mra.py:1093: in forward
    outputs = self.mra(
/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/mra/modeling_mra.py:1028: in forward
    encoder_outputs = self.encoder(
/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/mra/modeling_mra.py:782: in forward
    layer_outputs = layer_module(hidden_states, attention_mask)
/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/mra/modeling_mra.py:729: in forward
    self_attention_outputs = self.attention(hidden_states, attention_mask)
/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/mra/modeling_mra.py:681: in forward
    self_outputs = self.self(hidden_states, attention_mask)
/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/mra/modeling_mra.py:615: in forward
    context_layer = mra2_attention(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

query = tensor([[[[ 0.0500, -0.0523, -0.0260,  ...,  0.0000,  0.0000,  0.0000],
          [-0.1339,  0.0844,  0.0287,  ...,  0...       [ 0.0293,  0.1609,  0.0547,  ...,  0.0000,  0.0000,  0.0000]]]],
       device='cuda:0', grad_fn=<CatBackward0>)
key = tensor([[[[ 0.0185, -0.0316,  0.0150,  ...,  0.0000,  0.0000,  0.0000],
          [-0.0575, -0.1123,  0.0832,  ...,  0...       [ 0.0608,  0.0932, -0.0973,  ...,  0.0000,  0.0000,  0.0000]]]],
       device='cuda:0', grad_fn=<CatBackward0>)
value = tensor([[[[ 0.0131,  0.1242,  0.0672,  ...,  0.0000,  0.0000,  0.0000],
          [-0.0212,  0.0600,  0.0269,  ...,  0...       [-0.1005, -0.0048,  0.0561,  ...,  0.0000,  0.0000,  0.0000]]]],
       device='cuda:0', grad_fn=<CatBackward0>)
mask = tensor([[-2.1475e+09,  1.0000e+00,  1.0000e+00, -2.1475e+09,  1.0000e+00,
         -2.1475e+09, -2.1475e+09,  1.0000e+...  1.0000e+00,  1.0000e+00, -2.1475e+09,  1.0000e+00,
         -2.1475e+09, -2.1475e+09,  1.0000e+00]], device='cuda:0')
num_blocks = 64, approx_mode = 'full', block_size = 32, initial_prior_first_n_blocks = 0, initial_prior_diagonal_n_blocks = 0

    def mra2_attention(
        query,
        key,
        value,
        mask,
        num_blocks,
        approx_mode,
        block_size=32,
        initial_prior_first_n_blocks=0,
        initial_prior_diagonal_n_blocks=0,
    ):
        """
        Use Mra to approximate self-attention.
        """
        if cuda_kernel is None:
            return torch.zeros_like(query).requires_grad_()
    
        batch_size, num_head, seq_len, head_dim = query.size()
        meta_batch = batch_size * num_head
    
        if seq_len % block_size != 0:
>           raise ValueError("sequence length must be divisible by the block_size.")
E           ValueError: sequence length must be divisible by the block_size.

src/transformers/models/mra/modeling_mra.py:403: ValueError
```<|||||>Hello @ydshieh, thanks for bringing this up. We will likely have to use larger values for seq_len and hidden_size. Can you please try with the values [here](https://github.com/novice03/transformers/blob/1612188d6b6d094c81cc34a77641936687b8f7b3/tests/models/mra/test_modeling_mra.py)?