now it is working!

run.sh

@@ -1,7 +1,8 @@
 python run_clm_mp.py \
-    	--model_name_or_path "." \
-  	--tokenizer_name gpt2 \
-   	--dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
+    	--model_name_or_path /mnt/disks/flaxdisk/norwegian-gptneo-red/ \
+  	--tokenizer_name /mnt/disks/flaxdisk/norwegian-gptneo-red/ \
+   	--train_file /mnt/disks/flaxdisk/corpus/social_train.json \
+       	--validation_file /mnt/disks/flaxdisk/corpus/social_validation.json \
      	--do_train \
        	--do_eval  \
      	--block_size 1024 \
@@ -12,6 +13,8 @@ python run_clm_mp.py \
     	--overwrite_output_dir \
 	--output_dir /mnt/disks/flaxdisk/output/ \
 	--cache_dir /mnt/disks/flaxdisk/cache/ \
-       	--dtype bfloat16 \
-   	--logging_steps 97 \
-       	--eval_steps 96
+	--preprocessing_num_workers 96 \
+	--dtype bfloat16 \
+   	--logging_steps 5000 \
+       	--eval_steps 5000 \
+	--push_to_hub

run_bak.sh → run_wiki.sh

@@ -1,19 +1,17 @@
 python run_clm_mp.py \
-    	--model_name_or_path /mnt/disks/flaxdisk/norwegian-gptneo-red/ \
-  	--tokenizer_name /mnt/disks/flaxdisk/norwegian-gptneo-red/ \
-   	--train_file /mnt/disks/flaxdisk/corpus/social_train.json \
-       	--validation_file /mnt/disks/flaxdisk/corpus/social_validation.json \
+    	--model_name_or_path "." \
+  	--tokenizer_name gpt2 \
+   	--dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
      	--do_train \
        	--do_eval  \
      	--block_size 1024 \
    	--num_train_epochs 5 \
    	--learning_rate 4e-6 \
-	--per_device_train_batch_size 3 \
-       	--per_device_eval_batch_size 3 \
+	--per_device_train_batch_size 2 \
+       	--per_device_eval_batch_size 2 \
     	--overwrite_output_dir \
 	--output_dir /mnt/disks/flaxdisk/output/ \
 	--cache_dir /mnt/disks/flaxdisk/cache/ \
        	--dtype bfloat16 \
    	--logging_steps 97 \
-       	--eval_steps 96\
-	--push_to_hub
+       	--eval_steps 96