Merge pull request #2501 from AI-Hypercomputer:anisha-clean-path

PiperOrigin-RevId: 820032273

Author: Google-ML-Automation

src/MaxText/examples/grpo_llama3_1_70b_demo_pw.py

@@ -141,18 +141,21 @@
 TRAIN_FRACTION = 1.0
 
 
+# ====== Input Checkpoint directory =====
+MODEL_CHECKPOINT_PATH = "/path/to/scanned/model/ckpt_load_dir/"
+
 # ====== Checkpoint directory =====
 LOG_DIR = f"{HOME}/content/tensorboard/grpo/logs_llama3/"
 if not os.path.exists(LOG_DIR):
   os.makedirs(LOG_DIR)
 
 # ===== Profiling =====
-PROFILE_DIR = f"gs://mazumdera-test-bucket-europe-west4/rl-tuning/grpo/anisha-{run_id}/profiles_llama3/"
+PROFILE_DIR = f"/path/to/profile_dir/{run_id}/profiles_llama3/"
 if not epath.Path(PROFILE_DIR).exists():
   epath.Path(PROFILE_DIR).mkdir(parents=True)
 
 # ====== Checkpoint saving ======
-CKPT_DIR = f"gs://mazumdera-test-bucket-europe-west4/rl-tuning/grpo/anisha-{run_id}/ckpts_llama3/"
+CKPT_DIR = f"/path/to/ckpt_save_dir/{run_id}/ckpts_llama3/"
 
 if not epath.Path(CKPT_DIR).exists():
   epath.Path(CKPT_DIR).mkdir(parents=True)
@@ -396,8 +399,7 @@ def get_ref_maxtext_model(config, devices=None):
     run_name="test-tunix-maxtext-llama3.1-70b",
     tokenizer_type="tiktoken",
     tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer_llama3.tiktoken"),
-    load_parameters_path="gs://mazumdera-test-bucket-europe-west4/llama3.1-70b-Instruct/scanned-pathways/0/items",
-    # load_parameters_path="path/to/scanned/checkpoint",
+    load_parameters_path=MODEL_CHECKPOINT_PATH,
     per_device_batch_size=1,
     max_prefill_predict_length=4,
     max_target_length=1024,
@@ -470,8 +472,7 @@ def get_ref_maxtext_model(config, devices=None):
     run_name="test-tunix-maxtext-llama3.1-70b",  # This is not used in Tunix.
     tokenizer_type="tiktoken",
     tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer_llama3.tiktoken"),
-    load_parameters_path="gs://mazumdera-test-bucket-europe-west4/llama3.1-70b-Instruct/scanned-pathways/0/items",
-    # load_parameters_path="path/to/scanned/checkpoint",
+    load_parameters_path=MODEL_CHECKPOINT_PATH,
     per_device_batch_size=1,
     max_prefill_predict_length=4,
     max_target_length=1024,

src/MaxText/examples/grpo_llama3_1_8b_demo.ipynb

@@ -104,7 +104,7 @@
     "\n",
     "DEBUG = False  # set to True to run in debug mode, for more print statements\n",
     "# set this to the path of the checkpoint you want to load, gs://<bucket> supported\n",
