Faster TensorFlow Training: Mixed Precision, XLA, and the TensorBoard Profiler

Once your input pipeline is fast enough to keep the accelerator fed, the next question is whether the accelerator itself is doing useful work. Most TensorFlow training jobs we profile are leaving somewhere between 30% and 60% of the available throughput on the table - not because the model is wrong, but because it runs in float32 with unfused kernels on hardware that was designed for mixed precision and compiler fusion.

This tutorial covers the two changes that give the largest speedup for the least code: mixed precision and XLA compilation. It assumes you have already ruled out an input bottleneck; if model.fit is waiting on data, start with our tf.data tutorial instead, because neither of these techniques will help you.

Step 1: measure before you change anything

Never enable a performance flag you cannot measure. Capture a profile with the TensorBoard Profiler callback:

import keras, tensorflow as tf

tb = keras.callbacks.TensorBoard(
    log_dir="logs/baseline",
    profile_batch=(20, 40),   # skip warmup, capture 20 steady-state steps
)
model.fit(train_ds, epochs=1, callbacks=[tb])

pip install -U tensorboard-plugin-profile
tensorboard --logdir logs

In the Profile tab, the Overview Page classifies your step time. Two numbers matter:

• Device (accelerator) time vs host time. If the device is idle waiting on the host, you are input-bound - stop here and fix the pipeline.
• Kernel utilisation / Tensor Core eligibility. The TensorFlow Stats and Kernel Stats pages show which ops dominate and whether matmul and convolution kernels are running in reduced precision. A job with 0% Tensor Core usage on an NVIDIA A100 or H100 is the classic candidate for this tutorial.

Also record a hard baseline number you can compare against: steady-state steps per second, and validation metric after a fixed number of steps. Speed that costs accuracy is not speed.

Step 2: turn on mixed precision

Mixed precision keeps variables in float32 but runs the heavy math in a 16-bit type. On NVIDIA GPUs from Volta onwards this activates Tensor Cores; on TPUs and on recent GPUs bfloat16 is the preferred type because its exponent range matches float32 and it needs no loss scaling.

import keras

# NVIDIA GPU (Volta / Turing / Ampere / Hopper)
keras.mixed_precision.set_global_policy("mixed_float16")

# TPU, or Ampere+ where you prefer range over precision
# keras.mixed_precision.set_global_policy("mixed_bfloat16")

Set the policy before you build the model. Every layer created afterwards computes in 16-bit and stores its weights in float32.

Keep the output layer in float32

The single most common bug is a softmax or loss computed in float16, which overflows or saturates. Force the final activation back to float32:

outputs = keras.layers.Dense(num_classes, dtype="float32")(x)
# or, if you want the logits in 16-bit:
# x = keras.layers.Dense(num_classes)(x)
# outputs = keras.layers.Activation("softmax", dtype="float32")(x)

Verify after building:

print(model.layers[-1].dtype_policy)   # expect <DTypePolicy "float32">
print(model.layers[-2].dtype_policy)   # expect <DTypePolicy "mixed_float16">

Loss scaling

With mixed_float16, small gradients underflow to zero. Keras handles this for you when you use model.fit: the policy wraps your optimizer in a LossScaleOptimizer with dynamic scaling. If you write a custom training loop, you must do it yourself:

optimizer = keras.mixed_precision.LossScaleOptimizer(
    keras.optimizers.AdamW(learning_rate=1e-3)
)

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        preds = model(x, training=True)
        loss = loss_fn(y, preds)
        scaled_loss = optimizer.scale_loss(loss)
    scaled_grads = tape.gradient(scaled_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(scaled_grads, model.trainable_variables))
    return loss

mixed_bfloat16 does not need loss scaling at all - one reason we default to it on TPU and on Ampere-class GPUs when the model is numerically touchy.

Sizing rule

Tensor Cores prefer dimensions that are multiples of 8 (float16) or 16. A Dense(1000) and a batch size of 60 will silently fall back to slower paths; Dense(1024) with batch size 64 will not. When you are already rewriting the model head, round the numbers.

