Keras Implementation of Memory-Saving Recomputation Technique is Now Available · English (unofficial) translations of posts at kexue.fm

Many readers may have recently noticed the blog post "BERT Recomputation: Saving 5x Memory with 22.5% Training Time (with Code)". It introduces a technique called "recomputation." Simply put, it is a method used to save VRAM (Video RAM), allowing the batch_size to increase several times at the cost of a slightly slower average training speed. This technique was first published in the paper "Training Deep Nets with Sublinear Memory Cost" in 2016, though it hasn’t become particularly popular until recently.

Exploration

The aforementioned post mentioned that this technique has native implementations in PyTorch and PaddlePaddle, but not yet in TensorFlow. However, in reality, TensorFlow has included this feature since version 1.8, initially listed in the tf.contrib sub-library. Starting from TensorFlow 1.15, it was built-in as one of the main functions: tf.recompute_grad.

After finding tf.recompute_grad, I spent some time exploring its usage. After some effort, I successfully put it into practice and managed to increase the batch_size from 48 to 144! However, during further testing and organization, I discovered that this feature is broken in TensorFlow 2.x... Consequently, I spent another two days searching through various resources and debugging repeatedly. Finally, I successfully addressed this deficiency.

The following is my own open-source implementation:

Github Address: https://github.com/bojone/keras_recompute

This implementation has been integrated into bert4keras. Users of bert4keras can upgrade to the latest version (0.7.5+) to test this feature.

Usage

My implementation is also named recompute_grad. It is a decorator used to wrap the call function of a custom Keras layer, for example:

from recompute import recompute_grad

class MyLayer(Layer):
    @recompute_grad
    def call(self, inputs):
        return inputs * 2

For existing layers, you can decorate them through inheritance:

from recompute import recompute_grad

class MyDense(Dense):
    @recompute_grad
    def call(self, inputs):
        return super(MyDense, self).call(inputs)

After defining the custom layers, embed them into your code and set the environment variable RECOMPUTE=1 before execution to enable recomputation.

Note: Simply inserting @recompute_grad into the overall model will not achieve the goal of saving memory. Instead, you should insert @recompute_grad into each individual layer to better save VRAM. In short, the more @recompute_grad decorators you insert, the more VRAM you save. For specific reasons, please study the principles of recomputation carefully.

Effect

bert4keras 0.7.5+ has built-in recomputation support. Passing the environment variable RECOMPUTE=1 will enable it. Readers can try it themselves; the approximate effects are:

In the BERT Base version, the batch_size can be increased to about 3 times the original;
In the BERT Large version, the batch_size can be increased to about 4 times the original;
The average training time per sample increases by approximately 25%;
Theoretically, the more layers there are, the larger the multiplier for increasing batch_size.

Environment

Tests passed under the following environments:

tensorflow 1.14 + keras 2.3.1
tensorflow 1.15 + keras 2.3.1
tensorflow 2.0 + keras 2.3.1
tensorflow 2.1 + keras 2.3.1
tensorflow 2.0 + built-in tf.keras
tensorflow 2.1 + built-in tf.keras

Confirmed unsupported environment:

tensorflow 1.x + built-in tf.keras

Reports of more test results are welcome.

By the way, it is strongly recommended to use Keras 2.3.1 in conjunction with TensorFlow 1.x/2.x, and strongly discouraged to use the tf.keras that comes with TensorFlow 2.x.

References

Finally, my implementation mainly refers to the following two source codes, for which I express my gratitude: