Merge pull request #1562 from AI-Hypercomputer:bvandermoon-xpk-path

PiperOrigin-RevId: 746215624

Author: maxtext authors

benchmarks/globals.py

@@ -21,8 +21,6 @@
 from tempfile import gettempdir
 from benchmarks import xla_flags_library
 
-from benchmarks.globals import PKG_DIR
-
 # TODO(vbarr@) Abstract software features like checkpointing,
 # real data / synthetic data out of this config
 # TODO(vbarr@) Make slice dependent configurations to allow for a model's tuning
@@ -629,7 +627,7 @@ def _add_to_model_dictionary(
             "profiler": "xplane",
             "dataset_path": "gs://max-datasets-rogue",
             "dataset_type": "tfds",
-            "tokenizer_path": os.path.join(os.path.dirname(PKG_DIR), "assets", "tokenizer.llama2"),
+            "tokenizer_path": os.path.join("assets", "tokenizer.llama2"),
             "sa_block_q": 1024,
             "sa_block_q_dkv": 2048,
             "sa_block_q_dq": 2048,
@@ -1605,7 +1603,7 @@ def _add_to_model_dictionary(
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
-            "tokenizer_path": os.path.join(os.path.dirname(PKG_DIR), "assets", "tokenizer.llama2"),
+            "tokenizer_path": os.path.join("assets", "tokenizer.llama2"),
             "sa_block_q": 2048,
             "sa_block_q_dkv": 2048,
             "sa_block_q_dq": 2048,
@@ -1638,7 +1636,7 @@ def _add_to_model_dictionary(
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
-            "tokenizer_path": os.path.join(os.path.dirname(PKG_DIR), "assets", "tokenizer.llama2"),
+            "tokenizer_path": os.path.join("assets", "tokenizer.llama2"),
             "sa_block_q": 2048,
             "sa_block_q_dkv": 2048,
             "sa_block_q_dq": 2048,

benchmarks/maxtext_trillium_model_configs.py

@@ -16,7 +16,6 @@
 import os.path
 from benchmarks import xla_flags_library
 from benchmarks.maxtext_trillium_model_configs import MaxTextModel, _add_to_model_dictionary
-from benchmarks.globals import PKG_DIR
 
 
 v5e_model_dict = {}
@@ -162,7 +161,7 @@
         "remat_policy": "save_qkv_proj",
         "max_target_length": 2048,
         "use_iota_embed": True,
-        "tokenizer_path": os.path.join(os.path.dirname(PKG_DIR), "assets", "tokenizer.llama2"),
+        "tokenizer_path": os.path.join("assets", "tokenizer.llama2"),
         "dataset_path": "gs://max-datasets-rogue",
         "dataset_type": "synthetic",
         "reuse_example_batch": 1,
@@ -187,7 +186,7 @@
         "remat_policy": "qkv_proj_offloaded",
         "max_target_length": 2048,
         "use_iota_embed": True,
-        "tokenizer_path": os.path.join(os.path.dirname(PKG_DIR), "assets", "tokenizer.llama2"),
+        "tokenizer_path": os.path.join("assets", "tokenizer.llama2"),
         "dataset_path": "gs://max-datasets-rogue",
         "dataset_type": "synthetic",
         "reuse_example_batch": 1,

benchmarks/maxtext_v5e_model_configs.py

@@ -105,7 +105,7 @@ class WorkloadConfig:
   generate_metrics_and_upload_to_big_query: bool = True
   hardware_id: str = 'v6e'
   metrics_gcs_file: str = ''
-  base_config: str = os.path.join(PKG_DIR, "configs", "base.yml")
+  base_config: str = os.path.join("MaxText", "configs", "base.yml")
   topology: str = dataclasses.field(init=False)
   num_devices_per_slice: int = dataclasses.field(init=False)
   db_project: str = ""
@@ -349,7 +349,7 @@ def _build_args_from_config(wl_config: WorkloadConfig) -> dict:
   args["xla_flags"] = f"'{xla_flags_str}'"
   args["dataset"] = dataset
   args["run_type"] = "maxtext-xpk"
-  args["config_file"] = os.path.join(PKG_DIR, "configs", "base.yml")
+  args["config_file"] = os.path.join("MaxText", "configs", "base.yml")
   args["topology"] = wl_config.topology
   args["tuning_params"] = f"'{tuning_params_str}'"
   args["db_project"] = wl_config.db_project
@@ -413,7 +413,7 @@ def build_user_command(
       'export ENABLE_PATHWAYS_PERSISTENCE=1 &&',
       f'export JAX_PLATFORMS={jax_platforms} &&',
       'export ENABLE_PJRT_COMPATIBILITY=true &&',
-      f'python3 -m MaxText.train {os.path.join(PKG_DIR, "configs", "base.yml")}',
+      f'python3 -m MaxText.train {os.path.join("MaxText", "configs", "base.yml")}',
       f'{config_tuning_params}',
       f'steps={wl_config.num_steps}',
       f'model_name={wl_config.model.model_type}',