Add 50k train step model

config.gin

@@ -0,0 +1,148 @@
+from __gin__ import dynamic_registration
+import __main__ as train_script
+import seqio
+from t5x import adafactor
+from t5x.examples.scalable_t5 import network
+from t5x import gin_utils
+from t5x import models
+from t5x import partitioning
+from t5x import trainer
+from t5x import utils
+import tasks
+
+# Macros:
+# ==============================================================================
+BATCH_SIZE = 256
+DROPOUT_RATE = 0.0
+LABEL_SMOOTHING = 0.0
+LOSS_NORMALIZING_FACTOR = None
+MIXTURE_OR_TASK_MODULE = None
+MIXTURE_OR_TASK_NAME = 'pretrain_finnish'
+MODEL = @models.EncoderDecoderModel()
+MODEL_DIR = '/researchdisk/t5-small-nl16-finnish'
+OPTIMIZER = @adafactor.Adafactor()
+RANDOM_SEED = None
+SHUFFLE_TRAIN_EXAMPLES = True
+TASK_FEATURE_LENGTHS = {'inputs': 512, 'targets': 512}
+TRAIN_STEPS = 500000
+USE_CACHED_TASKS = False
+USE_HARDWARE_RNG = False
+VOCABULARY = @seqio.SentencePieceVocabulary()
+Z_LOSS = 0.0001
+
+# Parameters for adafactor.Adafactor:
+# ==============================================================================
+adafactor.Adafactor.decay_rate = 0.8
+adafactor.Adafactor.logical_factor_rules = \
+    @adafactor.standard_logical_factor_rules()
+adafactor.Adafactor.step_offset = 0
+
+# Parameters for utils.CheckpointConfig:
+# ==============================================================================
+utils.CheckpointConfig.restore = @utils.RestoreCheckpointConfig()
+utils.CheckpointConfig.save = @utils.SaveCheckpointConfig()
+
+# Parameters for utils.create_learning_rate_scheduler:
+# ==============================================================================
+utils.create_learning_rate_scheduler.base_learning_rate = 1.0
+utils.create_learning_rate_scheduler.factors = 'constant * rsqrt_decay'
+utils.create_learning_rate_scheduler.warmup_steps = 10000
+
+# Parameters for train/utils.DatasetConfig:
+# ==============================================================================
+train/utils.DatasetConfig.batch_size = %BATCH_SIZE
+train/utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+train/utils.DatasetConfig.module = %MIXTURE_OR_TASK_MODULE
+train/utils.DatasetConfig.pack = True
+train/utils.DatasetConfig.seed = None
+train/utils.DatasetConfig.shuffle = %SHUFFLE_TRAIN_EXAMPLES
+train/utils.DatasetConfig.split = 'train'
+train/utils.DatasetConfig.task_feature_lengths = %TASK_FEATURE_LENGTHS
+train/utils.DatasetConfig.use_cached = %USE_CACHED_TASKS
+
+# Parameters for train_eval/utils.DatasetConfig:
+# ==============================================================================
+train_eval/utils.DatasetConfig.batch_size = %BATCH_SIZE
+train_eval/utils.DatasetConfig.mixture_or_task_name = %MIXTURE_OR_TASK_NAME
+train_eval/utils.DatasetConfig.module = %MIXTURE_OR_TASK_MODULE
+train_eval/utils.DatasetConfig.pack = True
+train_eval/utils.DatasetConfig.seed = 42
+train_eval/utils.DatasetConfig.shuffle = False
+train_eval/utils.DatasetConfig.split = 'validation'
+train_eval/utils.DatasetConfig.task_feature_lengths = %TASK_FEATURE_LENGTHS
+train_eval/utils.DatasetConfig.use_cached = %USE_CACHED_TASKS
+
+# Parameters for models.EncoderDecoderModel:
+# ==============================================================================
+models.EncoderDecoderModel.input_vocabulary = %VOCABULARY
+models.EncoderDecoderModel.label_smoothing = %LABEL_SMOOTHING
+models.EncoderDecoderModel.loss_normalizing_factor = %LOSS_NORMALIZING_FACTOR
+models.EncoderDecoderModel.module = @network.Transformer()
+models.EncoderDecoderModel.optimizer_def = %OPTIMIZER
+models.EncoderDecoderModel.output_vocabulary = %VOCABULARY
+models.EncoderDecoderModel.z_loss = %Z_LOSS
+
+# Parameters for partitioning.PjitPartitioner:
+# ==============================================================================
+partitioning.PjitPartitioner.logical_axis_rules = \
+    @partitioning.standard_logical_axis_rules()
+partitioning.PjitPartitioner.model_parallel_submesh = None
+partitioning.PjitPartitioner.num_partitions = 1
+
+# Parameters for utils.RestoreCheckpointConfig:
+# ==============================================================================
+utils.RestoreCheckpointConfig.path = []
+
+# Parameters for utils.SaveCheckpointConfig:
+# ==============================================================================
+utils.SaveCheckpointConfig.dtype = 'float32'
+utils.SaveCheckpointConfig.keep = 10
+utils.SaveCheckpointConfig.period = 10000
+utils.SaveCheckpointConfig.save_dataset = False
+
+# Parameters for seqio.SentencePieceVocabulary:
+# ==============================================================================
+seqio.SentencePieceVocabulary.sentencepiece_model_file = 'spiece.model'
+
+# Parameters for network.T5Config:
+# ==============================================================================
+network.T5Config.dropout_rate = %DROPOUT_RATE
+network.T5Config.dtype = 'bfloat16'
+network.T5Config.emb_dim = 512
+network.T5Config.head_dim = 64
+network.T5Config.logits_via_embedding = False
+network.T5Config.mlp_activations = ('gelu', 'linear')
+network.T5Config.mlp_dim = 2048
+network.T5Config.num_decoder_layers = 16
+network.T5Config.num_encoder_layers = 16
+network.T5Config.num_heads = 8
+network.T5Config.remat_policy = 'minimal'
+network.T5Config.scan_layers = True
+network.T5Config.vocab_size = 32128
+
+# Parameters for train_script.train:
+# ==============================================================================
+train_script.train.checkpoint_cfg = @utils.CheckpointConfig()
+train_script.train.eval_period = 10000
+train_script.train.eval_steps = 20
+train_script.train.infer_eval_dataset_cfg = None
+train_script.train.model = %MODEL
+train_script.train.model_dir = %MODEL_DIR
+train_script.train.partitioner = @partitioning.PjitPartitioner()
+train_script.train.random_seed = %RANDOM_SEED
+train_script.train.summarize_config_fn = @gin_utils.summarize_gin_config
+train_script.train.total_steps = %TRAIN_STEPS
+train_script.train.train_dataset_cfg = @train/utils.DatasetConfig()
+train_script.train.train_eval_dataset_cfg = @train_eval/utils.DatasetConfig()
+train_script.train.trainer_cls = @trainer.Trainer
+train_script.train.use_gda = False
+train_script.train.use_hardware_rng = %USE_HARDWARE_RNG
+
+# Parameters for trainer.Trainer:
+# ==============================================================================
+trainer.Trainer.learning_rate_fn = @utils.create_learning_rate_scheduler()
+trainer.Trainer.num_microbatches = None
+
+# Parameters for network.Transformer:
+# ==============================================================================
+network.Transformer.config = @network.T5Config()

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "Finnish-NLP/t5-small-nl16-finnish",
+  "_name_or_path": "/researchdisk/t5-small-nl16-finnish",
   "architectures": [
     "T5ForConditionalGeneration"
   ],
