Repository: keras-team/keras-nlp
Pull request #1242 (2023-09-13 02:17:58)

Add a mixed precision test and fix mixed precision errors for layers

mattdangerw: /gcbrun
mattdangerw: Ah I think this needs a keras-core release to pass. Will sync with Francois.
jbischof: Is mixed precision working in Keras Core now?
mattdangerw: > Is mixed precision working in Keras Core now?

Yes, or at least it should be largely with the latest release. Landed the loss scaling optimizer, which was the main piece we were missing, as well as a few other fixes.
mattdangerw: /gcbrun

Patch:
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index 0e114691..2ea8f5b1 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -1,2 +1,2 @@
-FROM mcr.microsoft.com/vscode/devcontainers/python:3.8
+FROM mcr.microsoft.com/vscode/devcontainers/python:3.9
 COPY setup.sh /setup.sh
diff --git a/keras_nlp/layers/modeling/f_net_encoder.py b/keras_nlp/layers/modeling/f_net_encoder.py
index f8de353e..a5370d96 100644
--- a/keras_nlp/layers/modeling/f_net_encoder.py
+++ b/keras_nlp/layers/modeling/f_net_encoder.py
@@ -97,11 +97,13 @@ class FNetEncoder(keras.layers.Layer):
         # Layer Norm layers.
         self._mixing_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="mixing_layer_norm",
         )
         self._mixing_layer_norm.build(inputs_shape)
         self._output_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="output_layer_norm",
         )
         self._output_layer_norm.build(inputs_shape)
@@ -112,6 +114,7 @@ class FNetEncoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="intermediate_dense",
         )
         self._intermediate_dense.build(inputs_shape)
@@ -119,12 +122,17 @@ class FNetEncoder(keras.layers.Layer):
             feature_size,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="output_dense",
         )
         self._output_dense.build(
             self._intermediate_dense.compute_output_shape(inputs_shape)
         )
-        self._output_dropout = keras.layers.Dropout(rate=self.dropout)
+        self._output_dropout = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="output_dropout",
+        )
         self.built = True
 
     def call(self, inputs):
@@ -140,9 +148,12 @@ class FNetEncoder(keras.layers.Layer):
 
         def fourier_transform(input):
             # Apply FFT on the input and take the real part.
+            input_dtype = input.dtype
+            # FFT transforms do not support float16.
+            input = ops.cast(input, "float32")
             real_in, imaginary_in = (input, ops.zeros_like(input))
             real_out, _ = ops.fft2((real_in, imaginary_in))
-            return real_out
+            return ops.cast(real_out, input_dtype)
 
         def add_and_norm(input1, input2, norm_layer):
             return norm_layer(input1 + input2)
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
index 5bbcb491..eacee7e8 100644
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -141,10 +141,12 @@ class MaskedLMHead(keras.layers.Layer):
             activation=self.intermediate_activation,
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
+            dtype=self.dtype_policy,
             name="intermediate_dense",
         )
         self._intermediate_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="intermediate_layer_norm",
         )
         # The gather length does not affect any of our built variables, so
@@ -185,6 +187,7 @@ class MaskedLMHead(keras.layers.Layer):
             outputs = self.token_embedding(x, reverse=True)
         else:
             outputs = ops.matmul(x, self._kernel)
+        outputs = ops.cast(outputs, self.compute_dtype)
         outputs = outputs + self._bias
 
         # Apply a final activation.
diff --git a/keras_nlp/layers/modeling/reversible_embedding.py b/keras_nlp/layers/modeling/reversible_embedding.py
index 51de8dfd..e01d6698 100644
--- a/keras_nlp/layers/modeling/reversible_embedding.py
+++ b/keras_nlp/layers/modeling/reversible_embedding.py
@@ -49,6 +49,9 @@ class ReversibleEmbedding(keras.layers.Embedding):
             the `embeddings` matrix (see `keras.constraints`).
         mask_zero: Boolean, whether or not the input value 0 is a special
             "padding" value that should be masked out.
