Add 10k train step pytorch model

config.json

@@ -32,6 +32,7 @@
       "max_length": 50
     }
   },
+  "torch_dtype": "float32",
   "transformers_version": "4.16.0.dev0",
   "use_cache": true,
   "vocab_size": 50257

flax_model_to_pytorch.py

@@ -0,0 +1,27 @@
+from transformers import AutoModelForCausalLM, FlaxAutoModelForCausalLM, 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 = FlaxAutoModelForCausalLM.from_pretrained(MODEL_PATH)
+model.params = to_f32(model.params)
+model.save_pretrained(MODEL_PATH)
+
+pt_model = AutoModelForCausalLM.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

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

runs/events.out.tfevents.1642099734.t1v-n-42145f73-w-0.2317757.0.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3310027638bb05951334514395d3e81a9d7eb8ce109681141532c4e59751cb88
-size 1471449
+oid sha256:78cfa1f897e391e235a903d1ff19d56c36817328c8b3c8b76f575958a16fdf68
+size 1912863