[TFTrainer] gradient accumulation error 
Author: maurice-g
Pull request: no

## Environment info
     
- `transformers` version: master (#9a8c168)
- Tensorflow version: 2.3.0

### Who can help 

 Trainer: @sgugger
 tensorflow: @jplu 

## Information

When using >1 `gradient_accumulation_steps` with TFTrainer and model inputs which are *not* simple tensors (for example dicts) the trainer fails. Also, to me it looks like there are logic issues in the way the `reduced_features` are computed for the gradient accumulation (not sure though).

Issue is here: https://github.com/huggingface/transformers/blob/9a8c168f56fe3c0e21d554a577ac03beb004ef89/src/transformers/trainer_tf.py#L602

## To reproduce

```python
import tensorflow as tf
from transformers import TFT5ForConditionalGeneration, TFTrainer, TFTrainingArguments

input_ids = [[1, 2, 3], [1, 2, 3]]
labels = [1, 2, 3]

dataset = tf.data.Dataset.from_tensor_slices(({'input_ids': input_ids}, labels))
model = TFT5ForConditionalGeneration.from_pretrained("t5-base")

training_args = TFTrainingArguments(
    output_dir='./results',  # output directory
    logging_steps=100,
    max_steps=2,
    save_steps=2000,
    per_device_train_batch_size=2,  # batch size per device during training
    per_device_eval_batch_size=8,  # batch size for evaluation
    warmup_steps=0,  # number of warmup steps for learning rate scheduler
    weight_decay=0.0,  # strength of weight decay
    learning_rate=5e-5,
    gradient_accumulation_steps=2
)

with training_args.strategy.scope():
    model = TFT5ForConditionalGeneration.from_pretrained("t5-base")

trainer = TFTrainer(
    model=model,  # the instantiated 🤗 Transformers model to be trained
    args=training_args,  # training arguments, defined above
    train_dataset=dataset,  # training dataset
)

trainer.train()
```

## Issues

### Error
Produces error `TypeError: unhashable type: 'slice'`. Also `/` produce floats on python3, which I guess is not intended here.
Solution in same spirit could be conditional use of
```
reduced_features = {
 ft: features[ft][:self.args.train_batch_size // self.args.n_replicas]
 for ft in features
}
```

Already mentioned here https://github.com/huggingface/transformers/pull/6038#issuecomment-664706046

### Logic issue
I don't understand what the `n_replicas` has to do with the gradient accumulation here? Shouldn't the denominator rather be `gradient_accumulation_steps`? And shouldn't it actually use the different slices of the features, and not always the "first" slice? 
Might be totally misunderstanding this.

Also this line doesn't seem to have any purpose: https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py#L607

Happy to provide PR if someone can give me a hint on the logic issue.

Comments:
We should wait for @jplu to come back from holiday for this, since he wrote that part of the code.

Good catch @maurice-g!!

It is my fault, I did that part in a hury and I should have been more careful. This will be fixed in a next PR (currently doing it)

`n_replicas` is important here because we have to get the number of tuple (features, labels) corresponding to the batch size per GPU and `self.args.train_batch_size` gives the total batch size (batch size per GPU * number of replicas).

Should be fixed in https://github.com/huggingface/transformers/pull/6713 :+1: 

thanks for looking into this @jplu

One further remark on your PR #6713: The code now works _iff_ the features are a dict, but does not anymore if the features are a raw tensor (which worked before). IMO this should work for both, therefore there needs to be a conditional check on the type and then both situations should be handled. Or do you think that's not a relevant case?

Keeping open until @jplu answers your question @maurice-g 

You are right, it is not working anymore with list/tuple/raw tensors. This is on purpose because I'm gonna push the usage of dictionaries only in TF at some point. Is it a big issue for you to use dictionaries?

Ok, works for me, just wanted to point it out.