Repository: keras-team/keras-nlp
Pull request #1429 (2024-02-08 19:54:55)

Add a Causal LM model for Mistral
This PR adds a Causal LM for Mistral called `MistralCausalLM` and a preprocessor for it called `MistralCausalLMPreprocessor`. Presets are not added yet but can done in a follow-up PR.

Note that I removed sliding window attention cache from the attention layer for Mistral. This is because JAX was complaining about dynamic slicing which is required to make the caching work. More explaination in this commit: 19b0b89911b19d34f2d67290d77468d9e46df496

I am in the process of testing if this model matches the outputs of the original model after the weights transfer. Once that's done, I can open the PR up for reviews.

tirthasheshpatel: > I am in the process of testing if this model matches the outputs of the original model after the weights transfer. Once that's done, I can open the PR up for reviews.

@mattdangerw Tested it with the 7B preset. The outputs of both the backbone and the generator match up. This is ready from my side! I can share the preset with you once this is merged.

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 7b898138..8fd6a70a 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -93,6 +93,12 @@ from keras_nlp.models.gpt_neo_x.gpt_neo_x_preprocessor import (
 from keras_nlp.models.gpt_neo_x.gpt_neo_x_tokenizer import GPTNeoXTokenizer
 from keras_nlp.models.llama.llama_backbone import LlamaBackbone
 from keras_nlp.models.mistral.mistral_backbone import MistralBackbone
+from keras_nlp.models.mistral.mistral_causal_lm import MistralCausalLM
+from keras_nlp.models.mistral.mistral_causal_lm_preprocessor import (
+    MistralCausalLMPreprocessor,
+)
+from keras_nlp.models.mistral.mistral_preprocessor import MistralPreprocessor
+from keras_nlp.models.mistral.mistral_tokenizer import MistralTokenizer
 from keras_nlp.models.opt.opt_backbone import OPTBackbone
 from keras_nlp.models.opt.opt_causal_lm import OPTCausalLM
 from keras_nlp.models.opt.opt_causal_lm_preprocessor import (
diff --git a/keras_nlp/models/albert/albert_backbone.py b/keras_nlp/models/albert/albert_backbone.py
index 99a7cc97..1e342e79 100644
--- a/keras_nlp/models/albert/albert_backbone.py
+++ b/keras_nlp/models/albert/albert_backbone.py
@@ -72,6 +72,10 @@ class AlbertBackbone(Backbone):
             embeddings.
         num_segments: int. The number of types that the 'segment_ids' input can
             take.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -110,6 +114,7 @@ class AlbertBackbone(Backbone):
         dropout=0.0,
         max_sequence_length=512,
         num_segments=2,
+        dtype=None,
         **kwargs,
     ):
         if num_layers % num_groups != 0:
@@ -123,35 +128,41 @@ class AlbertBackbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=embedding_dim,
             embeddings_initializer=albert_kernel_initializer(),
+            dtype=dtype,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             initializer=albert_kernel_initializer(),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="position_embedding",
         )
         self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=embedding_dim,
             embeddings_initializer=albert_kernel_initializer(),
+            dtype=dtype,
             name="segment_embedding",
         )
         self.embeddings_add = keras.layers.Add(
+            dtype=dtype,
             name="embeddings_add",
         )
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
-            name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
-            dtype="float32",
+            dtype=dtype,
+            name="embeddings_layer_norm",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.embeddings_projection = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=albert_kernel_initializer(),
+            dtype=dtype,
             name="embedding_projection",
         )
         self.transformer_layers = []
@@ -165,6 +176,7 @@ class AlbertBackbone(Backbone):
                     dropout=dropout,
                     layer_norm_epsilon=1e-12,
                     kernel_initializer=albert_kernel_initializer(),
+                    dtype=dtype,
                     name=f"group_{group_idx}_inner_layer_{inner_idx}",
                 )
                 inner_layers.append(layer)
@@ -173,6 +185,7 @@ class AlbertBackbone(Backbone):
             hidden_dim,
             kernel_initializer=albert_kernel_initializer(),
             activation="tanh",
+            dtype=dtype,
             name="pooled_dense",
         )
 
diff --git a/keras_nlp/models/albert/albert_classifier.py b/keras_nlp/models/albert/albert_classifier.py
index e8621cf6..32a4e084 100644
--- a/keras_nlp/models/albert/albert_classifier.py
+++ b/keras_nlp/models/albert/albert_classifier.py
@@ -162,10 +162,12 @@ class AlbertClassifier(Task):
             num_classes,
             kernel_initializer=albert_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
         self.output_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="output_dropout",
         )
 
diff --git a/keras_nlp/models/albert/albert_masked_lm.py b/keras_nlp/models/albert/albert_masked_lm.py
index 139be0a7..1958713b 100644
--- a/keras_nlp/models/albert/albert_masked_lm.py
+++ b/keras_nlp/models/albert/albert_masked_lm.py
@@ -105,6 +105,7 @@ class AlbertMaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation=gelu_approximate,
             kernel_initializer=albert_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/backbone.py b/keras_nlp/models/backbone.py
index 9e22be3f..6fccf601 100644
--- a/keras_nlp/models/backbone.py
+++ b/keras_nlp/models/backbone.py
@@ -22,7 +22,7 @@ from keras_nlp.utils.python_utils import format_docstring
 
 @keras.saving.register_keras_serializable(package="keras_nlp")
 class Backbone(keras.Model):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, dtype=None, **kwargs):
         super().__init__(*args, **kwargs)
         self._functional_layer_ids = set(
             id(layer) for layer in self._flatten_layers()
@@ -74,8 +74,8 @@ class Backbone(keras.Model):
         self._token_embedding = value
 
     def get_config(self):
-        # Don't chain to super here. The default `get_config()` for functional
-        # models is nested and cannot be passed to our Backbone constructors.
+        # Don't chain to super here. `get_config()` for functional models is
+        # a nested layer config and cannot be passed to Backbone constructors.
         return {
             "name": self.name,
             "trainable": self.trainable,
diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
index fdb8e5df..803d5a2a 100644
--- a/keras_nlp/models/bart/bart_backbone.py
+++ b/keras_nlp/models/bart/bart_backbone.py
@@ -60,6 +60,10 @@ class BartBackbone(Backbone):
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -100,6 +104,7 @@ class BartBackbone(Backbone):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=1024,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -107,24 +112,28 @@ class BartBackbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=bart_kernel_initializer(),
+            dtype=dtype,
             name="token_embedding",
         )
         self.encoder_position_embedding = PositionEmbedding(
             initializer=bart_kernel_initializer(),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="encoder_position_embedding",
         )
         self.encoder_embeddings_add = keras.layers.Add(
+            dtype=dtype,
             name="encoder_embeddings_add",
         )
         self.encoder_embeddings_layer_norm = keras.layers.LayerNormalization(
-            name="encoder_embeddings_layer_norm",
             axis=-1,
             epsilon=1e-5,
-            dtype="float32",
+            dtype=dtype,
+            name="encoder_embeddings_layer_norm",
         )
         self.encoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="encoder_embeddings_dropout",
         )
         self.encoder_transformer_layers = []
@@ -136,25 +145,29 @@ class BartBackbone(Backbone):
                 dropout=dropout,
                 layer_norm_epsilon=1e-5,
                 kernel_initializer=bart_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_encoder_layer_{i}",
             )
             self.encoder_transformer_layers.append(layer)
         self.decoder_position_embedding = PositionEmbedding(
             initializer=bart_kernel_initializer(),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="decoder_position_embedding",
         )
         self.decoder_embeddings_add = keras.layers.Add(
+            dtype=dtype,
             name="decoder_embeddings_add",
         )
         self.decoder_embeddings_layer_norm = keras.layers.LayerNormalization(
-            name="decoder_embeddings_layer_norm",
             axis=-1,
             epsilon=1e-5,
-            dtype="float32",
+            dtype=dtype,
+            name="decoder_embeddings_layer_norm",
         )
         self.decoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="decoder_embeddings_dropout",
         )
         self.decoder_transformer_layers = []
@@ -166,6 +179,7 @@ class BartBackbone(Backbone):
                 activation=keras.activations.gelu,
                 layer_norm_epsilon=1e-5,
                 kernel_initializer=bart_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_decoder_layer_{i}",
             )
             self.decoder_transformer_layers.append(layer)
diff --git a/keras_nlp/models/bert/bert_backbone.py b/keras_nlp/models/bert/bert_backbone.py
index f511de36..2248260d 100644
--- a/keras_nlp/models/bert/bert_backbone.py
+++ b/keras_nlp/models/bert/bert_backbone.py
@@ -61,6 +61,10 @@ class BertBackbone(Backbone):
             embeddings.
         num_segments: int. The number of types that the 'segment_ids' input can
             take.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -97,6 +101,7 @@ class BertBackbone(Backbone):
         dropout=0.1,
         max_sequence_length=512,
         num_segments=2,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -104,30 +109,35 @@ class BertBackbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=bert_kernel_initializer(),
+            dtype=dtype,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             initializer=bert_kernel_initializer(),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="position_embedding",
         )
         self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=hidden_dim,
             embeddings_initializer=bert_kernel_initializer(),
+            dtype=dtype,
             name="segment_embedding",
         )
         self.embeddings_add = keras.layers.Add(
+            dtype=dtype,
             name="embeddings_add",
         )
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
-            name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
-            dtype="float32",
+            dtype=dtype,
+            name="embeddings_layer_norm",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.transformer_layers = []
@@ -139,6 +149,7 @@ class BertBackbone(Backbone):
                 dropout=dropout,
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=bert_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
@@ -146,6 +157,7 @@ class BertBackbone(Backbone):
             hidden_dim,
             kernel_initializer=bert_kernel_initializer(),
             activation="tanh",
+            dtype=dtype,
             name="pooled_dense",
         )
 
diff --git a/keras_nlp/models/bert/bert_classifier.py b/keras_nlp/models/bert/bert_classifier.py
index 3ddb90a9..09d2b881 100644
--- a/keras_nlp/models/bert/bert_classifier.py
+++ b/keras_nlp/models/bert/bert_classifier.py
@@ -145,12 +145,14 @@ class BertClassifier(Task):
         self.preprocessor = preprocessor
         self.output_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="classifier_dropout",
         )
         self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=bert_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
 
diff --git a/keras_nlp/models/bert/bert_masked_lm.py b/keras_nlp/models/bert/bert_masked_lm.py
index 555c562f..17b96696 100644
--- a/keras_nlp/models/bert/bert_masked_lm.py
+++ b/keras_nlp/models/bert/bert_masked_lm.py
@@ -109,6 +109,7 @@ class BertMaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=bert_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/bloom/bloom_attention.py b/keras_nlp/models/bloom/bloom_attention.py
index 366169d1..e36c6fac 100644
--- a/keras_nlp/models/bloom/bloom_attention.py
+++ b/keras_nlp/models/bloom/bloom_attention.py
@@ -75,7 +75,9 @@ class BloomAttention(keras.layers.Layer):
         )
         self._value_dense.build(inputs_shape)
 
-        self._alibi_layer = AlibiBias()
+        self._alibi_layer = AlibiBias(
+            dtype=self.dtype_policy,
+        )
 
         self._output_dense = keras.layers.Dense(
             hidden_dim,
@@ -87,10 +89,13 @@ class BloomAttention(keras.layers.Layer):
         self._output_dense.build(inputs_shape)
 
         self._dropout_layer = keras.layers.Dropout(
-            rate=self.dropout, dtype=self.dtype_policy, name="dropout"
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="dropout",
         )
         self._softmax = keras.layers.Softmax(
-            dtype=self.dtype_policy, name="softmax"
+            dtype="float32",
+            name="softmax",
         )
 
         self.built = True
diff --git a/keras_nlp/models/bloom/bloom_backbone.py b/keras_nlp/models/bloom/bloom_backbone.py
index 4f12bacd..5737dcc8 100644
--- a/keras_nlp/models/bloom/bloom_backbone.py
+++ b/keras_nlp/models/bloom/bloom_backbone.py
@@ -55,6 +55,10 @@ class BloomBackbone(Backbone):
             the transformer decoder.
         max_sequence_length: int. The maximum sequence length that this decoder
             can consume.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -93,6 +97,7 @@ class BloomBackbone(Backbone):
         dropout=0.0,
         layer_norm_epsilon=1e-5,
         max_sequence_length=2048,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -101,10 +106,12 @@ class BloomBackbone(Backbone):
             output_dim=hidden_dim,
             embeddings_initializer=_bloom_kernel_initializer(stddev=0.02),
             tie_weights=False,
+            dtype=dtype,
             name="token_embedding",
         )
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
             epsilon=layer_norm_epsilon,
