keep memory

run.sh

@@ -6,18 +6,17 @@ python3 run_mlm_flax.py \
     --config_name="./configs/base" \
     --tokenizer_name="./" \
     --dataset_name="munggok/KoPI" \
-    --cache_dir="/data/cache" \
     --dataset_config_name="full" \
     --max_seq_length="512" \
     --pad_to_max_length  \
     --per_device_train_batch_size="64" \
     --per_device_eval_batch_size="64" \
-    --preprocessing_num_workers="96" \
+    --preprocessing_num_workers="64" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
     --adam_epsilon="1e-6" \
     --learning_rate="8e-5" \
-    --num_train_epochs="15" \
+    --num_train_epochs="20" \
     --weight_decay="0.01" \
     --save_strategy="steps" \
     --save_steps="10000" \

run_mlm_flax.py

@@ -648,6 +648,7 @@ def main():
             tokenize_function,
             input_columns=[text_column_name],
             batched=True,
+            keep_in_memory=True,
             num_proc=data_args.preprocessing_num_workers,
             remove_columns=column_names,
             load_from_cache_file=not data_args.overwrite_cache,
@@ -663,6 +664,7 @@ def main():
         tokenized_datasets = datasets.map(
             tokenize_function,
             batched=True,
+            keep_in_memory=True,
             num_proc=data_args.preprocessing_num_workers,
             remove_columns=column_names,
             load_from_cache_file=not data_args.overwrite_cache,
@@ -694,6 +696,7 @@ def main():
         tokenized_datasets = tokenized_datasets.map(
             group_texts,
             batched=True,
+            keep_in_memory=True,
             num_proc=data_args.preprocessing_num_workers,
             load_from_cache_file=not data_args.overwrite_cache,
         )