Diagnosing NaN Losses After Enabling TensorFlow Mixed‑Precision

Problem Statement

When you enable tf.keras.mixed_precision.set_global_policy(\"mixed_float16\"), training may start normally but after a handful of steps the loss turns NaN or Inf. GPU utilization stays high and the logs show no obvious errors. This symptom indicates that mixed‑precision arithmetic is producing numerical instability, typically during gradient scaling or an unsupported operation.

Typical Causes

Cause	Why It Happens	Diagnostic Check
Missing or mis‑configured loss scaling	Gradients underflow to zero in float16 and become NaN when accumulated.	Check optimizer.loss_scale and use tf.debugging.check_numerics on losses.
Optimizer without loss scaling support	Custom optimizers that don’t call get_scaled_loss() produce unscaled gradients.	Verify optimizer is wrapped in tf.keras.mixed_precision.LossScaleOptimizer.
Unsupported ops in float16	Operations like tf.math.divide or tf.nn.l2_loss can produce Inf/NaN in float16.	Search for such ops in the model graph.
Dynamic loss scale too aggressive	The auto‑adjusted scale jumps to Inf or zero.	Log optimizer.loss_scale per step.
Shape‑cast issues with experimental_relax_shapes	TensorFlow may cast to float16 incorrectly before a layer.	Disable the flag or cast explicitly to float32.

Diagnostic Checklist

1. Run a short training loop with policy='float32' and confirm loss remains finite. If it does, the issue is precision‑related.
2. Set the global policy to mixed_float16 and wrap the optimizer:

   policy = tf.keras.mixed_precision.Policy(\"mixed_float16\")
    tf.keras.mixed_precision.set_global_policy(policy)
    base_optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
    optimizer = tf.keras.mixed_precision.LossScaleOptimizer(base_optimizer, loss_scale=\"dynamic\")
   

3. Insert tf.debugging.check_numerics(loss, \"Loss\") before and after scaling to catch NaNs early.
4. Log optimizer.loss_scale each step:

   tf.print(\"Step\", step, \"Loss scale\", optimizer.loss_scale)
   

5. Search the model for unsupported ops. Replace them with float32 equivalents or cast inputs:

   x = tf.cast(x, tf.float32)
   # compute op
   x = tf.cast(x, tf.float16)
   

6. If using tf.distribute.Strategy, ensure loss scaling is applied inside the strategy scope.

Step‑by‑Step Fixes

• Wrap the optimizer: If you’re using a custom optimizer, add

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

• Use a fixed loss scale when dynamic scaling misbehaves:

  optimizer = tf.keras.mixed_precision.LossScaleOptimizer(base_optimizer, loss_scale=256)
  

• Cast problematic ops to float32:

  y = tf.nn.l2_loss(tf.cast(x, tf.float32))
  

• Disable experimental_relax_shapes if you suspect shape‑related casts:

  tf.config.experimental_run_functions_eagerly(True)
  

Concrete Example: MNIST with LossScaleOptimizer

# train_mnist.py
import tensorflow as tf

# 1. Set mixed‑precision policy
policy = tf.keras.mixed_precision.Policy(\"mixed_float16\")
tf.keras.mixed_precision.set_global_policy(policy)

# 2. Build model
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation=\"relu\"),
    tf.keras.layers.Dense(10, activation=\"softmax\")
])

# 3. Wrap optimizer
base_opt = tf.keras.optimizers.Adam(learning_rate=1e-3)
optimizer = tf.keras.mixed_precision.LossScaleOptimizer(base_opt, loss_scale=\"dynamic\")

# 4. Compile
model.compile(optimizer=optimizer,
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=[\"accuracy\"])

# 5. Load data
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()

# 6. Train
model.fit(x_train, y_train, epochs=5, batch_size=256)

Run the script on a GPU:

python train_mnist.py

Check the console for optimizer.loss_scale values. If they stay in a reasonable range and no NaNs appear, the configuration is correct.

Monitoring & Escalation

• Use tf.debugging.check_numerics to automatically abort on NaN/Inf.
• If loss remains NaN after applying all fixes, consider downgrading to a static loss scale or reverting the mixed‑precision policy for the offending layer.
• For distributed training, enable tf.keras.mixed_precision.experimental.LossScaleOptimizer inside the strategy scope and verify that gradient aggregation preserves the scale.

Summary

NaN losses after enabling mixed‑precision usually stem from missing or incorrect loss scaling, unsupported ops, or aggressive dynamic scaling. By systematically verifying the policy, wrapping the optimizer, checking for unsupported ops, and monitoring the loss scale, you can isolate the root cause and apply the appropriate fix. Always validate against a float32 baseline to confirm that the issue is precision‑related.