[`Deberta/Deberta-v2`] Refactor code base to support compile, export, and fix LLM
Author: ArthurZucker
Pull request: yes
Labels: name: WIP

# What does this PR do?
Refactor both Deberta and DebertaV2 to make them more compatible with the overall transformers library

Should fix a bunch of issues related to torch-scripting with Deberta:
fixes #15216
fixes #15673
fixes #16456
fixes #18659
fixes #21300 
fixes #20815
fixes #12436 
fixes #18674 


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_22105). All of your documentation changes will be reflected on that endpoint.

Hey @ArthurZucker any updates on this? ETA for when it will be merged into main?

Hey! Just got back from holidays, this should be my main focus in the coming days! 

Sorry! Seem like I had to postpone this! If anyone want to take over feel free to do it, otherwise will be my priority once https://github.com/huggingface/transformers/pull/23909 is merge!

Regarding the `z_steps` in `DebertaV2Model`: I think this code is relevant for the [enhanced mask decoder of the generator model](https://github.com/microsoft/DeBERTa/blob/master/DeBERTa/apps/models/masked_language_model.py#L51)

```python
if attention_mask.dim() <= 2:
    extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
    att_mask = extended_attention_mask.byte()
    attention_mask = att_mask * att_mask.squeeze(-2).unsqueeze(-1)
elif attention_mask.dim() == 3:
    attention_mask = attention_mask.unsqueeze(1)
target_mask = target_ids > 0
hidden_states = encoder_layers[-2]
if not self.position_biased_input:
    layers = [encoder.layer[-1] for _ in range(2)]
    z_states += hidden_states
    query_states = z_states
    query_mask = attention_mask
    outputs = []
    rel_embeddings = encoder.get_rel_embedding()

    for layer in layers:
        # TODO: pass relative pos ids
        output = layer(hidden_states, query_mask, return_att=False, query_states=query_states,
                       relative_pos=relative_pos, rel_embeddings=rel_embeddings)
        query_states = output
        outputs.append(query_states)
else:
    outputs = [encoder_layers[-1]]
```

As far as I can tell, they hardcoded z_steps to 2 here. Although it should still be left as 0 for the discriminator. Adding the z_steps to the config seems like a good idea. 

`z_states` represents [the position embeddings](https://github.com/microsoft/DeBERTa/blob/master/DeBERTa/apps/models/masked_language_model.py#L111), which are non-zero if `position_biased_input` is set to `True`. They are passed from the [embedding layer](https://github.com/microsoft/DeBERTa/blob/master/DeBERTa/deberta/bert.py#L269). So in order to properly implement this, I think we need to return the position embeddings here:

```python
class DebertaV2Embeddings(nn.Module):
    def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=None, inputs_embeds=None):
        ...

        return embeddings, position_embeddings
```

and implement the `z_steps` like this:

```python
class DebertaV2Model(DebertaV2PreTrainedModel):
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        ...

        embedding_output, position_embedding_output = self.embeddings(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            mask=attention_mask,
            inputs_embeds=inputs_embeds,
        )
        ...

        if self.z_steps > 0:
            hidden_states = encoded_layers[-2]
            layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
            position_embedding_output += hidden_states
            query_states = position_embedding_output
            query_mask = self.encoder.get_attention_mask(attention_mask)
            rel_embeddings = self.encoder.get_rel_embedding()
            rel_pos = self.encoder.get_rel_pos(embedding_output)
            for layer in layers:
                query_states = layer(
                    hidden_states,
                    query_mask,
                    output_attentions=False,
                    query_states=query_states,
                    relative_pos=rel_pos,
                    rel_embeddings=rel_embeddings,
                )
                encoded_layers = encoded_layers + (query_states,)
```

What is the status?
The logs of the checks are expired.

#27734 should help with some of the issues in the mean time

Any update on this? We're blocked by some of the issues that this is meant to fix. Sepcifically:
* https://github.com/huggingface/transformers/issues/22790
* https://github.com/huggingface/transformers/issues/27586
* https://github.com/huggingface/transformers/issues/20815
* https://github.com/huggingface/transformers/pull/27734

Hi there, just checking in @ArthurZucker on whether there's any progress here please? 

Hey hey! Sorry I ended up dropping this, let me get back to you next week!

I reviewed #27734, will take it over this weekend if possible. 

Anyone up for the task can try to do it as well 🤗 

Hey @ArthurZucker any update on the status of this and related issues. Curious if this will also cover DeBERTa v3 now as well.

Hey @ArthurZucker !
Thank you so much for your work. I'm reviewing your changes as it does impact my code on some project I have. I saw that you swap `StableDropout` with `Dropout` in various places. I'm not entirely sure of the impact but it seems to diverge from microsoft implementation. Did you face an issue with `StableDropout` and tracing ? Or is it something else? I'm not super familiar with this part of the code, but if I'll try to help you if I can.

Thanks again

Yep, this decision is based on the fact that `StableDropout` is supposed to be just dropout but more efficient. 
This function is hard to trace / compile and is 4 years old. 
The idea is that if we remove this and XSoftmax, performance gains are so much better and code so much simpler, but equivalent, that is is worth it!

LGTM! Every test we have on our end work with this PR. Thanks @ArthurZucker 

cc @michaelbenayoun I am not sure I know how to avoid the HFProxy issue, do you have pointers? 

Had to skip some pt-tf equivalence tests. The slow tests ran for me and are passing. 
If anyone has a problem will be quick to fix!

Nice one, thanks @ArthurZucker !