+            dtype=dtype,
             name="token_embedding_layernorm",
         )
         self.transformer_layers = []
@@ -114,11 +121,13 @@ class BloomBackbone(Backbone):
                 intermediate_dim=intermediate_dim,
                 dropout=dropout,
                 layer_norm_epsilon=layer_norm_epsilon,
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
         self.layer_norm = keras.layers.LayerNormalization(
             epsilon=layer_norm_epsilon,
+            dtype=dtype,
             name="final_layernorm",
         )
 
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
index 87531aef..e7bd8ca2 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
@@ -67,6 +67,10 @@ class DebertaV3Backbone(Backbone):
             `max_sequence_length`.
         bucket_size: int. The size of the relative position buckets. Generally
             equal to `max_sequence_length // 2`.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Example:
     ```python
@@ -106,6 +110,7 @@ class DebertaV3Backbone(Backbone):
         dropout=0.1,
         max_sequence_length=512,
         bucket_size=256,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -113,15 +118,17 @@ class DebertaV3Backbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=deberta_kernel_initializer(),
+            dtype=dtype,
             name="token_embedding",
         )
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
             epsilon=1e-7,
-            dtype="float32",
+            dtype=dtype,
             name="embeddings_layer_norm",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.relative_embeddings = RelativeEmbedding(
@@ -129,6 +136,7 @@ class DebertaV3Backbone(Backbone):
             bucket_size=bucket_size,
             layer_norm_epsilon=1e-7,
             kernel_initializer=deberta_kernel_initializer(),
+            dtype=dtype,
             name="rel_embedding",
         )
         self.transformer_layers = []
@@ -142,6 +150,7 @@ class DebertaV3Backbone(Backbone):
                 activation=keras.activations.gelu,
                 layer_norm_epsilon=1e-7,
                 kernel_initializer=deberta_kernel_initializer(),
+                dtype=dtype,
                 name=f"disentangled_attention_encoder_layer_{i}",
             )
             self.transformer_layers.append(layer)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
index f5249cb3..d6eea636 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
@@ -168,22 +168,26 @@ class DebertaV3Classifier(Task):
         self.preprocessor = preprocessor
         self.pooled_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="pooled_dropout",
         )
         hidden_dim = hidden_dim or backbone.hidden_dim
         self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             activation=keras.activations.gelu,
+            dtype=backbone.dtype_policy,
             name="pooled_dense",
         )
         self.output_dropout = keras.layers.Dropout(
             backbone.dropout,
+            dtype=backbone.dtype_policy,
             name="classifier_dropout",
         )
         self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=deberta_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
 
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
index fadb4c0e..d050dde6 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
@@ -112,6 +112,7 @@ class DebertaV3MaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation=keras.activations.gelu,
             kernel_initializer=deberta_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py b/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py
index 081c345c..be79bc98 100644
--- a/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py
+++ b/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py
@@ -99,22 +99,26 @@ class DisentangledAttentionEncoder(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",
         )
         self._self_attention_layer.build(inputs_shape)
         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)
@@ -123,6 +127,7 @@ class DisentangledAttentionEncoder(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)
@@ -130,6 +135,7 @@ class DisentangledAttentionEncoder(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)
@@ -137,6 +143,7 @@ class DisentangledAttentionEncoder(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/models/deberta_v3/disentangled_self_attention.py b/keras_nlp/models/deberta_v3/disentangled_self_attention.py
index 1c9ae569..48730f2a 100644
--- a/keras_nlp/models/deberta_v3/disentangled_self_attention.py
+++ b/keras_nlp/models/deberta_v3/disentangled_self_attention.py
@@ -86,6 +86,7 @@ class DisentangledSelfAttention(keras.layers.Layer):
             output_shape=(None, self.num_heads, self.attn_head_size),
             bias_axes="de",
             **self._get_common_kwargs_for_sublayer(use_bias=True),
+            dtype=self.dtype_policy,
             name="query",
         )
         self._query_dense.build(inputs_shape)
@@ -94,6 +95,7 @@ class DisentangledSelfAttention(keras.layers.Layer):
             output_shape=(None, self.num_heads, self.attn_head_size),
             bias_axes="de",
             **self._get_common_kwargs_for_sublayer(use_bias=True),
+            dtype=self.dtype_policy,
             name="key",
         )
         self._key_dense.build(inputs_shape)
@@ -102,17 +104,27 @@ class DisentangledSelfAttention(keras.layers.Layer):
             output_shape=(None, self.num_heads, self.attn_head_size),
             bias_axes="de",
             **self._get_common_kwargs_for_sublayer(use_bias=True),
+            dtype=self.dtype_policy,
             name="value",
         )
         self._value_dense.build(inputs_shape)
 
         # Relative attention.
-        self._position_dropout_layer = keras.layers.Dropout(self.dropout)
+        self._position_dropout_layer = keras.layers.Dropout(
+            self.dropout,
+            dtype=self.dtype_policy,
+        )
 
         self._attn_dropout_layer = keras.layers.Dropout(
-            self.dropout, name="attention_dropout"
+            self.dropout,
+            dtype=self.dtype_policy,
+            name="attention_dropout",
+        )
+        self._softmax = keras.layers.Softmax(
+            axis=-1,
+            dtype="float32",
+            name="attention_softmax",
         )
-        self._softmax = keras.layers.Softmax(axis=-1, name="attention_softmax")
 
         # Output.
         self._output_dense = keras.layers.EinsumDense(
@@ -120,6 +132,7 @@ class DisentangledSelfAttention(keras.layers.Layer):
             output_shape=(None, self.hidden_dim),
             bias_axes="d",
             **self._get_common_kwargs_for_sublayer(use_bias=True),
+            dtype=self.dtype_policy,
             name="attention_output",
         )
         self._output_dense.build(inputs_shape)
diff --git a/keras_nlp/models/deberta_v3/relative_embedding.py b/keras_nlp/models/deberta_v3/relative_embedding.py
index 6ae29a5f..f727ce05 100644
--- a/keras_nlp/models/deberta_v3/relative_embedding.py
+++ b/keras_nlp/models/deberta_v3/relative_embedding.py
@@ -57,7 +57,9 @@ class RelativeEmbedding(keras.layers.Layer):
             name="rel_embedding",
         )
         self.layer_norm = keras.layers.LayerNormalization(
-            epsilon=layer_norm_epsilon, name="rel_embeddings_layer_norm"
+            epsilon=layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="rel_embeddings_layer_norm",
         )
 
     def call(self, inputs):
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone.py b/keras_nlp/models/distil_bert/distil_bert_backbone.py
index f97e24f5..1ae0840e 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone.py
@@ -62,6 +62,10 @@ class DistilBertBackbone(Backbone):
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -98,6 +102,7 @@ class DistilBertBackbone(Backbone):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=512,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -106,6 +111,7 @@ class DistilBertBackbone(Backbone):
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=distilbert_kernel_initializer(),
+            dtype=dtype,
             name="token_and_position_embedding",
         )
         # Keep the token_embedding property for consistency across models.
@@ -113,11 +119,12 @@ class DistilBertBackbone(Backbone):
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
             axis=-1,
             epsilon=1e-12,
-            dtype="float32",
+            dtype=dtype,
             name="embeddings_layer_norm",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.transformer_layers = []
@@ -129,6 +136,7 @@ class DistilBertBackbone(Backbone):
                 dropout=dropout,
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=distilbert_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier.py b/keras_nlp/models/distil_bert/distil_bert_classifier.py
index cf0db978..e82aaf27 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier.py
@@ -158,16 +158,19 @@ class DistilBertClassifier(Task):
             hidden_dim,
             activation="relu",
             kernel_initializer=distilbert_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="pooled_dense",
         )
         self.output_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="output_dropout",
         )
         self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=distilbert_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
 
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
index 9be43f8a..fcf54e01 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
@@ -112,6 +112,7 @@ class DistilBertMaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=distilbert_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/electra/electra_backbone.py b/keras_nlp/models/electra/electra_backbone.py
index 2e2a0197..13be2d8e 100644
--- a/keras_nlp/models/electra/electra_backbone.py
+++ b/keras_nlp/models/electra/electra_backbone.py
@@ -58,6 +58,10 @@ class ElectraBackbone(Backbone):
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
         ```python
@@ -92,6 +96,7 @@ class ElectraBackbone(Backbone):
         dropout=0.1,
         max_sequence_length=512,
         num_segments=2,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -99,34 +104,42 @@ class ElectraBackbone(Backbone):
             input_dim=vocab_size,
             output_dim=embedding_dim,
             embeddings_initializer=electra_kernel_initializer(),
+            dtype=dtype,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             initializer=electra_kernel_initializer(),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="position_embedding",
         )
         self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=embedding_dim,
             embeddings_initializer=electra_kernel_initializer(),
+            dtype=dtype,
             name="segment_embedding",
         )
-        self.embeddings_add = keras.layers.Add()
+        self.embeddings_add = keras.layers.Add(
+            dtype=dtype,
+            name="embeddings_add",
+        )
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
-            name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
-            dtype="float32",
+            dtype=dtype,
+            name="embeddings_layer_norm",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         if hidden_dim != embedding_dim:
             self.embeddings_projection = keras.layers.Dense(
                 hidden_dim,
                 kernel_initializer=electra_kernel_initializer(),
+                dtype=dtype,
                 name="embeddings_projection",
             )
         self.transformer_layers = []
@@ -138,6 +151,7 @@ class ElectraBackbone(Backbone):
                 dropout=dropout,
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=electra_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
@@ -145,6 +159,7 @@ class ElectraBackbone(Backbone):
             hidden_dim,
             kernel_initializer=electra_kernel_initializer(),
             activation="tanh",
+            dtype=dtype,
             name="pooled_dense",
         )
 
diff --git a/keras_nlp/models/f_net/f_net_backbone.py b/keras_nlp/models/f_net/f_net_backbone.py
index 9103a10d..309f312a 100644
--- a/keras_nlp/models/f_net/f_net_backbone.py
+++ b/keras_nlp/models/f_net/f_net_backbone.py
@@ -66,6 +66,10 @@ class FNetBackbone(Backbone):
             embeddings.
         num_segments: int. The number of types that the 'segment_ids' input can
             take.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -99,6 +103,7 @@ class FNetBackbone(Backbone):
         dropout=0.1,
         max_sequence_length=512,
         num_segments=4,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -106,34 +111,42 @@ class FNetBackbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=f_net_kernel_initializer(),
+            dtype=dtype,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             initializer=f_net_kernel_initializer(),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="position_embedding",
         )
         self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=hidden_dim,
             embeddings_initializer=f_net_kernel_initializer(),
+            dtype=dtype,
             name="segment_embedding",
         )
-        self.embeddings_add = keras.layers.Add()
+        self.embeddings_add = keras.layers.Add(
+            dtype=dtype,
+            name="embeddings_add",
+        )
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
-            name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
-            dtype="float32",
+            dtype=dtype,
+            name="embeddings_layer_norm",
         )
         self.embedding_projection = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=f_net_kernel_initializer(),
             bias_initializer=f_net_bias_initializer(),
+            dtype=dtype,
             name="embedding_projection",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.transformer_layers = []
@@ -145,6 +158,7 @@ class FNetBackbone(Backbone):
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=f_net_kernel_initializer(),
                 bias_initializer=f_net_bias_initializer(),
+                dtype=dtype,
                 name=f"f_net_layer_{i}",
             )
             self.transformer_layers.append(layer)
@@ -153,6 +167,7 @@ class FNetBackbone(Backbone):
             kernel_initializer=f_net_kernel_initializer(),
             bias_initializer=f_net_bias_initializer(),
             activation="tanh",
+            dtype=dtype,
             name="pooled_dense",
         )
 
diff --git a/keras_nlp/models/f_net/f_net_classifier.py b/keras_nlp/models/f_net/f_net_classifier.py
index f4ee31d1..512182d2 100644
--- a/keras_nlp/models/f_net/f_net_classifier.py
+++ b/keras_nlp/models/f_net/f_net_classifier.py
@@ -114,12 +114,14 @@ class FNetClassifier(Task):
         self.preprocessor = preprocessor
         self.output_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="output_dropout",
         )
         self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=f_net_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
 
