ProphetNet inconsistent with changing batch ordering
Author: mikkelfo
Pull request: no
Labels: name: Good Second Issue

name: Good Difficult Issue

name: bug

### System Info

```shell
- `transformers` version: 4.15.0
- Platform: Windows-10-10.0.19044-SP0
- Python version: 3.9.7
- PyTorch version (GPU?): 1.10.1+cpu (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?

@patrickvonplaten

### 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


``` python
from transformers import ProphetNetForConditionalGeneration, ProphetNetTokenizer

model = ProphetNetForConditionalGeneration.from_pretrained('microsoft/prophetnet-large-uncased')
tokenizer = ProphetNetTokenizer.from_pretrained('microsoft/prophetnet-large-uncased')

input_string = ['Hello, my dog is cute', 'I have a cat that is not cute']
labels = ['My dog is cute', 'My cat is not cute']

inputs = tokenizer(input_string, return_tensors="pt", padding=True, truncation=True)
targets = tokenizer(labels, return_tensors="pt", padding=True, truncation=True)

# Inverse the ordering of the input and labels using [::-1]
inputs_inv = tokenizer(input_string[::-1], return_tensors="pt", padding=True, truncation=True)
targets_inv = tokenizer(labels[::-1], return_tensors="pt", padding=True, truncation=True)

output = model(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, labels=targets.input_ids)
output_inv = model(input_ids=inputs_inv.input_ids, attention_mask=inputs_inv.attention_mask, labels=targets_inv.input_ids)

print(output.loss.item())
print(output_inv.loss.item())
```
Given two different forward passes, where one of them is the inverse of the other, the two losses are different.

output.loss.item() = 5.023777484893799
output_inv.loss.item() = 6.5036540031433105

When comparing their encoder output (last hidden states), they are again not equal
```Python
enc_output = model.prophetnet.encoder(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask)
enc_output_inv = model.prophetnet.encoder(input_ids=inputs_inv.input_ids, attention_mask=inputs_inv.attention_mask)
# We flip the inverse encoder output to have the same order
print((enc_output.last_hidden_state == enc_output_inv.last_hidden_state.flip(0)).all())
``` 
Also equals `False`


### Expected behavior

```shell
I expect the order of the batch to not have an influence on the model output (aside from the ordering of the output)
```


**Edit** The issue persistent regardless of padding and also causes model.generate to generate two different sets of text. 


Comments:
Hey @mikkelfo,

This indeed seems like a bug! I sadly won't be able to look into this any time soon - I'll mark the issue as a "good second issue" in case someone wants to give it a try :-). Also cc @patil-suraj 

Thanks for the answer @patrickvonplaten. I suspected so, although I did question my sanity for a moment. I might have some time to look into this myself as I'm using it for my thesis.

I've done a little digging and the issue appears to start during the self-attention of the encoder, more specifically after [line 700](https://github.com/huggingface/transformers/blob/v4.19.2/src/transformers/models/prophetnet/modeling_prophetnet.py#L700). It appears to me that the reshaping into the attention heads (with `.view(*proj_shape)`)causes the issue, as I am unable to invert things back to the original form after that point. 

Sorry for confusing you with this. To be honest, we've ported a relatively messy Prophetnet implementation and never fully tested it. If you have some time to dive into it, I'd suggest to add a test that ensures that changing the order of sequences in a batch still gives the same results (that's a great test btw) and then trying to correct this test while making sure that all other tests run correctly.

Happy to help you if you're stuck. Note the original implementation is here:
- https://github.com/microsoft/ProphetNet

The original author is @qiweizhen I think, we could maybe also ask him if you run into problems :-) 

After some more digging, the issue appears to be centered around ProphetNetAttention and ProphetNetNgramSelfAttention (so both the encoder and the decoder suffers from this). The fix essentially boils down to changing the `proj_shape = (batch_size * self.num_attn_heads, -1, self.head_dim)` to `proj_shape = (batch_size, self.num_attn_heads, -1, self.head_dim)` and adapting the rest of the code to fit this. Keeping the dimensions seperate is also utilised in other transformer models (Took inspiration from BERT) and avoids potential mishaps with reshapes. 

For the encoder, this is relatively simple as it only constitues changing the ProphetNetAttention [forward pass](https://github.com/huggingface/transformers/blob/v4.19.2/src/transformers/models/prophetnet/modeling_prophetnet.py#L655) and adapting the [attention mask](https://github.com/huggingface/transformers/blob/v4.19.2/src/transformers/models/prophetnet/modeling_prophetnet.py#L1317) to these new dimensions. I have done this, but it is only a fix for the encoder, so I don't know if that's worthy of a PR @patrickvonplaten? It greatly reduces the inconsistency to a few decimals (for this particular example), from a difference of 1.5 to 0.06. 

I am trying to do the same for the decoder part (ProphetNetNgramSelfAttention), but it is quite a bit more cumbersome. I'll keep trying for a bit, otherwise I'll simply offer my encoder fix for now.  

Hello again. I believe I have managed to fix both the encoder and decoder part, such that the hidden states are consistent. However, the loss computations of `ProphetNetForConditionalGeneration` and `ProphetNetForCasualLM` still differs (with a very slight amount). 

When the log_softmax is taken ([line 2014](https://github.com/huggingface/transformers/blob/6e535425feae20ca61a8b10ae5e8a7fab4d394ba/src/transformers/models/prophetnet/modeling_prophetnet.py#L2014)) the logits are equal when inversed, but the probabilities are not.

```
# logits = (n_gram, batch_size, seq_length, vocab_size)
print((logits == logits_inv.flip(1)).all()) # True (Flip on batch_size)

lprobs = torch.nn.functional.log_softmax(logits.view(-1, logits.size(-1)),dim=-1,dtype=torch.float32,)
lprobs_inv = torch.nn.functional.log_softmax(logits_inv.view(-1, logits.size(-1)),dim=-1,dtype=torch.float32,)

print(lprobs - lprobs_inv) # -4.7684e-07
```
I believe the .view part is cause the issue due to the collapsing of dimensions. As for the case of the other fixes, merging dimensions are most likely the cause, so seperating from (n_gram\*batch_size\*tokens) to (batch_size, n_gram\*tokens) should fix the issue. It requires a bit more than changing `logits.view(-1, logits.size(-1))` to `logits.view(2, -1, logits.size(-1))`. I'm currently stuck on this, but I'm wondering if any of you can be of help? @patrickvonplaten @patil-suraj @qiweizhen

It is only these 3 lines that are left to fix. How can this be computed while keeping the batch dimension intact?
```
# [batch_size, n_gram, tokens, vocab_size] -> # [n_gram, batch_size, tokens, vocab_size]
logits = logits.transpose(0, 1).contiguous() 
lprobs = nn.functional.log_softmax(
    logits.view(-1, logits.size(-1)),
    dim=-1,
    dtype=torch.float32,
)

loss = nn.functional.nll_loss(lprobs, expend_targets.view(-1), reduction="mean")
```

the difference of `-4.7684e-07` is very small so IMO it can be ignored. `1e-4` is the threshold that we use in `Transformers`. So if the difference is less than that it should be fine I think. wdyt @patrickvonplaten 

I checked a bit further (with my implemented fixes). 

Given two model outputs, where one is a batch and the other is single element (i.e. the first element in input_string), we expect that the first element of the batched output is equal to the single element output. Below is a code snippet of how it would look. 

```
outputs = model(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, labels=labels,  output_hidden_states=True)

batch_size = inputs.input_ids.shape[0]
for i in range(batch_size):
    single_outputs = model(input_ids=inputs.input_ids[i:i+1], attention_mask=inputs.attention_mask[i:i+1], labels=labels[i:i+1], output_hidden_states=True)
    print((single_outputs.encoder_last_hidden_state - outputs.encoder_last_hidden_state[i:i+1]).abs().mean())
    print((single_outputs.decoder_hidden_states[-1] - outputs.decoder_hidden_states[-1][i:i+1]).abs().mean())
```
We expect the differences between the two outputs to be 0, but there is a difference between `1e-6` to `1e-7` for each element in the hidden_states (was `1e` to `1e-1` previously) when using some examples from the CNN/DailyMail dataset. My previous example had equal hidden states, but a slight variation in loss, whereas this example does not have equal hidden states. As @patil-suraj mentioned, it is within the threshold used. 

I don't have time to look into this anymore, but I am happy to create a pull request for now. While I dont believe the problem is fully fixed, my fixes makes the differences significantly less. If this is enough for a pull request, let me know and I'll create one @patrickvonplaten.


Hi @patil-suraj , @patrickvonplaten.
If it's ok I would like to help and try to solve this issue.

If you want @kiansierra, here's a [link](https://github.com/mikkelfo/multi-document-abstractive-summarization/blob/main/src/models/prophetnet_fixes.py#L9) for my ad-hoc fixes I did. 

I basically copy-pasted the source code and implemented the fixes, so you must figure out yourself where they are if you want. The function [prophetnet_fixes](https://github.com/mikkelfo/multi-document-abstractive-summarization/blob/main/src/models/prophetnet_fixes.py#L9) updates the affected functions for both the encoder and the decoder.