Exception encountered when calling layer "tf_bert_for_pre_training" (type TFBertForPreTraining)
### System Info

```shell
`transformers` version: 4.20.0
- Platform: Linux-5.13.0-48-generic-x86_64-with-glibc2.31
- Python version: 3.7.13
- PyTorch version (GPU?): 1.11.0 (False)
- Tensorflow version (GPU?): 2.9.1 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
```


### Who can help?

@Rocketknight1 @gante 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Colab link to reproduce: https://colab.research.google.com/drive/1tusV1pNe7sV2To9y7tep4l2LTPunEd5h?usp=sharing

```
InvalidArgumentError: Exception encountered when calling layer "tf_bert_for_pre_training" (type TFBertForPreTraining).

Input to reshape is a tensor with 61 values, but the requested shape requires a multiple of 2 [Op:Reshape]

Call arguments received by layer "tf_bert_for_pre_training" (type TFBertForPreTraining):
  • input_ids={'input_ids': 'tf.Tensor(shape=(2, 512), dtype=int64)', 'token_type_ids': 'tf.Tensor(shape=(2, 512), dtype=int64)', 'attention_mask': 'tf.Tensor(shape=(2, 512), dtype=int64)', 'next_sentence_label': 'tf.Tensor(shape=(2,), dtype=int64)', 'labels': 'tf.Tensor(shape=(2, 512), dtype=int64)'}
  • attention_mask=None
  • token_type_ids=None
  • position_ids=None
  • head_mask=None
  • inputs_embeds=None
  • output_attentions=None
  • output_hidden_states=None
  • return_dict=None
  • labels=None
  • next_sentence_label=None
  • training=True
```

### Expected behavior

```shell
The pre-training should start. I suspect the problem occurs when the number of masked tokens is not divisible by the batch size because of the reshape operation here: https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_tf_bert.py#L146.

The solution doesn't look trivial as it would need loss to be reduced before addition, because the current implementation adds loss for each item/sample separately.
```

This line seems strange to me

https://github.com/huggingface/transformers/blob/401fcca6c561d61db6ce25d9b1cebb75325a034f/src/transformers/models/bert/modeling_tf_bert.py#L146

I don't think it makes sense to reshape `masked_lm_loss ` using `next_sentence_loss`. cc @Rocketknight1 <|||||>@ydshieh This is true. I suspect this is done because 1) During loss calculation no "reduction" is being done 2) Since no reduction is being done the code is trying to add loss sample/instance wise (for example `masked_lm_loss` + `next_sentence_loss` for each sentence in a batch).

The fix might be to do reduction but then I see no loss calculation with reduction anywhere in HF.<|||||>Investigating this now - I think this bug is real, but does not occur for most of our models, and might be specific to `BertForPreTraining()` and the next sentence prediction loss. As a workaround for now, you can use a language model that doesn't have a next sentence prediction loss, like `TFBertForMaskedLM` or `TFRobertaForMaskedLM` - the current consensus is that this loss isn't that helpful for training a language model anyway, and models more recent than BERT generally don't use it.<|||||>Hi @Rocketknight1 ,

Thank You for your reply. Actually, I am planning to make a community notebook for Tensorflow BERT pre-training (it's been a problem to figure out according to the [discussion](https://discuss.huggingface.co/t/pre-train-bert-from-scratch/1245/30)) since BERT still serves as a baseline for a lot of the research and very recent [work](https://arxiv.org/pdf/2203.15827.pdf) also shows variations of NSP to help in pre-training. 

So I thought this might be a nice feature to have. Thank You for the help!<|||||>Bug post-mortem: The bug is in the line that @ydshieh identified. The code here is quite old and was obviously trying to reshape the masked LM loss before reduction so that a per-sample loss tensor would result. However, the loss vector does not reshape cleanly after masking, because random positions are removed from each sample. I rewrote everything with static shapes to fix the issue, and add XLA compilation as a bonus!<|||||>@Sreyan88 We have now pushed a fix, so you can try installing from `main` and see if this fixes your problem. If it doesn't, please feel free to post the new error and reopen this issue!<|||||>Hi @Rocketknight1 ,

The code works perfectly fine on colab now! However, in my personal server, it's giving me `nan` loss since the beginning of training. Do you think there is a reason for this? I have the same tf version (2.8.0) on both and the same hf version too. The only difference is GPUs (Tesla T4 on Colab and RTX 3090 on person system). Any clues?<|||||>Hi @Sreyan88, we're in the process of rewriting some loss functions in preparation for our next release, so things are changing quite quickly on `main`. Can you try updating to the most recent commit on your personal server and let me know if you still get the error? Use
`pip install --upgrade git+https://github.com/huggingface/transformers.git`<|||||>The problem persists :( .

Is there anything more I should do beyond checking that both are on `4.21.0.dev0`?<|||||>It finally worked, I had to force-reinstall. Thank You!<|||||>Hi @Rocketknight1 ,

Just a question, can  `prediction_logits` keyword in [this](https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_tf_bert.py#L1247) line be converted to `logits`? a.k.a the `prediction_logits` in `TFBertForPreTrainingOutput` be converted to `logits`. This way the model gets compatible with `pipeline("fill-mask")` which will be useful since Pre-training also has MLM as a task!

P.S. - `pipeline("fill-mask")` currently errors out with `TFBertForPreTraining` because it expects `logits`.

Thank You! If you suggest doing this would be correct I can create a PR!<|||||>Hi @Sreyan88, I'm not sure - like I said, `TFBertForPreTraining` is mostly not used anymore because the next sentence loss doesn't seem to be helpful! If you'd like to use a model you trained with `TFBertForPreTraining` with the `fill_mask` pipeline, I suggest loading the checkpoint with `TFBertForMaskedLM.from_pretrained()` - this will give you a model without the next sentence prediction head, which `fill_mask` doesn't use anyway. 