diff --git a/keras_nlp/models/f_net/f_net_masked_lm.py b/keras_nlp/models/f_net/f_net_masked_lm.py
index c0eb231d..c715a708 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm.py
@@ -109,6 +109,7 @@ class FNetMaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=f_net_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/gpt2/gpt2_backbone.py b/keras_nlp/models/gpt2/gpt2_backbone.py
index d3f4a415..d93b2199 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone.py
@@ -61,6 +61,10 @@ class GPT2Backbone(Backbone):
             can consume. If `None`, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for the models computations and weights. Note that some
+            computations, such as softmax and layer normalization will always
+            be done a float32 precision regardless of dtype.
 
     Example:
     ```python
@@ -95,6 +99,7 @@ class GPT2Backbone(Backbone):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=1024,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -102,18 +107,22 @@ class GPT2Backbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_gpt_2_kernel_initializer(stddev=0.01),
+            dtype=dtype,
             name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             initializer=_gpt_2_kernel_initializer(stddev=0.02),
             sequence_length=max_sequence_length,
+            dtype=dtype,
             name="position_embedding",
         )
         self.embeddings_add = keras.layers.Add(
+            dtype=dtype,
             name="embeddings_add",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.transformer_layers = []
@@ -127,14 +136,15 @@ class GPT2Backbone(Backbone):
                     activation=gelu_approximate,
                     kernel_initializer=_gpt_2_kernel_initializer(stddev=0.02),
                     normalize_first=True,
+                    dtype=dtype,
                     name=f"transformer_layer_{i}",
                 )
             )
         self.layer_norm = keras.layers.LayerNormalization(
-            name="layer_norm",
             axis=-1,
             epsilon=1e-05,
-            dtype="float32",
+            dtype=dtype,
+            name="layer_norm",
         )
 
         # === Functional Model ===
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py
index 40cdc0d5..dee8addf 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_attention.py
@@ -64,7 +64,8 @@ class GPTNeoXAttention(keras.layers.Layer):
         self.rotary_max_wavelength = rotary_max_wavelength
         self.rotary_dim = int(self.attn_head_size * rotary_percentage)
         self.rotary_embedding_layer = RotaryEmbedding(
-            max_wavelength=rotary_max_wavelength
+            max_wavelength=rotary_max_wavelength,
+            dtype=self.dtype_policy,
         )
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
@@ -76,15 +77,22 @@ class GPTNeoXAttention(keras.layers.Layer):
             output_shape=(None, self.num_heads, 3 * self.attn_head_size),
             bias_axes="de",
             **self._get_common_kwargs_for_sublayer(use_bias=True),
+            dtype=self.dtype_policy,
             name="query_key_value",
         )
         self._qkv_dense.build(input_shape)
 
         self._attn_dropout_layer = keras.layers.Dropout(
-            self.dropout, name="attention_dropout"
+            self.dropout,
+            dtype=self.dtype_policy,
+            name="attention_dropout",
         )
 
-        self._softmax = keras.layers.Softmax(axis=-1, name="attention_softmax")
+        self._softmax = keras.layers.Softmax(
+            axis=-1,
+            dtype="float32",
+            name="attention_softmax",
+        )
 
         # Output.
         self._output_dense = keras.layers.EinsumDense(
@@ -92,6 +100,7 @@ class GPTNeoXAttention(keras.layers.Layer):
             output_shape=(None, self.hidden_dim),
             bias_axes="d",
             **self._get_common_kwargs_for_sublayer(use_bias=True),
+            dtype=self.dtype_policy,
             name="attention_output",
         )
 
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
index 5bbc11af..1955ed58 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
@@ -61,6 +61,10 @@ class GPTNeoXBackbone(Backbone):
             can consume. If `None`, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
     """
 
     def __init__(
@@ -75,6 +79,7 @@ class GPTNeoXBackbone(Backbone):
         rotary_max_wavelength=10000,
         layer_norm_epsilon=1e-5,
         max_sequence_length=512,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -82,10 +87,12 @@ class GPTNeoXBackbone(Backbone):
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_gpt_neo_x_kernel_initializer(stddev=0.01),
+            dtype=dtype,
             name="token_embedding",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.transformer_layers = []
@@ -100,14 +107,15 @@ class GPTNeoXBackbone(Backbone):
                 layer_norm_epsilon=layer_norm_epsilon,
                 activation=gelu_approximate,
                 kernel_initializer=_gpt_neo_x_kernel_initializer(stddev=0.02),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
         self.layer_norm = keras.layers.LayerNormalization(
-            name="layer_norm",
             axis=-1,
             epsilon=layer_norm_epsilon,
-            dtype="float32",
+            dtype=dtype,
+            name="layer_norm",
         )
 
         # === Functional Model ===
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_decoder.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_decoder.py
index f80cafd5..0a7bad7c 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_decoder.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_decoder.py
@@ -99,18 +99,21 @@ class GPTNeoXDecoder(keras.layers.Layer):
             max_sequence_length=self.max_sequence_length,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
             name="self_attention",
         )
         self._self_attention_layer.build(decoder_sequence_shape)
 
         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",
         )
 
@@ -120,6 +123,7 @@ class GPTNeoXDecoder(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(decoder_sequence_shape)
@@ -128,6 +132,7 @@ class GPTNeoXDecoder(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",
         )
 
@@ -137,12 +142,14 @@ class GPTNeoXDecoder(keras.layers.Layer):
 
         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(decoder_sequence_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/models/llama/llama_attention.py b/keras_nlp/models/llama/llama_attention.py
index a2604e53..529e73b0 100644
--- a/keras_nlp/models/llama/llama_attention.py
+++ b/keras_nlp/models/llama/llama_attention.py
@@ -58,6 +58,7 @@ class LlamaAttention(keras.layers.Layer):
             equation="bqm,muh->bquh",
             output_shape=(None, self.num_query_heads, self.attn_head_size),
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
             name="query",
         )
         self._query_dense.build(inputs_shape)
@@ -65,6 +66,7 @@ class LlamaAttention(keras.layers.Layer):
             equation="bkm,mvh->bkvh",
             output_shape=(None, self.num_key_value_heads, self.attn_head_size),
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
             name="key",
         )
         self._key_dense.build(inputs_shape)
@@ -73,16 +75,22 @@ class LlamaAttention(keras.layers.Layer):
             equation="bkm,mvh->bkvh",
             output_shape=(None, self.num_key_value_heads, self.attn_head_size),
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
             name="value",
         )
         self._value_dense.build(inputs_shape)
 
-        self._softmax = keras.layers.Softmax(axis=-1, name="attention_softmax")
+        self._softmax = keras.layers.Softmax(
+            axis=-1,
+            dtype="float32",
+            name="attention_softmax",
+        )
 
         self._output_dense = keras.layers.EinsumDense(
             equation="bqm,mh->bqh",
             output_shape=(None, self.hidden_dim),
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
             name="attention_output",
         )
         self._output_dense.build(inputs_shape)
@@ -90,6 +98,7 @@ class LlamaAttention(keras.layers.Layer):
         self._rotary_embedding_layer = RotaryEmbedding(
             max_wavelength=self.rope_max_wavelength,
             scaling_factor=self.rope_scaling_factor,
+            dtype=self.dtype_policy,
         )
         self._rotary_embedding_layer.build(inputs_shape)
 
@@ -173,10 +182,10 @@ class LlamaAttention(keras.layers.Layer):
         )
 
         attention_scores /= norm_factor
-
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )
+        attention_scores = ops.cast(attention_scores, self.compute_dtype)
         attention_output = ops.einsum(
             "acbe,aecd->abcd", attention_scores, value
         )
diff --git a/keras_nlp/models/llama/llama_backbone.py b/keras_nlp/models/llama/llama_backbone.py
index 46cfdc37..cc628ad7 100644
--- a/keras_nlp/models/llama/llama_backbone.py
+++ b/keras_nlp/models/llama/llama_backbone.py
@@ -58,7 +58,10 @@ class LlamaBackbone(Backbone):
             can consume. If `None`, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
-
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
     """
 
     def __init__(
@@ -73,6 +76,7 @@ class LlamaBackbone(Backbone):
         rope_max_wavelength=10000,
         layer_norm_epsilon=1e-5,
         max_sequence_length=4096,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -81,6 +85,7 @@ class LlamaBackbone(Backbone):
             output_dim=hidden_dim,
             embeddings_initializer=_llama_kernel_initializer(stddev=0.01),
             tie_weights=False,
+            dtype=dtype,
             name="token_embedding",
         )
         self.transformer_layers = []
@@ -95,12 +100,14 @@ class LlamaBackbone(Backbone):
                 layer_norm_epsilon=layer_norm_epsilon,
                 activation=ops.silu,
                 kernel_initializer=_llama_kernel_initializer(stddev=0.02),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
         self.layer_norm = LlamaLayerNorm(
-            name="layer_norm",
+            dtype=dtype,
             epsilon=layer_norm_epsilon,
+            name="layer_norm",
         )
 
         # === Functional Model ===
diff --git a/keras_nlp/models/llama/llama_decoder.py b/keras_nlp/models/llama/llama_decoder.py
index 2137831b..3b9d6906 100644
--- a/keras_nlp/models/llama/llama_decoder.py
+++ b/keras_nlp/models/llama/llama_decoder.py
@@ -64,11 +64,13 @@ class LlamaDecoder(keras.layers.Layer):
             max_sequence_length=self.max_sequence_length,
             rope_scaling_factor=self.rope_scaling_factor,
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
         )
         self._self_attention_layer.build(decoder_sequence_shape)
 
         self._self_attention_layernorm = LlamaLayerNorm(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
         )
         self._self_attention_layernorm.build(decoder_sequence_shape)
 
@@ -76,6 +78,7 @@ class LlamaDecoder(keras.layers.Layer):
         self._feedforward_intermediate_dense = keras.layers.Dense(
             self.intermediate_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
         )
         self._feedforward_intermediate_dense.build(decoder_sequence_shape)
 
@@ -83,12 +86,14 @@ class LlamaDecoder(keras.layers.Layer):
             self.intermediate_dim,
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
         )
         self._feedforward_gate_dense.build(decoder_sequence_shape)
 
         self._feedforward_output_dense = keras.layers.Dense(
             self.hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
+            dtype=self.dtype_policy,
         )
 
         intermediate_shape = list(decoder_sequence_shape)
@@ -97,6 +102,7 @@ class LlamaDecoder(keras.layers.Layer):
 
         self._feedforward_layernorm = LlamaLayerNorm(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
         )
         self._feedforward_layernorm.build(decoder_sequence_shape)
 
diff --git a/keras_nlp/models/mistral/mistral_attention.py b/keras_nlp/models/mistral/mistral_attention.py
index 680f1f6d..6511a224 100644
--- a/keras_nlp/models/mistral/mistral_attention.py
+++ b/keras_nlp/models/mistral/mistral_attention.py
@@ -69,7 +69,7 @@ class CachedMistralAttention(keras.layers.Layer):
             equation="bqm,muh->bquh",
             output_shape=(None, self._num_query_heads, self._head_dim),
             kernel_initializer=self._kernel_initializer,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="query",
         )
         self._query_dense.build(inputs_shape)
@@ -82,7 +82,7 @@ class CachedMistralAttention(keras.layers.Layer):
                 self._head_dim,
             ),
             kernel_initializer=self._kernel_initializer,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="key",
         )
         self._key_dense.build(inputs_shape)
@@ -95,22 +95,27 @@ class CachedMistralAttention(keras.layers.Layer):
                 self._head_dim,
             ),
             kernel_initializer=self._kernel_initializer,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="value",
         )
         self._value_dense.build(inputs_shape)
 
-        self._softmax = keras.layers.Softmax(axis=-1, name="attention_softmax")
+        self._softmax = keras.layers.Softmax(
+            axis=-1,
+            dtype="float32",
+            name="attention_softmax",
+        )
 
         self._dropout_layer = keras.layers.Dropout(
-            rate=self._dropout, dtype=self.compute_dtype
+            rate=self._dropout,
+            dtype=self.dtype_policy,
         )
 
         self._output_dense = keras.layers.EinsumDense(
             equation="bquh,uhm->bqm",
             output_shape=(None, self._hidden_dim),
             kernel_initializer=self._kernel_initializer,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="attention_output",
         )
         self._output_dense.build(
@@ -120,7 +125,7 @@ class CachedMistralAttention(keras.layers.Layer):
         self.rotary_embedding_layer = RotaryEmbedding(
             max_wavelength=self._rope_max_wavelength,
             scaling_factor=self._rope_scaling_factor,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
         )
 
         self._dot_product_equation = "bquh,bkuh->buqk"
@@ -136,7 +141,6 @@ class CachedMistralAttention(keras.layers.Layer):
         cache_update_index=None,
         training=None,
     ):
-        seq_len = ops.shape(hidden_states)[1]
         start_index = (
             cache_update_index if cache_update_index is not None else 0
         )
@@ -148,89 +152,34 @@ class CachedMistralAttention(keras.layers.Layer):
 
         query = self._query_dense(hidden_states)
 
-        # Note that the original PyTorch implementation uses
-        # view_as_complex/view_as_real while we use split/concatenate to
-        # convert to/from complex numbers. The transformations below make
-        # the rope computation numerically equivalent to the original
-        # implementation.
-        def _mistral_rope(x):
-            x = ops.concatenate([x[..., ::2], x[..., 1::2]], axis=-1)
-            x = self.rotary_embedding_layer(x, start_index=start_index)
-            x = ops.reshape(
-                ops.stack(ops.split(x, 2, axis=-1), axis=-1), ops.shape(x)
-            )
-            return x
-
         # Compute RoPE for queries
-        query = _mistral_rope(query)
+        query = self.rotary_embedding_layer(query, start_index=start_index)
 
         def _compute_key_value(x):
             key, value = self._key_dense(x), self._value_dense(x)
-            key = _mistral_rope(key)
+            # Compute RoPE for keys
+            key = self.rotary_embedding_layer(key, start_index=start_index)
             return key, value
 
         if cache is not None:
-            cache_k = cache[:, 0, ...]
-            cache_v = cache[:, 1, ...]
-
+            key_cache = cache[:, 0, ...]
+            value_cache = cache[:, 1, ...]
+            if cache_update_index is None:
+                key = key_cache
+                value = value_cache
+            else:
+                key_update, value_update = _compute_key_value(hidden_states)
+                start = [0, cache_update_index, 0, 0]
+                key = ops.slice_update(key_cache, start, key_update)
+                value = ops.slice_update(value_cache, start, value_update)
+                cache = ops.stack((key, value), axis=1)
+        else:
             if cache_update_index is not None:
-                # Compute the new keys and values
-                key, value = _compute_key_value(hidden_states)
-
-                # Cache is a rotating buffer, we want to warp around if
-                # the sequence length exceeds the sliding window.
-                update_end_index = (
-                    cache_update_index + seq_len - 1
-                ) % self._sliding_window + 1
-                update_end_index = ops.cast(update_end_index, "int32")
-                cache_update_index = cache_update_index % self._sliding_window
-                update_start_index = ops.cond(
-                    update_end_index > cache_update_index,
-                    lambda: ops.cast(cache_update_index, "int32"),
-                    lambda: ops.cast(0, "int32"),
-                )
-                # Also note that the update step below assumes that the
-                # sequence length is always one when `cache_update_index != 0`.
-                # This is necessary to support XLA compilation. Ideally, we
-                # would want to use
-                # `key[:, -(update_end_index - update_start_index):, ...]`
-                # as the update but updating using a dynamic slice gives an
-                # XLA compilation error in TensorFlow.
-                # Passing a sequence of length > 1 with cache update might give
-                # incorrect results (since there is no way to determine how
-                # many most recent tokens are to be saved if the tokens exceed
-                # the sliding window length).
-                cache_k = ops.slice_update(
-                    cache_k,
-                    [0, update_start_index, 0, 0],
-                    # We slice the keys and values since if the user has passed
-                    # a sequence of length > `self._sliding_window`. We want to
-                    # prefill the cache using just the most recent values in the
-                    # sliding window.
-                    ops.cast(
-                        key[:, -self._sliding_window :, ...], cache_k.dtype
-                    ),
+                raise ValueError(
+                    "`cache_update_index` should not be set if `cache` is "
+                    f"`None`. Received: cache={cache}, "
+                    f"cache_update_index={cache_update_index}"
                 )
-                cache_v = ops.slice_update(
-                    cache_v,
-                    [0, update_start_index, 0, 0],
-                    ops.cast(
-                        value[:, -self._sliding_window :, ...], cache_v.dtype
-                    ),
-                )
-                cache = ops.stack([cache_k, cache_v], axis=1)
-
-                # Get the required keys and values from the cache.
-                # Since we expect the user to pass a fixed-size cache, we just
-                # pick the first few slices up-to and including the newly computed
-                # keys and values.
-                cache_k = cache_k[:, :update_end_index, ...]
-                cache_v = cache_v[:, :update_end_index, ...]
-
-            key = ops.cast(cache_k, dtype=self.compute_dtype)
-            value = ops.cast(cache_v, dtype=self.compute_dtype)
-        else:
-            # Compute keys and values
             key, value = _compute_key_value(hidden_states)
 
         # [batch_shape, seq_len, num_key_value_heads, head_dim]
@@ -260,15 +209,15 @@ class CachedMistralAttention(keras.layers.Layer):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        attention_scores = ops.einsum(self._dot_product_equation, key, query)
+        attention_scores = ops.einsum(self._dot_product_equation, query, key)
 
         norm_factor = ops.sqrt(ops.cast(self._head_dim, self.compute_dtype))
 
         attention_scores = attention_scores / norm_factor
-
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )
+        attention_scores = ops.cast(attention_scores, self.compute_dtype)
         attention_output = ops.einsum(
             self._combine_equation, attention_scores, value
         )
diff --git a/keras_nlp/models/mistral/mistral_backbone.py b/keras_nlp/models/mistral/mistral_backbone.py
index 107e5699..375d3c54 100644
--- a/keras_nlp/models/mistral/mistral_backbone.py
+++ b/keras_nlp/models/mistral/mistral_backbone.py
@@ -64,7 +64,10 @@ class MistralBackbone(Backbone):
             layers in each transformer decoder. Only `sliding_window` number of tokens
             are saved in the cache and used to generate the next token.
             Defaults to `512`.
-        dtype (str, optional): The dtype policy for the mistral model.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
 
@@ -107,10 +110,10 @@ class MistralBackbone(Backbone):
         layer_norm_epsilon=1e-6,
         sliding_window=512,
         dropout=0,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
-        dtype = kwargs.pop("dtype", keras.backend.floatx())
         self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
diff --git a/keras_nlp/models/mistral/mistral_causal_lm.py b/keras_nlp/models/mistral/mistral_causal_lm.py
new file mode 100644
index 00000000..22defbc4
--- /dev/null
+++ b/keras_nlp/models/mistral/mistral_causal_lm.py
@@ -0,0 +1,213 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
+from keras_nlp.backend import ops
+from keras_nlp.models.generative_task import GenerativeTask
+from keras_nlp.models.mistral.mistral_backbone import MistralBackbone
+from keras_nlp.models.mistral.mistral_causal_lm_preprocessor import (
+    MistralCausalLMPreprocessor,
+)
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras_nlp_export("keras_nlp.models.MistralCausalLM")
+class MistralCausalLM(GenerativeTask):
+    """An end-to-end Mistral model for causal language modeling.
+
+    A causal language model (LM) predicts the next token based on previous
+    tokens. This task setup can be used to train the model unsupervised on
+    plain text input, or to autoregressively generate plain text similar to
+    the data used for training. This task can be used for pre-training or
+    fine-tuning a GPT-NeoX model, simply by calling `fit()`.
+
+    This model has a `generate()` method, which generates text based on a
+    prompt. The generation strategy used is controlled by an additional
+    `sampler` argument on `compile()`. You can recompile the model with
+    different `keras_nlp.samplers` objects to control the generation. By
+    default, `"top_k"` sampling will be used.
+
+    Args:
+        backbone: A `keras_nlp.models.MistralBackbone` instance.
+        preprocessor: A `keras_nlp.models.MistralCausalLMPreprocessor` or `None`.
+            If `None`, this model will not apply preprocessing, and inputs
+            should be preprocessed before calling the model.
+    """
+
+    def __init__(self, backbone, preprocessor=None, **kwargs):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
+        inputs = backbone.inputs
+        hidden_states = backbone(inputs)
+        outputs = backbone.token_embedding(hidden_states, reverse=True)
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === Default compilation ===
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(2e-5),
+            metrics=[keras.metrics.SparseCategoricalAccuracy()],
+            jit_compile=True,
+        )
+
+    @classproperty
+    def backbone_cls(cls):
+        return MistralBackbone
+
+    @classproperty
+    def preprocessor_cls(cls):
+        return MistralCausalLMPreprocessor
+
+    def call_with_cache(
+        self,
+        token_ids,
+        cache,
+        cache_update_index,
+    ):
+        """Forward pass of `MistralCausalLM` with cache.
+
+        `call_with_cache` adds an additional forward pass for the model for
+        autoregressive inference. Unlike calling the model directly, this method
+        allows caching previous key/value Tensors in multi-head attention layer,
+        and avoids recomputing the outputs of seen tokens.
+
+        Args:
+            token_ids: a dense int Tensor with shape `(batch_size, max_length)`.
+            cache: a dense float Tensor, the cache of key and value.
+            cache_update_index: int, or int Tensor. The index of current inputs
+            in the whole sequence.
+
+        Returns:
+            A (logits, hidden_states, cache) tuple. Where `logits` is the
+            language model logits for the input token_ids, `hidden_states` is
+            the final hidden representation of the input tokens, and `cache` is
+            the decoding cache.
+        """
+        x = self.backbone.token_embedding(token_ids)
+        # Each decoder layer has a cache; we update them separately.
+        updated_cache = []
+        for i in range(self.backbone.num_layers):
+            current_cache = cache[:, i, ...]
+            x, next_cache = self.backbone.transformer_layers[i](
+                x,
+                self_attention_cache=current_cache,
+                self_attention_cache_update_index=cache_update_index,
+            )
+            updated_cache.append(next_cache)
+        cache = ops.stack(updated_cache, axis=1)
+        hidden_states = x = self.backbone.layer_norm(x)
+        logits = self.backbone.token_embedding(x, reverse=True)
+        return logits, hidden_states, cache
+
+    def _build_cache(self, token_ids):
+        """Build an empty cache for use with `call_with_cache()`."""
+        batch_size = ops.shape(token_ids)[0]
+        max_length = ops.shape(token_ids)[1]
+        num_layers = self.backbone.num_layers
+        num_key_value_heads = self.backbone.num_key_value_heads
+        head_dim = self.backbone.hidden_dim // self.backbone.num_query_heads
+        shape = [
+            batch_size,
+            num_layers,
+            2,
+            max_length,
+            num_key_value_heads,
+            head_dim,
+        ]
+        cache = ops.zeros(shape, dtype=self.compute_dtype)
+        # Seed the cache.
+        _, hidden_states, cache = self.call_with_cache(token_ids, cache, 0)
+        return hidden_states, cache
+
+    def generate_step(
+        self,
+        inputs,
+        end_token_id=None,
+    ):
+        """A compilable generation function for a single batch of inputs.
+
+        This function represents the inner, XLA-compilable, generation function
+        for a single batch of inputs. Inputs should have the same structure as
+        model inputs, a dictionary with keys `"token_ids"` and `"padding_mask"`.
+
+        Args:
+            inputs: A dictionary with two keys `"token_ids"` and
+                `"padding_mask"` and batched tensor values.
+            end_token_id: The id of the end token to stop on. If all
+                sequences have produced a new `end_token_id`, generation
+                will stop.
+        """
+        token_ids, padding_mask = inputs["token_ids"], inputs["padding_mask"]
+        # Create and seed cache with a single forward pass.
+        hidden_states, cache = self._build_cache(token_ids)
+        # Compute the lengths of all user inputted tokens ids.
+        row_lengths = ops.sum(ops.cast(padding_mask, "int32"), axis=-1)
+        # Start at the first index that has no user inputted id.
+        index = ops.min(row_lengths)
+
+        def next(prompt, cache, index):
+            # The cache index is the index of our previous token.
+            cache_update_index = index - 1
+            batch_size = ops.shape(prompt)[0]
+            prompt = ops.slice(prompt, [0, cache_update_index], [batch_size, 1])
+            logits, hidden_states, cache = self.call_with_cache(
+                prompt,
+                cache,
+                cache_update_index,
+            )
+            return (
+                ops.squeeze(logits, axis=1),
+                ops.squeeze(hidden_states, axis=1),
+                cache,
+            )
+
+        token_ids = self._sampler(
+            next=next,
+            prompt=token_ids,
+            cache=cache,
+            index=index,
+            mask=padding_mask,
+            end_token_id=end_token_id,
+            hidden_states=hidden_states,
+        )
+
+        # Compute an output padding mask with the token ids we updated.
+        if end_token_id is not None:
+            # Build a mask of `end_token_id` locations not in the original
+            # prompt (not in locations where `padding_mask` is True).
+            end_locations = ops.logical_and(
+                ops.equal(token_ids, end_token_id),
+                ops.logical_not(padding_mask),
+            )
+            end_locations = ops.cast(end_locations, "int32")
+            # Use cumsum to get ones in all locations after end_locations.
+            cumsum = ops.cast(ops.cumsum(end_locations, axis=-1), "int32")
+            overflow = cumsum - end_locations
+            # Our padding mask is the inverse of these overflow locations.
+            padding_mask = ops.logical_not(ops.cast(overflow, "bool"))
+        else:
+            # Without early stopping, all locations will have been updated.
+            padding_mask = ops.ones_like(token_ids, dtype="bool")
+        return {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
diff --git a/keras_nlp/models/mistral/mistral_causal_lm_preprocessor.py b/keras_nlp/models/mistral/mistral_causal_lm_preprocessor.py
new file mode 100644
index 00000000..c8a08217
--- /dev/null
+++ b/keras_nlp/models/mistral/mistral_causal_lm_preprocessor.py
@@ -0,0 +1,171 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+from absl import logging
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import ops
+from keras_nlp.models.mistral.mistral_preprocessor import MistralPreprocessor
+from keras_nlp.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+
+
+@keras_nlp_export("keras_nlp.models.MistralCausalLMPreprocessor")
+class MistralCausalLMPreprocessor(MistralPreprocessor):
+    """Mistral Causal LM preprocessor.
+
+    This preprocessing layer is meant for use with
+    `keras_nlp.models.MistralCausalLM`. By default, it will take in batches of
+    strings, and return outputs in a `(x, y, sample_weight)` format, where the
+    `y` label is the next token id in the `x` sequence.
+
+    For use with generation, the layer also exposes two methods
+    `generate_preprocess()` and `generate_postprocess()`. When this preprocessor
+    is attached to a `keras_nlp.models.MistralCausalLM` instance, these methods
+    will be called implicitly in `generate()`. They can also be called
+    standalone (e.g. to precompute preprocessing inputs for generation in a
+    separate process).
+
+    Args:
+        tokenizer: A `keras_nlp.models.MistralTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        add_start_token: If `True`, the preprocessor will prepend the tokenizer
+            start token to each input sequence. Default is `True`.
+        add_end_token: If `True`, the preprocessor will append the tokenizer
+            end token to each input sequence. Default is `False`.
+
+    Call arguments:
+        x: A string, `tf.Tensor` or list of python strings.
+        y: Label data. Should always be `None` as the layer generates labels.
+        sample_weight: Label weights. Should always be `None` as the layer
+            generates label weights.
+        sequence_length: Pass to override the configured `sequence_length` of
+            the layer.
+
+    Examples:
+    ```python
+    # Load the preprocessor from a preset.
+    preprocessor = keras_nlp.models.MistralCausalLMPreprocessor.from_preset(
+        "mistral_base_en"
+    )
+
+    # Tokenize and pack a single sentence.
+    sentence = tf.constant("League of legends")
+    preprocessor(sentence)
+    # Same output.
+    preprocessor("League of legends")
+
+    # Tokenize a batch of sentences.
+    sentences = tf.constant(["Taco tuesday", "Fish taco please!"])
+    preprocessor(sentences)
+    # Same output.
+    preprocessor(["Taco tuesday", "Fish taco please!"])
+
+    # Map a dataset to preprocess a single sentence.
+    features = tf.constant(
+        [
+            "Avatar 2 is amazing!",
+            "Well, I am not sure.",
+        ]
+    )
+    labels = tf.constant([1, 0])
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess unlabled sentences.
+    ds = tf.data.Dataset.from_tensor_slices(features)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        if y is not None or sample_weight is not None:
+            logging.warning(
+                "`MistralCausalLMPreprocessor` generates `y` and "
+                "`sample_weight` based on your input data, but your data "
+                "already contains `y` or `sample_weight`. Your `y` and "
+                "`sample_weight` will be ignored."
+            )
+        sequence_length = sequence_length or self.sequence_length
+
+        x = convert_inputs_to_list_of_tensor_segments(x)[0]
+        x = self.tokenizer(x)
+        # Pad with one extra token to account for the truncation below.
+        token_ids, padding_mask = self.packer(
+            x,
+            sequence_length=sequence_length + 1,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        # The last token does not have a next token, so we truncate it out.
+        x = {
+            "token_ids": token_ids[..., :-1],
+            "padding_mask": padding_mask[..., :-1],
+        }
+        # Target `y` will be the next token.
+        y, sample_weight = token_ids[..., 1:], padding_mask[..., 1:]
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    def generate_preprocess(
+        self,
+        x,
+        sequence_length=None,
+    ):
+        x = convert_inputs_to_list_of_tensor_segments(x)[0]
+        x = self.tokenizer(x)
+        token_ids, padding_mask = self.packer(
+            x, sequence_length=sequence_length, add_end_value=False
+        )
+        return {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+
+    def generate_postprocess(
+        self,
+        x,
+    ):
+        """Covert integer token output to strings for generation.
+
+        This method reverses `generate_preprocess()`, by first removing all
+        padding and start/end tokens, and then converting the integer sequence
+        back to a string.
+        """
+        token_ids, padding_mask = x["token_ids"], x["padding_mask"]
+        # Convert the inputs to numpy arrays if they aren't a tensor already.
+        if not isinstance(token_ids, tf.Tensor):
+            token_ids = ops.convert_to_numpy(token_ids)
+            # Make sure the numpy array has type `int32` since
+            # `SentencePieceProcessor.detokenize` only accepts `int32` arrays.
+            token_ids = token_ids.astype("int32")
+        if not isinstance(padding_mask, tf.Tensor):
+            padding_mask = ops.convert_to_numpy(padding_mask)
+            padding_mask = padding_mask.astype("bool")
+        # Strip any special tokens during detokenization (e.g. the start and
+        # end markers). In the future we could make this configurable.
+        padding_mask = padding_mask & (token_ids != self.tokenizer.end_token_id)
+        padding_mask = padding_mask & (
+            token_ids != self.tokenizer.start_token_id
+        )
+        token_ids = tf.ragged.boolean_mask(token_ids, padding_mask)
+        return self.tokenizer.detokenize(token_ids)
diff --git a/keras_nlp/models/mistral/mistral_layer_norm.py b/keras_nlp/models/mistral/mistral_layer_norm.py
index 9f9ddf26..e714a854 100644
--- a/keras_nlp/models/mistral/mistral_layer_norm.py
+++ b/keras_nlp/models/mistral/mistral_layer_norm.py
@@ -32,7 +32,6 @@ class MistralLayerNormalization(keras.layers.Layer):
             trainable=True,
             shape=(self._dim,),
             initializer="ones",
-            dtype=self.compute_dtype,
         )
         self.built = True
 
diff --git a/keras_nlp/models/mistral/mistral_transformer_decoder.py b/keras_nlp/models/mistral/mistral_transformer_decoder.py
index 9b6f7fdb..7c90ab91 100644
--- a/keras_nlp/models/mistral/mistral_transformer_decoder.py
+++ b/keras_nlp/models/mistral/mistral_transformer_decoder.py
@@ -36,7 +36,7 @@ class MistralTransformerDecoder(keras.layers.Layer):
         num_key_value_heads,
         rope_max_wavelength=10000,
         rope_scaling_factor=1.0,
-        activation="relu",
+        activation="silu",
         layer_norm_epsilon=1e-5,
         kernel_initializer="glorot_uniform",
         sliding_window=512,
@@ -73,20 +73,20 @@ class MistralTransformerDecoder(keras.layers.Layer):
             sliding_window=self.sliding_window,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             dropout=self.dropout,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="self_attention",
         )
         self._self_attention_layer.build(decoder_sequence_shape)
 
         self._self_attention_layernorm = MistralLayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="self_attention_layernorm",
-            dtype=self.compute_dtype,
         )
         self._self_attention_layernorm.build(decoder_sequence_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="self_attention_dropout",
         )
 
@@ -95,7 +95,7 @@ class MistralTransformerDecoder(keras.layers.Layer):
             self.intermediate_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             use_bias=False,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(decoder_sequence_shape)
@@ -105,6 +105,7 @@ class MistralTransformerDecoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             use_bias=False,
+            dtype=self.dtype_policy,
             name="feedforward_gate_dense",
         )
         self._feedforward_gate_dense.build(decoder_sequence_shape)
@@ -113,7 +114,7 @@ class MistralTransformerDecoder(keras.layers.Layer):
             self.hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             use_bias=False,
-            dtype=self.compute_dtype,
+            dtype=self.dtype_policy,
             name="feedforward_output_dense",
         )
 
@@ -125,8 +126,8 @@ class MistralTransformerDecoder(keras.layers.Layer):
 
         self._feedforward_layernorm = MistralLayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
             name="feedforward_layernorm",
-            dtype=self.compute_dtype,
         )
         self._feedforward_layernorm.build(decoder_sequence_shape)
 
@@ -145,6 +146,8 @@ class MistralTransformerDecoder(keras.layers.Layer):
             decoder_sequence=decoder_sequence,
             decoder_padding_mask=decoder_padding_mask,
             decoder_attention_mask=decoder_attention_mask,
+            self_attention_cache=self_attention_cache,
+            self_attention_cache_update_index=self_attention_cache_update_index,
         )
         residual = decoder_sequence
 
@@ -184,23 +187,36 @@ class MistralTransformerDecoder(keras.layers.Layer):
         decoder_sequence,
         decoder_padding_mask,
         decoder_attention_mask,
+        self_attention_cache,
+        self_attention_cache_update_index,
     ):
         decoder_mask = merge_padding_and_attention_mask(
             decoder_sequence, decoder_padding_mask, decoder_attention_mask
         )
         batch_size = ops.shape(decoder_sequence)[0]
         input_length = output_length = ops.shape(decoder_sequence)[1]
+        # We need to handle a rectangular causal mask when doing cached
+        # decoding. For generative inference, `decoder_sequence` will
+        # generally be length 1, and `cache` will be the full generation length.
+        if self_attention_cache is not None:
+            input_length = ops.shape(self_attention_cache)[2]
+
+        cache_update_index = (
+            0
+            if self_attention_cache_update_index is None
+            else self_attention_cache_update_index
+        )
 
         # Mistral uses a banded attention mask
         causal_mask_lower = compute_causal_mask(
-            batch_size, input_length, output_length, 0
+            batch_size, input_length, output_length, cache_update_index
         )
         # Below is a workaround for `ops.triu` for Keras 2.
         # TODO(tirthasheshpatel): Use `ops.triu` once Keras 2 support is removed.
         # causal_mask = ops.triu(causal_mask_lower, k=-self.sliding_window)
-        i = ops.arange(output_length)[:, None]
+        i = ops.arange(output_length)[:, None] + cache_update_index
         j = ops.arange(input_length)[None, :]
-        causal_mask_upper = ops.cast(i <= j + self.sliding_window, "int32")
+        causal_mask_upper = ops.cast(i < j + self.sliding_window, "int32")
         causal_mask = ops.minimum(causal_mask_lower, causal_mask_upper)
 
         return (
diff --git a/keras_nlp/models/opt/opt_backbone.py b/keras_nlp/models/opt/opt_backbone.py
index d04f4b57..0b98a6c6 100644
--- a/keras_nlp/models/opt/opt_backbone.py
+++ b/keras_nlp/models/opt/opt_backbone.py
@@ -57,6 +57,10 @@ class OPTBackbone(Backbone):
             can consume. If `None`, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -91,6 +95,7 @@ class OPTBackbone(Backbone):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=2048,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -99,6 +104,7 @@ class OPTBackbone(Backbone):
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=opt_kernel_initializer(),
+            dtype=dtype,
             name="embeddings",
         )
         self.token_embedding = self.embeddings.token_embedding
@@ -112,13 +118,14 @@ class OPTBackbone(Backbone):
                 layer_norm_epsilon=1e-5,
                 normalize_first=True,
                 kernel_initializer=opt_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
         self.layer_norm = keras.layers.LayerNormalization(
             axis=-1,
             epsilon=1e-5,
-            dtype="float32",
+            dtype=dtype,
             name="layer_norm",
         )
 
diff --git a/keras_nlp/models/roberta/roberta_backbone.py b/keras_nlp/models/roberta/roberta_backbone.py
index 614104d8..1ab61eee 100644
--- a/keras_nlp/models/roberta/roberta_backbone.py
+++ b/keras_nlp/models/roberta/roberta_backbone.py
@@ -61,6 +61,10 @@ class RobertaBackbone(Backbone):
             consume. The sequence length of the input must be less than
             `max_sequence_length` default value. This determines the variable
             shape for positional embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
@@ -96,6 +100,7 @@ class RobertaBackbone(Backbone):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=512,
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -104,17 +109,19 @@ class RobertaBackbone(Backbone):
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=roberta_kernel_initializer(),
+            dtype=dtype,
             name="embeddings",
         )
         self.token_embedding = self.embeddings.token_embedding
         self.embeddings_layer_norm = keras.layers.LayerNormalization(
             axis=-1,
             epsilon=1e-5,  # Original paper uses this epsilon value
-            dtype="float32",
+            dtype=dtype,
             name="embeddings_layer_norm",
         )
         self.embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="embeddings_dropout",
         )
         self.transformer_layers = []
@@ -126,6 +133,7 @@ class RobertaBackbone(Backbone):
                 dropout=dropout,
                 layer_norm_epsilon=1e-5,
                 kernel_initializer=roberta_kernel_initializer(),
+                dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
diff --git a/keras_nlp/models/roberta/roberta_classifier.py b/keras_nlp/models/roberta/roberta_classifier.py
index e3d7666f..887bc657 100644
--- a/keras_nlp/models/roberta/roberta_classifier.py
+++ b/keras_nlp/models/roberta/roberta_classifier.py
@@ -149,22 +149,26 @@ class RobertaClassifier(Task):
         self.preprocessor = preprocessor
         self.pooled_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="pooled_dropout",
         )
         hidden_dim = hidden_dim or backbone.hidden_dim
         self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             activation="tanh",
+            dtype=backbone.dtype_policy,
             name="pooled_dense",
         )
         self.output_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="output_dropout",
         )
         self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=roberta_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
 
diff --git a/keras_nlp/models/roberta/roberta_masked_lm.py b/keras_nlp/models/roberta/roberta_masked_lm.py
index 0e62f4cf..bf961898 100644
--- a/keras_nlp/models/roberta/roberta_masked_lm.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm.py
@@ -111,6 +111,7 @@ class RobertaMaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=roberta_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/t5/t5_backbone.py b/keras_nlp/models/t5/t5_backbone.py
index 3fdf69ff..cf747c50 100644
--- a/keras_nlp/models/t5/t5_backbone.py
+++ b/keras_nlp/models/t5/t5_backbone.py
@@ -67,7 +67,11 @@ class T5Backbone(Backbone):
             layer normalization layers in the Transformer layers.
         tie_embedding_weights: boolean. If `True`, the weights of the token
             embedding and the weights projecting language model outputs from
-            `hidden_dim`
+            `hidden_dim`.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
     """
 
     def __init__(
@@ -83,6 +87,7 @@ class T5Backbone(Backbone):
         use_gated_activation=True,
         layer_norm_epsilon=1e-06,
         tie_embedding_weights=True,
+        dtype=None,
         **kwargs,
     ):
         # Token embedding layer. This layer is shared by encoder and decoder.
@@ -91,10 +96,12 @@ class T5Backbone(Backbone):
             output_dim=hidden_dim,
             tie_weights=tie_embedding_weights,
             embeddings_initializer=keras.initializers.TruncatedNormal(1.0),
+            dtype=dtype,
             name="token_embedding",
         )
         self.encoder_embedding_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="encoder_embedding_dropout",
         )
         self.encoder_transformer_layers = []
@@ -110,19 +117,23 @@ class T5Backbone(Backbone):
                 num_heads=num_heads,
                 use_gated_activation=use_gated_activation,
                 use_relative_attention_bias=bool(i == 0),
+                dtype=dtype,
                 name=f"transformer_encoder_layer_{i}",
             )
             self.encoder_transformer_layers.append(layer)
         self.encoder_layer_norm = T5LayerNorm(
             epsilon=layer_norm_epsilon,
+            dtype=dtype,
             name="encoder_output_layer_norm",
         )
         self.encoder_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="encoder_output_dropout",
         )
         self.decoder_embedding_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="decoder_embedding_dropout",
         )
         self.decoder_transformer_layers = []
@@ -138,15 +149,18 @@ class T5Backbone(Backbone):
                 num_heads=num_heads,
                 use_gated_activation=use_gated_activation,
                 use_relative_attention_bias=bool(i == 0),
+                dtype=dtype,
                 name=f"transformer_decoder_layer_{i}",
             )
             self.decoder_transformer_layers.append(layer)
         self.decoder_layer_norm = T5LayerNorm(
             epsilon=layer_norm_epsilon,
+            dtype=dtype,
             name="decoder_output_layer_norm",
         )
         self.decoder_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="decoder_output_dropout",
         )
 
diff --git a/keras_nlp/models/t5/t5_multi_head_attention.py b/keras_nlp/models/t5/t5_multi_head_attention.py
index 77e7109e..2e3647b1 100644
--- a/keras_nlp/models/t5/t5_multi_head_attention.py
+++ b/keras_nlp/models/t5/t5_multi_head_attention.py
@@ -45,36 +45,43 @@ class T5MultiHeadAttention(keras.layers.Layer):
         self.query_projector = keras.layers.Dense(
             self.inner_dim,
             use_bias=False,
-            name="query_projector",
             kernel_initializer=keras.initializers.RandomNormal(
                 mean=0, stddev=(self.inner_dim * self.key_value_dim) ** -0.5
             ),
+            dtype=self.dtype_policy,
+            name="query_projector",
         )
         self.key_projector = keras.layers.Dense(
             self.inner_dim,
             use_bias=False,
-            name="key_projector",
             kernel_initializer=keras.initializers.RandomNormal(
                 mean=0, stddev=self.inner_dim**-0.5
             ),
+            dtype=self.dtype_policy,
+            name="key_projector",
         )
         self.value_projector = keras.layers.Dense(
             self.inner_dim,
             use_bias=False,
-            name="value_projector",
             kernel_initializer=keras.initializers.RandomNormal(
                 mean=0, stddev=self.inner_dim**-0.5
             ),
+            dtype=self.dtype_policy,
+            name="value_projector",
         )
         self.output_projector = keras.layers.Dense(
             self.hidden_dim,
             use_bias=False,
-            name="output_projector",
             kernel_initializer=keras.initializers.RandomNormal(
                 mean=0, stddev=self.inner_dim**-0.5
             ),
+            dtype=self.dtype_policy,
+            name="output_projector",
+        )
+        self.dropout_layer = keras.layers.Dropout(
+            dropout,
+            dtype=self.dtype_policy,
         )
-        self.dropout_layer = keras.layers.Dropout(dropout)
 
         if self.use_relative_attention_bias:
             self.relative_attention_bias = self.add_weight(
@@ -298,7 +305,7 @@ class T5MultiHeadAttention(keras.layers.Layer):
                 mask = (1.0 - ops.cast(mask, position_bias.dtype)) * -1e9
                 position_bias = position_bias + mask
 
-        scores += position_bias
+        scores += ops.cast(position_bias, scores.dtype)
         weights = ops.nn.softmax(
             scores, axis=-1
         )  # (batch_size, num_heads, query_length, key_length)
diff --git a/keras_nlp/models/t5/t5_transformer_layer.py b/keras_nlp/models/t5/t5_transformer_layer.py
index 27b4c989..697af208 100644
--- a/keras_nlp/models/t5/t5_transformer_layer.py
+++ b/keras_nlp/models/t5/t5_transformer_layer.py
@@ -47,10 +47,17 @@ class T5TransformerLayer(keras.layers.Layer):
             num_heads=num_heads,
             dropout=dropout,
             use_relative_attention_bias=use_relative_attention_bias,
+            dtype=self.dtype_policy,
             name="self_attention",
         )
-        self.self_attention_layer_norm = T5LayerNorm(layer_norm_epsilon)
-        self.self_attention_dropout = keras.layers.Dropout(dropout)
+        self.self_attention_layer_norm = T5LayerNorm(
+            layer_norm_epsilon,
+            dtype=self.dtype_policy,
+        )
+        self.self_attention_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=self.dtype_policy,
+        )
 
         if self.is_decoder:
             self.cross_attention = T5MultiHeadAttention(
@@ -60,39 +67,55 @@ class T5TransformerLayer(keras.layers.Layer):
                 num_heads=num_heads,
                 dropout=dropout,
                 use_relative_attention_bias=False,
+                dtype=self.dtype_policy,
                 name="cross_attention",
             )
-            self.cross_attention_layer_norm = T5LayerNorm(layer_norm_epsilon)
-            self.cross_attention_dropout = keras.layers.Dropout(dropout)
+            self.cross_attention_layer_norm = T5LayerNorm(
+                layer_norm_epsilon,
+                dtype=self.dtype_policy,
+            )
+            self.cross_attention_dropout = keras.layers.Dropout(
+                dropout,
+                dtype=self.dtype_policy,
+            )
 
         self.input_projector = keras.layers.Dense(
             intermediate_dim,
             use_bias=False,
-            name="input_projector",
             activation=keras.activations.get(activation),
             kernel_initializer=keras.initializers.RandomNormal(
                 mean=0, stddev=hidden_dim**-0.5
             ),
+            dtype=self.dtype_policy,
+            name="input_projector",
         )
         if self.use_gated_activation:
             self.gate_projector = keras.layers.Dense(
                 intermediate_dim,
                 use_bias=False,
-                name="gate_projector",
                 kernel_initializer=keras.initializers.RandomNormal(
                     mean=0, stddev=hidden_dim**-0.5
                 ),
+                dtype=self.dtype_policy,
+                name="gate_projector",
             )
         self.output_projector = keras.layers.Dense(
             hidden_dim,
             use_bias=False,
-            name="output_projector",
             kernel_initializer=keras.initializers.RandomNormal(
                 mean=0, stddev=intermediate_dim**-0.5
             ),
+            dtype=self.dtype_policy,
+            name="output_projector",
+        )
+        self.layer_norm = T5LayerNorm(
+            epsilon=layer_norm_epsilon,
+            dtype=self.dtype_policy,
+        )
+        self.dropout_layer = keras.layers.Dropout(
+            dropout,
+            dtype=self.dtype_policy,
         )
-        self.layer_norm = T5LayerNorm(epsilon=layer_norm_epsilon)
-        self.dropout_layer = keras.layers.Dropout(dropout)
 
     def call(
         self,
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
index 1fe8d0b7..783cc0b4 100644
--- a/keras_nlp/models/task.py
+++ b/keras_nlp/models/task.py
@@ -220,9 +220,14 @@ class Task(PipelineModel):
 
         # Backbone case.
         if preset_cls == cls.backbone_cls:
+            # Forward dtype to the backbone.
+            config_overrides = {}
+            if "dtype" in kwargs:
+                config_overrides["dtype"] = kwargs.pop("dtype")
             backbone = load_from_preset(
                 preset,
                 load_weights=load_weights,
+                config_overrides=config_overrides,
             )
             if "preprocessor" in kwargs:
                 preprocessor = kwargs.pop("preprocessor")
diff --git a/keras_nlp/models/whisper/whisper_backbone.py b/keras_nlp/models/whisper/whisper_backbone.py
index 2e842190..c66a61d4 100644
--- a/keras_nlp/models/whisper/whisper_backbone.py
+++ b/keras_nlp/models/whisper/whisper_backbone.py
@@ -75,6 +75,10 @@ class WhisperBackbone(Backbone):
             positional embedding layer.
         max_decoder_sequence_length: int. The maximum sequence length that the
             text decoder can consume.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
 
@@ -112,6 +116,7 @@ class WhisperBackbone(Backbone):
         dropout=0.0,
         max_encoder_sequence_length=3000,
         max_decoder_sequence_length=448,
+        dtype=None,
         **kwargs,
     ):
         assert_tf_backend(self.__class__.__name__)
@@ -122,6 +127,7 @@ class WhisperBackbone(Backbone):
             kernel_size=3,
             strides=1,
             padding="same",
+            dtype=dtype,
             name="encoder_token_embedding_conv_layer_1",
         )
         self.encoder_conv_layer_2 = keras.layers.Conv1D(
@@ -129,22 +135,27 @@ class WhisperBackbone(Backbone):
             kernel_size=3,
             strides=2,
             padding="valid",
+            dtype=dtype,
             name="encoder_token_embedding_conv_layer_2",
         )
         self.encoder_padder = Padder(
+            dtype=dtype,
             name="encoder_padder",
         )
         self.encoder_position_embedding = PositionEmbedding(
             initializer=whisper_kernel_initializer(),
             sequence_length=max_encoder_sequence_length // 2,
+            dtype=dtype,
             name="encoder_position_embedding",
             trainable=False,
         )
         self.encoder_embeddings_add = keras.layers.Add(
+            dtype=dtype,
             name="encoder_embeddings_add",
         )
         self.encoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="encoder_embeddings_dropout",
         )
         self.encoder_transformer_layers = []
@@ -157,25 +168,28 @@ class WhisperBackbone(Backbone):
                 dropout=dropout,
                 kernel_initializer=whisper_kernel_initializer(),
                 normalize_first=True,
+                dtype=dtype,
                 name=f"transformer_encoder_layer_{i}",
             )
             self.encoder_transformer_layers.append(layer)
         self.encoder_layer_norm = keras.layers.LayerNormalization(
-            name="encoder_layer_norm",
             axis=-1,
             epsilon=1e-5,
-            dtype="float32",
+            dtype=dtype,
+            name="encoder_layer_norm",
         )
         self.decoder_embeddings = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_decoder_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=whisper_kernel_initializer(),
+            dtype=dtype,
             name="decoder_token_and_position_embedding",
         )
         self.token_embedding = self.decoder_embeddings.token_embedding
         self.decoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="decoder_embeddings_dropout",
         )
         self.decoder_transformer_layers = []
@@ -188,14 +202,15 @@ class WhisperBackbone(Backbone):
                 layer_norm_epsilon=1e-5,
                 kernel_initializer=whisper_kernel_initializer(),
                 normalize_first=True,
+                dtype=dtype,
                 name=f"transformer_decoder_layer_{i}",
             )
             self.decoder_transformer_layers.append(layer)
         self.decoder_layer_norm = keras.layers.LayerNormalization(
-            name="decoder_layer_norm",
             axis=-1,
             epsilon=1e-5,
-            dtype="float32",
+            dtype=dtype,
+            name="decoder_layer_norm",
         )
 
         # === Functional Model ===
@@ -222,7 +237,7 @@ class WhisperBackbone(Backbone):
         # For the second conv. layer, we cannot use `padding="same"` since
         # that corresponds to a padding size of 1.5 (since stride is 2). Hence,
         # we will manually pad the input.
-        embedded_features = Padder()(embedded_features)
+        embedded_features = self.encoder_padder(embedded_features)
         embedded_features = keras.activations.gelu(
             self.encoder_conv_layer_2(embedded_features),
             approximate=False,
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
index b83b4596..c74a0fd6 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
@@ -52,6 +52,10 @@ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
             consume. The sequence length of the input must be less than
             `max_sequence_length` default value. This determines the variable
             shape for positional embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Examples:
     ```python
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index 45d79eb3..fcd8bfe9 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -162,22 +162,26 @@ class XLMRobertaClassifier(Task):
         self.preprocessor = preprocessor
         self.pooled_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="pooled_dropout",
         )
         hidden_dim = hidden_dim or backbone.hidden_dim
         self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             activation="tanh",
+            dtype=backbone.dtype_policy,
             name="pooled_dense",
         )
         self.output_dropout = keras.layers.Dropout(
             dropout,
+            dtype=backbone.dtype_policy,
             name="output_dropout",
         )
         self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=roberta_kernel_initializer(),
             activation=activation,
+            dtype=backbone.dtype_policy,
             name="logits",
         )
 
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py b/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
index b29aa30d..e231f3dc 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
@@ -114,6 +114,7 @@ class XLMRobertaMaskedLM(Task):
             token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=roberta_kernel_initializer(),
+            dtype=backbone.dtype_policy,
             name="mlm_head",
         )
 
diff --git a/keras_nlp/models/xlnet/relative_attention.py b/keras_nlp/models/xlnet/relative_attention.py
index a11ae3fd..6ae56d14 100644
--- a/keras_nlp/models/xlnet/relative_attention.py
+++ b/keras_nlp/models/xlnet/relative_attention.py
@@ -154,6 +154,7 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
                 output_rank - 1, [self._num_heads, self._key_dim]
             ),
             bias_axes=bias_axes if self._use_bias else None,
+            dtype=self.dtype_policy,
             name="query",
             **self._get_common_kwargs_for_sublayer(),
         )
@@ -168,6 +169,7 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
                 output_rank - 1, [self._num_heads, self._key_dim]
             ),
             bias_axes=bias_axes if self._use_bias else None,
