Cannot use `transformers.GradientAccumulator` with `tf.function`
Author: jarednielsen
Pull request: no
Labels: name: wontfix

name: TensorFlow

# 🐛 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


Comments:
`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.