Updated code to have different seed and reduced lr

run.sh

@@ -11,7 +11,7 @@
     --preprocessing_num_workers="64" \
     --per_device_train_batch_size="64" \
     --per_device_eval_batch_size="64" \
-    --learning_rate="3e-4" \
+    --learning_rate="2e-4" \
     --warmup_steps="1000" \
     --overwrite_output_dir \
     --num_train_epochs="8" \

run_mlm_flax.py

@@ -324,6 +324,7 @@ if __name__ == "__main__":
     logger.info(f"Training/evaluation parameters {training_args}")
 
     # Set seed before initializing model.
+    training_args.seed = 42
     set_seed(training_args.seed)
 
     # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
@@ -587,6 +588,7 @@ if __name__ == "__main__":
 
     train_time = 0
     epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
+    save_checkpoint=True
     for epoch in epochs:
         # ======================== Training ================================
         train_start = time.time()
@@ -609,6 +611,14 @@ if __name__ == "__main__":
             model_inputs = shard(model_inputs.data)
             state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
             train_metrics.append(train_metric)
+            if save_checkpoint and (train_metric['loss'] < 5.).all():
+                params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
+                model.save_pretrained(
+                '/home/khandelia1000/checkpoints/',
+                params=params,
+                push_to_hub=False
+                )
+                save_checkpoint = False
 
         train_time += time.time() - train_start