diff --git a/keras_nlp/layers/f_net_encoder_test.py b/keras_nlp/layers/f_net_encoder_test.py
index b154e1845f..e6f889d453 100644
--- a/keras_nlp/layers/f_net_encoder_test.py
+++ b/keras_nlp/layers/f_net_encoder_test.py
@@ -133,7 +133,9 @@ def test_saved_model(self, save_format, filename):
         data = tf.random.uniform(shape=[2, 4, 6])
         model(data)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         loaded_model = keras.models.load_model(path)
 
         model_output = model(data)
diff --git a/keras_nlp/layers/masked_lm_head_test.py b/keras_nlp/layers/masked_lm_head_test.py
index f66f0b1918..ff266c75f5 100644
--- a/keras_nlp/layers/masked_lm_head_test.py
+++ b/keras_nlp/layers/masked_lm_head_test.py
@@ -16,12 +16,13 @@
 import os
 
 import tensorflow as tf
+from absl.testing import parameterized
 from tensorflow import keras
 
 from keras_nlp.layers import masked_lm_head
 
 
-class MaskedLMHeadTest(tf.test.TestCase):
+class MaskedLMHeadTest(tf.test.TestCase, parameterized.TestCase):
     def test_valid_call(self):
         head = masked_lm_head.MaskedLMHead(
             vocabulary_size=100,
@@ -156,7 +157,11 @@ def test_checkpointing(self):
         head2_output = head2(token_data, mask_positions=position_data)
         self.assertAllClose(head1_output, head2_output)
 
-    def test_saving_model(self):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         head = masked_lm_head.MaskedLMHead(
             vocabulary_size=100,
             activation="softmax",
@@ -171,9 +176,11 @@ def test_saving_model(self):
             shape=(4, 5), maxval=10, dtype="int32"
         )
         model_output = model((token_data, position_data))
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path)
-        restored = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
-        restored_output = restored((token_data, position_data))
+        restored_output = restored_model((token_data, position_data))
         self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/layers/multi_segment_packer_test.py b/keras_nlp/layers/multi_segment_packer_test.py
index f271850c3f..4434486a77 100644
--- a/keras_nlp/layers/multi_segment_packer_test.py
+++ b/keras_nlp/layers/multi_segment_packer_test.py
@@ -178,7 +178,9 @@ def test_saved_model(self, save_format, filename):
         outputs = packer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model((seq1, seq2)),
diff --git a/keras_nlp/layers/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
index f627cfde9c..0dbd7ff341 100644
--- a/keras_nlp/layers/position_embedding_test.py
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -309,7 +309,9 @@ def test_saved_model(self, save_format, filename):
         model(data)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         loaded_model = keras.models.load_model(path)
 
         model_output = model.predict(data)
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
index 7b9558864d..455d2899d6 100644
--- a/keras_nlp/layers/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -148,7 +148,9 @@ def test_saved_model(self, save_format, filename):
         model(data)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         loaded_model = keras.models.load_model(path)
 
         model_output = model.predict(data)
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 6b1929b237..b3bf1a612f 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -332,7 +332,9 @@ def test_saved_model(self, save_format, filename):
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model([decoder_sequence, encoder_sequence])
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         loaded_model = keras.models.load_model(path)
         model_output = model([decoder_sequence, encoder_sequence])
@@ -358,7 +360,9 @@ def test_saved_model_without_cross_attention(self, save_format, filename):
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model(decoder_sequence)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         loaded_model = keras.models.load_model(path)
 
         model_output = model(decoder_sequence)
diff --git a/keras_nlp/layers/transformer_encoder_test.py b/keras_nlp/layers/transformer_encoder_test.py
index d68dd308e0..20a2922e99 100644
--- a/keras_nlp/layers/transformer_encoder_test.py
+++ b/keras_nlp/layers/transformer_encoder_test.py
@@ -176,7 +176,9 @@ def test_saved_model(self, save_format, filename):
         data = tf.random.uniform(shape=[2, 4, 6])
         model_output = model(data)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         loaded_model = keras.models.load_model(path)
         loaded_model_output = loaded_model(data)