+        reverse_dtype: The dtype for the reverse projection computation.
+            For stability, it is usually best to use full precision even when
+            working with half or mixed precision training.
 
     Call args:
         inputs: The tensor inputs to the layer.
@@ -87,6 +90,7 @@ class ReversibleEmbedding(keras.layers.Embedding):
         embeddings_regularizer=None,
         embeddings_constraint=None,
         mask_zero=False,
+        reverse_dtype="float32",
         **kwargs,
     ):
         super().__init__(
@@ -99,6 +103,7 @@ class ReversibleEmbedding(keras.layers.Embedding):
             **kwargs,
         )
         self.tie_weights = tie_weights
+        self.reverse_dtype = reverse_dtype
 
     def build(self, inputs_shape=None):
         super().build(inputs_shape)
@@ -114,12 +119,12 @@ class ReversibleEmbedding(keras.layers.Embedding):
     def call(self, inputs, reverse=False):
         if reverse:
             if self.tie_weights:
-                reverse_embeddings = ops.transpose(
-                    ops.convert_to_tensor(self.embeddings)
-                )
+                kernel = ops.transpose(ops.convert_to_tensor(self.embeddings))
             else:
-                reverse_embeddings = self.reverse_embeddings
-            return ops.matmul(inputs, reverse_embeddings)
+                kernel = self.reverse_embeddings
+            inputs = ops.cast(inputs, self.reverse_dtype)
+            kernel = ops.cast(kernel, self.reverse_dtype)
+            return ops.matmul(inputs, kernel)
 
         return super().call(inputs)
 
@@ -128,6 +133,7 @@ class ReversibleEmbedding(keras.layers.Embedding):
         config.update(
             {
                 "tie_weights": self.tie_weights,
+                "reverse_dtype": self.reverse_dtype,
             }
         )
         return config
diff --git a/keras_nlp/layers/modeling/token_and_position_embedding.py b/keras_nlp/layers/modeling/token_and_position_embedding.py
index ad89af05..f3dffe34 100644
--- a/keras_nlp/layers/modeling/token_and_position_embedding.py
+++ b/keras_nlp/layers/modeling/token_and_position_embedding.py
@@ -93,11 +93,13 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
                 self.embeddings_initializer
             ),
             mask_zero=mask_zero,
+            dtype=self.dtype_policy,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             sequence_length=sequence_length,
             initializer=clone_initializer(self.embeddings_initializer),
+            dtype=self.dtype_policy,
             name="position_embedding",
         )
         self.supports_masking = self.token_embedding.supports_masking
diff --git a/keras_nlp/layers/modeling/transformer_decoder.py b/keras_nlp/layers/modeling/transformer_decoder.py
index 4117a82e..92734fca 100644
--- a/keras_nlp/layers/modeling/transformer_decoder.py
+++ b/keras_nlp/layers/modeling/transformer_decoder.py
@@ -111,7 +111,6 @@ class TransformerDecoder(keras.layers.Layer):
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
         normalize_first=False,
-        name=None,
         **kwargs,
     ):
         # Work around for model saving, we need to ensure our model is built
@@ -119,7 +118,7 @@ class TransformerDecoder(keras.layers.Layer):
         decoder_sequence_shape = kwargs.pop("decoder_sequence_shape", None)
         encoder_sequence_shape = kwargs.pop("encoder_sequence_shape", None)
 
-        super().__init__(name=name, **kwargs)
+        super().__init__(**kwargs)
         self.intermediate_dim = intermediate_dim
         self.num_heads = num_heads
         self.dropout = dropout
@@ -160,6 +159,7 @@ class TransformerDecoder(keras.layers.Layer):
             dropout=self.dropout,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="self_attention",
         )
         if hasattr(self._self_attention_layer, "_build_from_signature"):
