Merge pull request #3044 from AI-Hypercomputer:mohit/move_configs

PiperOrigin-RevId: 867740632

Author: Google-ML-Automation

.github/workflows/run_jupyter_notebooks.yml

@@ -90,8 +90,11 @@ jobs:
           PYTHONPATH: "${{ github.workspace }}/src"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          MAXTEXT_REPO_ROOT=$(pwd)
-          MAXTEXT_NOTEBOOKS_ROOT="$MAXTEXT_REPO_ROOT/src/maxtext/examples"
+          source .venv/bin/activate
+
+          export MAXTEXT_REPO_ROOT=$(pwd)
+          export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
+          export MAXTEXT_NOTEBOOKS_ROOT="$MAXTEXT_REPO_ROOT/src/maxtext/examples"
 
           for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl}*.ipynb; do
             filename=$(basename "$notebook")
@@ -101,7 +104,7 @@ jobs:
             echo "Running $filename ..."
             echo "------------------------------------------------------"
 
-            .venv/bin/papermill "$notebook" "$output_name" -k maxtext_venv
+            papermill "$notebook" "$output_name" -k maxtext_venv
           done
       - name: Record Commit IDs
         shell: bash

.github/workflows/run_pathways_tests.yml

@@ -100,7 +100,7 @@ jobs:
           export MAXTEXT_REPO_ROOT=$(pwd)
           export MAXTEXT_ASSETS_ROOT=$(pwd)/src/maxtext/assets
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
-          export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
+          export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
           .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0
         env:

.github/workflows/run_tests_against_package.yml

@@ -110,7 +110,7 @@ jobs:
           export MAXTEXT_REPO_ROOT=$(pwd)
           export MAXTEXT_ASSETS_ROOT=$(pwd)/src/maxtext/assets
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
-          export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
+          export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           # omit this libtpu init args for gpu tests
           if [ "${{ inputs.device_type }}" != "cuda12" ]; then
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'

.vscode/launch.json

@@ -9,7 +9,7 @@
       "justMyCode": false,
       "python": "python3",
       "module": "maxtext.decode",
