Saving weights and logs of step 2500

events.out.tfevents.1733914848.t1v-n-53cd541d-w-35.928604.0.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7c188c2540a746555d61e89dc5c37f08447e2ef1ff92679381718754ca42f56e
-size 15496782
+oid sha256:2db3f0432de2d811e3e7bae521e572f24fd21674f4cad549ed9e3a35e0260974
+size 15571352

events.out.tfevents.1734608156.t1v-n-53cd541d-w-35.1559154.0.v2

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

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88ad4d4c86901db3d51130b4a8b0c5e72b2c4ff0eff5c681352672231aaaaf4c
+oid sha256:cfc030a0b79e3275c47d1d1f0f5d207f753f01ebe904a7fcf12243ee4459dc47
 size 1421662309

run_mlm_flax.py

@@ -830,7 +830,7 @@ def main():
             state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
             train_metrics.append(train_metric)
 
-            cur_step = epoch * (num_train_samples // train_batch_size) + step
+            cur_step = epoch * (num_train_samples // (train_batch_size * jax.process_count())) + step
 
             if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics