Using Mixed Precision and XLA to Speed Up Training in bert4keras

在bert4keras中使用混合精度和XLA加速训练 苏剑林 · 2022-04-28 · Machine Learning model optimization gradient

Up until now I've mostly focused on model design and implementation, and rarely paid much attention to training acceleration techniques. Things like mixed precision and XLA — I'd heard of them, but never actually put them into practice. Over the past couple of days I did some tinkering and successfully got mixed precision and XLA working together in bert4keras to speed up training. Here's a short write-up for reference.

Most of the takeaways in this post aren't specific to bert4keras. The reason I highlight bert4keras in the title is simply that the model implementations in bert4keras are relatively well-organized, so getting these acceleration tricks working requires fewer modifications there.

Experimental Setup

The GPU used in this post's experiments is a 3090, running the Docker image nvcr.io/nvidia/tensorflow:21.09-tf1-py3, which comes with TensorFlow 1.15.5. The bert4keras version used in the experiments is 0.11.3. You can set up other environments by analogy, but do keep an adventurous spirit — don't expect things to just work without any tinkering.

By the way, cards like the 3090 and A100 only work with CUDA 11, and the official TensorFlow 1.15 release doesn't support CUDA 11. So if you still want to use TensorFlow 1.x, your only option is nvidia-tensorflow, maintained directly by NVIDIA, or the Docker image built from it. Using the NVIDIA-maintained version instead of Google's not only lets you run TF 1.x on the latest GPUs, but also gives you some extra optimizations that NVIDIA specifically implemented — see the details here. more

Don't be the person who says "TensorFlow is already up to 2.8, why are you still using 1.15." Your GPU is made by NVIDIA, so which version of TensorFlow works best isn't up to you or me to decide — not even Google gets a say. NVIDIA gets the final say, and NVIDIA is still maintaining 1.15, which tells you that 1.15 is the real yyds (goat).

Mixed Precision

Let's start with mixed-precision training. Simply put, this means computing with FP16 while doing parameter updates and storage in FP32. The representable range of FP16 is roughly $6\times 10^{-8}\sim 65504$, and both bounds are ones we can realistically run into when implementing models. So the biggest problem introduced by FP16 is overflow and precision loss. For the detailed theory, please search around on your own — this post focuses mainly on how to actually use it.

The nvidia-tensorflow documentation introduces mixed-precision training here. The simplest way to enable it is to add an environment variable at the top of your script:

import os
os.environ['TF_KERAS'] = '1'  # 必须使用tf.keras
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE'] = '1'  # 混合精度训练

You may have noticed that most tutorials mention TF_ENABLE_AUTO_MIXED_PRECISION, whereas here I use TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE. The difference is that the former automatically adds "dynamic loss scaling," while the latter doesn't. However, in my testing I found that dynamic loss scaling can't fully replace manual loss adjustment, so I just skip that feature altogether.

After adding the environment variable, restart your training script and see what happens. If you get NaNs right from the start, try adjusting the infinity and epsilon values:

from bert4keras.backend import K
K.set_infinity(1e4)
K.set_epsilon(1e-5)

After this adjustment, you usually won't get NaNs immediately (if you still do, check whether some other part of the model uses an infinity or epsilon value not controlled by these two functions, and fix that too). But you might still see the loss decrease at first, then rise, and eventually go to NaN. This happens because of poor initialization, or because — as in DeepNet — the model is deliberately designed so that some parameters have extremely small gradients (smaller than $10^{-8}$). Within FP16 precision, such gradients round straight down to zero, so those parameters never get updated — or equivalently, the gradients become inaccurate. Training for a long time with inaccurate gradients tends to lead to non-convergence.

The solution here is "loss scaling." We can simply multiply the loss function by a scaling factor (say, 1000 — feel free to tune this, and as large as possible without causing NaNs is generally better). This amplifies gradients that would otherwise be too small, bringing them back into the representable FP16 range instead of being rounded to zero, thereby avoiding precision loss in the gradients. For optimizers we commonly use, like Adam or LAMB, multiplying the loss by a constant doesn't change their training dynamics at all — in other words, they are fully compatible with loss scaling.

In fact, I found that the "loss scaling" trick is useful not only in mixed-precision training, but also gives some benefit even when training in full FP32 precision: when training in full FP32 without loss scaling, the model tends to get stuck at some loss value for a while at the start before slowly decreasing; but with loss scaling applied, the model keeps a steady downward trend from the very beginning, converging comparatively faster.

Algebraic Acceleration

Now let's look at XLA, short for "Accelerated Linear Algebra" — as the name suggests, it's specifically designed to speed up linear algebra computations. In simple terms, XLA compiles and optimizes the computation graph ahead of time: it fuses operators that can be merged (reducing the number of cached intermediate variables to save memory) and parallelizes operators that can run concurrently (improving compute speed).

In nvidia-tensorflow, the simplest way to enable XLA is again via an environment variable:

import os
os.environ['TF_KERAS'] = '1'  # 必须使用tf.keras
os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=1'  # 启用XLA

But be aware that XLA doesn't guarantee a speedup. As mentioned, XLA tries to parallelize as many operators as it can, which is clearly a strategy that trades space for time — so enabling XLA might consume more GPU memory and lead to OOM errors, or even cause a performance regression if the parallel clusters get too large. The official documentation gives a fairly thorough analysis of the possible issues and offers corresponding recommendations. My personal recommendation is to add the --tf_xla_enable_lazy_compilation=false parameter:

import os
os.environ['TF_KERAS'] = '1'  # 必须使用tf.keras
os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=1'  # 启用XLA
os.environ['TF_XLA_FLAGS'] += ' --tf_xla_enable_lazy_compilation=false'  # 优化XLA

If that still doesn't solve it, try switching to XLA Lite:

import os
os.environ['TF_KERAS'] = '1'  # 必须使用tf.keras
os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=fusible'  # 启用XLA Lite

If switching to XLA Lite still doesn't fix things, that basically means XLA isn't a good fit for your model.

Performance Comparison

On the 3090, enabling mixed-precision training gives a speedup of a little over 10%. This might be less than you'd expect, and my guess is that on newer cards like the 3090 and A100, the default FP32 format is actually implemented internally using a format called TF32 (see here). In a sense, TF32 is itself already a kind of "half-precision format," and it's faster than plain FP32. In other words, FP32 on the 3090 is already somewhat half-precision-optimized under the hood, so it's already fast to begin with — which is why the additional gain from switching to mixed precision is comparatively smaller.

As for the gain from XLA, it's roughly 15%. In my training script, setting the TF_XLA_FLAGS environment variable directly to --tf_xla_auto_jit=1 caused an OOM; adding --tf_xla_enable_lazy_compilation=false still didn't help; but switching to --tf_xla_auto_jit=fusible let training proceed normally.

Finally — and most importantly — mixed precision and XLA can be stacked together! Using both at once gives a combined speedup of roughly 30%, and adding mixed precision essentially cancels out the extra GPU memory consumption introduced by XLA. The two really do complement each other nicely.

Summary

This post described my experience using mixed precision and XLA to accelerate training in bert4keras. Enabling both together gives roughly a 30% speedup on a 3090.