-      "args": ["src/MaxText/configs/base.yml",
+      "args": ["src/maxtext/configs/base.yml",
                "run_name=runner_$(date +%Y-%m-%d-%H-%M)",
                "base_output_directory=gs://test-maxtext-output",
                "dataset_path=gs://test-maxtext-dataset",
@@ -36,7 +36,7 @@
       "justMyCode": false,
       "python": "python3",
       "module": "maxtext.decode",
-      "args": ["src/MaxText/configs/base.yml",
+      "args": ["src/maxtext/configs/base.yml",
                "run_name=runner_$(date +%Y-%m-%d-%H-%M)",
                "base_output_directory=gs://test-maxtext-output",
                "dataset_path=gs://test-maxtext-dataset",
@@ -52,7 +52,7 @@
       "justMyCode": false,
       "python": "python3",
       "module": "MaxText.train",
-      "args": ["src/MaxText/configs/base.yml",
+      "args": ["src/maxtext/configs/base.yml",
                "run_name=runner_$(date +%Y-%m-%d-%H-%M)",
                "base_output_directory=gs://test-maxtext-output",
                "dataset_path=gs://test-maxtext-dataset",
@@ -68,7 +68,7 @@
       "python": "python3",
       "module": "maxtext.inference.inference_microbenchmark",
       "args": [
-        "src/MaxText/configs/base.yml",
+        "src/maxtext/configs/base.yml",
         "model_name=llama2-7b",
         "tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.llama2",
         "weight_dtype=bfloat16",

PREFLIGHT.md

@@ -7,12 +7,12 @@ Before you run ML workload on Multihost with GCE or GKE, simply apply `bash pref
 
 Here is an example for GCE:
 ```
-bash preflight.sh PLATFORM=GCE && python3 -m MaxText.train src/MaxText/configs/base.yml run_name=$YOUR_JOB_NAME
+bash preflight.sh PLATFORM=GCE && python3 -m MaxText.train src/maxtext/configs/base.yml run_name=$YOUR_JOB_NAME
 ```
 
 Here is an example for GKE:
 ```
-bash preflight.sh PLATFORM=GKE && python3 -m MaxText.train src/MaxText/configs/base.yml run_name=$YOUR_JOB_NAME
+bash preflight.sh PLATFORM=GKE && python3 -m MaxText.train src/maxtext/configs/base.yml run_name=$YOUR_JOB_NAME
 ```
 
 # Optimization 2: Numa binding (You can only apply this to v4 and v5p)
@@ -22,14 +22,14 @@ For GCE,
 [preflight.sh](https://github.com/google/maxtext/blob/main/preflight.sh) will help you install `numactl` dependency, so you can use it directly, here is an example:
 
 ```
-bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m MaxText.train src/MaxText/configs/base.yml run_name=$YOUR_JOB_NAME
+bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m MaxText.train src/maxtext/configs/base.yml run_name=$YOUR_JOB_NAME
 ```
 
 For GKE,
 `numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example
 
 ```
-bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m MaxText.train src/MaxText/configs/base.yml run_name=$YOUR_JOB_NAME
+bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m MaxText.train src/maxtext/configs/base.yml run_name=$YOUR_JOB_NAME
 ```
 
 1. `numactl`: This is the command-line tool used for controlling NUMA policy for processes or shared memory. It's particularly useful on multi-socket systems where memory locality can impact performance.

README.md

@@ -42,7 +42,7 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies
 ## 🔥 Latest news 🔥
 
 * \[December 22, 2025\] [Muon optimizer](https://kellerjordan.github.io/posts/muon) is now supported.
-* \[December 10, 2025\] DeepSeek V3.1 is now supported. Use existing configs for [DeepSeek V3 671B](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/configs/models/deepseek3-671b.yml) and load in V3.1 checkpoint to use model.
+* \[December 10, 2025\] DeepSeek V3.1 is now supported. Use existing configs for [DeepSeek V3 671B](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/models/deepseek3-671b.yml) and load in V3.1 checkpoint to use model.
 * \[December 9, 2025\] [New RL and SFT Notebook tutorials](https://github.com/AI-Hypercomputer/maxtext/tree/main/src/maxtext/examples) are available.
 * \[December 4, 2025\] The [ReadTheDocs documentation site](https://maxtext.readthedocs.io/en/latest/index.html) has been reorganized.
 * \[December 3, 2025\] Multi-host support for GSPO and GRPO is now available via [new RL tutorials](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_on_multi_host.html).

benchmarks/api_server/README.md

@@ -33,7 +33,7 @@ export HF_TOKEN=<your_hugging_face_token>
 
 The primary way to launch the API server is by using the `start_server.sh` script. This script ensures that the server is run from the project's root directory, which is necessary for the Python interpreter to find all the required modules.
 
-The script takes the path to a base configuration file (e.g., `MaxText/configs/base.yml`) followed by any number of model-specific configuration overrides.
+The script takes the path to a base configuration file (e.g., `maxtext/configs/base.yml`) followed by any number of model-specific configuration overrides.
 
 ### Benchmarking Configuration
 
@@ -56,7 +56,7 @@ Here is an example of how to launch the server with a `qwen3-30b-a3b` model, con
 # Make sure you are in the root directory of the maxtext project.
 
 bash benchmarks/api_server/start_server.sh \
-    MaxText/configs/base.yml \
+    maxtext/configs/base.yml \
     model_name="qwen3-30b-a3b" \
     tokenizer_path="Qwen/Qwen3-30B-A3B-Thinking-2507" \
     load_parameters_path="<path_to_your_checkpoint>" \
@@ -135,7 +135,7 @@ CMD="export HF_TOKEN=${HF_TOKEN} && \
      pip install --upgrade pip && \
      pip install -r benchmarks/api_server/requirements.txt && \
      bash benchmarks/api_server/start_server.sh \
-        MaxText/configs/base.yml \
+        maxtext/configs/base.yml \
         model_name="${MODEL_NAME}" \
         tokenizer_path="${TOKENIZER_PATH}" \
         load_parameters_path="${LOAD_PARAMETERS_PATH}" \

benchmarks/api_server/launch_gke_server.sh.template

@@ -53,7 +53,7 @@ CMD="export HF_TOKEN=${HF_TOKEN} && \
      pip install --upgrade pip && \
      pip install -r benchmarks/api_server/requirements.txt && \
      bash benchmarks/api_server/start_server.sh \
-        MaxText/configs/base.yml \
+        maxtext/configs/base.yml \
         model_name=\"${MODEL_NAME}\" \
         tokenizer_path=\"${TOKENIZER_PATH}\" \
         load_parameters_path=\"${LOAD_PARAMETERS_PATH}\" \

benchmarks/api_server/start_server.sh

@@ -20,7 +20,7 @@
 #
 # Example:
 # bash benchmarks/api_server/start_server.sh \
-#     MaxText/configs/base.yml \
+#     maxtext/configs/base.yml \
 #     model_name="qwen3-30b-a3b" \
 #     tokenizer_path="Qwen/Qwen3-30B-A3B-Thinking-2507" \
 #     load_parameters_path="<path_to_your_checkpoint>" \

benchmarks/globals.py

@@ -25,7 +25,10 @@
     r if os.path.isdir(os.path.join(r := os.path.dirname(os.path.dirname(__file__)), ".git")) else MAXTEXT_PKG_DIR,
 )
 
+# This is the configs root: with "base.yml"; "models/"; &etc.
+MAXTEXT_CONFIGS_DIR = os.environ.get("MAXTEXT_CONFIGS_DIR", os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs"))
+
 # This is the assets root: with "tokenizers/"; &etc.
 MAXTEXT_ASSETS_ROOT = os.environ.get("MAXTEXT_ASSETS_ROOT", os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "assets"))
 
-__all__ = ["MAXTEXT_ASSETS_ROOT", "MAXTEXT_PKG_DIR", "MAXTEXT_REPO_ROOT"]
+__all__ = ["MAXTEXT_ASSETS_ROOT", "MAXTEXT_CONFIGS_DIR", "MAXTEXT_PKG_DIR", "MAXTEXT_REPO_ROOT"]