@@ -174,11 +174,14 @@ class TransformerDecoder(keras.layers.Layer):
             )
         self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="self_attention_layer_norm",
         )
         self._self_attention_layer_norm.build(decoder_sequence_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="self_attention_dropout",
         )
 
         # Cross attention layers are optional.
@@ -191,6 +194,7 @@ class TransformerDecoder(keras.layers.Layer):
                 dropout=self.dropout,
                 kernel_initializer=clone_initializer(self.kernel_initializer),
                 bias_initializer=clone_initializer(self.bias_initializer),
+                dtype=self.dtype_policy,
                 name="cross_attention",
             )
             if hasattr(self._cross_attention_layer, "_build_from_signature"):
@@ -205,11 +209,14 @@ class TransformerDecoder(keras.layers.Layer):
                 )
             self._cross_attention_layer_norm = keras.layers.LayerNormalization(
                 epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
                 name="cross_attention_layer_norm",
             )
             self._cross_attention_layer_norm.build(encoder_sequence_shape)
             self._cross_attention_dropout = keras.layers.Dropout(
                 rate=self.dropout,
+                dtype=self.dtype_policy,
+                name="cross_attention_dropout",
             )
 
         # Feedforward layers.
@@ -218,25 +225,30 @@ class TransformerDecoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
-            name="intermediate_dense",
+            dtype=self.dtype_policy,
+            name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(decoder_sequence_shape)
         self._feedforward_output_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
-            name="output_dense",
+            dtype=self.dtype_policy,
+            name="feedforward_output_dense",
         )
         intermediate_shape = list(decoder_sequence_shape)
         intermediate_shape[-1] = self.intermediate_dim
         self._feedforward_output_dense.build(tuple(intermediate_shape))
         self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
-            name="output_layer_norm",
+            dtype=self.dtype_policy,
+            name="feedforward_layer_norm",
         )
         self._feedforward_layer_norm.build(decoder_sequence_shape)
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="feedforward_dropout",
         )
         # Create layers based on input shape.
         self.built = True
diff --git a/keras_nlp/layers/modeling/transformer_encoder.py b/keras_nlp/layers/modeling/transformer_encoder.py
index 6f8f2ad9..32cdd355 100644
--- a/keras_nlp/layers/modeling/transformer_encoder.py
+++ b/keras_nlp/layers/modeling/transformer_encoder.py
@@ -92,10 +92,9 @@ class TransformerEncoder(keras.layers.Layer):
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
         normalize_first=False,
-        name=None,
         **kwargs,
     ):
-        super().__init__(name=name, **kwargs)
+        super().__init__(**kwargs)
         self.intermediate_dim = intermediate_dim
         self.num_heads = num_heads
         self.dropout = dropout
@@ -125,6 +124,7 @@ class TransformerEncoder(keras.layers.Layer):
             dropout=self.dropout,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="self_attention_layer",
         )
         if hasattr(self._self_attention_layer, "_build_from_signature"):
@@ -139,17 +139,21 @@ class TransformerEncoder(keras.layers.Layer):
             )
         self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="self_attention_layer_norm",
         )
         self._self_attention_layer_norm.build(inputs_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
             name="self_attention_dropout",
         )
 
         # Feedforward layers.
         self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="feedforward_layer_norm",
         )
         self._feedforward_layer_norm.build(inputs_shape)
         self._feedforward_intermediate_dense = keras.layers.Dense(
@@ -157,6 +161,7 @@ class TransformerEncoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(inputs_shape)
@@ -164,6 +169,7 @@ class TransformerEncoder(keras.layers.Layer):
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="feedforward_output_dense",
         )
         intermediate_shape = list(inputs_shape)
@@ -171,6 +177,8 @@ class TransformerEncoder(keras.layers.Layer):
         self._feedforward_output_dense.build(tuple(intermediate_shape))
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="feedforward_dropout",
         )
         self.built = True
 
