Try fix git

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "gpt2-large",
+  "_name_or_path": "./",
   "activation_function": "gelu_new",
   "architectures": [
     "GPT2LMHeadModel"

run_clm_flax.py

@@ -61,7 +61,7 @@ from transformers import (
 from transformers.file_utils import get_full_repo_name
 from transformers.testing_utils import CaptureLogger
 
-from distributed_shampoo import distributed_shampoo, GraftingType
+# from distributed_shampoo import distributed_shampoo, GraftingType
 
 
 logger = logging.getLogger(__name__)

start_train.sh

@@ -10,19 +10,19 @@ python3 run_clm_flax.py \
     --dataset_filepath="/researchdisk/training_dataset_full_deduplicated" \
     --do_train --do_eval \
     --block_size="512" \
-    --per_device_train_batch_size="16" \
-    --per_device_eval_batch_size="16" \
-    --preprocessing_num_workers="1" \
+    --per_device_train_batch_size="8" \
+    --per_device_eval_batch_size="8" \
+    --preprocessing_num_workers="96" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
-    --learning_rate="1e-5" \
+    --learning_rate="4e-5" \
     --weight_decay="0.01" \
     --warmup_steps="4000" \
     --cosine_decay \
     --overwrite_output_dir \
     --logging_steps="500" \
-    --eval_steps="10000" \
-    --save_steps="50000" \
+    --eval_steps="100" \
+    --save_steps="100" \
     --num_train_epochs="5" \
     --dtype="bfloat16" \
     --push_to_hub \