@@ -26,7 +26,7 @@
   "relative_attention_num_buckets": 32,
   "tie_word_embeddings": false,
   "torch_dtype": "float32",
-  "transformers_version": "4.20.1",
+  "transformers_version": "4.21.2",
   "use_cache": true,
   "vocab_size": 32128
 }

convert_t5x_checkpoint_to_flax.py

@@ -3,7 +3,8 @@
 import argparse
 
 from t5x import checkpoints
-from transformers import T5Config, FlaxT5ForConditionalGeneration
+from transformers import T5Config, FlaxT5ForConditionalGeneration, AutoModelForSeq2SeqLM
+import torch
 
 
 def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_folder_path):
@@ -11,37 +12,36 @@ def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_f
     flax_model = FlaxT5ForConditionalGeneration(config=config)
     t5x_model = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)
 
-    split_mlp_wi = "wi_0" in t5x_model["target"]["encoder"]["layers_0"]["mlp"]
+    split_mlp_wi = "wi_0" in t5x_model["target"]["encoder"]["encoder"]["mlp"]
 
     # Encoder
     for layer_index in range(config.num_layers):
-        layer_name = f"layers_{str(layer_index)}"
         
         # Self-Attention
-        t5x_attention_key = t5x_model["target"]["encoder"][layer_name]["attention"]["key"]["kernel"]
-        t5x_attention_out = t5x_model["target"]["encoder"][layer_name]["attention"]["out"]["kernel"]
-        t5x_attention_query = t5x_model["target"]["encoder"][layer_name]["attention"]["query"]["kernel"]
-        t5x_attention_value = t5x_model["target"]["encoder"][layer_name]["attention"]["value"]["kernel"]
+        t5x_attention_key = t5x_model["target"]["encoder"]["encoder"]["attention"]["key"]["kernel"][:, layer_index, :, :]
+        t5x_attention_out = t5x_model["target"]["encoder"]["encoder"]["attention"]["out"]["kernel"][:, layer_index, :, :]
+        t5x_attention_query = t5x_model["target"]["encoder"]["encoder"]["attention"]["query"]["kernel"][:, layer_index, :, :]
+        t5x_attention_value = t5x_model["target"]["encoder"]["encoder"]["attention"]["value"]["kernel"][:, layer_index, :, :]
 
         ## Layer Normalization
