Diagnosing and Fixing NaN Loss in Keras Mixed‑Precision Training

Problem Overview

When you enable tf.keras.mixed_precision and start training, you may see the loss drop to nan after a few steps. The training stalls, metrics diverge, and the model never converges. This guide walks through the most common culprits, a quick diagnostic checklist, and targeted fixes.

Root Causes in One Table

Cause	Typical Symptoms	Quick Check
Unsupported layers (e.g., Lambda, custom ops)	Loss becomes NaN immediately after the first batch.	Replace the layer with a Keras equivalent or cast to float32.
Missing or mis‑configured loss scaling	Gradients explode, loss spikes to inf then NaN.	Print tf.keras.mixed_precision.global_policy().loss_scale.
Optimizer incompatibility (e.g., AdamW)	Training stops after a few steps, loss stays at NaN.	Switch to Adam or add manual scaling.
Premature float16 data cast in the pipeline	Gradients are NaN from the start of training.	Ensure tf.data pipelines output float32 tensors.
Inconsistent dtype in loss/metrics	NaNs appear only when evaluating metrics.	Force float32 dtype in loss function.

Step‑by‑Step Diagnostic Flow

1. Confirm Mixed‑Precision Policy

   import tensorflow as tf
   from tensorflow.keras import mixed_precision
   mixed_precision.set_global_policy('mixed_float16')
   print(mixed_precision.global_policy())
   # Expected output: mixed_float16 with loss_scale=Dynamic
   

   Running this on a machine without a compatible GPU will silently fall back to float32, so verify you’re on a Volta or newer NVIDIA card.

2. Check Loss Scaling State

   policy = mixed_precision.global_policy()
   print('Loss scale:', policy.loss_scale)
   # Should be "Dynamic" or a numeric value, not None
   

   If loss_scale is None, loss scaling is disabled and NaNs are likely due to overflow.

3. Run a Minimal Model

   Use a tiny model on a small dataset to isolate the problem.

   import numpy as np
   from tensorflow.keras import layers, models
   
   x = np.random.rand(32, 28, 28, 1).astype('float32')
   y = np.random.randint(0, 10, 32)
   
   def build_model():
       return models.Sequential([
           layers.Conv2D(32, 3, activation='relu'),
           layers.Flatten(),
           layers.Dense(10)
       ])
   
   model = build_model()
   model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))
   model.fit(x, y, epochs=1)
   

   If this runs without NaNs, the issue is in your custom layers or data pipeline.

4. Inspect Custom Layers

   Wrap any Lambda or user‑defined op with explicit tf.cast to float32 before computation, then cast back.

   def my_layer(x):
       x32 = tf.cast(x, tf.float32)
       # compute in float32
       y = tf.math.sin(x32)
       return tf.cast(y, tf.float16)
   

5. Validate Data Pipeline Dtypes

   Check that tf.data mapping functions do not cast to float16 before the model consumes them.

   dataset = tf.data.Dataset.from_tensor_slices((x, y))
   # Ensure dtype remains float32
   for batch in dataset.take(1):
       print(batch[0].dtype)
   

6. Test with mixed_float32

   Temporarily switch to mixed_float32 to see if NaNs disappear. If they do, the problem is definitely related to precision.

   mixed_precision.set_global_policy('mixed_float32')
   

7. Check Optimizer Compatibility

   Some optimizers, like AdamW from tensorflow_addons, do not apply loss scaling automatically. Replace with Adam or wrap the optimizer with mixed_precision.LossScaleOptimizer:

   optimizer = tf.keras.optimizers.Adam()
   optimizer = mixed_precision.LossScaleOptimizer(optimizer)
   model.compile(optimizer=optimizer, loss=...)
   

8. Monitor Gradient Norms

   Large batch sizes can still overflow even with scaling. Log the norm of gradients after each step:

   @tf.function
   def train_step(x, y):
       with tf.GradientTape() as tape:
           logits = model(x, training=True)
           loss = loss_fn(y, logits)
       grads = tape.gradient(loss, model.trainable_variables)
       norm = tf.linalg.global_norm(grads)
       print('Grad norm:', norm)
       optimizer.apply_gradients(zip(grads, model.trainable_variables))
   

9. Escalation Criteria
   1. If NaNs persist after all above checks, isolate by removing custom layers one by one.
   2. Run the same model on a CPU to confirm the issue is GPU‑specific.
   3. Submit a minimal reproducible example to the TensorFlow GitHub issue tracker with TensorFlow 2.16 and GPU details.

Practical Example: Fixing a NaN in a Custom Lambda Layer

Suppose you have the following layer in a model:

def custom_activation(x):
    return tf.nn.relu(x) * tf.math.exp(-x)

model = models.Sequential([
    layers.Dense(64, activation='relu'),
    layers.Lambda(custom_activation),
    layers.Dense(10)
])

When training with mixed_float16, the exponential can under‑flow, producing NaNs. Fix: cast to float32 inside the lambda and cast back.

def custom_activation(x):
    x32 = tf.cast(x, tf.float32)
    y = tf.nn.relu(x32) * tf.math.exp(-x32)
    return tf.cast(y, tf.float16)

Re‑run training; the loss stays finite.

Key Takeaways

• Always verify the global mixed‑precision policy and loss scaling before training.
• Unsupported ops and premature dtype changes are the most common NaN triggers.
• Use mixed_float32 as a quick sanity check; if NaNs vanish, precision is the culprit.
• Wrap incompatible optimizers with LossScaleOptimizer or switch to a supported one.
• Monitor gradient norms to catch hidden overflows early.

Final Checklist for Release

1. [ ] Mixed‑precision policy set to mixed_float16 and loss scaling active.
2. [ ] All custom layers cast to float32 internally.
3. [ ] Data pipeline outputs float32 tensors.
4. [ ] Optimizer is compatible or wrapped with LossScaleOptimizer.
5. [ ] Gradient norms stay below a reasonable threshold (e.g., 1e4).
6. [ ] NaN loss is not present after the first epoch.

Conclusion

Mixed‑precision can boost performance dramatically, but it introduces subtle dtype pitfalls. By following the ordered checks above, you can quickly isolate NaN loss issues and apply the appropriate fix—be it a dtype cast, loss‑scaling adjustment, or optimizer change.