not save so often

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4fe397b86f6a483191db26e2778ff9e645c14107b76758b09db78a9338e7a8a6
+oid sha256:b4af9ac77cf162bc4703db2637e08a4dc5861b3e5fcaefe2ef38d8edfda1898e
 size 5262371934

run_clm_mp.py

@@ -632,7 +632,7 @@ def main():
                         f"Step... ({cur_step} | Eval loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']}"
                     )
 
-                if cur_step % training_args.save_steps == 0 and cur_step > 0:
+                if cur_step % steps_per_epoch == 0 and cur_step > 0:
                     # save checkpoint after each epoch and push checkpoint to the hub
                     if jax.process_index() == 0:
                         params = jax.device_get(params)