Different inference results of a keras including transformer model on TPU vs CPU?
transformers version: 4.0.0
Platform: Linux-4.9.0-11-amd64-x86_64-with-debian-9.11
Python version: 3.7.9
PyTorch version (GPU?): 1.6.0a0+bf2bbd9 (False)
Tensorflow version (GPU?): 2.3.1 (False)
Using GPU in script?: No
Using distributed or parallel set-up in script?: Distributed

I am building a Keras model which consists of a TFRobertaModel with 2 custom heads on top. One is a QA head which outputs span predictions and the other is a binary classification head. I train the model on TPUs and everything works fine with prediction and inference on the TPU with great model performance. The issue I am having is loading the saved model and/or weights and doing inference on CPU. I am getting completely different results compared to inference on the TPU.

I save the model using model.save and the weights as well with models.save_weights and it doesn't matter which one I load, I get the same results (using tf.keras.models.load_model). I do get a warning that there was an error saving the state of the optimiser which is initialized at random on model loading. I figure this is not an issue with keras since I compile the model with a tf.keras.optimizers.Adam which should be saved with model.save with a keras only model. I have also tried building the model from scratch and only loading the saved weights but I get the same results.

This is a convoluted problem to reproduce but I was wondering if you had any pointers on how to debug this or if this was a known problem. Here is a sample of the model output on TPU vs CPU:

TPU - the first 10 binary answer predictions:

array([[9.9994159e-01, 5.8382138e-05],
       [9.9990284e-01, 9.7181561e-05],
       [9.9995410e-01, 4.5917721e-05],
       [9.9996519e-01, 3.4784229e-05],
       [9.9975628e-01, 2.4374224e-04],
       [9.9997389e-01, 2.6103005e-05],
       [9.9995828e-01, 4.1662061e-05],
       [9.9998319e-01, 1.6824890e-05],
       [7.0182599e-02, 9.2981732e-01],
       [9.9993420e-01, 6.5814391e-05]], dtype=float32)

CPU - the first 10 binary answer predictions:

array([[0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ],
       [0.06942184, 0.9305781 ]], dtype=float32)

The model building and compilation:

```
dropout=False
def create_model():
    
    input_shape = (None,)
    model = TFRobertaForQuestionAnswering.from_pretrained(model_path, from_pt=True, trainable=True)
    
    input_ids = tf.keras.layers.Input(shape=input_shape, dtype=np.int32, name='input_ids')
    attention_mask = tf.keras.layers.Input(shape=input_shape, dtype=np.int32, name='attention_mask')
    token_type_ids = tf.keras.layers.Input(shape=input_shape, dtype=np.int32, name='token_type_ids')
    
    outputs = model.roberta(input_ids, 
                            attention_mask=attention_mask, 
                            token_type_ids=token_type_ids, 
                            output_hidden_states=True)
    

    seq_output = outputs[0]
    logits = model.qa_outputs(seq_output)
    start_logits, end_logits = tf.split(logits, 2, axis=-1)
    start_logits = tf.squeeze(start_logits, axis=-1)
    end_logits = tf.squeeze(end_logits, axis=-1)
    
    #BINARY_ANSWER    
    concat_hidden_layers = tf.concat(tuple([outputs.hidden_states[i] for i in [-4, -3, -2, -1]]), axis=-1)
    pooled_output = concat_hidden_layers[:, 0, :]

    
    binary_answer_logits = tf.keras.layers.Dense(768,
            kernel_initializer=tf.keras.initializers.truncated_normal(stddev=0.02),
            activation="tanh",
            name="dense_tanh")(pooled_output)

    if dropout:
        binary_answer_logits = tf.keras.layers.Dropout(0.1)(binary_answer_logits)
        
    binary_answer_probs = tf.keras.layers.Dense(2, activation='softmax', name="binary_answer")(binary_answer_logits)
    
    keras_model = Model(inputs={'input_ids':input_ids, 
                                'attention_mask':attention_mask, 
                                'token_type_ids':token_type_ids}, 
                        outputs={'start_logits':start_logits, 
                                 'end_logits':end_logits, 
                                 'binary_answer_probs':binary_answer_probs})    
    
    return keras_model
    
with strategy.scope():
    keras_model = create_model()
    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
    keras_model.compile(loss={'start_logits':compute_loss,
                              'end_logits':compute_loss,
                              'binary_answer_probs':tf.keras.losses.binary_crossentropy}, 
                                optimizer=optimizer)
```

I also get completely different results for span predictions on TPU vs CPU which is not surprising, seeing how different the binary prediction is. Any help or pointers are appreciated. 

## Expected behavior

The same inference results on TPU vs CPU.

Hello, thanks for opening an issue! We try to keep the github issues for bugs/feature requests.
Could you ask your question on the [forum](https://discusss.huggingface.co) instead?

Thanks!