-        t5x_attention_layer_norm = t5x_model["target"]["encoder"][layer_name]["pre_attention_layer_norm"]["scale"]
+        t5x_attention_layer_norm = t5x_model["target"]["encoder"]["encoder"]["pre_attention_layer_norm"]["scale"][:, layer_index]
         
         if split_mlp_wi:
-            t5x_mlp_wi_0 = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi_0"]["kernel"]
-            t5x_mlp_wi_1 = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi_1"]["kernel"]
+            t5x_mlp_wi_0 = t5x_model["target"]["encoder"]["encoder"]["mlp"]["wi_0"]["kernel"][:, layer_index, :]
+            t5x_mlp_wi_1 = t5x_model["target"]["encoder"]["encoder"]["mlp"]["wi_1"]["kernel"][:, layer_index, :]
         else:
-            t5x_mlp_wi = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi"]["kernel"]
+            t5x_mlp_wi = t5x_model["target"]["encoder"]["encoder"]["mlp"]["wi"]["kernel"][:, layer_index, :]
         
-        t5x_mlp_wo = t5x_model["target"]["encoder"][layer_name]["mlp"]["wo"]["kernel"]
+        t5x_mlp_wo = t5x_model["target"]["encoder"]["encoder"]["mlp"]["wo"]["kernel"][:, layer_index, :]
         
         ## Layer Normalization
-        t5x_mlp_layer_norm = t5x_model["target"]["encoder"][layer_name]["pre_mlp_layer_norm"]["scale"]
+        t5x_mlp_layer_norm = t5x_model["target"]["encoder"]["encoder"]["pre_mlp_layer_norm"]["scale"][:, layer_index]
         
         # Assigning
-        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key
-        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out
-        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query
-        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key.reshape(*t5x_attention_key.shape[:-2], -1)
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out.reshape(-1, t5x_attention_out.shape[-1])
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query.reshape(*t5x_attention_query.shape[:-2], -1)
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value.reshape(*t5x_attention_value.shape[:-2], -1)
         
         flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["layer_norm"]["weight"] = t5x_attention_layer_norm
         
@@ -55,59 +55,58 @@ def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_f
         flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["layer_norm"]["weight"] = t5x_mlp_layer_norm
 
     # Only for layer 0:
