The memory-saving recomputation trick now has a Keras version too

节省显存的重计算技巧也有了Keras版了 苏剑林 · 2020-04-29 · Machine Learning model deep learning Keras

Quite a few readers may recently have noticed the WeChat article "BERT Recomputation: Save 5x GPU Memory at the Cost of 22.5% More Training Time (with Code)", which introduced a technique called "recomputation." In short, it's a method for saving GPU memory by making the average training step a bit slower, in exchange for being able to increase the batch size several times over. This trick was first published in the paper "Training Deep Nets with Sublinear Memory Cost", and it was actually proposed back in 2016—it just doesn't seem to have become especially popular since.

Exploration

The WeChat article mentions that this trick already has native implementations in PyTorch and PaddlePaddle, but not in TensorFlow. In fact, though, TensorFlow has had this functionality built in since version 1.8, where it was originally included in the tf.contrib sub-library. Starting from TensorFlow 1.15, it was folded into TensorFlow's main namespace as one of its core functions—namely tf.recompute_grad.

After finding tf.recompute_grad, I spent some time figuring out how to use it, and after a fair bit of fiddling, I actually managed to get it working—successfully bumping batch_size from 48 all the way up to 144! However, while continuing to organize and test the code, I discovered that this thing is actually broken under TensorFlow 2.x... So I spent another two days digging through various resources and debugging repeatedly, and eventually managed to patch up this shortcoming.

Here is my own open-source implementation:

GitHub: https://github.com/bojone/keras_recompute

This implementation has already been built into bert4keras; readers who use bert4keras can upgrade to the latest version (0.7.5+) to try it out. more

Usage

My implementation is also named recompute_grad. It's 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 layers that already exist, you can apply the decorator via inheritance:

from recompute import recompute_grad

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

Once you've defined the custom layer, embed it in your code as usual, and before running your script, add the environment variable RECOMPUTE=1 to enable recomputation.

Note: simply inserting @recompute_grad somewhere in the overall model isn't enough to achieve memory savings—you need to insert @recompute_grad into every single layer to get the best memory savings. Put simply, the more instances of @recompute_grad you insert, the more memory you save. For the exact reason why, please make sure you understand how recomputation actually works under the hood.

Results

bert4keras 0.7.5+ already has recomputation built in—just pass in the environment variable RECOMPUTE=1 to enable it. Feel free to try it out yourself; roughly speaking, the effects are:

1. For BERT Base, the batch size can be increased to roughly 3x the original;
2. For BERT Large, the batch size can be increased to roughly 4x the original;
3. Average training time per sample increases by about 25%;
4. In theory, the more layers a model has, the larger the multiplier by which the batch size can be increased.

Environment

Tested and confirmed working in 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 NOT supported:

TensorFlow 1.x + built-in tf.keras

More test reports from readers are welcome.

By the way, I strongly recommend running Keras 2.3.1 together with TensorFlow 1.x/2.x, and strongly recommend against using the built-in tf.keras that ships with TensorFlow 2.x.

References

Finally, my implementation is mainly based on the following two pieces of source code—many thanks to their authors.

https://github.com/davisyoshida/tf2-gradient-checkpointing
https://github.com/tensorflow/tensorflow/blob/v2.1.0/tensorflow/python/ops/custom_gradient.py#L454-L499