LED for Seq2Seq output shape mismatch between tensorflow and pytorch
Author: federicoruggeri
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.4.2
- Platform: Linux-5.4.0-48-generic-x86_64-with-Ubuntu-20.04-focal
- Python version: 3.6.13
- PyTorch version (GPU?): 1.8.1+cu102 (True)
- Tensorflow version (GPU?): 2.3.0 (True)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
@patrickvonplaten
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.



Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): 'allenai/led-base-16384' via AutoModelForSeq2SeqLM

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1.  Run the tensorflow version of a simple test script. 
```python
from transformers import TFAutoModelForSeq2SeqLM

preloaded_name = 'allenai/led-base-16384'
led = TFAutoModelForSeq2SeqLM.from_pretrained(preloaded_name)

"""
In this example, we have the following shapes:
input_length --> 2234
output_length --> 70
"""
inputs = {...}

print('Inputs...')
for key, value in inputs.items():
    print('Key: {0} - Value: {1}'.format(key, value.shape))

"""
Prints:
input_ids - Value: (1, 2234)
attention_mask - Value: (1, 2234)
global_attention_mask - Value: (1, 2234)
labels - Value: (1, 70)
"""

led_output = led(
    input_ids=inputs['input_ids'],
    attention_mask=inputs['attention_mask'],
    labels=inputs['labels'],
    global_attention_mask=inputs['global_attention_mask'] if 'global_attention_mask' in inputs else None,
    training=False, use_cache=False, return_dict=True, output_hidden_states=True)

print('Outputs...')
for key, value in led_output.items():
    if type(value) != tuple:
        print('Key: {0} - Value: {1}'.format(key, value.shape))
    else:
        print('Key: {0} - Length: {1} - value shapes: {2}'.format(key, len(value), [item.shape for item in value]))

"""
Prints:
loss - Value: (17,)
logits - Value: (1, 70, 50265)
encoder_last_hidden_state - Value: (1, 70, 768)
encoder_hidden_states - Length 7 - value shapes : [TensorShape([1, 2234, 768]), ....]
decoder_hidden_states - Length: 7 - value shapes: [TensorShape([1, 70, 768]), ....]
```

2. Run the torch version of the same script.
```python
from transformers import AutoModelForSeq2SeqLM
import torch

preloaded_name = 'allenai/led-base-16384'
led = AutoModelForSeq2SeqLM.from_pretrained(preloaded_name)

"""
#NOTE: same inputs as in the tensorflow example!

In this example, we have the following shapes:
input_length --> 2234
output_length --> 70
"""
inputs = {...}

print('Inputs...')
for key, value in inputs.items():
    print('Key: {0} - Value: {1}'.format(key, value.shape))

"""
Prints:
input_ids - Value: (1, 2234)
attention_mask - Value: (1, 2234)
global_attention_mask - Value: (1, 2234)
labels - Value: (1, 70)
"""

led_output = led(
    input_ids=inputs['input_ids'],
    attention_mask=inputs['attention_mask'],
    labels=inputs['labels'],
    global_attention_mask=inputs['global_attention_mask'] if 'global_attention_mask' in inputs else None,
    use_cache=False, return_dict=True, output_hidden_states=True)

print('Outputs...')
for key, value in led_output.items():
    if type(value) != tuple:
        print('Key: {0} - Value: {1}'.format(key, value.shape))
    else:
        print('Key: {0} - Length: {1} - value shapes: {2}'.format(key, len(value), [item.shape for item in value]))

"""
Prints:
loss - Value: torch.Size([])
logits - Value: torch.Size([1, 70, 50265])
encoder_last_hidden_state - Value: torch.Size([1, 2234, 768])
encoder_hidden_states - Length 7 - value shapes : [torch.Size([1, 3072, 768]), ....]
decoder_hidden_states - Length: 7 - value shapes: [torch.Size([1, 70, 768]), ....]
"""
```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

If we compare both LED outputs (tensorflow and pytorch), there are some evident shape mismatches.

1. loss: loss seems to be already averaged (not a big problem)
2. encoder_last_hidden_state: it seems to be completely wrong in tensorflow version (it outputs a sequence of 70 instead of 2234 values)
3. encoder_hidden_states: all encoder hidden states from torch model have different sequence length compared to the ones given by the tensorflow model.

We would expect both models (tensorflow and pytorch) to output the same values.

<!-- A clear and concise description of what you would expect to happen. -->


Comments:
@Rocketknight1 - do you have some bandwidth for taking a look? Otherwise, I can take a look as well :-)

I'm investigating this today, will update!

So, firstly I've reproduced the issues above. They seem to stem from multiple issues within the model code.

1) States in PyTorch are expanded to 3072 because the model uses the `_pad_to_window_size` method to expand input sequences to the next largest multiple of `config.window_size`, which in this case is 1024 (3072 == 1024 * 3).

2) This is supposed to be reversed by the line `hidden_states = hidden_states[:, :-padding_len]`, but this doesn't seem to be working in PyTorch. In this case, I think the TF implementation is correct, and the output shape should be (1, 2234, 768)

3) The Tensorflow implementation of Layerdrop uses `random.uniform` and so will not survive compilation by TF (I don't know if it's the part of the problem here, but it is **a** problem)

4) `encoder_last_hidden_state` in the TF model seems to contain the decoder output states instead, I'm working on figuring that one out.

Hopefully these are just isolated bugs and can be resolved, but I'm worried that there seem to be separate, sizeable bugs in both frameworks' implementations. Will keep investigating!

@federicoruggeri Can you try the PR and confirm it's working for you now? You can install from the branch with the fix by doing `pip install git+https://github.com/huggingface/transformers.git@fix_led`

@Rocketknight1 I've tested the above scripts after mentioned PR and shapes are now fine!
Many thanks for the quick fix!

@federicoruggeri No problem! I'm going to reopen the issue though - it's linked to the PR, so when that PR is merged that will auto-close the issue and keep everything in sync for us.