T5 training with Keras: InvalidArgumentError:  logits and labels must have the same first dimension
Author: marton-avrios
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.3.0.dev0
- Platform: Linux version 4.19.0-14-cloud-amd64 (debian-kernel@lists.debian.org) (gcc version 8.3.0 (Debian 8.3.0-6)) #1 SMP Debian 4.19.171-2 (2021-01-30)
- Python version: 3.7
- PyTorch version (GPU?): 1.7.1, No
- Tensorflow version (GPU?): 2.3.1, No
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten, @patil-suraj
<!-- 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: @jplu

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- 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 ...): T5

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

See code below:

```
import numpy as np
import tensorflow as tf
from transformers import T5TokenizerFast, TFT5ForConditionalGeneration

MODEL_NAME = "t5-small"
INPUT_TEXTS = [
    "When Liana Barrientos was 23 years old, she got married in Westchester County, New York.",
    "Only 18 days after that marriage, she got hitched yet again.",
    "Then, Barrientos declared 'I do' five more times, sometimes only within two weeks of each other.",
    "In 2010, she married once more, this time in the Bronx.",
    "In an application for a marriage license, she stated it was her 'first and only' marriage.",
    "Prosecutors said the marriages were part of an immigration scam.",
    "In total, Barrientos has been married 10 times, with nine of her marriages occurring between 1999 and 2002.",
    "All occurred either in Westchester County, Long Island, New Jersey or the Bronx.",
    "Any divorces happened only after such filings were approved.",
    "It was unclear whether any of the men will be prosecuted.",
]
LABEL_TEXTS = ["Yes", "No", "Yes", "Yes", "No", "Yes", "No", "No", "Well, you never know, right?", "Yes"]

tokenizer = T5TokenizerFast.from_pretrained(MODEL_NAME)

tokenized_inputs = tokenizer(INPUT_TEXTS, padding="max_length", truncation=True, return_tensors="tf")
tokenized_labels = tokenizer(LABEL_TEXTS, padding="max_length", truncation=True, return_tensors="tf")

decoder_input_texts = ["<pad> " + _txt for _txt in LABEL_TEXTS]
tokenized_decoder_inputs = tokenizer(decoder_input_texts, padding="max_length", truncation=True, return_tensors="tf")


def add_dec_inp_ids(_features, _labels, _dec_inp_ids):
    _features["decoder_input_ids"] = _dec_inp_ids
    return (_features, _labels)


ds = tf.data.Dataset.from_tensor_slices(
    (tokenized_inputs.data, tokenized_decoder_inputs.input_ids, tokenized_labels.input_ids))\
    .map(add_dec_inp_ids)

batch_size = 2
steps_per_epoch = np.ceil(len(INPUT_TEXTS) / batch_size)

train_ds = ds.repeat().prefetch(tf.data.experimental.AUTOTUNE).batch(batch_size)

model = TFT5ForConditionalGeneration.from_pretrained(MODEL_NAME)
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer=optimizer, loss=loss)
model.fit(train_ds, epochs=2, steps_per_epoch=steps_per_epoch)
```

And what I get:
```
tensorflow.python.framework.errors_impl.InvalidArgumentError:  logits and labels must have the same first dimension, got logits shape [8192,64] and labels shape [1024]
	 [[node sparse_categorical_crossentropy_4/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits (defined at <ipython-input-2-0152c3165ef3>:51) ]] [Op:__inference_train_function_27365]
```
<!-- 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
Training starts and then finishes without error.
<!-- A clear and concise description of what you would expect to happen. -->


Comments:
Would like to help you here, I've created a [Colab notebook](https://colab.research.google.com/drive/1PtRxbK4oNUsm4lrsOvWoNYzA-BhOWwf2?usp=sharing) that illustrates how to fine-tune `TFT5ForConditionalGeneration` using Keras. However, I'm having the same issue as posted in #6817, namely:

`ValueError: No gradients provided for any variable: ['shared/shared/weight:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._0/SelfAttention/q/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._0/SelfAttention/k/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._0/SelfAttention/v/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._0/SelfAttention/o/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._0/SelfAttention/relative_attention_bias/embeddings:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._0/layer_norm/weight:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._1/DenseReluDense/wi/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._1/DenseReluDense/wo/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._0/layer_._1/layer_norm/weight:0', 'tf_t5for_conditional_generation_2/encoder/block_._1/layer_._0/SelfAttention/q/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._1/layer_._0/SelfAttention/k/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._1/layer_._0/SelfAttention/v/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._1/layer_._0/SelfAttention/o/kernel:0', 'tf_t5for_conditional_generation_2/encoder/block_._1/layer_._0/layer_norm/weight:0', 'tf_t5for_conditional_generation_2/encoder/block_._1/layer_._1/DenseReluDense/wi/kernel:0', 'tf_t5for_conditional_generation_2/encoder/bloc...`

