Saving weights and logs of step 10000

events.out.tfevents.1637523639.t1v-n-8eba1090-w-0.295438.0.v2 → events.out.tfevents.1637622530.t1v-n-8eba1090-w-0.18483.0.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ecd923c8c7884cefb47b293d21c23b3dae875ca57d3697a3891225499c03af29
-size 17857247
+oid sha256:46b1ed32532cb8c66a31d50c192133429b153d1fb9e6ce38e84ddc41202ef2d1
+size 1470757

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cfe37e1b1a4fd67ba61d17ae18ea79fee0b60886def1a5fd73f0ef9d720f260e
+oid sha256:416eebb5d6b9be6fb90af53e09925009426f2e9a53571890b22b61245d85d1a4
 size 1421662309

run_mlm_flax.py

@@ -508,14 +508,6 @@ if __name__ == "__main__":
 
     # save the tokenized dataset for future runs
     if data_args.save_tokenized_dataset_filepath is not None:
-        if data_args.dataset_filepath is not None:
-            try:
-                os.system(f"sudo rm {data_args.dataset_filepath}/train/cache*")
-                os.system(f"sudo rm {data_args.dataset_filepath}/validation/cache*")
-                os.system(f"sudo rm {data_args.dataset_filepath}/train/tmp*")
-                os.system(f"sudo rm {data_args.dataset_filepath}/validation/tmp*")
-            except:
-                pass
         tokenized_datasets.save_to_disk(data_args.save_tokenized_dataset_filepath)
 
 
@@ -592,6 +584,7 @@ if __name__ == "__main__":
         # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
         optimizer = optax.adafactor(
             learning_rate=linear_decay_lr_schedule_fn,
+            weight_decay_rate=training_args.weight_decay,
         )
     else:
         optimizer = optax.adamw(

start_train.sh

@@ -1,23 +1,24 @@
 # set train hyperparams
 unset LD_PRELOAD
 export HF_DATASETS_CACHE="/researchdisk/datasets_cache"
+export USE_TORCH=False
 python3 run_mlm_flax.py \
     --output_dir="./" \
     --model_type="roberta" \
     --config_name="./" \
     --tokenizer_name="./" \
     --dataset_filepath="/researchdisk/training_dataset_full" \
-    --save_tokenized_dataset_filepath="/researchdisk/training_dataset_full_tokenized_128" \
     --max_seq_length="128" \
     --pad_to_max_length \
     --preprocessing_num_workers="96" \
     --per_device_train_batch_size="64" \
     --per_device_eval_batch_size="64" \
+    --weight_decay="0.01" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
     --adam_epsilon="1e-6" \
     --learning_rate="2e-4" \
-    --warmup_steps="1500" \
+    --warmup_steps="25000" \
     --overwrite_output_dir \
     --num_train_epochs="2" \
     --save_strategy="steps" \