🐛 Weird learning rate with TPU Trainer
# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): **ELECTRA**

Language I am using the model on (English, Chinese ...): **English**

The problem arises when using:
* [ ] the official example scripts: 
* [x] my own modified scripts: 

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: 
* [x] my own task or dataset: CNN/DM

## To reproduce

I use `run_glue.py` as example to build a training script for TPU, with the Trainer API. My task is sequence classification with CNN/DM dataset.

I initialized the Trainer with following optimizer / scheduler :

```python
optimizer = AdamW(optimizer_grouped_parameters, lr=training_args.learning_rate, eps=training_args.adam_epsilon)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=training_args.warmup_steps, num_training_steps=287113)

trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=cnn_dm["train"] if training_args.do_train else None,
        eval_dataset=cnn_dm["validation"],
        data_collator=MyCollator(),
        prediction_loss_only=True,
        optimizers=(optimizer, scheduler),
    )
```

Now, the training procedure is working : the code run fine on 8 TPU core.

**But the loss is not decreasing.**

After looking into Tensorboards logs, I found the learning rate to be very weird :

 
![image](https://user-images.githubusercontent.com/43774355/82189972-3e888800-992b-11ea-8887-6715500f388d.png)

A few points to note :
* I specified a learning rate of **1e-4** (with the command argument `--learning_rate 1e-4`), but as you can see, the maximum value for the learning rate is **3.5e-6**.
* The shape of learning rate is not what I expected : After warmup, learning rate is supposed to decrease linearly, but instead, it stays fixed.

I don't know why the learning rate is being like this. Any idea what I might be doing wrong ?

_I can't share my notebook, but this seems to be the exact same issue with the official script `run_glue.py`, as described in #4358_ 

## 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: **2.9.1**
- Platform: **Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic**
- Python version: **3.6.9**
- PyTorch version (GPU?): **1.6.0a0+176174a (False)**
- Tensorflow version (GPU?): **2.2.0 (False)**
- Using GPU in script?: **No**
- Using distributed or parallel set-up in script?: **Yes : `xla_spawn.py`**

@julien-c 
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
