GPU Benchmarking + Accumulated Optimizer for TF2
Author: iliaschalkidis
Pull request: no
Labels: name: wontfix

## 🚀 Feature

- Create a GPU benchmarking section in Documentation (Wiki).
- Build and Include TF2 Optimizer with gradient accumulation.

```python
optimizer=AccumulatedOptimizer(Adam(lr=2e-5, clipnorm=1.0), accumulate_steps=4)
```

## Motivation

I experiment with transformers library on Tensorflow 2 for a week and what you offer to us the users seems really useful. It saves a ton of time and keeping us up-to-date with the latest advances on pre-trained models accelerate our research.

Although, this come with a great cost in computational resources (e.g., GPUs). For example the the use of BERT through Tensorflow Hub is much lighter than the one you offer in this great library. For example in a 12GB GPU (e.g., 1080Ti, RTX 2080Ti), someone can go up to batches of 8 of 512 tokens if he makes use of the Tensorflow Hub BERT-BASE module , while using transformers library will lead to 1/2 of the batch size (=4). 

Based on the above facts I think we need three crucial things:

- An explanation on why this happens? It's seem weird to me that for a model with same parameters, we cannot go up to the same batching based on different implementations. I think understanding this mismatch is very interesting from a theoretical and practical point of view.
- A table with benchmarks given different GPUs, different batches, different max sequence
lengths. This will helps us find possible limitations and also find bugs in our code, if we do not meet the benchmarking.
- Given this limitation, I propose the release of a new TF2 optimizer that used gradient accumulation, thus we can go up to our batch size limit and use accumulation in order to avoid one back-propagation step per forward step.




Comments:
I would also like to see benchmarks, however this is a computationally heavy task. It might be useful to provide a benchmark script and a benchmark table. Contributors can then run the script on their available hardware, and add their results to the table - highlighting the used parameters and hardware.

Here's my findings from testing the Transformers library on a Titan V (12GB), which I'm also using to run my dual displays (about 400MB VRAM there).

The VRAM usage below is the VRAM allocated by TensorFlow, not necessarily the exact VRAM needed to run the model. TF tends to allocate VRAM in chunks once a smaller chunk is not sufficient. All experiments below use a token length of 512 tokens, and the same batch size for train and eval, and Adam optimizer. The script used is a minimally modified `run_tf_glue.py`. 

| Batch Size | VRAM        | Mixed Precision |
| ----------- | ---------- | ---------------- |
| 4                  | 8723MB   |  No                       |
| 4                  | 8723MB   |  Yes                      |
| 8                  | 11265MB  |  No                       |
| 8                  | 11265MB  |  Yes                      |
| 9                  | 11265MB  |  No                       |
| 9                  | 11265MB  |  Yes                     |
| 10                | OOM         |  No                       |
| 10                | OOM         |  Yes                      |

On 1080 Ti (11GB), I managed to run it at batch size 8, with VRAM usage of 10753MB.

From the results I got, one should be able to run at batch size 8 on a 2080 Ti (11GB), but I don't have a 2080 Ti to test. Worth nothing that you might not be able to run a display AND the training at the same time, as every last bit of VRAM seems to be required.

The script can be found [here](https://gist.github.com/tlkh/d252abcb3a5b59a7b8c47660997fd390#file-tf_run_glue-py).

I will test with the TF Hub BERT module at a later date if I have time, but from memory the VRAM usage seems to be similar.

cc @thomwolf who asked about it on #1441

Hi @tlkh, I was able to rerun your script `tf_run_glue.py` successfully in a 1080Ti. Then I tried to pass the core elements that affect GPU acceleration and optimization in my own code:

```python
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

USE_XLA = False
USE_AMP = False

tf.config.optimizer.set_jit(USE_XLA)
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP})

# Compile t TFBertForSequenceClassification with Adam

# Load my data using a custom Keras Generator 
# that yields a list of numpy ndarrays of shape: [(8,512), (8,512), (8,512)]  to pass token_ids, mask_ids, segment_ids

# Call model.fit_generator(train_data=generator) on keras Model

```

This always lead me to an OOM error on specific steps of the network:

> tensorflow.python.framework.errors_impl.ResourceExhaustedError:  OOM when allocating tensor with shape[8,12,512,512] 

OR

> tensorflow.python.framework.errors_impl.ResourceExhaustedError:  OOM when allocating tensor with shape[6,512,768] 

OR

> tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[6,512,3072]

All of them are internal functions/computations in transformers...


- Is there any reason to suspect that tf.Dataset objects are optimized compared to a keras Generator?
- I tried to activate XLA, but this leads to the following error:

> tensorflow.python.framework.errors_impl.NotFoundError:  ./bin/ptxas not found
>          [[{{node cluster_0_1/xla_compile}}]] [Op:__inference_call_6748]

and I can't find ptxas path on server....

> Is there any reason to suspect that tf.Dataset objects are optimized compared to a keras Generator?

I don't think the VRAM usage will defer, but tf.Dataset objects *should* be more optimized. I believe Keras Generators are supposed to be deprecated eventually. 

> I tried to activate XLA

Your TensorFlow build needs to support XLA, although I believe it should be already built in by default.

If @iliaschalkidis don't mind sharing your code I could try running and see if I can replicate the problem.

@tlkh I just made this demo script that replicates the way I handle data and fit the model. In my server this script leads to OOM as well, as the actual project. I thought it's much easier than sharing the whole project, cover dependencies and load real datasets.

```python
import numpy as np
import tensorflow as tf
from transformers import TFBertForSequenceClassification

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)


class SampleGenerator(tf.keras.utils.Sequence):
    """Generates data for Keras"""

    def __len__(self):
        # 10 batches of samples each
        return 10

    def __getitem__(self, index):
        # Yield mock data batch
        token_ids = np.zeros((8, 512), dtype=np.int32)
        mask_ids = np.zeros((8, 512), dtype=np.int32)
        segment_ids = np.zeros((8, 512), dtype=np.int32)

        targets = np.zeros((8, 1000), dtype=np.int32)

        return [token_ids, mask_ids, segment_ids], targets


# script parameters
BATCH_SIZE = 8
EVAL_BATCH_SIZE = BATCH_SIZE
USE_XLA = False
USE_AMP = False

tf.config.optimizer.set_jit(USE_XLA)
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP})

# Load model from pretrained model/vocabulary
model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=1000)

# Prepare datasets as Keras generators
train_generator = SampleGenerator()
val_generator = SampleGenerator()

# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
opt = tf.keras.optimizers.Adam(learning_rate=3e-5)
if USE_AMP:
    # loss scaling is currently required when using mixed precision
    opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')

model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
model.fit_generator(generator=train_generator, epochs=2, validation_data=val_generator)

```

I replaced Keras generator with `tf.data.Dataset`:

```python
import numpy as np
import tensorflow as tf
from transformers import TFBertForSequenceClassification

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

# script parameters
BATCH_SIZE = 8
EVAL_BATCH_SIZE = BATCH_SIZE
USE_XLA = False
USE_AMP = False

tf.config.optimizer.set_jit(USE_XLA)
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": USE_AMP})

# Load model from pretrained model/vocabulary
model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=1000)


def gen():
    for x1, x2, x3, y in zip(np.zeros((80, 512), dtype=np.int32),
                          np.zeros((80, 512), dtype=np.int32),
                          np.zeros((80, 512), dtype=np.int32),
                          np.zeros((80, 1000), dtype=np.int32)):
        yield ({'input_ids': x1,
                'attention_mask': x2,
                'token_type_ids': x3}, y)


# Prepare dataset as tf.Dataset from generator
dataset = tf.data.Dataset.from_generator(gen,
            ({'input_ids': tf.int32,
              'attention_mask': tf.int32,
              'token_type_ids': tf.int32},
             tf.int32),
            ({'input_ids': tf.TensorShape([None]),
              'attention_mask': tf.TensorShape([None]),
              'token_type_ids': tf.TensorShape([None])},
             tf.TensorShape([None])))

train_dataset = dataset.shuffle(128).batch(BATCH_SIZE).repeat(-1)

# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
opt = tf.keras.optimizers.Adam(learning_rate=3e-5)
if USE_AMP:
    # loss scaling is currently required when using mixed precision
    opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')

model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])

# Train and evaluate using tf.keras.Model.fit()
model.fit(train_dataset, epochs=2, steps_per_epoch= 80//BATCH_SIZE)
```

The model know fits in the GPU... So it seems `tf.data.Dataset.from_generator()` is more memory efficient from Keras generators...

Interesting to know that the issue is Keras generators. Glad you have a way of running the model now!

> Interesting to know that the issue is Keras generators. Glad you have a way of running the model now!

Unfortunately, I still haven't. Because I need to refactor a lot aspects of my personal codebase in order to load the datasets in the same fashion as in this example... Let's hope this will not take more than a single day 😄 

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.