Update MaskedLMHead to support dtype=bfloat16/float16/float64.

keras_nlp/layers/modeling/masked_lm_head.py

@@ -153,9 +153,11 @@ def build(self, inputs_shape, masked_positions_shape=None):
  activation=self.intermediate_activation,
  kernel_initializer=self.kernel_initializer,
  bias_initializer=self.bias_initializer,
+ dtype=self._dtype_policy,
  )
  self._layer_norm = keras.layers.LayerNormalization(
  epsilon=self.layer_norm_epsilon,
+ dtype=self._dtype_policy,
  )
  if masked_positions_shape:
  gather_length = masked_positions_shape[1]
@@ -181,18 +183,25 @@ def call(self, inputs, masked_positions):
  # Gather the encoded tokens at the masked indices.
  masked_positions = ops.expand_dims(masked_positions, axis=-1)
  x = ops.take_along_axis(inputs, masked_positions, axis=1)
+ print("XXX/1", x.dtype)
 
  # Apply a trainable linear transformation and a layer norm.
  x = self._dense(x)
+ print("XXX/2", x.dtype)
  x = self._layer_norm(x)
+ print("XXX/3", x.dtype)
 
  # Transform encodings to vocabulary_size predictions.
  if self.embedding_weights is None:
  kernel = self._kernel
+ print("XXX/4", kernel)
  else:
  kernel = ops.cast(self.embedding_weights, self.compute_dtype)
+ print("XXX/5", kernel)
  kernel = ops.transpose(kernel)
+ print("XXX/6", kernel)
  outputs = ops.matmul(x, kernel)
+ print("XXX", outputs.dtype, self._bias.dtype)
  outputs = outputs + self._bias
 
  # Apply a final activation.

keras_nlp/layers/modeling/masked_lm_head_test.py

@@ -15,6 +15,9 @@
 
 import os
 
+import tensorflow as tf
+from absl.testing import parameterized
+
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.layers.modeling import masked_lm_head
@@ -36,6 +39,30 @@ def test_valid_call(self):
  position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
  model((token_data, position_data))
 
+ @parameterized.named_parameters(
+ ("bfloat16", tf.bfloat16),
+ ("float16", tf.float16),
+ ("float32", tf.float32),
+ ("float64", tf.float64),
+ )
+ def test_valid_call_with_dtype(self, dtype):
+ head = masked_lm_head.MaskedLMHead(
+ vocabulary_size=100,
+ activation="softmax",
+ dtype=dtype,
+ )
+ encoded_tokens = keras.Input(shape=(10, 16))
+ positions = keras.Input(shape=(5,), dtype="int32")
+ outputs = head(encoded_tokens, masked_positions=positions)
+ model = keras.Model((encoded_tokens, positions), outputs)
+
+ token_data = ops.random.uniform(shape=(4, 10, 16))
+ position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
+ model((token_data, position_data))
+
+ for w in head.weights:
+ self.assertEqual(w.dtype, dtype, "Wrong type: " + w.name)
+
  def test_valid_call_with_embedding_weights(self):
  embedding = keras.layers.Embedding(100, 16)
  embedding.build((4, 10))
@@ -119,6 +146,32 @@ def test_one_train_step(self):
  loss = model.train_on_batch(x=(token_data, position_data), y=label_data)
  self.assertGreater(loss, 0)
 
+ @parameterized.named_parameters(
+ ("bfloat16", tf.bfloat16),
+ ("float16", tf.float16),
+ ("float32", tf.float32),
+ ("float64", tf.float64),
+ )
+ def test_one_train_step_with_dtype(self, dtype):
+ head = masked_lm_head.MaskedLMHead(
+ vocabulary_size=100,
+ dtype=dtype,
+ )
+ encoded_tokens = keras.Input(shape=(10, 16))
+ positions = keras.Input(shape=(5,), dtype="int32")
+ outputs = head(encoded_tokens, masked_positions=positions)
+ model = keras.Model((encoded_tokens, positions), outputs)
+
+ token_data = ops.random.uniform(shape=(4, 10, 16))
+ position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
+ label_data = ops.random.randint(minval=0, maxval=2, shape=(4, 5, 1))
+
+ loss = keras.losses.SparseCategoricalCrossentropy(from_logits=False)
+ optimizer = keras.optimizers.Adam()
+ model.compile(loss=loss, optimizer=optimizer)
+ loss = model.train_on_batch(x=(token_data, position_data), y=label_data)
+ self.assertGreater(loss, 0)
+
  def test_saved_model(self):
  head = masked_lm_head.MaskedLMHead(
  vocabulary_size=100,