init checkpoint_state_dir as output_dir

Author: MohdElgaar

open_instruct/grpo_fast.py

@@ -2154,6 +2154,8 @@ def setup_runtime_variables(args: Args) -> Args:
     """Set up runtime variables for the experiment."""
     args.run_name = f"{args.exp_name}__{args.seed}__{int(time.time())}"
     args.output_dir = os.path.join(args.output_dir, args.run_name)
+    if args.checkpoint_state_dir is None:
+        args.checkpoint_state_dir = args.output_dir
     args.dataset_local_cache_dir = os.path.abspath(args.dataset_local_cache_dir)
     if is_beaker_job():
         args.dataset_local_cache_dir = "/weka/oe-adapt-default/allennlp/deletable_open_instruct_dataset_cache"