Repository: keras-team/keras-nlp
Pull request #1221 (2023-08-21 22:52:59)

Stop asserting key order in bart preprocessor


Patch:
diff --git a/keras_nlp/layers/modeling/f_net_encoder.py b/keras_nlp/layers/modeling/f_net_encoder.py
index 806a3dac..03ed33b3 100644
--- a/keras_nlp/layers/modeling/f_net_encoder.py
+++ b/keras_nlp/layers/modeling/f_net_encoder.py
@@ -99,11 +99,13 @@ class FNetEncoder(keras.layers.Layer):
 
         # Layer Norm layers.
         self._mixing_layer_norm = keras.layers.LayerNormalization(
-            epsilon=self.layer_norm_epsilon
+            epsilon=self.layer_norm_epsilon,
+            name="mixing_layer_norm",
         )
         self._mixing_layer_norm.build(inputs_shape)
         self._output_layer_norm = keras.layers.LayerNormalization(
-            epsilon=self.layer_norm_epsilon
+            epsilon=self.layer_norm_epsilon,
+            name="output_layer_norm",
         )
         self._output_layer_norm.build(inputs_shape)
 
@@ -113,12 +115,14 @@ class FNetEncoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="intermediate_dense",
         )
         self._intermediate_dense.build(inputs_shape)
         self._output_dense = keras.layers.Dense(
             feature_size,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="output_dense",
         )
         self._output_dense.build(
             self._intermediate_dense.compute_output_shape(inputs_shape)
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
index 46c15f02..eeea9bbd 100644
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -49,7 +49,7 @@ class MaskedLMHead(keras.layers.Layer):
         token_embedding: Optional. A `keras_nlp.layers.ReversibleEmbedding`
             instance. If passed, the layer will be used to project from the
             `hidden_dim` of the model to the output `vocabulary_size`.
-        intermediate_activation: The activation function of inner dense layer.
+        intermediate_activation: The activation function of intermediate dense layer.
         activation: The activation function for the outputs of the layer.
             Usually either `None` (return logits), or `"softmax"`
             (return probabilities).
@@ -138,22 +138,24 @@ class MaskedLMHead(keras.layers.Layer):
         else:
             feature_size = inputs_shape[-1]
 
-        self._dense = keras.layers.Dense(
+        self._intermediate_dense = keras.layers.Dense(
             feature_size,
             activation=self.intermediate_activation,
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
+            name="intermediate_dense",
         )
-        self._layer_norm = keras.layers.LayerNormalization(
+        self._intermediate_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            name="intermediate_layer_norm",
         )
         # The gather length does not affect any of our built variables, so
         # we can pass any value here.
         gather_length = None
         shape = (inputs_shape[0], gather_length, inputs_shape[-1])
-        self._dense.build(shape)
+        self._intermediate_dense.build(shape)
         shape = (inputs_shape[0], gather_length, feature_size)
-        self._layer_norm.build(shape)
+        self._intermediate_layer_norm.build(shape)
         if self.token_embedding is None:
             self._kernel = self.add_weight(
                 name="output_kernel",
@@ -174,8 +176,8 @@ class MaskedLMHead(keras.layers.Layer):
         x = ops.take_along_axis(inputs, mask_positions, axis=1)
 
         # Apply a trainable linear transformation and a layer norm.
-        x = self._dense(x)
-        x = self._layer_norm(x)
+        x = self._intermediate_dense(x)
+        x = self._intermediate_layer_norm(x)
 
         # Transform encodings to vocabulary_size predictions.
         if self.token_embedding:
diff --git a/keras_nlp/layers/modeling/token_and_position_embedding.py b/keras_nlp/layers/modeling/token_and_position_embedding.py
index aebf9090..3b72ff46 100644
--- a/keras_nlp/layers/modeling/token_and_position_embedding.py
+++ b/keras_nlp/layers/modeling/token_and_position_embedding.py
@@ -95,14 +95,12 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
                 self.embeddings_initializer
             ),
             mask_zero=mask_zero,
-            name="token_embedding"
-            + str(keras.backend.get_uid("token_embedding")),
+            name="token_embedding",
         )
         self.position_embedding = PositionEmbedding(
             sequence_length=sequence_length,
             initializer=clone_initializer(self.embeddings_initializer),
-            name="position_embedding"
-            + str(keras.backend.get_uid("position_embedding")),
+            name="position_embedding",
         )
         self.supports_masking = self.token_embedding.supports_masking
 
diff --git a/keras_nlp/layers/modeling/transformer_decoder.py b/keras_nlp/layers/modeling/transformer_decoder.py
index b608f6f4..7216a87a 100644
--- a/keras_nlp/layers/modeling/transformer_decoder.py
+++ b/keras_nlp/layers/modeling/transformer_decoder.py
@@ -173,10 +173,10 @@ class TransformerDecoder(keras.layers.Layer):
                 query_shape=decoder_sequence_shape,
                 value_shape=decoder_sequence_shape,
             )
-        self._self_attention_layernorm = keras.layers.LayerNormalization(
+        self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
-        self._self_attention_layernorm.build(decoder_sequence_shape)
+        self._self_attention_layer_norm.build(decoder_sequence_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
         )
@@ -202,10 +202,10 @@ class TransformerDecoder(keras.layers.Layer):
                     query_shape=encoder_sequence_shape,
                     value_shape=encoder_sequence_shape,
                 )
-            self._cross_attention_layernorm = keras.layers.LayerNormalization(
+            self._cross_attention_layer_norm = keras.layers.LayerNormalization(
                 epsilon=self.layer_norm_epsilon,
             )
-            self._cross_attention_layernorm.build(encoder_sequence_shape)
+            self._cross_attention_layer_norm.build(encoder_sequence_shape)
             self._cross_attention_dropout = keras.layers.Dropout(
                 rate=self.dropout,
             )
@@ -226,10 +226,10 @@ class TransformerDecoder(keras.layers.Layer):
         intermediate_shape = list(decoder_sequence_shape)
         intermediate_shape[-1] = self.intermediate_dim
         self._feedforward_output_dense.build(tuple(intermediate_shape))
-        self._feedforward_layernorm = keras.layers.LayerNormalization(
+        self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
-        self._feedforward_layernorm.build(decoder_sequence_shape)
+        self._feedforward_layer_norm.build(decoder_sequence_shape)
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
         )
@@ -364,7 +364,7 @@ class TransformerDecoder(keras.layers.Layer):
         # Self attention block.
         residual = x
         if self.normalize_first:
-            x = self._self_attention_layernorm(x)
+            x = self._self_attention_layer_norm(x)
         x, self_attention_cache = self._self_attention_layer(
             query=x,
             value=x,
@@ -375,7 +375,7 @@ class TransformerDecoder(keras.layers.Layer):
         x = self._self_attention_dropout(x)
         x = x + residual
         if not self.normalize_first:
-            x = self._self_attention_layernorm(x)
+            x = self._self_attention_layer_norm(x)
 
         # Cross attention is optional.
         if has_cross_attention:
@@ -387,7 +387,7 @@ class TransformerDecoder(keras.layers.Layer):
             # Cross attention block.
             residual = x
             if self.normalize_first:
-                x = self._cross_attention_layernorm(x)
+                x = self._cross_attention_layer_norm(x)
             x, cross_attention_cache = self._cross_attention_layer(
                 query=x,
                 value=encoder_sequence,
@@ -398,18 +398,18 @@ class TransformerDecoder(keras.layers.Layer):
             x = self._cross_attention_dropout(x)
             x = x + residual
             if not self.normalize_first:
-                x = self._cross_attention_layernorm(x)
+                x = self._cross_attention_layer_norm(x)
 
         # Feedforward block.
         residual = x
         if self.normalize_first:
-            x = self._feedforward_layernorm(x)
+            x = self._feedforward_layer_norm(x)
         x = self._feedforward_intermediate_dense(x)
         x = self._feedforward_output_dense(x)
         x = self._feedforward_dropout(x)
         x = x + residual
         if not self.normalize_first:
-            x = self._feedforward_layernorm(x)
+            x = self._feedforward_layer_norm(x)
 
         if self_attention_cache is not None:
             if has_cross_attention:
diff --git a/keras_nlp/layers/modeling/transformer_encoder.py b/keras_nlp/layers/modeling/transformer_encoder.py
index 92d933e9..78620dcb 100644
--- a/keras_nlp/layers/modeling/transformer_encoder.py
+++ b/keras_nlp/layers/modeling/transformer_encoder.py
@@ -127,6 +127,7 @@ class TransformerEncoder(keras.layers.Layer):
             dropout=self.dropout,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="self_attention_layer",
         )
         if hasattr(self._self_attention_layer, "_build_from_signature"):
             self._self_attention_layer._build_from_signature(
@@ -138,30 +139,34 @@ class TransformerEncoder(keras.layers.Layer):
                 query_shape=inputs_shape,
                 value_shape=inputs_shape,
             )
-        self._self_attention_layernorm = keras.layers.LayerNormalization(
+        self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            name="self_attention_layer_norm",
         )
-        self._self_attention_layernorm.build(inputs_shape)
+        self._self_attention_layer_norm.build(inputs_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            name="self_attention_dropout",
         )
 
         # Feedforward layers.
-        self._feedforward_layernorm = keras.layers.LayerNormalization(
+        self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
-        self._feedforward_layernorm.build(inputs_shape)
+        self._feedforward_layer_norm.build(inputs_shape)
         self._feedforward_intermediate_dense = keras.layers.Dense(
             self.intermediate_dim,
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(inputs_shape)
         self._feedforward_output_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="feedforward_output_dense",
         )
         intermediate_shape = list(inputs_shape)
         intermediate_shape[-1] = self.intermediate_dim
@@ -197,7 +202,7 @@ class TransformerEncoder(keras.layers.Layer):
         # Self attention block.
         residual = x
         if self.normalize_first:
-            x = self._self_attention_layernorm(x)
+            x = self._self_attention_layer_norm(x)
         x = self._self_attention_layer(
             query=x,
             value=x,
@@ -206,18 +211,18 @@ class TransformerEncoder(keras.layers.Layer):
         x = self._self_attention_dropout(x)
         x = x + residual
         if not self.normalize_first:
-            x = self._self_attention_layernorm(x)
+            x = self._self_attention_layer_norm(x)
 
         # Feedforward block.
         residual = x
         if self.normalize_first:
-            x = self._feedforward_layernorm(x)
+            x = self._feedforward_layer_norm(x)
         x = self._feedforward_intermediate_dense(x)
         x = self._feedforward_output_dense(x)
         x = self._feedforward_dropout(x)
         x = x + residual
         if not self.normalize_first:
-            x = self._feedforward_layernorm(x)
+            x = self._feedforward_layer_norm(x)
 
         return x
 
diff --git a/keras_nlp/models/bart/bart_preprocessor.py b/keras_nlp/models/bart/bart_preprocessor.py
index 04626ddb..0493c89c 100644
--- a/keras_nlp/models/bart/bart_preprocessor.py
+++ b/keras_nlp/models/bart/bart_preprocessor.py
@@ -179,7 +179,7 @@ class BartPreprocessor(Preprocessor):
     def call(self, x, y=None, sample_weight=None):
         if not (
             isinstance(x, dict)
-            and ["encoder_text", "decoder_text"] == list(x.keys())
+            and all(k in x for k in ("encoder_text", "decoder_text"))
         ):
             raise ValueError(
                 '`x` must be a dictionary, containing the keys `"encoder_text"`'
diff --git a/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py b/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py
index 594e7188..b192350a 100644
--- a/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py
+++ b/keras_nlp/models/deberta_v3/disentangled_attention_encoder.py
@@ -101,38 +101,45 @@ class DisentangledAttentionEncoder(keras.layers.Layer):
             dropout=self.dropout,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="self_attention_layer",
         )
         self._self_attention_layer.build(inputs_shape)
-        self._self_attention_layernorm = keras.layers.LayerNormalization(
+        self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            name="self_attention_layer_norm",
         )
-        self._self_attention_layernorm.build(inputs_shape)
+        self._self_attention_layer_norm.build(inputs_shape)
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            name="self_attention_dropout",
         )
 
         # Feedforward layers.
-        self._feedforward_layernorm = keras.layers.LayerNormalization(
+        self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            name="feedforward_layer_norm",
         )
-        self._feedforward_layernorm.build(inputs_shape)
+        self._feedforward_layer_norm.build(inputs_shape)
         self._feedforward_intermediate_dense = keras.layers.Dense(
             self.intermediate_dim,
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(inputs_shape)
         self._feedforward_output_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="feedforward_output_dense",
         )
         intermediate_shape = list(inputs_shape)
         intermediate_shape[-1] = self.intermediate_dim
         self._feedforward_output_dense.build(tuple(intermediate_shape))
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            name="feedforward_dropout",
         )
         self.built = True
 
@@ -177,7 +184,7 @@ class DisentangledAttentionEncoder(keras.layers.Layer):
         )
         x = self._self_attention_dropout(x)
         x = x + residual
-        x = self._self_attention_layernorm(x)
+        x = self._self_attention_layer_norm(x)
 
         # Feedforward block.
         residual = x
@@ -185,7 +192,7 @@ class DisentangledAttentionEncoder(keras.layers.Layer):
         x = self._feedforward_output_dense(x)
         x = self._feedforward_dropout(x)
         x = x + residual
-        x = self._feedforward_layernorm(x)
+        x = self._feedforward_layer_norm(x)
 
         return x
 
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 1ad2fd26..f0f66f4d 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
@@ -98,16 +98,19 @@ 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),
+            name="self_attention",
         )
         self._self_attention_layer.build(decoder_sequence_shape)
 
-        self._self_attention_layernorm = keras.layers.LayerNormalization(
+        self._self_attention_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            name="self_attention_layer_norm",
         )
-        self._self_attention_layernorm.build(decoder_sequence_shape)
+        self._self_attention_layer_norm.build(decoder_sequence_shape)
 
         self._self_attention_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            name="self_attention_dropout",
         )
 
         # Feedforward layers.
@@ -116,6 +119,7 @@ class GPTNeoXDecoder(keras.layers.Layer):
             activation=self.activation,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="feedforward_intermediate_dense",
         )
         self._feedforward_intermediate_dense.build(decoder_sequence_shape)
 
@@ -123,19 +127,22 @@ class GPTNeoXDecoder(keras.layers.Layer):
             hidden_dim,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             bias_initializer=clone_initializer(self.bias_initializer),
+            name="feedforward_output_dense",
         )
 
         intermediate_shape = list(decoder_sequence_shape)
         intermediate_shape[-1] = self.intermediate_dim
         self._feedforward_output_dense.build(tuple(intermediate_shape))
 
-        self._feedforward_layernorm = keras.layers.LayerNormalization(
+        self._feedforward_layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
+            name="feedforward_layer_norm",
         )
-        self._feedforward_layernorm.build(decoder_sequence_shape)
+        self._feedforward_layer_norm.build(decoder_sequence_shape)
 
         self._feedforward_dropout = keras.layers.Dropout(
             rate=self.dropout,
+            name="feedforward_dropout",
         )
         self.built = True
 
@@ -157,7 +164,7 @@ class GPTNeoXDecoder(keras.layers.Layer):
 
         residual = decoder_sequence
 
-        x = self._self_attention_layernorm(decoder_sequence)
+        x = self._self_attention_layer_norm(decoder_sequence)
 
         # Self attention block.
         x, self_attention_cache = self._self_attention_layer(
@@ -169,7 +176,7 @@ class GPTNeoXDecoder(keras.layers.Layer):
         x = self._self_attention_dropout(x)
         attention_output = x
 
-        x = self._feedforward_layernorm(decoder_sequence)
+        x = self._feedforward_layer_norm(decoder_sequence)
         x = self._feedforward_intermediate_dense(x)
         x = self._feedforward_output_dense(x)
         feedforward_output = x
diff --git a/keras_nlp/models/t5/t5_transformer_layer.py b/keras_nlp/models/t5/t5_transformer_layer.py
index 8e516b8b..1e8ccec5 100644
--- a/keras_nlp/models/t5/t5_transformer_layer.py
+++ b/keras_nlp/models/t5/t5_transformer_layer.py
@@ -45,8 +45,9 @@ class T5TransformerLayer(keras.layers.Layer):
             num_heads,
             dropout,
             use_relative_attention_bias=use_relative_attention_bias,
+            name="self_attention",
         )
-        self.self_attention_layernorm = T5LayerNorm(layer_norm_epsilon)
+        self.self_attention_layer_norm = T5LayerNorm(layer_norm_epsilon)
         self.self_attention_dropout = keras.layers.Dropout(dropout)
 
         if self.is_decoder:
@@ -56,8 +57,9 @@ class T5TransformerLayer(keras.layers.Layer):
                 num_heads,
                 dropout,
                 use_relative_attention_bias=False,
+                name="cross_attention",
             )
-            self.cross_attention_layernorm = T5LayerNorm(layer_norm_epsilon)
+            self.cross_attention_layer_norm = T5LayerNorm(layer_norm_epsilon)
             self.cross_attention_dropout = keras.layers.Dropout(dropout)
 
         self.input_projector = keras.layers.Dense(
@@ -109,7 +111,7 @@ class T5TransformerLayer(keras.layers.Layer):
         x = hidden_states  # Intermediate result.
 
         residual = x
-        x = self.self_attention_layernorm(x)
+        x = self.self_attention_layer_norm(x)
         x, position_bias = self.self_attention(
             x,
             mask=attention_mask,
@@ -121,7 +123,7 @@ class T5TransformerLayer(keras.layers.Layer):
 
         if self.is_decoder:
             residual = x
-            x = self.cross_attention_layernorm(x)
+            x = self.cross_attention_layer_norm(x)
             x, _ = self.cross_attention(
                 x,
                 key_value_states=encoder_hidden_states,
diff --git a/tools/checkpoint_conversion/bert_base_cased.ipynb b/tools/checkpoint_conversion/bert_base_cased.ipynb
index a9622211..a7fc9db3 100644
--- a/tools/checkpoint_conversion/bert_base_cased.ipynb
+++ b/tools/checkpoint_conversion/bert_base_cased.ipynb
@@ -600,14 +600,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
@@ -642,14 +642,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
diff --git a/tools/checkpoint_conversion/bert_base_multi_cased.ipynb b/tools/checkpoint_conversion/bert_base_multi_cased.ipynb
index 8788b01f..c87bbdbd 100644
--- a/tools/checkpoint_conversion/bert_base_multi_cased.ipynb
+++ b/tools/checkpoint_conversion/bert_base_multi_cased.ipynb
@@ -600,14 +600,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
@@ -642,14 +642,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
diff --git a/tools/checkpoint_conversion/bert_base_uncased.ipynb b/tools/checkpoint_conversion/bert_base_uncased.ipynb
index 0a199251..282998b7 100644
--- a/tools/checkpoint_conversion/bert_base_uncased.ipynb
+++ b/tools/checkpoint_conversion/bert_base_uncased.ipynb
@@ -600,14 +600,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
@@ -642,14 +642,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
diff --git a/tools/checkpoint_conversion/bert_base_zh.ipynb b/tools/checkpoint_conversion/bert_base_zh.ipynb
index 19a6d758..7dbf1dac 100644
--- a/tools/checkpoint_conversion/bert_base_zh.ipynb
+++ b/tools/checkpoint_conversion/bert_base_zh.ipynb
@@ -600,14 +600,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
@@ -642,14 +642,14 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[\n",
     "            f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"\n",
     "        ]\n",
diff --git a/tools/checkpoint_conversion/bert_large_cased_en.ipynb b/tools/checkpoint_conversion/bert_large_cased_en.ipynb
index 36cb0215..0c164a5b 100644
--- a/tools/checkpoint_conversion/bert_large_cased_en.ipynb
+++ b/tools/checkpoint_conversion/bert_large_cased_en.ipynb
@@ -824,12 +824,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/beta\"]\n",
     "    )\n",
     "    model.get_layer(\n",
@@ -854,12 +854,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/beta\"]\n",
     "    )\n",
     "\n",
diff --git a/tools/checkpoint_conversion/bert_large_uncased_en.ipynb b/tools/checkpoint_conversion/bert_large_uncased_en.ipynb
index 647b6d76..a80069e1 100644
--- a/tools/checkpoint_conversion/bert_large_uncased_en.ipynb
+++ b/tools/checkpoint_conversion/bert_large_uncased_en.ipynb
@@ -824,12 +824,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/beta\"]\n",
     "    )\n",
     "    model.get_layer(\n",
@@ -854,12 +854,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/beta\"]\n",
     "    )\n",
     "\n",
diff --git a/tools/checkpoint_conversion/bert_medium_uncased_en.ipynb b/tools/checkpoint_conversion/bert_medium_uncased_en.ipynb
index ea7aaf81..01e5169a 100644
--- a/tools/checkpoint_conversion/bert_medium_uncased_en.ipynb
+++ b/tools/checkpoint_conversion/bert_medium_uncased_en.ipynb
@@ -518,12 +518,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/beta\"]\n",
     "    )\n",
     "    model.get_layer(\n",
@@ -548,12 +548,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/beta\"]\n",
     "    )\n",
     "\n",
diff --git a/tools/checkpoint_conversion/bert_small_uncased_en.ipynb b/tools/checkpoint_conversion/bert_small_uncased_en.ipynb
index 36e6d9cc..ded288b2 100644
--- a/tools/checkpoint_conversion/bert_small_uncased_en.ipynb
+++ b/tools/checkpoint_conversion/bert_small_uncased_en.ipynb
@@ -442,12 +442,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/beta\"]\n",
     "    )\n",
     "    model.get_layer(\n",
@@ -472,12 +472,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/beta\"]\n",
     "    )\n",
     "\n",
diff --git a/tools/checkpoint_conversion/bert_tiny_uncased_en.ipynb b/tools/checkpoint_conversion/bert_tiny_uncased_en.ipynb
index d346e8a5..d1fb1119 100644
--- a/tools/checkpoint_conversion/bert_tiny_uncased_en.ipynb
+++ b/tools/checkpoint_conversion/bert_tiny_uncased_en.ipynb
@@ -404,12 +404,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.gamma.assign(\n",
+    "    )._self_attention_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._self_attention_layernorm.beta.assign(\n",
+    "    )._self_attention_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/attention/output/LayerNorm/beta\"]\n",
     "    )\n",
     "    model.get_layer(\n",
@@ -434,12 +434,12 @@
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.gamma.assign(\n",
+    "    )._feedforward_layer_norm.gamma.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/gamma\"]\n",
     "    )\n",
     "    model.get_layer(\n",
     "        f\"transformer_layer_{i}\"\n",
-    "    )._feedforward_layernorm.beta.assign(\n",
+    "    )._feedforward_layer_norm.beta.assign(\n",
     "        weights[f\"bert/encoder/layer_{i}/output/LayerNorm/beta\"]\n",
     "    )\n",
     "\n",
diff --git a/tools/checkpoint_conversion/convert_albert_checkpoints.py b/tools/checkpoint_conversion/convert_albert_checkpoints.py
index a089f800..049da603 100644
--- a/tools/checkpoint_conversion/convert_albert_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_albert_checkpoints.py
@@ -158,14 +158,14 @@ def convert_checkpoints(hf_model):
 
             keras_nlp_model.get_layer(
                 f"group_{i}_inner_layer_{j}"
-            )._self_attention_layernorm.gamma.assign(
+            )._self_attention_layer_norm.gamma.assign(
                 hf_wts[
                     f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.LayerNorm.weight"
                 ].numpy()
             )
             keras_nlp_model.get_layer(
                 f"group_{i}_inner_layer_{j}"
-            )._self_attention_layernorm.beta.assign(
+            )._self_attention_layer_norm.beta.assign(
                 hf_wts[
                     f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.LayerNorm.bias"
                 ].numpy()
@@ -207,14 +207,14 @@ def convert_checkpoints(hf_model):
 
             keras_nlp_model.get_layer(
                 f"group_{i}_inner_layer_{j}"
-            )._feedforward_layernorm.gamma.assign(
+            )._feedforward_layer_norm.gamma.assign(
                 hf_wts[
                     f"encoder.albert_layer_groups.{i}.albert_layers.{j}.full_layer_layer_norm.weight"
                 ].numpy()
             )
             keras_nlp_model.get_layer(
                 f"group_{i}_inner_layer_{j}"
-            )._feedforward_layernorm.beta.assign(
+            )._feedforward_layer_norm.beta.assign(
                 hf_wts[
                     f"encoder.albert_layer_groups.{i}.albert_layers.{j}.full_layer_layer_norm.bias"
                 ].numpy()
diff --git a/tools/checkpoint_conversion/convert_bart_checkpoints.py b/tools/checkpoint_conversion/convert_bart_checkpoints.py
index 1bbc464b..3beab9b1 100644
--- a/tools/checkpoint_conversion/convert_bart_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_bart_checkpoints.py
@@ -61,10 +61,10 @@ def convert_checkpoints(hf_model):
     ).position_embeddings.assign(hf_wts["encoder.embed_positions.weight"][2:])
 
     keras_nlp_model.get_layer("encoder_embeddings_layer_norm").gamma.assign(
-        hf_wts["encoder.layernorm_embedding.weight"]
+        hf_wts["encoder.layer_norm_embedding.weight"]
     )
     keras_nlp_model.get_layer("encoder_embeddings_layer_norm").beta.assign(
-        hf_wts["encoder.layernorm_embedding.bias"]
+        hf_wts["encoder.layer_norm_embedding.bias"]
     )
 
     for i in range(keras_nlp_model.num_layers):
@@ -133,12 +133,12 @@ def convert_checkpoints(hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_encoder_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             hf_wts[f"encoder.layers.{i}.self_attn_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_encoder_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             hf_wts[f"encoder.layers.{i}.self_attn_layer_norm.bias"].numpy()
         )
 
@@ -167,12 +167,12 @@ def convert_checkpoints(hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_encoder_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             hf_wts[f"encoder.layers.{i}.final_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_encoder_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             hf_wts[f"encoder.layers.{i}.final_layer_norm.bias"].numpy()
         )
 
@@ -183,10 +183,10 @@ def convert_checkpoints(hf_model):
     ).position_embeddings.assign(hf_wts["decoder.embed_positions.weight"][2:])
 
     keras_nlp_model.get_layer("decoder_embeddings_layer_norm").gamma.assign(
-        hf_wts["decoder.layernorm_embedding.weight"]
+        hf_wts["decoder.layer_norm_embedding.weight"]
     )
     keras_nlp_model.get_layer("decoder_embeddings_layer_norm").beta.assign(
-        hf_wts["decoder.layernorm_embedding.bias"]
+        hf_wts["decoder.layer_norm_embedding.bias"]
     )
 
     for i in range(keras_nlp_model.num_layers):
@@ -255,12 +255,12 @@ def convert_checkpoints(hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_decoder_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             hf_wts[f"decoder.layers.{i}.self_attn_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_decoder_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             hf_wts[f"decoder.layers.{i}.self_attn_layer_norm.bias"].numpy()
         )
 
@@ -329,12 +329,12 @@ def convert_checkpoints(hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_decoder_layer_{i}"
-        )._cross_attention_layernorm.gamma.assign(
+        )._cross_attention_layer_norm.gamma.assign(
             hf_wts[f"decoder.layers.{i}.encoder_attn_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_decoder_layer_{i}"
-        )._cross_attention_layernorm.beta.assign(
+        )._cross_attention_layer_norm.beta.assign(
             hf_wts[f"decoder.layers.{i}.encoder_attn_layer_norm.bias"].numpy()
         )
 
@@ -363,12 +363,12 @@ def convert_checkpoints(hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_decoder_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             hf_wts[f"decoder.layers.{i}.final_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_decoder_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             hf_wts[f"decoder.layers.{i}.final_layer_norm.bias"].numpy()
         )
 
diff --git a/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py b/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
index 82f35a8e..f4d6239f 100644
--- a/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
@@ -205,14 +205,14 @@ def convert_checkpoints(keras_nlp_model, hf_model):
 
         keras_nlp_model.get_layer(
             f"disentangled_attention_encoder_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             hf_wts[
                 f"encoder.layer.{i}.attention.output.LayerNorm.weight"
             ].numpy()
         )
         keras_nlp_model.get_layer(
             f"disentangled_attention_encoder_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             hf_wts[f"encoder.layer.{i}.attention.output.LayerNorm.bias"].numpy()
         )
 
@@ -244,12 +244,12 @@ def convert_checkpoints(keras_nlp_model, hf_model):
 
         keras_nlp_model.get_layer(
             f"disentangled_attention_encoder_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             hf_wts[f"encoder.layer.{i}.output.LayerNorm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"disentangled_attention_encoder_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             hf_wts[f"encoder.layer.{i}.output.LayerNorm.bias"].numpy()
         )
 
diff --git a/tools/checkpoint_conversion/convert_distilbert_checkpoints.py b/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
index 80af89eb..be190353 100644
--- a/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
@@ -197,12 +197,12 @@ def convert_checkpoints(keras_nlp_model, hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             hf_wts[f"transformer.layer.{i}.sa_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             hf_wts[f"transformer.layer.{i}.sa_layer_norm.bias"].numpy()
         )
 
@@ -234,12 +234,12 @@ def convert_checkpoints(keras_nlp_model, hf_model):
 
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             hf_wts[f"transformer.layer.{i}.output_layer_norm.weight"].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             hf_wts[f"transformer.layer.{i}.output_layer_norm.bias"].numpy()
         )
 
diff --git a/tools/checkpoint_conversion/convert_gpt2_checkpoints.py b/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
index 79e67eef..f9e37540 100644
--- a/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
@@ -148,10 +148,10 @@ def convert_checkpoints(num_params):
 
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(weights[f"model/h{i}/ln_1/g"])
+        )._self_attention_layer_norm.gamma.assign(weights[f"model/h{i}/ln_1/g"])
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.beta.assign(weights[f"model/h{i}/ln_1/b"])
+        )._self_attention_layer_norm.beta.assign(weights[f"model/h{i}/ln_1/b"])
 
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
@@ -177,10 +177,10 @@ def convert_checkpoints(num_params):
 
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(weights[f"model/h{i}/ln_2/g"])
+        )._feedforward_layer_norm.gamma.assign(weights[f"model/h{i}/ln_2/g"])
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.beta.assign(weights[f"model/h{i}/ln_2/b"])
+        )._feedforward_layer_norm.beta.assign(weights[f"model/h{i}/ln_2/b"])
 
     keras_nlp_model.get_layer("layer_norm").gamma.assign(
         weights["model/ln_f/g"]
diff --git a/tools/checkpoint_conversion/convert_gpt_neox_checkpoints.py b/tools/checkpoint_conversion/convert_gpt_neox_checkpoints.py
index 05800d71..0a047a4f 100644
--- a/tools/checkpoint_conversion/convert_gpt_neox_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_gpt_neox_checkpoints.py
@@ -108,25 +108,25 @@ for layer_index in range(cfg["num_layers"]):
     # LAYERNORM
     keras_model.get_layer(
         f"transformer_layer_{layer_index}"
-    )._self_attention_layernorm.gamma.assign(
+    )._self_attention_layer_norm.gamma.assign(
         hf_wts[f"layers.{layer_index}.input_layernorm.weight"]
     )
 
     keras_model.get_layer(
         f"transformer_layer_{layer_index}"
-    )._self_attention_layernorm.beta.assign(
+    )._self_attention_layer_norm.beta.assign(
         hf_wts[f"layers.{layer_index}.input_layernorm.bias"]
     )
 
     keras_model.get_layer(
         f"transformer_layer_{layer_index}"
-    )._feedforward_layernorm.gamma.assign(
+    )._feedforward_layer_norm.gamma.assign(
         hf_wts[f"layers.{layer_index}.post_attention_layernorm.weight"]
     )
 
     keras_model.get_layer(
         f"transformer_layer_{layer_index}"
-    )._feedforward_layernorm.beta.assign(
+    )._feedforward_layer_norm.beta.assign(
         hf_wts[f"layers.{layer_index}.post_attention_layernorm.bias"]
     )
 
diff --git a/tools/checkpoint_conversion/convert_opt_checkpoints.py b/tools/checkpoint_conversion/convert_opt_checkpoints.py
index 2f577bfb..fb1ea656 100644
--- a/tools/checkpoint_conversion/convert_opt_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_opt_checkpoints.py
@@ -129,12 +129,12 @@ def convert_weights(hf_model):
         # Attention LayerNorm
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             hf_model.model.decoder.layers[i].self_attn_layer_norm.gamma
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             hf_model.model.decoder.layers[i].self_attn_layer_norm.beta
         )
 
@@ -165,12 +165,12 @@ def convert_weights(hf_model):
         # FF LayerNorm
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             hf_model.model.decoder.layers[i].final_layer_norm.gamma
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             hf_model.model.decoder.layers[i].final_layer_norm.beta
         )
 
diff --git a/tools/checkpoint_conversion/convert_roberta_checkpoints.py b/tools/checkpoint_conversion/convert_roberta_checkpoints.py
index 092eb470..aaa9cac4 100644
--- a/tools/checkpoint_conversion/convert_roberta_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_roberta_checkpoints.py
@@ -206,14 +206,14 @@ def convert_checkpoints(size):
         # Attention LayerNorm
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.self_attn_layer_norm.weight"
             ].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.self_attn_layer_norm.bias"
             ].numpy()
@@ -250,14 +250,14 @@ def convert_checkpoints(size):
         # FF LayerNorm
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.final_layer_norm.weight"
             ].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.final_layer_norm.bias"
             ].numpy()
diff --git a/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py b/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
index 40cf3862..b1295b1e 100644
--- a/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
@@ -187,14 +187,14 @@ def convert_checkpoints(size):
         # Attention LayerNorm
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.gamma.assign(
+        )._self_attention_layer_norm.gamma.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.self_attn_layer_norm.weight"
             ].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._self_attention_layernorm.beta.assign(
+        )._self_attention_layer_norm.beta.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.self_attn_layer_norm.bias"
             ].numpy()
@@ -231,14 +231,14 @@ def convert_checkpoints(size):
         # FF LayerNorm
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.gamma.assign(
+        )._feedforward_layer_norm.gamma.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.final_layer_norm.weight"
             ].numpy()
         )
         keras_nlp_model.get_layer(
             f"transformer_layer_{i}"
-        )._feedforward_layernorm.beta.assign(
+        )._feedforward_layer_norm.beta.assign(
             pt_model[
                 f"decoder.sentence_encoder.layers.{i}.final_layer_norm.bias"
             ].numpy()


Test patch:
diff --git a/keras_nlp/models/bart/bart_preprocessor_test.py b/keras_nlp/models/bart/bart_preprocessor_test.py
index df569c5d..0b192989 100644
--- a/keras_nlp/models/bart/bart_preprocessor_test.py
+++ b/keras_nlp/models/bart/bart_preprocessor_test.py
@@ -74,6 +74,22 @@ class BartPreprocessorTest(TestCase):
             output["decoder_padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1, 0]
         )
 
+    def test_key_order(self):
+        self.assertAllClose(
+            self.preprocessor(
+                {
+                    "encoder_text": " airplane at airport",
+                    "decoder_text": " kohli is the best",
+                }
+            ),
+            self.preprocessor(
+                {
+                    "decoder_text": " kohli is the best",
+                    "encoder_text": " airplane at airport",
+                }
+            ),
+        )
+
     def test_tokenize_list_of_strings(self):
         input_data = {
             "encoder_text": [" airplane at airport"] * 4,