diff --git a/keras_nlp/samplers/beam_sampler.py b/keras_nlp/samplers/beam_sampler.py
index f6f4c8e4..9562f95d 100644
--- a/keras_nlp/samplers/beam_sampler.py
+++ b/keras_nlp/samplers/beam_sampler.py
@@ -16,7 +16,6 @@ import tensorflow as tf
 import tree
 
 from keras_nlp.api_export import keras_nlp_export
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.samplers.sampler import Sampler
 from keras_nlp.samplers.sampler import call_args_docstring
@@ -161,7 +160,7 @@ class BeamSampler(Sampler):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
             vocab_size = ops.shape(logits)[-1]
-            probs = keras.activations.softmax(logits / self.temperature)
+            probs = self.compute_probabilities(logits)
 
             # Compute the running log-likelihood of each new candidate.
             next_log_probs = ops.log(probs) + log_probs[..., None]
diff --git a/keras_nlp/samplers/contrastive_sampler.py b/keras_nlp/samplers/contrastive_sampler.py
index a5a91a79..bac65bcf 100644
--- a/keras_nlp/samplers/contrastive_sampler.py
+++ b/keras_nlp/samplers/contrastive_sampler.py
@@ -15,7 +15,6 @@
 import tree
 
 from keras_nlp.api_export import keras_nlp_export
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.samplers.sampler import Sampler
 from keras_nlp.samplers.sampler import call_args_docstring
@@ -131,7 +130,7 @@ class ContrastiveSampler(Sampler):
 
         def body(prompt, cache, index, logits, hidden_states):
             # Compute the softmax distribution for the next token.
-            probabilities = keras.activations.softmax(logits / self.temperature)
+            probabilities = self.compute_probabilities(logits)
 
             # Replicate for `self.k` times to find the best token in top-k
             # candidates.
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index f7f08322..e28fbe9d 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -143,7 +143,7 @@ class Sampler:
         def body(prompt, cache, index):
             # Compute the softmax distribution for the next token.
             logits, _, cache = next(prompt, cache, index)
-            probabilities = keras.activations.softmax(logits / self.temperature)
+            probabilities = self.compute_probabilities(logits)
             # Compute the next token.
             next_token = self.get_next_token(probabilities)
             # Don't overwrite anywhere mask is True.
@@ -164,6 +164,17 @@ class Sampler:
         )
         return prompt
 
+    def compute_probabilities(self, logits):
+        """Compute token probabilities from logits.
+
+        This will always be done in full precision, regardless of dtype, and
+        scale by `temperature`.
+        """
+        logits_dtype = logits.dtype
+        logits = ops.cast(logits, "float32")
+        probs = keras.activations.softmax(logits / self.temperature)
+        return ops.cast(probs, logits_dtype)
+
     def run_loop(self, cond, body, loop_vars=None, maximum_iterations=None):
         """Run ops.while_loops with a `StatelessScope` if necessary."""
         if config.backend() == "jax":
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
index bd6f568a..7da1e9d7 100644
--- a/keras_nlp/tokenizers/tokenizer.py
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -123,6 +123,3 @@ class Tokenizer(PreprocessingLayer):
 
     def call(self, inputs, *args, training=None, **kwargs):
         return self.tokenize(inputs, *args, **kwargs)
-
-    def compute_output_shape(self, inputs_shape):
-        return tuple(inputs_shape) + (self.sequence_length,)


Test patch:
diff --git a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
index 3f6ecbc7..fdaab606 100644
--- a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.backend import config
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.cached_multi_head_attention import (
     CachedMultiHeadAttention,
@@ -34,6 +35,9 @@ class CachedMultiHeadAttentionTest(TestCase):
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=8,
             expected_num_non_trainable_variables=1,
+            # tf.keras does not handle mixed precision correctly when not set
+            # globally.
+            run_mixed_precision_check=config.multi_backend(),
         )
 
     def test_cache_call_is_correct(self):
