LED for Seq2Seq output shape mismatch between tensorflow and pytorch

Environment info

• 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

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)
• my own modified scripts: (give details below)

The tasks I am working on is:

• an official GLUE/SQUaD task: (give the name)
• 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.

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]), ....]

1. Run the torch version of the same script.

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]), ....]
"""

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.

[Rocketknight1]

I’m investigating this today, will update!

[Rocketknight1]

@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.