UPDATE: this issue was resolved by providing the data in the correct format, namely a tuple of `(inputs, outputs)`. A forward pass on a random batch is now working. However, having the following error when calling `model.fit()`:

```
ValueError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:805 train_function  *
        return step_function(self, iterator)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:795 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:788 run_step  **
        outputs = model.train_step(data)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:758 train_step
        self.compiled_metrics.update_state(y, y_pred, sample_weight)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:387 update_state
        self.build(y_pred, y_true)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:318 build
        self._metrics, y_true, y_pred)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:1163 map_structure_up_to
        **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:1245 map_structure_with_tuple_paths_up_to
        expand_composites=expand_composites)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:878 assert_shallow_structure
        input_length=len(input_tree), shallow_length=len(shallow_tree)))

    ValueError: The two structures don't have the same sequence length. Input structure has length 4, while shallow structure has length 3.
```

Hello!

For now T5 cannot be trained with usual `.compile()` and `.fit()` methods (such as multiple other models but we are currently working on this). You have to either use the TFTrainer or to update yourself the behavior of the internal training loop of Keras. An example of how to deal with T5 and properly training it, is showed in this nice [Colab](https://colab.research.google.com/github/snapthat/TF-T5-text-to-text/blob/master/snapthatT5/notebooks/TF-T5-Datasets%20Training.ipynb).

Ok, I saw that Colab but now understand why the author defined his own `train_step` (this was failing for me). Will update my notebook to add this.

Thank you!

@marton-avrios  Here's an updated version of my Colab notebook, illustrating how to fine-tune `TFT5ForConditionalGeneration` on your data: https://colab.research.google.com/drive/1PtRxbK4oNUsm4lrsOvWoNYzA-BhOWwf2?usp=sharing

@NielsRogge Your colab looks much better but will be buggy for a few cases. Your train step must look something like this:

```python
def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compiled_loss(y, y_pred.logits, regularization_losses=self.losses)

        gradients = tape.gradient(loss, self.trainable_variables)

        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        self.compiled_metrics.update_state(y, y_pred.logits)

        return {m.name: m.result() for m in self.metrics}
```

Nice adaptation in your Colab BTW :)

Thank you guys, very useful resource! Will it work on TPU or with other distribution strategies? Keras handles that when I stick to using ```compile()``` and ```fit()``` even if I redefine ```train_step()```, right? I mean dividing loss with global batch size, etc.

I doubt you will be able to train a T5 on TPU because T5 is not entirely XLA compliant so that you might encounter some unexpected issues. Sorry for that, it is also something on which we are currently working :)

So PyTorch version won't work either on TPU? Any hints as to which parts? I might be able to look into it. I get this error:
```
Invalid argument: {{function_node __inference_distributed_training_steps_51234}} Compilation failure: Detected unsupported operations when trying to compile graph cluster_distributed_training_steps_2864851955408122598[] on XLA_TPU_JIT: StringFormat (No registered 'StringFormat' OpKernel for XLA_TPU_JIT devices compatible with node {{node tf_t5for_conditional_generation/encoder/StringFormat}}){{node tf_t5for_conditional_generation/encoder/StringFormat}}
```

There are still a lot of work to make TFT5 XLA compliant, so I don't suggest you to use it for this case. The Pytorch version is fully TPU compliant yes.

@marton-avrios turns out that the Tensorflow implementation of T5 already creates the `decoder_input_ids` for you as seen [here](https://github.com/huggingface/transformers/blob/587197dcd2b50ad9e96aedbfa389bf4fcc294c3c/src/transformers/models/t5/modeling_tf_t5.py#L1376), you don't need to prepare them yourself (I thought this was only supported in the PyTorch version for now). So I've updated my notebook, it's simpler now 



Thank you @NielsRogge ! I think you need to provide either label input ids as ```labels``` or shifted label input ids as ```decoder_input_ids``` in input dictionary. At least it failed for me with only ```input_ids``` and ```attention_mask```.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.