Add 128 pytorch model

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

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ce839b5e41f8765fb370f9d951ce8337ae17c2685c247f1273d66058f88f093b
-size 40215815
+oid sha256:c1ab87187bbc8b9b5e02916e5b331c99f73c0112cf4a39c025c8d103da252fad
+size 40812375

flax_model_to_pytorch.py

@@ -0,0 +1,27 @@
+from transformers import RobertaForMaskedLM, FlaxRobertaForMaskedLM, AutoTokenizer
+import torch
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+def to_f32(t):
+    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+jax.config.update('jax_platform_name', 'cpu')
+MODEL_PATH = "./"
+model = FlaxRobertaForMaskedLM.from_pretrained(MODEL_PATH)
+model.params = to_f32(model.params)
+model.save_pretrained(MODEL_PATH)
+
+pt_model = RobertaForMaskedLM.from_pretrained(
+    MODEL_PATH, from_flax=True).to('cpu')
+
+input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
+input_ids_pt = torch.tensor(input_ids)
+
+logits_pt = pt_model(input_ids_pt).logits
+print(logits_pt)
+logits_fx = model(input_ids).logits
+print(logits_fx)
+
+pt_model.save_pretrained(MODEL_PATH)

pytorch_model.bin

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

start_train.sh

@@ -5,15 +5,14 @@ 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_deduplicated" \
-    --max_seq_length="128" \
+    --max_seq_length="512" \
     --pad_to_max_length \
     --preprocessing_num_workers="64" \
-    --per_device_train_batch_size="64" \
-    --per_device_eval_batch_size="64" \
+    --per_device_train_batch_size="16" \
+    --per_device_eval_batch_size="16" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
     --adam_epsilon="1e-6" \