TensorFlow "predict" returns empty output with MirroredStrategy
Author: ZJaume
Pull request: no

I'm trying to use the `predict` method of the Keras TensorFlow API but it returns an empty output despite the input is being processed. Calling the model seems to work.

EDIT: the predict method works correctly if the model is loaded with single GPu strategy.

## 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.5.1`
- Platform: Linux CentOS 8.1
- Python version: `3.7.10`
- PyTorch version (GPU?): -
- Tensorflow version (GPU?): `2.3.2`(True)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: multi-gpu on a single machine

### Who can help
<!-- 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

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)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

```python
from transformers import BertTokenizerFast, TFBertForSequenceClassification
import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
#strategy = tf.distribute.OneDeviceStrategy("/gpu:0")
with strategy.scope():
    tf_model = TFBertForSequenceClassification.from_pretrained('bert-base-cased')

    tokenizer = BertTokenizerFast.from_pretrained('bert-base-cased')
    inputs = tokenizer('This is a test', 'Esto es una prueba',
                   return_tensors='tf', max_length=200,
                   padding='max_length', truncation=True,
                   return_attention_mask=True,
                   return_token_type_ids=False)
    print(tf_model.predict([inputs["input_ids"], inputs["attention_mask"]],
                           verbose=1))
    print(tf_model([inputs["input_ids"], inputs["attention_mask"]]))
```
```
All model checkpoint layers were used when initializing TFBertForSequenceClassification.

Some layers of TFBertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
WARNING:tensorflow:From /venv/lib/python3.7/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
1/1 [==============================] - 0s 241us/step
TFSequenceClassifierOutput(loss=None, logits=None, hidden_states=None, attentions=None)
TFSequenceClassifierOutput(loss=None, logits=<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[-0.47814545,  0.35146457]], dtype=float32)>, hidden_states=None, attentions=None)
```

## Expected behavior

Output should be the same as when model is being called.


Comments:
Tested on latest release and still present.

Pinging our TensorFlow expert, @Rocketknight1 

I've managed to reproduce this but I'm very confused about the cause, especially because I'm pretty sure I've used model.predict with MirroredStrategy in our codebase before.

I've tested your code snippet with a standard RNN instead of BERT and confirmed that it works fine, and I tried distilbert instead of BERT and the problem remained, so the problem does seem to be the combination of MirroredStrategy and our models.

I'm going to keep poking around at this, but if you discover anything else that might help me figure out what's going on, please let me know!

Update: This bug appears in our `run_text_classification.py` script too, again only when using predict(). I'm investigating.

Update 2: `fit()` and `evaluate()` seemed to work correctly in a MirroredStrategy context (which is good because I have a whole example that uses them). The issue is specific to `predict()`

Hi, just keeping this issue alive! I've traced the issue to the way we return our values from the `call()` methods - I think Keras doesn't like the thing we do with a subclassed OrderedDict. We're going to reach out to our contacts at Google in the next couple of days and figure out what the best approach is - whether we need to refactor that totally, or if there's an easy workaround.

Putting this here as a writeup of what we know so far:

