Mixed Precision Training

2-3× faster training, 40% less memory, ~free

Mixed precision training uses lower-precision arithmetic (float16 or bfloat16) for compute while keeping float32 for weight storage. Memory savings + throughput improvements on modern GPUs and TPUs with minimal accuracy impact.

Why this works: neural networks are surprisingly tolerant of reduced precision in forward/backward passes. Gradients computed in float16 are accurate enough for the optimizer. But weights stay float32 so small gradient updates (much smaller than float16 epsilon) don't get lost.

Policy	Compute dtype	Variable dtype	Best for
mixed_float16	float16	float32	NVIDIA GPU (V100, A100, H100)
mixed_bfloat16	bfloat16	float32	TPU (all versions)
float32	float32	float32	Default; or when stability matters

Critical detail for mixed_float16: the final activation must be float32 for numerical stability. Always set dtype='float32' on the output Activation layer (e.g., softmax). For TPUs (bfloat16) this isn't needed because bfloat16 has the same exponent range as float32.

Code

Mixed precision policy 설정·python

import tensorflow as tf
from tensorflow.keras import mixed_precision, layers, Input, Model

# Set BEFORE building the model
mixed_precision.set_global_policy('mixed_float16')   # GPU
# mixed_precision.set_global_policy('mixed_bfloat16')  # TPU

policy = mixed_precision.Policy('mixed_float16')
print(f"Compute dtype: {policy.compute_dtype}")    # float16
print(f"Variable dtype: {policy.variable_dtype}")  # float32

# Build model — layers auto-use float16 for compute
inputs  = Input(shape=(784,))
x       = layers.Dense(4096, activation='relu')(inputs)
x       = layers.Dense(4096, activation='relu')(x)
x       = layers.Dense(10)(x)

# IMPORTANT: final activation must be float32 for numerical stability
outputs = layers.Activation('softmax', dtype='float32')(x)

model = Model(inputs, outputs)
model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=tf.keras.optimizers.Adam(),
    metrics=['accuracy'],
)

# Larger batches benefit more from float16 throughput
history = model.fit(x_train, y_train, batch_size=8192, epochs=5,
                    validation_split=0.2)

Custom loop with LossScaleOptimizer·python

import tensorflow as tf
from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy('mixed_float16')

optimizer = tf.keras.optimizers.Adam()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)  # prevent fp16 underflow

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
        scaled_loss = optimizer.get_scaled_loss(loss)

    scaled_grads = tape.gradient(scaled_loss, model.trainable_variables)
    grads = optimizer.get_unscaled_gradients(scaled_grads)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# Note: bfloat16 (TPU) does NOT need LossScaleOptimizer
# because bfloat16 has the same exponent range as float32

2-3× faster training, 40% less memory, ~free

Code

External links

Progress

Comments 0