Add LMHeadModel similar to BertLMHeadModel to modeling_distilbert.py
Author: KMFODA
Pull request: yes

# What does this PR do?

Similar to the `BertLMHeadModel` this PR aims to add a `DistilBertForCausalLM` model in `modeling_distilbert.py`.

Fixes #7397 

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to the it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors which may be interested in your PR.

Comments:
I don't know if this PR is in its final state and ready for review now, but here are a few pointers:

The PR looks good, but there are a few things to do before merging.

- You should take a look at the failing tests and fix them if they're related to your PR (they should be)
- You should implement the tests related to the added classes
- You should add these classes to the auto models and init.py

Hi @LysandreJik. Thanks for the comments. Apologies though, this PR is still a WIP so no need for review atm. I plan on adding my classes to auto models and init.py and running the appropriate tests before committing my final version.

Great, thanks for letting me know. Will wait for your ping before reviewing again!

HI @LysandreJik & @patrickvonplaten, I'm hoping to get your help with the only remaining test that's failing for this PR. The `tests/test_modeling_distilbert.py` script fails due to an error in the `tests/test_modeling_common.py` in line 966. I've detailed the error message below but in short what I think is happening is that check_equivalence fails on the following check:

```
check_equivalence(
    model, tuple_inputs, dict_inputs, {"output_hidden_states": True, "output_attentions": True}
)
```

What's strange is that all the other equivalence checks pass and I haven't changed the hidden states or attention layers.

Error Message:

```
tests/test_modeling_common.py:1037: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/test_modeling_common.py:1001: in check_equivalence
    recursive_check(tuple_output, dict_output)
tests/test_modeling_common.py:990: in recursive_check
    recursive_check(tuple_iterable_value, dict_iterable_value)
tests/test_modeling_common.py:990: in recursive_check
    recursive_check(tuple_iterable_value, dict_iterable_value)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

tuple_object = tensor([[[[0.4205, 0.0000, 0.2311,  ..., 0.0000, 0.0000, 0.3484],
          [0.4074, 0.0000, 0.2666,  ..., 0.0000, 0.0...71, 0.0000, 0.1483,  ..., 0.1858, 0.1648, 0.1354],
          [0.1422, 0.0000, 0.2250,  ..., 0.1523, 0.1632, 0.1669]]]])
dict_object = tensor([[[ 1.3640e+00,  2.6618e-01,  5.9669e-01,  ..., -7.1322e-01,
          -3.5144e-02, -1.3287e+00],
         [ 1....1.0870e-01],
         [-9.9694e-01, -7.1065e-01, -5.7949e-01,  ...,  1.4234e+00,
          -2.1929e+00,  1.2390e-01]]])

    def recursive_check(tuple_object, dict_object):
        if isinstance(tuple_object, (List, Tuple)):
            for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object):
                recursive_check(tuple_iterable_value, dict_iterable_value)
        elif tuple_object is None:
            return
        else:
            self.assertTrue(
                torch.allclose(
>                   set_nan_tensor_to_zero(tuple_object), set_nan_tensor_to_zero(dict_object), atol=1e-5
                ),
                msg=f"Tuple and dict output are not equal. Difference: {torch.max(torch.abs(tuple_object - dict_object))}. Tuple has `nan`: {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}.",
            )
E           RuntimeError: The size of tensor a (7) must match the size of tensor b (768) at non-singleton dimension 3

tests/test_modeling_common.py:996: RuntimeError
```

Hey @KMFODA, 

Thanks for all the hard work already! We'll definitely help you get this merged :-) @patil-suraj, could you maybe help here since you recently added the caching mechanism for `BertForCausalLM`? Otherwise I'll try to take a look by the end of this week!

Thanks @patrickvonplaten. I'm working through Lysandre's comment which was very helpful and hoping this will fix the remaining error. Will await any further comments / advice on how to merge this from you / suraj as well in case I've missed something else out.

Hi @KMFODA 

Thanks for all the work! 

`DistillBertForCausalLM` can be used as a decoder in `EncoderDecoder` models. To use it as a decoder we should now implement `cross attention`  (or `encoder-decoder` attention) so that it can attend to `encoder_hidden_states` and caching the `past_key_values` to speed-up inference.

We should implement this feature analogous to how it is implemented in `Bert`. This means that we should

- Add cross-attention and caching mechanism `MultiHeadSelfAttention` as shown here in Bert
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/src/transformers/models/bert/modeling_bert.py#L259

- this means we need to pass `encoder_hidden_states`, `encoder_attention_mask`, `past_key_values` through the layers. So introduce these parameters in `forward` of `DistilBertForCausalLM`, `DistilBertModel`, `Transformer`, `TransformerBlock`  and `MultiHeadSelfAttention`.

- compute cross attention in `TransformerBlock` as shown here in  `BertLayer`
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/src/transformers/models/bert/modeling_bert.py#L478-L480

- Adjust the position embedding and attention masks accordingly as done in `BertEmbeddings` and `BertModel`
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/src/transformers/models/bert/modeling_bert.py#L194
and
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/src/transformers/models/bert/modeling_bert.py#L194

The important thing to note here is that, when `past_key_values` are enabled, the actual seq length is `cur_seq_length + pas_key_values_length`

- Add tests for `DistillBertForCausalLM` that checks that the decoder and caching mechanism works as expected:
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/tests/test_modeling_bert.py#L170
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/tests/test_modeling_bert.py#L203
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/tests/test_modeling_bert.py#L230
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/tests/test_modeling_bert.py#L263


All of this should be pretty similar to how it's implemented in Bert, so you could use that as a reference. Let me know if something is not clear. Happy to help you here :)



Apolgoies, this PR got too messy when I tried to rebase to master and it closed automatically when I tried to revert back to an old commit. I'm working on @patil-suraj's very useful recommendations and will post a new PR when it passes all tests. Thanks!

Hey @KMFODA 

Glad to know you're still working on this :) Don't hesitate to ping me if you're stuck or need some help.