Cannot use `transformers.GradientAccumulator` with `tf.function`
# 🐛 Bug

I'm using gradient accumulation to get larger batch sizes when pretraining BERT, and GradientAccumulator works wonderfully. Thanks for including that!

However, tf.AutoGraph fails when I wrap in the gradient accumulation step in tf.function. I have to write code like this:

```code
@tf.function
def batch_step():
    step_grads = ...
    return step_grads

@tf.function
def allreduce():
    grads = [hvd.allreduce(grad) for grad in gradient_accumulator.gradients]
    optimizer.apply_gradients(grads)
    gradient_accumulator.reset()

def step():
    for i in range(gradient_accumulation_steps):
        step_grads = batch_step(dataset[i])
        gradient_accumulator(step_grads)

    allreduce()
```

This code works. However, it is very inefficient because TensorFlow stores the GradientAccumulator on CPU rather than GPU. My PCI throughput is around 9 GiB/s (on 8-GPU single-node training) when using gradient accumulation, as opposed to in the kilobytes when everything is wrapped in a tf.function. GPU utilization also drops because the bottleneck is transfer of gradients between GPU and CPU. I also suspect that it's causing a memory leak somewhere, but that's another story :) Being able to wrap the GradientAccumulator in a tf.function (or somehow pin it to the GPU, using `tf.device()` doesn't seem to be working) would be wonderful. Any tips?

Sample PCI throughput, ideally RX (GPU-to-CPU) should be close to 0.
<img width="682" alt="Screen Shot 2020-02-26 at 11 21 47 AM" src="https://user-images.githubusercontent.com/4564897/75380114-5223d580-588b-11ea-9594-e85755ac0db0.png">

## Information

Model I am using (Bert, XLNet ...): TFAlbert

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


## Environment info
     
- `transformers` version: 2.5.0
- Platform: Ubuntu 18.04
- Python version: 3.6.9
- PyTorch version (GPU?): None
- Tensorflow version (GPU?): 2.1
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes, but it also applies to single-node

`transformers`'s implementation of the GradientAccumulator has been inspired by the OpenNMT implementation. Can you try their (updated) version and report the results? You can find it here:

https://github.com/OpenNMT/OpenNMT-tf/blob/0311e473d9788b0363ca663cafbdfaf7777e53f9/opennmt/optimizers/utils.py#L64

You can also see how they use it in their training script, e.g.

https://github.com/OpenNMT/OpenNMT-tf/blob/0311e473d9788b0363ca663cafbdfaf7777e53f9/opennmt/optimizers/utils.py#L64

Disclaimer: I'm a PyTorch guy, so I can't help with TF. It just seems like something to test.<|||||>@jarednielsen Can you try moving the call to `gradient_accumulator(step_grads)` directly inside `batch_step()` and avoid returning the gradients from the function?

The implementation over at OpenNMT-tf was indeed updated because the previous one (and included in `transformers`) used the `device_map` property which is removed in recent TensorFlow versions.<|||||>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.