+            dtype=self.dtype_policy,
             name="key",
             **self._get_common_kwargs_for_sublayer(),
         )
@@ -182,6 +184,7 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
                 output_rank - 1, [self._num_heads, self._value_dim]
             ),
             bias_axes=bias_axes if self._use_bias else None,
+            dtype=self.dtype_policy,
             name="value",
             **self._get_common_kwargs_for_sublayer(),
         )
@@ -197,6 +200,7 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
                 output_rank - 1, [self._query_shape[-1]]
             ),
             bias_axes=None,
+            dtype=self.dtype_policy,
             name="attention_output",
             **self._get_common_kwargs_for_sublayer(),
         )
@@ -213,6 +217,7 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
                 output_rank - 1, [self._num_heads, self._key_dim]
             ),
             bias_axes=None,
+            dtype=self.dtype_policy,
             name="encoding",
             **self._get_common_kwargs_for_sublayer(),
         )
diff --git a/keras_nlp/models/xlnet/xlnet_backbone.py b/keras_nlp/models/xlnet/xlnet_backbone.py
index 45a4b8f4..0d660bea 100644
--- a/keras_nlp/models/xlnet/xlnet_backbone.py
+++ b/keras_nlp/models/xlnet/xlnet_backbone.py
@@ -52,6 +52,10 @@ class XLNetBackbone(Backbone):
         bias_initializer: string or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for
             the dense and multiheaded relative attention layers.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
 
     Call arguments:
         token_ids: Indices of input sequence tokens in the vocabulary of shape
@@ -101,6 +105,7 @@ class XLNetBackbone(Backbone):
         activation="gelu",
         kernel_initializer_range=0.02,
         bias_initializer="zeros",
+        dtype=None,
         **kwargs,
     ):
         # === Layers ===
@@ -108,14 +113,17 @@ class XLNetBackbone(Backbone):
             vocabulary_size=vocabulary_size,
             hidden_dim=hidden_dim,
             dropout=dropout,
+            dtype=dtype,
             name="content_query_embedding",
         )
         self.attn_mask_layer = XLNetAttentionMaskLayer(
             hidden_dim=hidden_dim,
             kernel_initializer_range=kernel_initializer_range,
+            dtype=dtype,
             name="encoder_block_attn_mask_layer",
         )
         self.seg_mat_layer = XLNetSegmentMatrixLayer(
+            dtype=dtype,
             name="encoder_block_seg_mat_layer",
         )
         head_dim = hidden_dim // num_heads
@@ -131,11 +139,13 @@ class XLNetBackbone(Backbone):
                 layer_norm_epsilon=1e-12,
                 kernel_initializer_range=kernel_initializer_range,
                 bias_initializer=bias_initializer,
+                dtype=dtype,
                 name=f"xlnet_encoder_{i}",
             )
             self.transformer_layers.append(layer)
         self.dropout = keras.layers.Dropout(
             dropout,
+            dtype=dtype,
             name="dropout",
         )
 
diff --git a/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py b/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
index a2bc2d0d..2de2c317 100644
--- a/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
+++ b/keras_nlp/models/xlnet/xlnet_content_and_query_embedding.py
@@ -58,7 +58,8 @@ class ContentAndQueryEmbedding(keras.layers.Layer):
                     ops.shape(pos_emb)[0],
                     ops.shape(pos_emb)[1] * bsz,
                     ops.shape(pos_emb)[2],
-                ]
+                ],
+                dtype=self.compute_dtype,
             )
             * pos_emb
         )
@@ -67,12 +68,14 @@ class ContentAndQueryEmbedding(keras.layers.Layer):
 
     def relative_positional_encoding(self, qlen, klen, bsz=None, clamp_len=-1):
         """create relative positional encoding."""
-        freq_seq = ops.arange(0, self.hidden_dim, 2.0, dtype=self.compute_dtype)
+        freq_seq = ops.arange(0, self.hidden_dim, 2.0, dtype="float32")
+        freq_seq = ops.cast(freq_seq, self.compute_dtype)
         inv_freq = 1 / (10000 ** (freq_seq / self.hidden_dim))
 
         beg, end = klen, -qlen
 
-        fwd_pos_seq = ops.arange(beg, end, -1.0, dtype=self.compute_dtype)
+        fwd_pos_seq = ops.arange(beg, end, -1.0, dtype="float32")
+        fwd_pos_seq = ops.cast(fwd_pos_seq, self.compute_dtype)
         if clamp_len > 0:
             fwd_pos_seq = ops.clip(
                 fwd_pos_seq, x_min=-clamp_len, x_max=clamp_len
@@ -85,11 +88,14 @@ class ContentAndQueryEmbedding(keras.layers.Layer):
         self.word_embed = keras.layers.Embedding(
             input_dim=self.vocabulary_size,
             output_dim=self.hidden_dim,
+            dtype=self.dtype_policy,
             name="word_embedding",
         )
         self.word_embed.build(input_shape)
-        self.dropout_layer = keras.layers.Dropout(self.dropout)
-
+        self.dropout_layer = keras.layers.Dropout(
+            self.dropout,
+            dtype=self.dtype_policy,
+        )
         super().build(input_shape)
 
     def call(
diff --git a/keras_nlp/models/xlnet/xlnet_encoder.py b/keras_nlp/models/xlnet/xlnet_encoder.py
index bb8e56e4..6aa44b15 100644
--- a/keras_nlp/models/xlnet/xlnet_encoder.py
+++ b/keras_nlp/models/xlnet/xlnet_encoder.py
@@ -94,25 +94,34 @@ class XLNetEncoder(keras.layers.Layer):
             key_dim=self.head_dim,
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
+            dtype=self.dtype_policy,
             name="rel_attn",
         )
 
         self.layer_norm = keras.layers.LayerNormalization(
-            epsilon=self.layer_norm_epsilon, name="layer_norm_rel_attn"
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm_rel_attn",
         )
         self.layer_norm.build(input_shape)
 
-        self.dropout_attn = keras.layers.Dropout(self.dropout)
+        self.dropout_attn = keras.layers.Dropout(
+            self.dropout,
+            dtype=self.dtype_policy,
+        )
 
         # Feed-Forward Part
         self.layer_norm_ff = keras.layers.LayerNormalization(
-            epsilon=self.layer_norm_epsilon, name="layer_norm_ff"
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="layer_norm_ff",
         )
         self.layer_norm_ff.build(input_shape)
 
         self.feedforward_intermediate_dense = keras.layers.Dense(
             self.intermediate_dim,
             kernel_initializer=self.kernel_initializer,
+            dtype=self.dtype_policy,
             name="feedforward_intermediate_dense",
         )
         self.feedforward_intermediate_dense.build(input_shape)
@@ -120,6 +129,7 @@ class XLNetEncoder(keras.layers.Layer):
         self.feedforward_output_dense = keras.layers.Dense(
             self.hidden_dim,
             kernel_initializer=self.kernel_initializer,
+            dtype=self.dtype_policy,
             name="feedforward_output_dense",
         )
         self.feedforward_output_dense.build(
@@ -128,7 +138,10 @@ class XLNetEncoder(keras.layers.Layer):
             )
         )
 
-        self.dropout_ff = keras.layers.Dropout(self.dropout)
+        self.dropout_ff = keras.layers.Dropout(
+            self.dropout,
+            dtype=self.dtype_policy,
+        )
 
         self.activation_function_ff = keras.activations.get(self.activation)
 


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 16b423e0..6bf43114 100644
--- a/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/modeling/cached_multi_head_attention_test.py
@@ -39,7 +39,7 @@ class CachedMultiHeadAttentionTest(TestCase):
             expected_num_non_trainable_variables=1,
             # Keras 2 does not handle mixed precision correctly when not set
             # globally.
