Mix-precisioned training about athena

[tjadamlee]

@iou2much thanks, @leixiaoning please help check this

[leixiaoning]

@iou2much thanks, @leixiaoning please help check this

ok, i will check. @tjadamlee

[leixiaoning]

Hi. We want to use it in mix-precisioned mode, as our GPU don't have much memory, and we want to speed up the training.

I change the code to use mix-precisioned training feature in TF2. It works for MPC (stage 1).
But for the fine-tuning stage, the loss becomes nan at the very beginning.
I try to debug it, and find out the PositionalEncoding in speech_transformer.py is always returning NaN.

        input_labels = layers.Input(shape=data_descriptions.sample_shape["output"], dtype=tf.int32)
        inner = layers.Embedding(self.num_class, d_model)(input_labels)
        inner = PositionalEncoding(d_model, scale=True)(inner) #it returns NaN
        inner = layers.Dropout(self.hparams.rate)(inner)
        self.y_net = tf.keras.Model(inputs=input_labels, outputs=inner, name="y_net")

could anyone help? Thanks a lot

did you use os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1' for supporting mixed precision training？ @iou2much

[iou2much]

nope. I add this before the model initialization.
tf.keras.mixed_precision.experimental.set_policy('mixed_float16')

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[stale]

This issue is closed. You can also re-open it if needed.