diff --git a/keras_nlp/models/albert/albert_backbone_test.py b/keras_nlp/models/albert/albert_backbone_test.py
index 9472c95e87..da29944f51 100644
--- a/keras_nlp/models/albert/albert_backbone_test.py
+++ b/keras_nlp/models/albert/albert_backbone_test.py
@@ -93,9 +93,11 @@ def test_error_for_invalid_num_groups(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, AlbertBackbone)
diff --git a/keras_nlp/models/albert/albert_classifier_test.py b/keras_nlp/models/albert/albert_classifier_test.py
index c6b8ad6d17..af84805b26 100644
--- a/keras_nlp/models/albert/albert_classifier_test.py
+++ b/keras_nlp/models/albert/albert_classifier_test.py
@@ -125,9 +125,11 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saving_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back
         self.assertIsInstance(restored_model, AlbertClassifier)
diff --git a/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py b/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py
index 2692bd0340..01c37843b0 100644
--- a/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py
@@ -164,7 +164,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]
diff --git a/keras_nlp/models/albert/albert_masked_lm_test.py b/keras_nlp/models/albert/albert_masked_lm_test.py
index cf18c59e91..8af1845cda 100644
--- a/keras_nlp/models/albert/albert_masked_lm_test.py
+++ b/keras_nlp/models/albert/albert_masked_lm_test.py
@@ -130,9 +130,11 @@ def test_classifier_fit_no_xla(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.masked_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, AlbertMaskedLM)
diff --git a/keras_nlp/models/albert/albert_preprocessor_test.py b/keras_nlp/models/albert/albert_preprocessor_test.py
index afc21e1c2b..4367db794e 100644
--- a/keras_nlp/models/albert/albert_preprocessor_test.py
+++ b/keras_nlp/models/albert/albert_preprocessor_test.py
@@ -177,7 +177,9 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data)["token_ids"],
diff --git a/keras_nlp/models/albert/albert_tokenizer_test.py b/keras_nlp/models/albert/albert_tokenizer_test.py
index 40fb242904..a8965198de 100644
--- a/keras_nlp/models/albert/albert_tokenizer_test.py
+++ b/keras_nlp/models/albert/albert_tokenizer_test.py
@@ -104,7 +104,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/bart/bart_backbone_test.py b/keras_nlp/models/bart/bart_backbone_test.py
index 02279376f8..1d56a10565 100644
--- a/keras_nlp/models/bart/bart_backbone_test.py
+++ b/keras_nlp/models/bart/bart_backbone_test.py
@@ -97,9 +97,11 @@ def test_compile_batched_ds(self, jit_compile):
     )
     def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.model.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.model.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, BartBackbone)