-    t5x_encoder_rel_embedding = t5x_model["target"]["encoder"]["relpos_bias"]["rel_embedding"].T
-    flax_model.params["encoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_encoder_rel_embedding
+    t5x_encoder_rel_embedding = t5x_model["target"]["encoder"]["encoder"]["relpos_bias"]["rel_embedding"].T
+    flax_model.params["encoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_encoder_rel_embedding[:, 0, :]
 
     # Assigning
     t5x_encoder_norm = t5x_model["target"]["encoder"]["encoder_norm"]["scale"]
     flax_model.params["encoder"]["final_layer_norm"]["weight"] = t5x_encoder_norm
 
     # Decoder
-    for layer_index in range(config.num_layers):
-        layer_name = f"layers_{str(layer_index)}"
+    for layer_index in range(config.num_decoder_layers):
 
         # Self-Attention
-        t5x_attention_key = t5x_model["target"]["decoder"][layer_name]["self_attention"]["key"]["kernel"]
-        t5x_attention_out = t5x_model["target"]["decoder"][layer_name]["self_attention"]["out"]["kernel"]
-        t5x_attention_query = t5x_model["target"]["decoder"][layer_name]["self_attention"]["query"]["kernel"]
-        t5x_attention_value = t5x_model["target"]["decoder"][layer_name]["self_attention"]["value"]["kernel"]
+        t5x_attention_key = t5x_model["target"]["decoder"]["decoder"]["self_attention"]["key"]["kernel"][:, layer_index, :, :]
+        t5x_attention_out = t5x_model["target"]["decoder"]["decoder"]["self_attention"]["out"]["kernel"][:, layer_index, :, :]
+        t5x_attention_query = t5x_model["target"]["decoder"]["decoder"]["self_attention"]["query"]["kernel"][:, layer_index, :, :]
+        t5x_attention_value = t5x_model["target"]["decoder"]["decoder"]["self_attention"]["value"]["kernel"][:, layer_index, :, :]
 
         ## Layer Normalization
-        t5x_pre_attention_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_self_attention_layer_norm"]["scale"]
+        t5x_pre_attention_layer_norm = t5x_model["target"]["decoder"]["decoder"]["pre_self_attention_layer_norm"]["scale"][:, layer_index]
 
         # Encoder-Decoder-Attention
-        t5x_enc_dec_attention_key = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["key"]["kernel"]
-        t5x_enc_dec_attention_out = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["out"]["kernel"]
-        t5x_enc_dec_attention_query = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["query"]["kernel"]
-        t5x_enc_dec_attention_value = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["value"]["kernel"]
+        t5x_enc_dec_attention_key = t5x_model["target"]["decoder"]["decoder"]["encoder_decoder_attention"]["key"]["kernel"][:, layer_index, :, :]
+        t5x_enc_dec_attention_out = t5x_model["target"]["decoder"]["decoder"]["encoder_decoder_attention"]["out"]["kernel"][:, layer_index, :, :]
+        t5x_enc_dec_attention_query = t5x_model["target"]["decoder"]["decoder"]["encoder_decoder_attention"]["query"]["kernel"][:, layer_index, :, :]
+        t5x_enc_dec_attention_value = t5x_model["target"]["decoder"]["decoder"]["encoder_decoder_attention"]["value"]["kernel"][:, layer_index, :, :]
 
         ## Layer Normalization
-        t5x_cross_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_cross_attention_layer_norm"]["scale"]
+        t5x_cross_layer_norm = t5x_model["target"]["decoder"]["decoder"]["pre_cross_attention_layer_norm"]["scale"][:, layer_index]
 
         # MLP
         if split_mlp_wi:
-            t5x_mlp_wi_0 = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi_0"]["kernel"]
-            t5x_mlp_wi_1 = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi_1"]["kernel"]
+            t5x_mlp_wi_0 = t5x_model["target"]["decoder"]["decoder"]["mlp"]["wi_0"]["kernel"][:, layer_index, :]
+            t5x_mlp_wi_1 = t5x_model["target"]["decoder"]["decoder"]["mlp"]["wi_1"]["kernel"][:, layer_index, :]
         else:
-            t5x_mlp_wi = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi"]["kernel"]
+            t5x_mlp_wi = t5x_model["target"]["decoder"]["decoder"]["mlp"]["wi"]["kernel"][:, layer_index, :]
 
-        t5x_mlp_wo = t5x_model["target"]["decoder"][layer_name]["mlp"]["wo"]["kernel"]
+        t5x_mlp_wo = t5x_model["target"]["decoder"]["decoder"]["mlp"]["wo"]["kernel"][:, layer_index, :]
 
         ## Layer Normalization
-        tx5_mlp_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_mlp_layer_norm"]["scale"]
+        tx5_mlp_layer_norm = t5x_model["target"]["decoder"]["decoder"]["pre_mlp_layer_norm"]["scale"][:, layer_index]
         
         # Assigning
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key.reshape(*t5x_attention_key.shape[:-2], -1)
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out.reshape(-1, t5x_attention_out.shape[-1])
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query.reshape(*t5x_attention_query.shape[:-2], -1)
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value.reshape(*t5x_attention_value.shape[:-2], -1)
 
         flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["layer_norm"]["weight"] = t5x_pre_attention_layer_norm
         
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["k"]["kernel"] = t5x_enc_dec_attention_key
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["o"]["kernel"] = t5x_enc_dec_attention_out
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["q"]["kernel"] = t5x_enc_dec_attention_query
-        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["v"]["kernel"] = t5x_enc_dec_attention_value
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["k"]["kernel"] = t5x_enc_dec_attention_key.reshape(*t5x_enc_dec_attention_key.shape[:-2], -1)
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["o"]["kernel"] = t5x_enc_dec_attention_out.reshape(-1, t5x_enc_dec_attention_out.shape[-1])
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["q"]["kernel"] = t5x_enc_dec_attention_query.reshape(*t5x_enc_dec_attention_query.shape[:-2], -1)
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["v"]["kernel"] = t5x_enc_dec_attention_value.reshape(*t5x_enc_dec_attention_value.shape[:-2], -1)
         
         flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["layer_norm"]["weight"] = t5x_cross_layer_norm
         
@@ -126,8 +125,8 @@ def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_f
     flax_model.params["decoder"]["final_layer_norm"]["weight"] = tx5_decoder_norm
 
     # Only for layer 0:
-    t5x_decoder_rel_embedding = t5x_model["target"]["decoder"]["relpos_bias"]["rel_embedding"].T
-    flax_model.params["decoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_decoder_rel_embedding
+    t5x_decoder_rel_embedding = t5x_model["target"]["decoder"]["decoder"]["relpos_bias"]["rel_embedding"].T
+    flax_model.params["decoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_decoder_rel_embedding[:, 0, :]
 
     # Token Embeddings
     tx5_token_embeddings = t5x_model["target"]["token_embedder"]["embedding"]
@@ -139,6 +138,10 @@ def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_f
     flax_model.save_pretrained(flax_dump_folder_path)
     print("T5X Model was sucessfully converted!")
 
+def convert_flax_to_pytorch(flax_dump_folder_path, pytorch_dump_folder_path):
+    model = AutoModelForSeq2SeqLM.from_pretrained(flax_dump_folder_path, from_flax=True, torch_dtype=torch.float32)
+    model.save_pretrained(pytorch_dump_folder_path)
+    print("Flax model was sucessfully converted to Pytorch!")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -154,4 +157,6 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
     convert_t5x_checkpoint_to_flax(args.t5x_checkpoint_path, args.config_name, args.flax_dump_folder_path)
+    convert_flax_to_pytorch(args.flax_dump_folder_path, args.flax_dump_folder_path)
+
     

flax_model.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af1dfe073774d3fd1413eb3ab376f16e6a8448dd3f979197a1297c2538377b89
+size 735762207

model-info.txt

@@ -0,0 +1,148 @@
+Variable decoder/decoder/encoder_decoder_attention/key/kernel                             size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable decoder/decoder/encoder_decoder_attention/out/kernel                             size 4194304      shape (heads=8, layers=16, kv=64, embed=512)   partition spec ('model', None, None, None)
+Variable decoder/decoder/encoder_decoder_attention/query/kernel                           size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable decoder/decoder/encoder_decoder_attention/value/kernel                           size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable decoder/decoder/mlp/wi_0/kernel                                                  size 16777216     shape (embed=512, layers=16, mlp=2048)         partition spec (None, None, 'model')
+Variable decoder/decoder/mlp/wi_1/kernel                                                  size 16777216     shape (embed=512, layers=16, mlp=2048)         partition spec (None, None, 'model')
+Variable decoder/decoder/mlp/wo/kernel                                                    size 16777216     shape (mlp=2048, layers=16, embed=512)         partition spec ('model', None, None)
+Variable decoder/decoder/pre_cross_attention_layer_norm/scale                             size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable decoder/decoder/pre_mlp_layer_norm/scale                                         size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable decoder/decoder/pre_self_attention_layer_norm/scale                              size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable decoder/decoder/relpos_bias/rel_embedding                                        size 4096         shape (heads=8, layers=16, relpos_buckets=32)  partition spec ('model', None, None)
+Variable decoder/decoder/self_attention/key/kernel                                        size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable decoder/decoder/self_attention/out/kernel                                        size 4194304      shape (heads=8, layers=16, kv=64, embed=512)   partition spec ('model', None, None, None)
+Variable decoder/decoder/self_attention/query/kernel                                      size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable decoder/decoder/self_attention/value/kernel                                      size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable decoder/decoder_norm/scale                                                       size 512          shape (embed=512)                              partition spec (None,)
+Variable decoder/logits_dense/kernel                                                      size 16449536     shape (embed=512, vocab=32128)                 partition spec (None, 'model')
+Variable encoder/encoder/attention/key/kernel                                             size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable encoder/encoder/attention/out/kernel                                             size 4194304      shape (heads=8, layers=16, kv=64, embed=512)   partition spec ('model', None, None, None)
+Variable encoder/encoder/attention/query/kernel                                           size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable encoder/encoder/attention/value/kernel                                           size 4194304      shape (embed=512, layers=16, heads=8, kv=64)   partition spec (None, None, 'model', None)
+Variable encoder/encoder/mlp/wi_0/kernel                                                  size 16777216     shape (embed=512, layers=16, mlp=2048)         partition spec (None, None, 'model')
+Variable encoder/encoder/mlp/wi_1/kernel                                                  size 16777216     shape (embed=512, layers=16, mlp=2048)         partition spec (None, None, 'model')
+Variable encoder/encoder/mlp/wo/kernel                                                    size 16777216     shape (mlp=2048, layers=16, embed=512)         partition spec ('model', None, None)
+Variable encoder/encoder/pre_attention_layer_norm/scale                                   size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable encoder/encoder/pre_mlp_layer_norm/scale                                         size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable encoder/encoder/relpos_bias/rel_embedding                                        size 4096         shape (heads=8, layers=16, relpos_buckets=32)  partition spec ('model', None, None)
+Variable encoder/encoder_norm/scale                                                       size 512          shape (embed=512)                              partition spec (None,)
+Variable token_embedder/embedding                                                         size 16449536     shape (vocab=32128, embed=512)                 partition spec ('model', None)
+Total number of parameters: 183944192
+
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/m              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/v              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/v_col          size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/key/kernel/v_row          size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/m              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/v              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/v_col          size 8192         shape (8, 16, 64)                              partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/out/kernel/v_row          size 8192         shape (16, 512)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/m            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/v            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/v_col        size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/query/kernel/v_row        size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/m            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/v            size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/v_col        size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/decoder/decoder/encoder_decoder_attention/value/kernel/v_row        size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/v_col                               size 32768        shape (16, 2048)                               partition spec None
+Variable param_states/decoder/decoder/mlp/wi_0/kernel/v_row                               size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/v_col                               size 32768        shape (16, 2048)                               partition spec None
+Variable param_states/decoder/decoder/mlp/wi_1/kernel/v_row                               size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/m                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/v                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/v_col                                 size 32768        shape (2048, 16)                               partition spec None
+Variable param_states/decoder/decoder/mlp/wo/kernel/v_row                                 size 8192         shape (16, 512)                                partition spec None
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/m              size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/v              size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/v_col          size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_cross_attention_layer_norm/scale/v_row          size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/m                          size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/v                          size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/v_col                      size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_mlp_layer_norm/scale/v_row                      size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/m               size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/v               size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/v_col           size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/pre_self_attention_layer_norm/scale/v_row           size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/v                         size 4096         shape (heads=8, layers=16, relpos_buckets=32)  partition spec ('model', None, None)
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/v_col                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/relpos_bias/rel_embedding/v_row                     size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/v                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/v_col                     size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/decoder/decoder/self_attention/key/kernel/v_row                     size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/v                         size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/v_col                     size 8192         shape (8, 16, 64)                              partition spec None
+Variable param_states/decoder/decoder/self_attention/out/kernel/v_row                     size 8192         shape (16, 512)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/m                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/v                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/v_col                   size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/decoder/decoder/self_attention/query/kernel/v_row                   size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/m                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/v                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/v_col                   size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/decoder/decoder/self_attention/value/kernel/v_row                   size 8192         shape (512, 16)                                partition spec None
+Variable param_states/decoder/decoder_norm/scale/m                                        size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder_norm/scale/v                                        size 512          shape (embed=512)                              partition spec (None,)
+Variable param_states/decoder/decoder_norm/scale/v_col                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/decoder_norm/scale/v_row                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/logits_dense/kernel/m                                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/logits_dense/kernel/v                                       size 1            shape (1,)                                     partition spec None
+Variable param_states/decoder/logits_dense/kernel/v_col                                   size 32128        shape (32128,)                                 partition spec None
+Variable param_states/decoder/logits_dense/kernel/v_row                                   size 512          shape (512,)                                   partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/m                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/v                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/v_col                          size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/encoder/encoder/attention/key/kernel/v_row                          size 8192         shape (512, 16)                                partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/m                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/v                              size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/v_col                          size 8192         shape (8, 16, 64)                              partition spec None
+Variable param_states/encoder/encoder/attention/out/kernel/v_row                          size 8192         shape (16, 512)                                partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/m                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/v                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/v_col                        size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/encoder/encoder/attention/query/kernel/v_row                        size 8192         shape (512, 16)                                partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/m                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/v                            size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/v_col                        size 8192         shape (16, 8, 64)                              partition spec None
+Variable param_states/encoder/encoder/attention/value/kernel/v_row                        size 8192         shape (512, 16)                                partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/v_col                               size 32768        shape (16, 2048)                               partition spec None
+Variable param_states/encoder/encoder/mlp/wi_0/kernel/v_row                               size 8192         shape (512, 16)                                partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/m                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/v                                   size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/v_col                               size 32768        shape (16, 2048)                               partition spec None
+Variable param_states/encoder/encoder/mlp/wi_1/kernel/v_row                               size 8192         shape (512, 16)                                partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/m                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/v                                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/v_col                                 size 32768        shape (2048, 16)                               partition spec None
+Variable param_states/encoder/encoder/mlp/wo/kernel/v_row                                 size 8192         shape (16, 512)                                partition spec None
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/m                    size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/v                    size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/v_col                size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_attention_layer_norm/scale/v_row                size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/m                          size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/v                          size 8192         shape (embed=512, layers=16)                   partition spec (None, None)
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/v_col                      size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/pre_mlp_layer_norm/scale/v_row                      size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/m                         size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/v                         size 4096         shape (heads=8, layers=16, relpos_buckets=32)  partition spec ('model', None, None)
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/v_col                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder/relpos_bias/rel_embedding/v_row                     size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder_norm/scale/m                                        size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder_norm/scale/v                                        size 512          shape (embed=512)                              partition spec (None,)
+Variable param_states/encoder/encoder_norm/scale/v_col                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/encoder/encoder_norm/scale/v_row                                    size 1            shape (1,)                                     partition spec None
+Variable param_states/token_embedder/embedding/m                                          size 1            shape (1,)                                     partition spec None
+Variable param_states/token_embedder/embedding/v                                          size 1            shape (1,)                                     partition spec None
+Variable param_states/token_embedder/embedding/v_col                                      size 32128        shape (32128,)                                 partition spec None
+Variable param_states/token_embedder/embedding/v_row                                      size 512          shape (512,)                                   partition spec None
+Variable step                                                                             size 1            shape ()                                       partition spec None

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05fe4c98850f026cd8154f4358131f3fe9f8538fb692a4621d31a316ac620c80
+size 735867349

train/events.out.tfevents.1661710468.t1v-n-12f94ad0-w-0.60675.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee5357b199786bd136f34c89c093f98ec5417d1cf220340749fc2496418fc60c
+size 16868

training_eval/pretrain_finnish/events.out.tfevents.1661710468.t1v-n-12f94ad0-w-0.60675.1.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:776cdaf0c0ff210e9e367110778093f7d42c2d9c7836a1a8a4667fb780f2e758
+size 9244