-    "MODEL_CHECKPOINT_PATH = \"path/to/scanned/checkpoint\""
+    "MODEL_CHECKPOINT_PATH = \"/path/to/scanned/model/ckpt_load_dir/\""
    ]
   },
   {

src/MaxText/examples/grpo_llama3_1_8b_demo.py

@@ -109,6 +109,9 @@
 TRAIN_FRACTION = 1.0
 
 
+# ====== Input Checkpoint directory =====
+MODEL_CHECKPOINT_PATH = "/path/to/scanned/model/ckpt_load_dir/"
+
 # ====== Checkpoint directory =====
 LOG_DIR = os.path.join(HOME, "content", "tensorboard", "grpo", "logs_llama3", "")
 if not os.path.exists(LOG_DIR):
@@ -404,8 +407,7 @@ def get_ref_maxtext_model(config):
     run_name="test-tunix-maxtext-llama3.1-8b",
     tokenizer_type="tiktoken",
     tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer_llama3.tiktoken"),
-    load_parameters_path="gs://yixuannwang-maxtext-logs/llama3.1-8b-Instruct/scanned/0/items",
-    # load_parameters_path="path/to/scanned/checkpoint",
+    load_parameters_path=MODEL_CHECKPOINT_PATH,
     per_device_batch_size=1,
     max_prefill_predict_length=4,
     max_target_length=1024,
@@ -464,8 +466,7 @@ def get_ref_maxtext_model(config):
     run_name="test-tunix-maxtext-llama3.1-8b",  # This is not used in Tunix.
     tokenizer_type="tiktoken",
     tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer_llama3.tiktoken"),
-    load_parameters_path="gs://yixuannwang-maxtext-logs/llama3.1-8b-Instruct/scanned/0/items",
-    # load_parameters_path="path/to/scanned/checkpoint",
+    load_parameters_path=MODEL_CHECKPOINT_PATH,
     per_device_batch_size=1,
     max_prefill_predict_length=4,
     max_target_length=1024,

src/MaxText/examples/grpo_llama3_1_8b_demo_pw.py

@@ -141,18 +141,21 @@
 TRAIN_FRACTION = 1.0
 
 
+# ====== Input Checkpoint directory =====
+MODEL_CHECKPOINT_PATH = "/path/to/scanned/model/ckpt_load_dir/"
+
 # ====== Checkpoint directory =====
 LOG_DIR = f"{HOME}/content/tensorboard/grpo/logs_llama3/"
 if not os.path.exists(LOG_DIR):
   os.makedirs(LOG_DIR)
 
 # ===== Profiling =====
-PROFILE_DIR = f"gs://mazumdera-test-bucket-europe-west4/rl-tuning/grpo/{run_id}/profiles_llama3/"
+PROFILE_DIR = f"/path/to/profile_dir/{run_id}/profiles_llama3/"
 if not epath.Path(PROFILE_DIR).exists():
   epath.Path(PROFILE_DIR).mkdir(parents=True)
 
 # ====== Checkpoint saving ======
-CKPT_DIR = f"gs://mazumdera-test-bucket-europe-west4/rl-tuning/grpo/{run_id}/ckpts_llama3/"
+CKPT_DIR = f"/path/to/ckpt_save_dir/{run_id}/ckpts_llama3/"
 
 if not epath.Path(CKPT_DIR).exists():
   epath.Path(CKPT_DIR).mkdir(parents=True)
@@ -394,8 +397,7 @@ def get_ref_maxtext_model(config, devices=None):
     run_name="test-tunix-maxtext-llama3.1-8b",
     tokenizer_type="tiktoken",
     tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer_llama3.tiktoken"),
-    load_parameters_path="gs://mazumdera-test-bucket-europe-west4/llama3.1-8b-Instruct/scanned-pathways/0/items",
-    # load_parameters_path="path/to/scanned/checkpoint",
+    load_parameters_path=MODEL_CHECKPOINT_PATH,
     per_device_batch_size=1,
     max_prefill_predict_length=4,
     max_target_length=1024,
@@ -469,8 +471,7 @@ def get_ref_maxtext_model(config, devices=None):
     run_name="test-tunix-maxtext-llama3.1-8b",  # This is not used in Tunix.
     tokenizer_type="tiktoken",
     tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizer_llama3.tiktoken"),
-    load_parameters_path="gs://mazumdera-test-bucket-europe-west4/llama3.1-8b-Instruct/scanned-pathways/0/items",
-    # load_parameters_path="path/to/scanned/checkpoint",
+    load_parameters_path=MODEL_CHECKPOINT_PATH,
     per_device_batch_size=1,
     max_prefill_predict_length=4,
     max_target_length=1024,