Saving weights and logs of step 35000

convert.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+import tempfile
+
+import jax
+from jax import numpy as jnp
+from transformers import AutoTokenizer, FlaxBertForMaskedLM, BertForMaskedLM
+
+
+def to_f32(t):
+    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+
+def main():
+    # Saving extra files from config.json and tokenizer.json files
+    tokenizer = AutoTokenizer.from_pretrained("./")
+    tokenizer.save_pretrained("./")
+
+    # Temporary saving bfloat16 Flax model into float32
+    tmp = tempfile.mkdtemp()
+    flax_model = FlaxBertForMaskedLM.from_pretrained("./")
+    flax_model.params = to_f32(flax_model.params)
+    flax_model.save_pretrained(tmp)
+    # Converting float32 Flax to PyTorch
+    model = BertForMaskedLM.from_pretrained(tmp, from_flax=True)
+    model.save_pretrained("./", save_config=False)
+
+
+if __name__ == "__main__":
+    main()

events.out.tfevents.1649146360.t1v-n-eedfb410-w-0.232238.0.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd9a16b3f835bda324ec54c6d055ddbef0ff69304226141f833adf5e119b9ee8
-size 5061593
+oid sha256:b0b0a97b939b6c937442a77a2798640d608b058f3eca8584ff7cff0ddc53098d
+size 5211435

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9e221febfb233514ab4836cc21f9905fffada2463fa5f5db7562aa39eec7e96a
+oid sha256:ca4841024e3ac944e0ca9eb204cea434ed7ec5ccda7f8b9c8028ad1ba103acd7
 size 711905363