tf2.4 mixed_precision with float16 return 0 gradient

Question

This question was posted here before, and I re-opened it here to draw more attentions.

The main issue is that when testing in normal float32 env, the tensorflow returns gradient like, but after I shift to mixed_precision.set_global_policy('mixed_float16') with float16, the returned gradient is always 0.

Below is a toy code that can reproduce the error.

System information

OS Platform and Distribution: linux TensorFlow version (use command below): tf2.4.1

Reproducing code

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import mixed_precision
import numpy as np
from tqdm import tqdm

gpus = tf.config.experimental.list_physical_devices('GPU')

for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)


mixed_precision.set_global_policy('mixed_float16')


def forward_conv(x, filters, kernels, name='forward', padding='same'):
    i = 0
    for flt, kernel in zip(filters, kernels):
        x = layers.Conv3D(flt, kernel, activation='relu', padding=padding, dilation_rate=(1, 1, 1),
                          use_bias=False, name=str(i) + '_' + name)(x)
        x = layers.BatchNormalization(name=str(i) + '_bn_' + name)(x)
        i += 1
    return x


def part_one(ipt):
    l1 = forward_conv(ipt, (4, 4), (3, 3), name='enc1')
    d2 = layers.MaxPool3D(pool_size=(2, 2, 2))(l1)
    l2 = forward_conv(d2, (4, 4), (3, 3), name='enc2')
    return l1, l2


def part_inner(ipt1, ipt2):
    l1 = forward_conv(ipt1, (4, 4), (3, 3), name='enc1')
    l2 = forward_conv(ipt2, (4, 4), (3, 3), name='enc2')
    return l1, l2


def part_two(ipt1, ipt2):
    l2 = forward_conv(ipt2, (4, 4), (3, 3), name='dec2')
    u1 = layers.UpSampling3D(size=(2, 2, 2))(l2)
    r1 = forward_conv(ipt1 + u1, (4, 4), (3, 3), name='dec1')
    return r1


initial = tf.ones([1, 256, 368, 368, 1], dtype=tf.float16)

tf.random.set_seed(1)

with tf.GradientTape() as g:
    g.watch(initial)
    l1_, l2_ = part_one(initial)
    for _ in range(2):
        l1_, l2_ = part_inner(l1_, l2_)
    opt_ = part_two(l1_, l2_)
    loss = tf.reduce_mean(l1_) + tf.reduce_mean(opt_)
    gd = g.gradient(loss, initial)
    print('-' * 100)
    print(f'loss is {loss} and grad is {np.sum(gd)} with ckpt= {ckpt}')

Behavior description

When using tf.float32 setting, the result of gradient is reasonable with value around 0.6, however, when shift to tf.float16 with mixed_precision, the gradient is constantly 0. Should we expect that the computed gradient is so different between normal float32 mode and mixed_precision float16 mode? Thank you!

Answer 1

In the Tensorflow documentation regarding mixed_precision they talk about using loss scaling in order to deal with this problem.

Since documentation often becomes obsolete with tensorflow, here's the suggested code :

loss_scale = 1024
loss = model(inputs)
loss *= loss_scale
# Assume `grads` are float32. You do not want to divide float16 gradients.
grads = compute_gradient(loss, model.trainable_variables)
grads /= loss_scale

This shouuld fix the problem.