Step 3: compile with XLA

XLA fuses chains of elementwise ops into single kernels, removes intermediate tensors, and picks better layouts. In Keras 3 you enable it per model:

model.compile(
    optimizer=keras.optimizers.AdamW(1e-3),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
    jit_compile=True,
)

For a custom tf.function, use @tf.function(jit_compile=True).

What to expect:

• Big wins on models dominated by many small elementwise ops - transformers with layer norms and activations between matmuls, RNNs, anything with heavy per-token arithmetic.
• Modest or no wins on models that are already one big convolution stack calling into cuDNN.
• Slower first steps. Compilation happens on the first call for each input shape. Budget 30-120 seconds of warmup, and exclude those steps from your benchmark.

The dynamic shape trap

XLA recompiles for every distinct input shape. A dataset whose last batch is ragged, or variable-length sequences, will trigger recompilation on nearly every step and can make the job slower than eager. Fix it at the pipeline level:

train_ds = train_ds.batch(BATCH, drop_remainder=True)   # fixed batch dimension
# and pad sequences to fixed buckets rather than to per-batch max length

If you see repeated XlaCompile entries throughout the profile rather than only at the start, this is what is happening.

Step 4: re-measure, then check accuracy

Profile again into logs/tuned and compare the two runs side by side in TensorBoard. On a typical transformer fine-tune on an A100 we expect something in this range:

Configuration	Relative step time	Peak memory
float32, no XLA	1.00x	1.00x
mixed_float16	0.55-0.70x	~0.6x
mixed_float16 + XLA	0.40-0.60x	~0.55x

Your numbers will differ - the point is that you now have your numbers, not ours.

Then validate quality. Train the baseline and the tuned configuration for the same number of steps on the same seed and compare the validation curve, not just the final number. Warning signs:

• Loss becomes NaN in the first few hundred steps -> an op that should be float32 is running in float16, or your learning rate was already at the edge. Check the head dtype first.
• Loss plateaus higher than baseline -> try mixed_bfloat16, or cast a specific layer with dtype="float32" on the layer that is losing precision (normalisation layers and attention logits are the usual suspects).
• Metrics look fine but occasionally spike -> dynamic loss scaling is skipping steps. That is normal in small numbers; check optimizer.loss_scale is not collapsing.

The freed memory is itself a result: mixed precision typically lets you raise the batch size 1.5-2x, which compounds with the per-step gain. Re-tune the learning rate when you do, and remember that larger batches interact with the sharding decisions in our multi-GPU scaling guide.

Step 5: keep the settings honest

Two operational notes we insist on with clients:

1. Record the precision policy and jit_compile flag with the run. A model trained in mixed_float16 and later re-trained in float32 will produce slightly different weights; if you cannot say which produced the artifact in production, you cannot reproduce a regression.
2. Inference is a separate decision. Mixed precision at training time does not change your SavedModel's serving behaviour - weights are still float32. If you want faster inference, that is post-training quantization, which is a different tool with different accuracy risk. See our serving tutorial and the Edge AI & LiteRT work.

Checklist

1. Profile first; confirm the job is compute-bound, not input-bound.
2. Record baseline steps/second and a validation curve.
3. set_global_policy("mixed_float16") (or mixed_bfloat16) before building the model.
4. Force the output layer or final activation to float32.
5. Use LossScaleOptimizer if you have a custom training loop.
6. Round layer widths and batch sizes to multiples of 8 or 16.
7. Add jit_compile=True; fix dynamic shapes with drop_remainder=True and bucketed padding.
8. Re-profile, compare accuracy, then raise the batch size with the memory you freed.
9. Log the precision policy and compile flags alongside the checkpoint.

Reference material: the Keras mixed precision guide and the TensorFlow XLA documentation.

If your training runs are expensive and you want a second opinion before buying more GPUs, this is usually a one- or two-day engagement: we profile, apply the changes above, and hand back a before/after report. Contact us with your model type and hardware, or read more about how our TensorFlow consultants work.