diff --git a/keras_nlp/layers/modeling/reversible_embedding_test.py b/keras_nlp/layers/modeling/reversible_embedding_test.py
index 16cc3f8f..ceb04578 100644
--- a/keras_nlp/layers/modeling/reversible_embedding_test.py
+++ b/keras_nlp/layers/modeling/reversible_embedding_test.py
@@ -17,6 +17,7 @@ import os
 import numpy as np
 from absl.testing import parameterized
 
+from keras_nlp.backend import config
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
@@ -73,3 +74,22 @@ class ReversibleEmbeddingTest(TestCase):
 
         input_data = ops.ones(shape=(4, 10), dtype="int32")
         self.assertAllClose(untied_model(input_data), tied_model(input_data))
+
+    def test_reverse_dtype(self):
+        embedding = ReversibleEmbedding(100, 16, reverse_dtype="float32")
+        input_data = ops.ones(shape=(4, 10, 16))
+        output_data = embedding(input_data, reverse=True)
+        self.assertEqual(output_data.shape, (4, 10, 100))
+        self.assertDTypeEqual(output_data, "float32")
+
+        if config.backend() == "torch":
+            import torch
+
+            if not torch.cuda.is_available():
+                self.skipTest("Torch CPU does not support float16")
+
+        embedding = ReversibleEmbedding(100, 16, reverse_dtype="float16")
+        input_data = ops.ones(shape=(4, 10, 16))
+        output_data = embedding(input_data, reverse=True)
+        self.assertEqual(output_data.shape, (4, 10, 100))
+        self.assertDTypeEqual(output_data, "float16")
diff --git a/keras_nlp/tests/test_case.py b/keras_nlp/tests/test_case.py
index a53c29a6..d0af5a47 100644
--- a/keras_nlp/tests/test_case.py
+++ b/keras_nlp/tests/test_case.py
@@ -17,7 +17,8 @@ import json
 import tensorflow as tf
 import tree
 from absl.testing import parameterized
-from keras_core.src.backend import standardize_dtype
+from keras_core.backend import is_float_dtype
+from keras_core.backend import standardize_dtype
 
 from keras_nlp.backend import config
 from keras_nlp.backend import keras
@@ -67,6 +68,10 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         x2 = tree.map_structure(convert_to_comparible_type, x2)
         super().assertAllEqual(x1, x2, msg=msg)
 
+    def assertDTypeEqual(self, x, expected_dtype, msg=None):
+        input_dtype = standardize_dtype(x.dtype)
+        super().assertEqual(input_dtype, expected_dtype, msg=msg)
+
     def run_layer_test(
         self,
         layer_cls,
@@ -78,6 +83,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         expected_num_non_trainable_weights=0,
         expected_num_non_trainable_variables=0,
         run_training_check=True,
+        run_mixed_precision_check=True,
     ):
         # Serialization test.
         layer = layer_cls(**init_kwargs)
@@ -168,6 +174,25 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         if run_training_check:
             run_training_step(layer, input_data, output_data)
 
+        # Never test mixed precision on torch CPU. Torch lacks support.
+        if run_mixed_precision_check and config.backend() == "torch":
+            import torch
+
+            run_mixed_precision_check = torch.cuda.is_available()
+
+        if run_mixed_precision_check:
+            layer = layer_cls(**{**init_kwargs, "dtype": "mixed_float16"})
+            if isinstance(input_data, dict):
+                output_data = layer(**input_data)
+            else:
+                output_data = layer(input_data)
+            for tensor in tree.flatten(output_data):
+                if is_float_dtype(tensor.dtype):
+                    self.assertDTypeEqual(tensor, "float16")
+            for weight in layer.weights:
+                if is_float_dtype(weight.dtype):
+                    self.assertDTypeEqual(weight, "float32")
+
     def run_class_serialization_test(self, instance):
         # get_config roundtrip
         cls = instance.__class__