diff --git a/keras_nlp/models/bart/bart_tokenizer_test.py b/keras_nlp/models/bart/bart_tokenizer_test.py
index 1ec896eb1a..f4ec9d7835 100644
--- a/keras_nlp/models/bart/bart_tokenizer_test.py
+++ b/keras_nlp/models/bart/bart_tokenizer_test.py
@@ -81,7 +81,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/bert/bert_backbone_test.py b/keras_nlp/models/bert/bert_backbone_test.py
index f548378067..c0cae4645f 100644
--- a/keras_nlp/models/bert/bert_backbone_test.py
+++ b/keras_nlp/models/bert/bert_backbone_test.py
@@ -79,9 +79,11 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, BertBackbone)
diff --git a/keras_nlp/models/bert/bert_classifier_test.py b/keras_nlp/models/bert/bert_classifier_test.py
index 83c6616555..11c254315f 100644
--- a/keras_nlp/models/bert/bert_classifier_test.py
+++ b/keras_nlp/models/bert/bert_classifier_test.py
@@ -98,9 +98,11 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, BertClassifier)
diff --git a/keras_nlp/models/bert/bert_masked_lm_preprocessor_test.py b/keras_nlp/models/bert/bert_masked_lm_preprocessor_test.py
index 27cbaa2cea..29938d3b34 100644
--- a/keras_nlp/models/bert/bert_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/bert/bert_masked_lm_preprocessor_test.py
@@ -148,7 +148,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]
diff --git a/keras_nlp/models/bert/bert_masked_lm_test.py b/keras_nlp/models/bert/bert_masked_lm_test.py
index 59f3dd9f76..2efa268e0d 100644
--- a/keras_nlp/models/bert/bert_masked_lm_test.py
+++ b/keras_nlp/models/bert/bert_masked_lm_test.py
@@ -104,13 +104,15 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.masked_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, BertMaskedLM)
 
         # Check that output matches.
         restored_output = restored_model.predict(self.raw_batch)
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
diff --git a/keras_nlp/models/bert/bert_preprocessor_test.py b/keras_nlp/models/bert/bert_preprocessor_test.py
index d0053ff058..ebd35b937c 100644
--- a/keras_nlp/models/bert/bert_preprocessor_test.py
+++ b/keras_nlp/models/bert/bert_preprocessor_test.py
@@ -131,7 +131,9 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data)["token_ids"],
diff --git a/keras_nlp/models/bert/bert_tokenizer_test.py b/keras_nlp/models/bert/bert_tokenizer_test.py
index 5bca0d366b..c36646909c 100644
--- a/keras_nlp/models/bert/bert_tokenizer_test.py
+++ b/keras_nlp/models/bert/bert_tokenizer_test.py
@@ -78,7 +78,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
index 931410a5be..cfa89fbe22 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
@@ -83,9 +83,11 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DebertaV3Backbone)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
index 79cdb6c63e..9c255d2103 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
@@ -120,9 +120,11 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saving_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DebertaV3Classifier)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py
index f84c403781..38f4847e3f 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py
@@ -161,7 +161,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
index 406fc616af..37e54e7209 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
@@ -118,13 +118,15 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.masked_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DebertaV3MaskedLM)
 
         # Check that output matches.
         restored_output = restored_model.predict(self.raw_batch)
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py b/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
index ec02a6ee58..1bec64da06 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
@@ -156,7 +156,9 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data)["token_ids"],
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
index 6672068e57..3f0884c7e2 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
@@ -116,7 +116,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone_test.py b/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
index d4877d7725..0360a3e4de 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
@@ -83,9 +83,11 @@ def test_distilbert_base_compile_batched_ds(self, jit_compile):
     )
     def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.model.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.model.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DistilBertBackbone)
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
index a9a0fc06d1..369dd0a343 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
@@ -118,9 +118,11 @@ def test_distilbert_classifier_fit_default_compile(self):
     )
     def test_saving_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DistilBertClassifier)
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py
index 62eadb5e81..4d8a1b18de 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor_test.py
@@ -117,8 +117,10 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
 
-        path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py
index 995a993123..30aa0c789a 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py
@@ -117,9 +117,11 @@ def test_distilbert_masked_lm_fit_no_preprocessing(self, jit_compile):
         ("keras_format", "keras_v3", "model.keras"),
     )
     def test_saved_model(self, save_format, filename):
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.masked_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DistilBertMaskedLM)
@@ -127,4 +129,4 @@ def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm(self.preprocessed_batch)
         restored_output = restored_model(self.preprocessed_batch)
 
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
diff --git a/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py b/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py
index 36ba4492e5..c7ebbb7d36 100644
--- a/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py
@@ -111,7 +111,9 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data)["token_ids"],
diff --git a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
index b51c224723..692180db76 100644
--- a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
@@ -70,7 +70,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/models/f_net/f_net_backbone_test.py b/keras_nlp/models/f_net/f_net_backbone_test.py
index 47b0ec9c1c..9e33d56377 100644
--- a/keras_nlp/models/f_net/f_net_backbone_test.py
+++ b/keras_nlp/models/f_net/f_net_backbone_test.py
@@ -73,9 +73,11 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, FNetBackbone)
diff --git a/keras_nlp/models/f_net/f_net_classifier_test.py b/keras_nlp/models/f_net/f_net_classifier_test.py
index 6f3e01d023..d6dbaf36f5 100644
--- a/keras_nlp/models/f_net/f_net_classifier_test.py
+++ b/keras_nlp/models/f_net/f_net_classifier_test.py
@@ -120,9 +120,11 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, FNetClassifier)
diff --git a/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py b/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py
index 600afc09c5..4830c6c81a 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py
@@ -142,7 +142,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]
diff --git a/keras_nlp/models/f_net/f_net_masked_lm_test.py b/keras_nlp/models/f_net/f_net_masked_lm_test.py
index 508b3f3e52..999262040c 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm_test.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm_test.py
@@ -116,13 +116,15 @@ def test_serialization(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.masked_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, FNetMaskedLM)
 
         # Check that output matches.
         restored_output = restored_model.predict(self.raw_batch)
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
diff --git a/keras_nlp/models/f_net/f_net_preprocessor_test.py b/keras_nlp/models/f_net/f_net_preprocessor_test.py
index 230247e880..10c0a3aac6 100644
--- a/keras_nlp/models/f_net/f_net_preprocessor_test.py
+++ b/keras_nlp/models/f_net/f_net_preprocessor_test.py
@@ -159,7 +159,9 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data)["token_ids"],
diff --git a/keras_nlp/models/f_net/f_net_tokenizer_test.py b/keras_nlp/models/f_net/f_net_tokenizer_test.py
index 15673a046c..33e26cc8cd 100644
--- a/keras_nlp/models/f_net/f_net_tokenizer_test.py
+++ b/keras_nlp/models/f_net/f_net_tokenizer_test.py
@@ -104,7 +104,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/gpt2/gpt2_backbone_test.py b/keras_nlp/models/gpt2/gpt2_backbone_test.py
index 4115a982f4..eb92022943 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone_test.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone_test.py
@@ -85,9 +85,11 @@ def test_gpt2_compile_batched_ds(self, jit_compile):
     )
     def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.model.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.model.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, GPT2Backbone)
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py b/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py
index 43b576fb0d..33354ca517 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py
@@ -140,7 +140,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm_test.py b/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
index ccf447532c..7e9bb6add2 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
@@ -153,9 +153,11 @@ def test_generate(self, jit_compile):
     def test_saved_model(self, save_format, filename):
         keras.utils.set_random_seed(42)
         model_output = self.causal_lm.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.causal_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.causal_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, GPT2CausalLM)
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessor_test.py b/keras_nlp/models/gpt2/gpt2_preprocessor_test.py
index 2e5b7311d4..58790dfe3f 100644
--- a/keras_nlp/models/gpt2/gpt2_preprocessor_test.py
+++ b/keras_nlp/models/gpt2/gpt2_preprocessor_test.py
@@ -136,7 +136,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
index bbc11cfca1..9caea75958 100644
--- a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
+++ b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
@@ -82,7 +82,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/opt/opt_backbone_test.py b/keras_nlp/models/opt/opt_backbone_test.py
index cb476591b2..85e80d3350 100644
--- a/keras_nlp/models/opt/opt_backbone_test.py
+++ b/keras_nlp/models/opt/opt_backbone_test.py
@@ -78,9 +78,11 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, OPTBackbone)
diff --git a/keras_nlp/models/opt/opt_tokenizer_test.py b/keras_nlp/models/opt/opt_tokenizer_test.py
index 3940bbea79..6babc65111 100644
--- a/keras_nlp/models/opt/opt_tokenizer_test.py
+++ b/keras_nlp/models/opt/opt_tokenizer_test.py
@@ -92,7 +92,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/roberta/roberta_backbone_test.py b/keras_nlp/models/roberta/roberta_backbone_test.py
index 203d3474e3..ac0777a8d4 100644
--- a/keras_nlp/models/roberta/roberta_backbone_test.py
+++ b/keras_nlp/models/roberta/roberta_backbone_test.py
@@ -83,9 +83,11 @@ def test_variable_sequence_length_call_roberta(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, RobertaBackbone)
diff --git a/keras_nlp/models/roberta/roberta_classifier_test.py b/keras_nlp/models/roberta/roberta_classifier_test.py
index d0fbd67364..bc1b0ca9c9 100644
--- a/keras_nlp/models/roberta/roberta_classifier_test.py
+++ b/keras_nlp/models/roberta/roberta_classifier_test.py
@@ -116,9 +116,11 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, RobertaClassifier)
diff --git a/keras_nlp/models/roberta/roberta_masked_lm_preprocessor_test.py b/keras_nlp/models/roberta/roberta_masked_lm_preprocessor_test.py
index 97719d63f5..f34f4b22e6 100644
--- a/keras_nlp/models/roberta/roberta_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm_preprocessor_test.py
@@ -162,7 +162,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         outputs = model(input_data)[0]["token_ids"]
diff --git a/keras_nlp/models/roberta/roberta_masked_lm_test.py b/keras_nlp/models/roberta/roberta_masked_lm_test.py
index 315a22e72d..8334ad42af 100644
--- a/keras_nlp/models/roberta/roberta_masked_lm_test.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm_test.py
@@ -121,13 +121,15 @@ def test_classifier_fit_no_xla(self):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.masked_lm.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.masked_lm.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.masked_lm.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, RobertaMaskedLM)
 
         # Check that output matches.
         restored_output = restored_model.predict(self.raw_batch)
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=0.01, rtol=0.01)
diff --git a/keras_nlp/models/roberta/roberta_preprocessor_test.py b/keras_nlp/models/roberta/roberta_preprocessor_test.py
index b12a344324..d98644f2f6 100644
--- a/keras_nlp/models/roberta/roberta_preprocessor_test.py
+++ b/keras_nlp/models/roberta/roberta_preprocessor_test.py
@@ -161,7 +161,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/roberta/roberta_tokenizer_test.py b/keras_nlp/models/roberta/roberta_tokenizer_test.py
index 0e73285d88..07401a403d 100644
--- a/keras_nlp/models/roberta/roberta_tokenizer_test.py
+++ b/keras_nlp/models/roberta/roberta_tokenizer_test.py
@@ -92,7 +92,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/t5/t5_backbone_test.py b/keras_nlp/models/t5/t5_backbone_test.py
index 813052b11e..4d89336dff 100644
--- a/keras_nlp/models/t5/t5_backbone_test.py
+++ b/keras_nlp/models/t5/t5_backbone_test.py
@@ -102,9 +102,11 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         outputs = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, T5Backbone)
diff --git a/keras_nlp/models/t5/t5_tokenizer_test.py b/keras_nlp/models/t5/t5_tokenizer_test.py
index a4b62b75ac..7386872dcf 100644
--- a/keras_nlp/models/t5/t5_tokenizer_test.py
+++ b/keras_nlp/models/t5/t5_tokenizer_test.py
@@ -103,7 +103,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/whisper/whisper_backbone_test.py b/keras_nlp/models/whisper/whisper_backbone_test.py
index 456cb7b564..bb4ee39ee7 100644
--- a/keras_nlp/models/whisper/whisper_backbone_test.py
+++ b/keras_nlp/models/whisper/whisper_backbone_test.py
@@ -104,9 +104,11 @@ def test_key_projection_bias_absence(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, WhisperBackbone)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
index d91330a0c0..6112c37984 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
@@ -81,9 +81,11 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saved_model(self, save_format, filename):
         model_output = self.backbone(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.backbone.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.backbone.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, XLMRobertaBackbone)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
index bf47672032..082b2eea9c 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
@@ -120,9 +120,11 @@ def test_serialization(self):
     @pytest.mark.large  # Saving is slow, so mark these large.
     def test_saving_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        self.classifier.save(save_path, save_format=save_format)
-        restored_model = keras.models.load_model(save_path)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        self.classifier.save(path, save_format=save_format, **kwargs)
+        restored_model = keras.models.load_model(path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, XLMRobertaClassifier)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py
index d983f8bcec..a06490c298 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py
@@ -152,7 +152,9 @@ def test_saved_model(self, save_format, filename):
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data)["token_ids"],
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py
index 42e1fe1877..912403fc4e 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py
@@ -118,7 +118,9 @@ def test_saved_model(self, save_format, filename):
         model = keras.Model(inputs, outputs)
 
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
 
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
index d3cd6eb886..d208207985 100644
--- a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
@@ -164,7 +164,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/byte_tokenizer_test.py b/keras_nlp/tokenizers/byte_tokenizer_test.py
index 0520ce616c..500de32666 100644
--- a/keras_nlp/tokenizers/byte_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_tokenizer_test.py
@@ -261,7 +261,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
index 8571eed1a3..8c06bedaa8 100644
--- a/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
@@ -201,7 +201,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py b/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py
index b40794bc26..c4e6317cee 100644
--- a/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py
+++ b/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py
@@ -363,7 +363,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index fa1b8bd95b..a81f875083 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -217,7 +217,9 @@ def test_saved_model(self, save_format, filename):
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), filename)
-        model.save(path, save_format=save_format)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/utils/pipeline_model_test.py b/keras_nlp/utils/pipeline_model_test.py
index a5a995166b..3a433a38ad 100644
--- a/keras_nlp/utils/pipeline_model_test.py
+++ b/keras_nlp/utils/pipeline_model_test.py
@@ -140,10 +140,12 @@ def test_saved_model(self, save_format, filename):
         model = NoopPipeline()
         x = tf.random.uniform((8, 5))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(
-            save_path, custom_objects={"NoopPipeline": NoopPipeline}
+            path, custom_objects={"NoopPipeline": NoopPipeline}
         )
 
         # Check we got the real object back.
@@ -256,10 +258,12 @@ def test_saved_model(self, save_format, filename):
         model = FeaturePipeline()
         x = tf.strings.as_string(tf.random.uniform((8, 5)))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(
-            save_path, custom_objects={"FeaturePipeline": FeaturePipeline}
+            path, custom_objects={"FeaturePipeline": FeaturePipeline}
         )
 
         # Check we got the real object back.
@@ -365,10 +369,12 @@ def test_saved_model(self, save_format, filename):
         model = LabelPipeline()
         x = tf.random.uniform((8, 5))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(
-            save_path, custom_objects={"LabelPipeline": LabelPipeline}
+            path, custom_objects={"LabelPipeline": LabelPipeline}
         )
 
         # Check we got the real object back.
@@ -457,10 +463,12 @@ def test_saved_model(self, save_format, filename):
         model = DataPipeline()
         data = tf.strings.as_string(tf.random.uniform((8, 1)))
         model_output = model.predict(data)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(
-            save_path, custom_objects={"DataPipeline": DataPipeline}
+            path, custom_objects={"DataPipeline": DataPipeline}
         )
 
         # Check we got the real object back.
@@ -506,10 +514,12 @@ def test_saved_model(self, save_format, filename):
         model = FunctionalPipeline()
         x = tf.strings.as_string(tf.random.uniform((8, 5)))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
+        kwargs = {"save_traces": False} if save_format == "tf" else {}
+        model.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(
-            save_path, custom_objects={"FunctionalPipeline": FunctionalPipeline}
+            path, custom_objects={"FunctionalPipeline": FunctionalPipeline}
         )
 
         # Check we got the real object back.