The issue is not caused by returning an `OrderedDict`, but instead because we return a `TFBaseModelOutput`, which is a subclass of `OrderedDict` decorated with dataclass. Refer to the code [here](https://github.com/huggingface/transformers/blob/38a716cd41f22f6a7d5ff3dc081903090198803a/src/transformers/modeling_tf_outputs.py#L24-L46).

If we just return a dict, `OrderedDict` or `ModelOutput` (the parent class for `TFBaseModelOutput`, subclassed from `OrderedDict`), everything works okay. Therefore the central issue is this data class, which will probably need to be removed. We're looking at how we can do that now!

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.

Hi,
Any updates about this issue? 

definitely looking forward to a fix for this. how can we help @Rocketknight1?

@jmwoloso @ayalaall 

Hey all! I'm going to reopen this issue, even though we're short on bandwidth for it right now. The current situation is that we know where the problem lies - it's in the fact that we're returning a `@dataclass` decorated object from our models, and that doesn't play nicely with Keras. We get away with it when we're not in a `Strategy` context, but it breaks inside of one, even though `fit()` still usually works correctly.

The problem is that even though the change needed is relatively small, it's finicky because we place a lot of value on maintaining a very consistent API for users, and changing the return class for every TF model on the entire hub is a big deal. So we need to find some way to make sure existing code is as unaffected as possible in the process, and that requires some engineering exploration.

The good news is the `@dataclass` decorator is really just for convenience rather than a critical part of the class - we just use it to ensure that certain keys are always present in the output dict, and set with default values, and it got ported over from the original PyTorch code. We could probably make some other subclass of `Dict` or `OrderedDict` and return that, and maybe that would play nicer with Keras, but I have a few other major things on my to do list, so I don't know if I'll be able to get to that for a month or two. If anyone wants to experiment and file a PR, feel free to ask any questions you want here. If not, I'll do my best to get to it as soon as I can.

Say no more @Rocketknight1! I'll take a look and get familiar with the components involved and see if I can devise a minimally-invasive solution. Thanks for re-opening!

You can assign this to me if you like as well.

@jmwoloso Sure, if you'd like! If you have any questions along the way, feel free to ask.

@ZJaume @ayalaall @Rocketknight1 
An update for the group. I'm still doing some testing, but this is fixed in both `master` and `transformers==4.10.0`!

Using a single VM (4 v100 GPUs) with `MirroredStrategy` works out of the box. `transformers==4.9.2` (the version I happened to be using) it does not work.

```
from transformers import TFDistilBertForSequenceClassification, DistilBertTokenizerFast
import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    tf_model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')

    tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
    inputs = tokenizer('This is a test', 'Esto es una prueba',
                   return_tensors='tf', max_length=200,
                   padding='max_length', truncation=True,
                   return_attention_mask=True,
                   return_token_type_ids=False)

    print(tf_model.predict([inputs["input_ids"], inputs["attention_mask"]], verbose=1))
    print(tf_model([inputs["input_ids"], inputs["attention_mask"]]))

```

```
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
Downloading: 100%|██████████| 483/483 [00:00<00:00, 551kB/s]
Downloading: 100%|██████████| 363M/363M [00:04<00:00, 79.6MB/s] 
Some layers from the model checkpoint at distilbert-base-uncased were not used when initializing TFDistilBertForSequenceClassification: ['vocab_projector', 'vocab_layer_norm', 'vocab_transform', 'activation_13']
- This IS expected if you are initializing TFDistilBertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing TFDistilBertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some layers of TFDistilBertForSequenceClassification were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['dropout_19', 'classifier', 'pre_classifier']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Downloading: 100%|██████████| 232k/232k [00:00<00:00, 1.03MB/s]
Downloading: 100%|██████████| 466k/466k [00:00<00:00, 1.52MB/s]
Downloading: 100%|██████████| 28.0/28.0 [00:00<00:00, 28.1kB/s]
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:From /databricks/python/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:5043: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.
Instructions for updating:
The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
1/1 [==============================] - 10s 10s/step
TFSequenceClassifierOutput(loss=None, logits=array([[ 0.03777119, -0.12381434]], dtype=float32), hidden_states=None, attentions=None)
TFSequenceClassifierOutput(loss=None, logits=<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[ 0.0377712 , -0.12381432]], dtype=float32)>, hidden_states=None, attentions=None)
```

@jmwoloso That's really fascinating! I didn't think I touched any relevant code between those releases, but possibly one of the other engineers did. Can you try a couple of other models, say BERT or RoBERTa, to see if you see the same pattern with both?

I tried with Roberta and DistilBert with the new version and it doesn't give empty output any more. Thank you!

Hi @jmwoloso @ZJaume this is great, thank you! Can you confirm it still works with an input array larger than the batch size? (to ensure the work is getting distributed to multiple GPUs and then merged correctly)

@Rocketknight1 yeah i'll take a look at doing that today and posting confirmation in here and then we can close this out!

Working with 1024 samples and 8 batch size per gpu.

I'm still trying to test it out but databricks is having issues spinning up gpu clusters today :roll_eyes: 

I think we're good to close this out @Rocketknight1 unless there are other scenarios you want us to check out.

So I noticed I had the same problem when I do this with basic Tensorflow. I found that the Tokenizer() function from tensorflow.keras.preprocessing.text seems to be an empty when you load the model. Which is understandable because you are not loading any sort of data to the Tokenizer. 
How I was able to solve it was

```
import pickle

# saving
with open('tokenizer.pickle', 'wb') as handle:
    pickle.dump(tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL)

# loading
with open('tokenizer.pickle', 'rb') as handle:
    tokenizer = pickle.load(handle)
```

@jmwoloso @ZJaume Thank you for all your help! I'm gonna mark this as resolved now, since the problem doesn't seem to have recurred.

@JithLord I think that's a separate problem, unless it's also unique to `MirroredStrategy` contexts, so I'm gonna close this issue anyway. If you think you've found a bugs in the repo, though, please feel free to file a separate issue (or file it with Tensorflow upstream if you think the bug is there).