push-to-hub fix

Author: arampacha

run_clm_streaming.sh

@@ -1,6 +1,6 @@
 #! /bin/bash
 ./run_clm_streaming_flax_v2.py \
-    --output_dir $HOME/gpt-neo-125M-code-clippy \
+    --output_dir $HOME/gpt-neo-125M-test \
     --model_name_or_path="EleutherAI/gpt-neo-125M" \
     --dataset_name $HOME/gpt-code-clippy/code_clippy.py \
     --data_dir /home/shared/code-clippy-dataset/merged-data \
@@ -11,27 +11,27 @@
     --per_device_eval_batch_size="16" \
     --preprocessing_num_workers="8" \
     --learning_rate="6e-4" \
-    --adafactor \
-    --max_steps 10000 \
-    --warmup_steps 3000 \
-    --decay_steps 5000 \
+    --max_steps 500 \
+    --warmup_steps 150 \
+    --decay_steps 250 \
     --adam_beta1="0.9" \
     --adam_beta2="0.95" \
     --weight_decay="0.01" \
     --overwrite_output_dir \
-    --logging_steps="100" \
-    --eval_steps="100" \
-    --push_to_hub="False" \
+    --logging_steps="10" \
+    --eval_steps="50" \
+    --push_to_hub="True" \
     --report_to="all" \
     --dtype="bfloat16" \
     --skip_memory_metrics="False" \
-    --save_steps="100" \
+    --save_steps="50" \
     --save_total_limit 2 \
     --gradient_accumulation_steps 8 \
     --report_to="wandb" \
-    --run_name="testing" \
+    --run_name="testing-mini" \
     --max_eval_samples 100 \
     --save_optimizer true \
+    # --adafactor \
     # --resume_from_checkpoint $HOME/gpt-neo-125M-code-clippy/ \
     # --max_train_samples="10000" \
 

run_clm_streaming_flax_v2.py

@@ -807,7 +807,7 @@ def eval_step(params, batch):
                                       push_to_hub=training_args.push_to_hub)
                 if model_args.save_optimizer:
                     # this saves full state including optimizer
-                    save_checkpoint(training_args.output_dir, jax_utils.unreplicate(state), cur_step, keep=training_args.save_total_limit, overwrite=False)
+                    save_checkpoint(training_args.output_dir, jax_utils.unreplicate(state), cur_step, keep=training_args.save_total_limit, overwrite=True)
                 if training_args.save_total_limit is not None:
                     rotate_checkpoints(training_args.output_dir, training_args.save_total_limit)