Saving weights and logs of step 10000

config.json

@@ -1,12 +1,13 @@
 {
-  "_name_or_path": "/researchdisk/roberta-large-finnish-wechsel",
+  "_name_or_path": "./",
   "architectures": [
-    "RobertaModel"
+    "RobertaForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.1,
   "bos_token_id": 0,
   "classifier_dropout": null,
   "eos_token_id": 2,
+  "gradient_checkpointing": false,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
   "hidden_size": 1024,
@@ -19,7 +20,7 @@
   "num_hidden_layers": 24,
   "pad_token_id": 1,
   "position_embedding_type": "absolute",
-  "torch_dtype": "float64",
+  "torch_dtype": "float32",
   "transformers_version": "4.13.0.dev0",
   "type_vocab_size": 1,
   "use_cache": true,

events.out.tfevents.1639865567.t1v-n-8eba1090-w-0.1317510.0.v2

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

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c9a123577826ae147f24d257b4f877eaa05fd6b67d294bef5786cd5b174f7eb7
-size 1421452955
+oid sha256:ed2451f3cd9fff0108476a448868d5478995e86c8bac4935e33645d62109b5de
+size 1421662309

pytorch_model.bin

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

run_mlm_flax.py

@@ -164,6 +164,10 @@ class ModelArguments:
             "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
         },
     )
+    from_pytorch: bool = field(
+        default=False,
+        metadata={"help": "Whether to use Pytorch model checkpoint for weight initialization or not."},
+    )
 
 
 @dataclass
@@ -608,7 +612,7 @@ def main():
 
     if model_args.model_name_or_path:
         model = FlaxAutoModelForMaskedLM.from_pretrained(
-            model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
+            model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype), from_pt=model_args.from_pytorch
         )
     else:
         model = FlaxAutoModelForMaskedLM.from_config(

run_wechsel.py

@@ -1,10 +1,10 @@
 import torch
-from transformers import AutoModel, AutoTokenizer, FlaxAutoModel
+from transformers import AutoModelForMaskedLM, AutoTokenizer, FlaxAutoModelForMaskedLM
 from datasets import load_dataset
 from wechsel import WECHSEL, load_embeddings
 
 source_tokenizer = AutoTokenizer.from_pretrained("roberta-large")
-model = AutoModel.from_pretrained("roberta-large")
+model = AutoModelForMaskedLM.from_pretrained("roberta-large")
 
 target_tokenizer = AutoTokenizer.from_pretrained("./")
 
@@ -20,9 +20,8 @@ target_embeddings, info = wechsel.apply(
     model.get_input_embeddings().weight.detach().numpy(),
 )
 
-model.get_input_embeddings().weight.data = torch.from_numpy(target_embeddings)
-
+model.get_input_embeddings().weight.data = torch.from_numpy(target_embeddings).to(torch.float32)
 model.save_pretrained("./")
 
-flax_model = FlaxAutoModel.from_pretrained("./", from_pt=True)
-flax_model.save_pretrained("./")
+# flax_model = FlaxAutoModelForMaskedLM.from_pretrained("./", from_pt=True)
+# flax_model.save_pretrained("./")

start_train.sh

@@ -5,9 +5,10 @@ export USE_TORCH=0
 python3 run_mlm_flax.py \
     --output_dir="./" \
     --model_name_or_path="./" \
+    --from_pytorch \
     --config_name="./" \
     --tokenizer_name="./" \
-    --dataset_filepath="/researchdisk/training_dataset_full" \
+    --dataset_filepath="/researchdisk/training_dataset_full_deduplicated" \
     --max_seq_length="128" \
     --pad_to_max_length \
     --preprocessing_num_workers="64" \
@@ -25,5 +26,6 @@ python3 run_mlm_flax.py \
     --eval_steps="10000" \
     --logging_steps="500" \
     --dtype="bfloat16" \
+    --adafactor \
     --push_to_hub \
     --hub_model_id="Finnish-NLP/roberta-large-wechsel-finnish"