-            run_mixed_precision_check=config.keras_3(),
+            run_precision_checks=config.keras_3(),
         )
 
     def test_cache_call_is_correct(self):
diff --git a/keras_nlp/layers/modeling/masked_lm_head_test.py b/keras_nlp/layers/modeling/masked_lm_head_test.py
index 8d22ea03..69a62889 100644
--- a/keras_nlp/layers/modeling/masked_lm_head_test.py
+++ b/keras_nlp/layers/modeling/masked_lm_head_test.py
@@ -58,6 +58,7 @@ class MaskedLMHeadTest(TestCase):
             },
             expected_output_shape=(4, 5, 100),
             expected_num_trainable_weights=6,
+            run_precision_checks=False,
         )
 
     def test_value_error_when_neither_embedding_or_vocab_size_set(self):
diff --git a/keras_nlp/layers/modeling/sine_position_encoding_test.py b/keras_nlp/layers/modeling/sine_position_encoding_test.py
index 80dad26c..e7ea4b1b 100644
--- a/keras_nlp/layers/modeling/sine_position_encoding_test.py
+++ b/keras_nlp/layers/modeling/sine_position_encoding_test.py
@@ -107,13 +107,3 @@ class SinePositionEncodingTest(TestCase):
                 sequential_output, (0, i, 0), parial_output
             )
         self.assertAllClose(full_output, sequential_output)
-
-    def test_float16_dtype(self):
-        pos_encoding = SinePositionEncoding(dtype="float16")
-        seq_length = 100
-        hidden_size = 32
-        inputs = keras.Input(shape=(seq_length, hidden_size))
-        outputs = pos_encoding(inputs)
-
-        # output dtype for this layer should be tf.float16.
-        self.assertEqual(outputs.dtype, "float16")
diff --git a/keras_nlp/layers/modeling/transformer_decoder_test.py b/keras_nlp/layers/modeling/transformer_decoder_test.py
index fda2ea7e..4757a0a4 100644
--- a/keras_nlp/layers/modeling/transformer_decoder_test.py
+++ b/keras_nlp/layers/modeling/transformer_decoder_test.py
@@ -49,7 +49,6 @@ class TransformerDecoderTest(TestCase):
         ("with_norm_first", True),
     )
     def test_layer_behaviors_with_cross_attention(self, normalize_first):
-        pass
         self.run_layer_test(
             cls=TransformerDecoder,
             init_kwargs={
diff --git a/keras_nlp/models/mistral/mistral_causal_lm_preprocessor_test.py b/keras_nlp/models/mistral/mistral_causal_lm_preprocessor_test.py
new file mode 100644
index 00000000..42099501
--- /dev/null
+++ b/keras_nlp/models/mistral/mistral_causal_lm_preprocessor_test.py
@@ -0,0 +1,81 @@
+# Copyright 2024 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from keras_nlp.models.mistral.mistral_causal_lm_preprocessor import (
+    MistralCausalLMPreprocessor,
+)
+from keras_nlp.models.mistral.mistral_tokenizer import MistralTokenizer
+from keras_nlp.tests.test_case import TestCase
+
+
+class MistralCausalLMPreprocessorTest(TestCase):
+    def setUp(self):
+        self.tokenizer = MistralTokenizer(
+            # Generated using create_mistral_test_proto.py
+            proto=os.path.join(
+                self.get_test_data_dir(), "mistral_test_vocab.spm"
+            )
+        )
+        self.init_kwargs = {
+            "tokenizer": self.tokenizer,
+            "sequence_length": 8,
+        }
+        self.input_data = (["the quick brown fox"],)
+
+    def test_preprocessor_basics(self):
+        self.run_preprocessing_layer_test(
+            cls=MistralCausalLMPreprocessor,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output=(
+                {
+                    "token_ids": [[1, 3, 8, 4, 6, 0, 0, 0]],
+                    "padding_mask": [[1, 1, 1, 1, 1, 0, 0, 0]],
+                },
+                [[3, 8, 4, 6, 0, 0, 0, 0]],  # Pass through labels.
+                [[1, 1, 1, 1, 0, 0, 0, 0]],  # Pass through sample_weights.
+            ),
+        )
+
+    def test_no_start_end_token(self):
+        input_data = ["the quick brown fox"] * 4
+
+        preprocessor = MistralCausalLMPreprocessor(
+            **self.init_kwargs,
+            add_start_token=False,
+            add_end_token=False,
+        )
+        x, y, sw = preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [[3, 8, 4, 6, 0, 0, 0, 0]] * 4)
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 0, 0, 0, 0]] * 4)
+        self.assertAllEqual(y, [[8, 4, 6, 0, 0, 0, 0, 0]] * 4)
+        self.assertAllEqual(sw, [[1, 1, 1, 0, 0, 0, 0, 0]] * 4)
+
+    def test_generate_preprocess(self):
+        input_data = "the quick brown fox"
+        preprocessor = MistralCausalLMPreprocessor(**self.init_kwargs)
+        x = preprocessor.generate_preprocess(input_data)
+        self.assertAllEqual(x["token_ids"], [1, 3, 8, 4, 6, 0, 0, 0])
+        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 0, 0, 0])
+
+    def test_generate_postprocess(self):
+        input_data = {
+            "token_ids": [1, 3, 8, 4, 6, 0, 0, 0],
+            "padding_mask": [1, 1, 1, 1, 1, 0, 0, 0],
+        }
+        preprocessor = MistralCausalLMPreprocessor(**self.init_kwargs)
+        x = preprocessor.generate_postprocess(input_data)
+        self.assertAllEqual(x, "the quick brown fox")
diff --git a/keras_nlp/models/mistral/mistral_causal_lm_test.py b/keras_nlp/models/mistral/mistral_causal_lm_test.py
new file mode 100644
index 00000000..3f9d7fab
--- /dev/null
+++ b/keras_nlp/models/mistral/mistral_causal_lm_test.py
@@ -0,0 +1,130 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from unittest.mock import patch
+
+import pytest
+
+from keras_nlp.backend import ops
+from keras_nlp.models.mistral.mistral_backbone import MistralBackbone
+from keras_nlp.models.mistral.mistral_causal_lm import MistralCausalLM
+from keras_nlp.models.mistral.mistral_causal_lm_preprocessor import (
+    MistralCausalLMPreprocessor,
+)
+from keras_nlp.models.mistral.mistral_tokenizer import MistralTokenizer
+from keras_nlp.tests.test_case import TestCase
+
+
+class MistralCausalLMTest(TestCase):
+    def setUp(self):
+        self.preprocessor = MistralCausalLMPreprocessor(
+            MistralTokenizer(
+                # Generated using create_mistral_test_proto.py
+                proto=os.path.join(
+                    self.get_test_data_dir(), "mistral_test_vocab.spm"
+                )
+            ),
+            sequence_length=8,
+        )
+        self.backbone = MistralBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            num_query_heads=4,
+            num_key_value_heads=2,
+            hidden_dim=8,
+            intermediate_dim=16,
+        )
+        self.init_kwargs = {
+            "preprocessor": self.preprocessor,
+            "backbone": self.backbone,
+        }
+        self.train_data = (["the quick brown fox", "the earth is round"],)
+        self.input_data = self.preprocessor(*self.train_data)[0]
+
+    def test_causal_lm_basics(self):
+        self.run_task_test(
+            cls=MistralCausalLM,
+            init_kwargs=self.init_kwargs,
+            train_data=self.train_data,
+            expected_output_shape=(2, 8, 10),
+        )
+
+    def test_generate(self):
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        # String input.
+        prompt = "the quick brown fox"
+        output = causal_lm.generate(prompt)
+        self.assertTrue(prompt in output)
+        # Int tensor input.
+        prompt_ids = self.preprocessor.generate_preprocess([prompt])
+        causal_lm.preprocessor = None
+        outputs = causal_lm.generate(prompt_ids)
+        # Assert prompt is in output in token id space.
+        self.assertAllEqual(
+            outputs["token_ids"][:, :5],
+            prompt_ids["token_ids"][:, :5],
+        )
+        self.assertAllEqual(
+            outputs["padding_mask"][:, :5],
+            prompt_ids["padding_mask"][:, :5],
+        )
+
+    def test_early_stopping(self):
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        call_with_cache = causal_lm.call_with_cache
+
+        def wrapper(*args, **kwargs):
+            """Modify output logits to always favor end_token_id"""
+            logits, hidden_states, cache = call_with_cache(*args, **kwargs)
+            index = self.preprocessor.tokenizer.end_token_id
+            update = ops.ones_like(logits)[:, :, index] * 1.0e9
+            update = ops.expand_dims(update, axis=-1)
+            logits = ops.slice_update(logits, (0, 0, index), update)
+            return logits, hidden_states, cache
+
+        with patch.object(causal_lm, "call_with_cache", wraps=wrapper):
+            prompt = ["the quick brown fox", "the earth"]
+            output = causal_lm.generate(prompt)
+            # We should immediately abort and output the prompt.
+            self.assertEqual(prompt, output)
+
+    def test_generate_compilation(self):
+        causal_lm = MistralCausalLM(**self.init_kwargs)
+        # Assert we do not recompile with successive calls.
+        causal_lm.generate("the quick brown fox")
+        first_fn = causal_lm.generate_function
+        causal_lm.generate("the quick brown fox")
+        second_fn = causal_lm.generate_function
+        self.assertEqual(first_fn, second_fn)
+        # Assert we do recompile after compile is called.
+        causal_lm.compile(sampler="greedy")
+        self.assertIsNone(causal_lm.generate_function)
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=MistralCausalLM,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in MistralCausalLM.presets:
+            self.run_preset_test(
+                cls=MistralCausalLM,
+                preset=preset,
+                input_data=self.input_data,
+            )
diff --git a/keras_nlp/tests/test_case.py b/keras_nlp/tests/test_case.py
index 10102628..50e77335 100644
--- a/keras_nlp/tests/test_case.py
+++ b/keras_nlp/tests/test_case.py
@@ -87,7 +87,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,
+        run_precision_checks=True,
     ):
         """Run basic tests for a modeling layer."""
         # Serialization test.
@@ -187,24 +187,8 @@ 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 = 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")
+        if run_precision_checks:
+            self.run_precision_test(cls, init_kwargs, input_data)
 
     def run_preprocessing_layer_test(
         self,
@@ -283,6 +267,40 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
                     lst.remove("__annotations__")
             self.assertEqual(set(ref_dir), set(new_dir))
 
+    def run_precision_test(self, cls, init_kwargs, input_data):
+        # Keras 2 has some errors as non-float32 precision.
+        if not config.keras_3():
+            return
+        # Never test mixed precision on torch CPU. Torch lacks support.
+        if config.backend() == "torch":
+            import torch
+
+            if not torch.cuda.is_available():
+                return
+
+        for policy in ["mixed_float16", "mixed_bfloat16", "bfloat16"]:
+            policy = keras.mixed_precision.Policy(policy)
+            layer = cls(**{**init_kwargs, "dtype": policy})
+            if isinstance(layer, keras.Model):
+                output_data = layer(input_data)
+            elif 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, policy.compute_dtype)
+            for weight in layer.weights:
+                if is_float_dtype(weight.dtype):
+                    self.assertDTypeEqual(weight, policy.variable_dtype)
+            for sublayer in layer._flatten_layers(include_self=False):
+                if isinstance(
+                    sublayer, (keras.layers.Softmax, keras.layers.InputLayer)
+                ):
+                    continue
+                self.assertEqual(policy.compute_dtype, sublayer.compute_dtype)
+                self.assertEqual(policy.variable_dtype, sublayer.variable_dtype)
+
     def run_model_saving_test(
         self,
         cls,
@@ -310,6 +328,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         input_data,
         expected_output_shape,
         variable_length_data=None,
+        run_mixed_precision_check=True,
     ):
         """Run basic tests for a backbone, including compilation."""
         backbone = cls(**init_kwargs)
@@ -351,6 +370,8 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         name = re.sub("([a-z])([A-Z])", r"\1_\2", name).lower()
         self.assertRegexpMatches(backbone.name, name)
 
+        self.run_precision_test(cls, init_kwargs, input_data)
+
     def run_task_test(
         self,
         cls,