Upload RWForCausalLM

modelling_RW.py

@@ -363,14 +363,12 @@ class DecoderLayer(nn.Module):
         super().__init__()
         hidden_size = config.hidden_size
 
-        self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.ln_attn = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.ln_mlp = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+
         self.num_heads = config.n_head
         self.self_attention = Attention(config)
 
-        if not config.parallel_attn:
-            # unused if parallel attn
-            self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-
         self.mlp = MLP(config)
 
         self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
@@ -389,12 +387,14 @@ class DecoderLayer(nn.Module):
         output_attentions: bool = False,
     ):
 
-        layernorm_output = self.input_layernorm(hidden_states)
+        ln_attn = self.ln_attn(hidden_states)
+        ln_mlp = self.ln_mlp(hidden_states)
+
         residual = hidden_states
 
         # Self attention.
         attn_outputs = self.self_attention(
-            layernorm_output,
+            ln_attn,
             layer_past=layer_past,
             attention_mask=attention_mask,
             alibi=alibi,
@@ -405,19 +405,14 @@ class DecoderLayer(nn.Module):
 
         attention_output = attn_outputs[0]
 
-        if not self.config.parallel_attn:
-            residual = dropout_add(attention_output, residual, self.config.attention_dropout, training=self.training)
-            layernorm_output = self.post_attention_layernorm(residual)
-
         outputs = attn_outputs[1:]
 
         # MLP.
-        mlp_output = self.mlp(layernorm_output)
-
-        if self.config.parallel_attn:
-            mlp_output += attention_output
+        mlp_output = self.mlp(ln_mlp)
 
-        output = dropout_add(mlp_output, residual, self.config.hidden_dropout, training=self.training)
+        output = dropout_add(
+            mlp_output + attention_output, residual, self.config.hidden_dropout, training=self.training
+        )
 
         if use_cache:
             outputs = (output,) + outputs

pytorch_model-00001-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cf8d79dce91486c166640f2dbde437985109617072c8747abc875ff7b35f4937
-size 9504536357
+oid sha256:384bd92e6d1d0b6a133b1d81bad915a10d954d1bfcdf4fdc8b6d0d81b4f5e9e8
+size 9504770141

pytorch_model-00002-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f2bda814cd53c53b6e7c9c168dce4e48371714fcbbe1ee273b6a35173867942a
-size 9512925303
+oid sha256:638cb048ee6ad66f85a176e980332b0d61cbddb2d5dc9b5ad1863d87e020350c
+size 9513159151

pytorch_model-00003-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:29ea68447ced396b92bae50150340fffcb55166ee72835b8e0fb7c0a6002d0a7
-size 9512925367
+oid sha256:b2125ceee3992a093f9808efc9a720643a44e59b17baa77eddb0fbec6965ce5a
+size 9513159151

pytorch_model-00004-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e99fbdaf8f48474b55c23cbaea749790bedf0d52e0d1d7a522f9b4bceed6a4c3
-size 9512925367
+oid sha256:2fc011e39d4a9c5958630b1bb946dec9ca54adbccc52804fed545490f995b20b
+size 9513159151

pytorch_model-00005-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c9398b3a9b7b758f333073f15db1d3ab958c9c9d30556221c0e51a07afa167ca
-size 9512925367
+oid sha256:ae277c7c5aef0d0fea9e9e6b13113dfef0078a36c78805bebe14d32e9ab887d1
+size 9513159151

pytorch_model-00006-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2f158d1250afe86b2e455cdba83a5ba183c7329a5b9f82c53d418878bf6e4b53
-size 9512925367
+oid sha256:d6376a884bc80a6ea95da83f963508e57d8508333c424bed7beda565e4aaa0f3
+size 9513159151

pytorch_model-00007-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:72b644f83bf1b351942b6801cb1da27933cd16c9d4bea0b009d48ca8b864bce2
-size 9512925367
+oid sha256:c3f07c8ec31edebe5ff80ed32d175475a044a0517d84c5eff6f6247ad0a4432d
+size 9513159151

pytorch_model-00008-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4bc35d185c86958ac48722c41434bc9bd4b268e15a17e91b9f24d3cb1f4f2eb1
-size 9512925367
+oid sha256:7af5021fe9afb9b4186a59a91e75aa60570c66fba43390b634387ee54bd73223
+size 9513159151

pytorch_model-00009-of-00009.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:063b41bbd55c371cac7d7cd5a5a93bb1709f54a661c6e4703a2849fc836069b1
-size 7575086769
+oid sha256:c5c794d3106945f514b5501904506485377c95ce2cef030373c7cbe7459ef610
+size 7575220277

pytorch_model.bin.index.json

@@ -1,365 +1,485 @@
 {
   "metadata": {
-    "total_size": 83669975040
+    "total_size": 83671941120
   },
   "weight_map": {
     "lm_head.weight": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.0.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.0.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.0.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.0.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.0.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.0.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.0.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.0.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.0.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.0.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.1.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.1.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.1.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.1.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.1.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.1.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.1.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.1.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.1.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.1.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.10.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.10.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.10.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.10.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.10.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.10.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.10.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.10.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.10.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.10.self_attention.query_key_value.weight": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.11.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.11.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.11.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.11.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.11.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.11.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.11.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.11.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.11.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.11.self_attention.query_key_value.weight": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.12.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.12.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.12.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.12.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.12.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.12.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.12.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.12.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.12.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.12.self_attention.query_key_value.weight": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.13.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.13.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.13.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.13.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.13.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.13.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.13.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.13.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.13.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.13.self_attention.query_key_value.weight": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.14.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.14.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.14.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.14.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.14.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.14.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.14.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.14.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.14.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.14.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.15.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.15.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.15.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.15.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.15.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.15.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.15.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.15.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.15.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.15.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.16.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.16.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.16.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.16.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.16.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.16.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.16.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.16.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.16.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.16.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.17.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.17.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.17.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.17.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.17.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.17.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.17.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.17.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.17.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.17.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.18.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.18.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.18.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.18.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.18.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.18.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.18.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.18.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.18.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.18.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.19.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.19.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.19.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.19.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.19.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.19.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.19.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.19.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.19.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.19.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.2.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.2.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.2.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.2.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.2.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.2.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.2.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.2.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.2.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.2.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.20.input_layernorm.bias": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.20.input_layernorm.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.20.ln_attn.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.20.ln_attn.weight": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.20.ln_mlp.bias": "pytorch_model-00003-of-00009.bin",
+    "transformer.h.20.ln_mlp.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.20.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.20.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.20.self_attention.dense.weight": "pytorch_model-00003-of-00009.bin",
     "transformer.h.20.self_attention.query_key_value.weight": "pytorch_model-00003-of-00009.bin",
-    "transformer.h.21.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.21.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.21.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.21.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.21.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.21.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.21.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.21.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.21.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.21.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.22.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.22.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.22.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.22.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.22.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.22.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.22.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.22.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.22.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.22.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.23.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.23.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.23.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.23.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.23.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.23.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.23.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.23.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.23.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.23.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.24.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.24.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.24.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.24.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.24.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.24.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.24.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.24.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.24.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.24.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.25.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.25.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.25.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.25.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.25.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.25.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.25.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.25.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.25.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.25.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.26.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.26.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.26.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.26.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.26.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.26.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.26.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.26.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.26.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.26.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.27.input_layernorm.bias": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.27.input_layernorm.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.27.ln_attn.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.27.ln_attn.weight": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.27.ln_mlp.bias": "pytorch_model-00004-of-00009.bin",
+    "transformer.h.27.ln_mlp.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.27.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.27.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.27.self_attention.dense.weight": "pytorch_model-00004-of-00009.bin",
     "transformer.h.27.self_attention.query_key_value.weight": "pytorch_model-00004-of-00009.bin",
-    "transformer.h.28.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.28.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.28.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.28.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.28.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.28.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.28.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.28.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.28.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.28.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.29.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.29.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.29.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.29.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.29.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.29.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.29.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.29.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.29.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.29.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.3.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.3.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.3.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.3.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.3.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.3.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.3.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.3.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.3.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.3.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.30.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.30.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.30.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.30.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.30.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.30.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.30.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.30.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.30.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.30.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.31.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.31.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.31.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.31.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.31.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.31.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.31.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.31.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.31.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.31.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.32.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.32.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.32.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.32.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.32.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.32.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.32.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.32.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.32.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.32.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.33.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.33.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.33.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.33.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.33.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.33.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.33.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.33.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.33.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.33.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.34.input_layernorm.bias": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.34.input_layernorm.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.34.ln_attn.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.34.ln_attn.weight": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.34.ln_mlp.bias": "pytorch_model-00005-of-00009.bin",
+    "transformer.h.34.ln_mlp.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.34.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.34.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.34.self_attention.dense.weight": "pytorch_model-00005-of-00009.bin",
     "transformer.h.34.self_attention.query_key_value.weight": "pytorch_model-00005-of-00009.bin",
-    "transformer.h.35.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.35.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.35.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.35.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.35.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.35.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.35.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.35.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.35.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.35.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.36.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.36.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.36.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.36.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.36.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.36.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.36.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.36.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.36.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.36.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.37.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.37.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.37.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.37.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.37.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.37.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.37.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.37.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.37.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.37.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.38.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.38.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.38.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.38.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.38.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.38.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.38.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.38.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.38.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.38.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.39.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.39.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.39.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.39.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.39.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.39.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.39.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.39.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.39.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.39.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.4.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.4.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.4.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.4.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.4.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.4.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.4.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.4.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.4.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.4.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.40.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.40.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.40.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.40.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.40.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.40.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.40.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.40.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.40.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.40.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.41.input_layernorm.bias": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.41.input_layernorm.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.41.ln_attn.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.41.ln_attn.weight": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.41.ln_mlp.bias": "pytorch_model-00006-of-00009.bin",
+    "transformer.h.41.ln_mlp.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.41.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.41.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.41.self_attention.dense.weight": "pytorch_model-00006-of-00009.bin",
     "transformer.h.41.self_attention.query_key_value.weight": "pytorch_model-00006-of-00009.bin",
-    "transformer.h.42.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.42.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.42.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.42.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.42.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.42.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.42.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.42.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.42.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.42.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.43.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.43.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.43.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.43.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.43.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.43.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.43.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.43.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.43.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.43.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.44.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.44.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.44.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.44.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.44.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.44.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.44.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.44.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.44.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.44.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.45.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.45.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.45.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.45.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.45.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.45.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.45.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.45.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.45.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.45.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.46.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.46.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.46.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.46.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.46.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.46.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.46.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.46.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.46.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.46.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.47.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.47.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.47.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.47.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.47.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.47.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.47.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.47.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.47.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.47.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.48.input_layernorm.bias": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.48.input_layernorm.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.48.ln_attn.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.48.ln_attn.weight": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.48.ln_mlp.bias": "pytorch_model-00007-of-00009.bin",
+    "transformer.h.48.ln_mlp.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.48.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.48.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.48.self_attention.dense.weight": "pytorch_model-00007-of-00009.bin",
     "transformer.h.48.self_attention.query_key_value.weight": "pytorch_model-00007-of-00009.bin",
-    "transformer.h.49.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.49.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.49.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.49.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.49.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.49.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.49.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.49.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.49.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.49.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.5.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.5.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.5.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.5.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.5.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.5.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.5.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.5.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.5.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.5.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.50.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.50.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.50.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.50.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.50.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.50.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.50.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.50.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.50.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.50.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.51.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.51.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.51.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.51.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.51.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.51.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.51.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.51.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.51.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.51.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.52.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.52.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.52.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.52.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.52.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.52.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.52.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.52.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.52.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.52.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.53.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.53.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.53.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.53.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.53.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.53.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.53.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.53.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.53.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.53.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.54.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.54.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.54.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.54.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.54.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.54.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.54.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.54.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.54.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.54.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.55.input_layernorm.bias": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.55.input_layernorm.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.55.ln_attn.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.55.ln_attn.weight": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.55.ln_mlp.bias": "pytorch_model-00008-of-00009.bin",
+    "transformer.h.55.ln_mlp.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.55.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.55.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.55.self_attention.dense.weight": "pytorch_model-00008-of-00009.bin",
     "transformer.h.55.self_attention.query_key_value.weight": "pytorch_model-00008-of-00009.bin",
-    "transformer.h.56.input_layernorm.bias": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.56.input_layernorm.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.56.ln_attn.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.56.ln_attn.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.56.ln_mlp.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.56.ln_mlp.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.56.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.56.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.56.self_attention.dense.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.56.self_attention.query_key_value.weight": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.57.input_layernorm.bias": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.57.input_layernorm.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.57.ln_attn.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.57.ln_attn.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.57.ln_mlp.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.57.ln_mlp.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.57.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.57.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.57.self_attention.dense.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.57.self_attention.query_key_value.weight": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.58.input_layernorm.bias": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.58.input_layernorm.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.58.ln_attn.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.58.ln_attn.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.58.ln_mlp.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.58.ln_mlp.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.58.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.58.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.58.self_attention.dense.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.58.self_attention.query_key_value.weight": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.59.input_layernorm.bias": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.59.input_layernorm.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.59.ln_attn.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.59.ln_attn.weight": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.59.ln_mlp.bias": "pytorch_model-00009-of-00009.bin",
+    "transformer.h.59.ln_mlp.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.59.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.59.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.59.self_attention.dense.weight": "pytorch_model-00009-of-00009.bin",
     "transformer.h.59.self_attention.query_key_value.weight": "pytorch_model-00009-of-00009.bin",
-    "transformer.h.6.input_layernorm.bias": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.6.input_layernorm.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.6.ln_attn.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.6.ln_attn.weight": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.6.ln_mlp.bias": "pytorch_model-00001-of-00009.bin",
+    "transformer.h.6.ln_mlp.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.6.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.6.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.6.self_attention.dense.weight": "pytorch_model-00001-of-00009.bin",
     "transformer.h.6.self_attention.query_key_value.weight": "pytorch_model-00001-of-00009.bin",
-    "transformer.h.7.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.7.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.7.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.7.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.7.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.7.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.7.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.7.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.7.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.7.self_attention.query_key_value.weight": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.8.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.8.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.8.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.8.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.8.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.8.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.8.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.8.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.8.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.8.self_attention.query_key_value.weight": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.9.input_layernorm.bias": "pytorch_model-00002-of-00009.bin",
-    "transformer.h.9.input_layernorm.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.9.ln_attn.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.9.ln_attn.weight": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.9.ln_mlp.bias": "pytorch_model-00002-of-00009.bin",
+    "transformer.h.9.ln_mlp.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.9.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.9.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00009.bin",
     "transformer.h.9.self_attention.dense.weight": "pytorch_model-00002-of-00009.bin",