Merge pull request #3356 from AI-Hypercomputer:fix-pypi-config-resolve

PiperOrigin-RevId: 881115051

Author: Google-ML-Automation

src/maxtext/configs/pyconfig.py

@@ -59,6 +59,12 @@ def resolve_config_path(param: str) -> str:
     lowercase_param = param.replace("MaxText", "maxtext")
     if os.path.isfile(lowercase_param):
       return lowercase_param
+  # For pip-installed packages, strip the src prefix and resolve against
+  # the installed configs directory (MAXTEXT_CONFIGS_DIR).
+  if param.startswith("src/maxtext/configs/"):
+    candidate = os.path.join(MAXTEXT_CONFIGS_DIR, param[len("src/maxtext/configs/"):])
+    if os.path.isfile(candidate):
+      return candidate
   return os.path.join("src", param)
 
 

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -23,7 +23,7 @@
 Usage Examples:
 
 # GRPO on Llama3.1-8B-Instruct
-python3 -m src.maxtext.trainers.post_train.rl.train_rl src/maxtext/configs/post_train/rl.yml \
+python3 -m maxtext.trainers.post_train.rl.train_rl src/maxtext/configs/post_train/rl.yml \
   model_name=llama3.1-8b \
   tokenizer_path=meta-llama/Llama-3.1-8B-Instruct \
   load_parameters_path=gs://path/to/checkpoint/0/items \
@@ -32,7 +32,7 @@
   hf_access_token=${HF_TOKEN?}
 
 # GSPO on Llama3.1-70B-Instruct
-python3 -m src.maxtext.trainers.post_train.rl.train_rl src/maxtext/configs/post_train/rl.yml \
+python3 -m maxtext.trainers.post_train.rl.train_rl src/maxtext/configs/post_train/rl.yml \
   model_name=llama3.1-70b \
   tokenizer_path=meta-llama/Llama-3.1-70B-Instruct \
   load_parameters_path=gs://path/to/checkpoint/0/items \

tests/unit/pyconfig_test.py

@@ -14,12 +14,13 @@
 
 """Tests for pyconfig."""
 
-import unittest
 import os.path
+import tempfile
+import unittest
 
 from maxtext.configs import pyconfig
 from maxtext.configs.pyconfig import resolve_config_path
-from maxtext.utils.globals import MAXTEXT_PKG_DIR
+from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR, MAXTEXT_PKG_DIR
 from tests.utils.test_helpers import get_test_config_path, get_post_train_test_config_path
 
 
@@ -101,6 +102,19 @@ def test_resolve_config_path(self):
     self.assertEqual(resolve_config_path("foo"), os.path.join("src", "foo"))
     self.assertEqual(resolve_config_path(__file__), __file__)
 
+  def test_resolve_config_path_pip_install(self):
+    """Simulates pip-installed env where cwd has no src/ folder."""
+    orig = os.getcwd()
+    with tempfile.TemporaryDirectory() as tmpdir:
+      try:
+        os.chdir(tmpdir)
+        result = resolve_config_path("src/maxtext/configs/base.yml")
+        self.assertEqual(result, os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml"))
+        result = resolve_config_path("src/maxtext/configs/post_train/rl.yml")
+        self.assertEqual(result, os.path.join(MAXTEXT_CONFIGS_DIR, "post_train/rl.yml"))
+      finally:
+        os.chdir(orig)
+
 
 if __name__ == "__main__":
   unittest.main()