Merge pull request #3615 from AI-Hypercomputer:bvandermoon-uxr

PiperOrigin-RevId: 900246240

Author: Google-ML-Automation

docs/run_maxtext/decoupled_mode.md

@@ -51,10 +51,10 @@ Optional environment variables:
 
 ## Centralized Decoupling API (`gcloud_stub.py`)
 
-MaxText exposes a single module `MaxText.gcloud_stub` to avoid scattering environment checks:
+MaxText exposes a single module `maxtext.common.gcloud_stub` to avoid scattering environment checks:
 
 ```python
-from MaxText.gcloud_stub import is_decoupled, cloud_diagnostics, jetstream
+from maxtext.common.gcloud_stub import is_decoupled, cloud_diagnostics, jetstream
 
 if is_decoupled():
     # Skip optional integrations or use local fallbacks

src/maxtext/checkpoint_conversion/standalone_scripts/llama4_ckpt_unscanned.py

@@ -21,8 +21,10 @@
 
 Example cmd:
 To save a ckpt
-JAX_PLATFORMS=CPU python -m  MaxText.utils.ckpt_scripts.llama4_ckpt_unscanned --base-model-path [CHKPT_DIR] \
- --maxtext-model-path [OUTPUT_CHKPT_DIR] --model-size llama4-17b-16e
+JAX_PLATFORMS=CPU python -m  maxtext.checkpoint_conversion.standalone_scripts.llama4_ckpt_unscanned \
+    --base-model-path [CHKPT_DIR] \
+    --maxtext-model-path [OUTPUT_CHKPT_DIR] \
+    --model-size llama4-17b-16e
 
 If using a PT checkpoint, the base model checkpoints should be in the format `{name}.{chkpt_idx}.pth`
 For example: `llama4-17b-16e.00.pth`

src/maxtext/examples/multimodal_gemma3_demo.ipynb

@@ -63,10 +63,10 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "import MaxText\n",
+    "import maxtext\n",
     "\n",
     "# Get the root directory of the MaxText\n",
-    "MAXTEXT_PKG_DIR = os.path.dirname(MaxText.__file__)\n",
+    "MAXTEXT_PKG_DIR = os.path.dirname(maxtext.__file__)\n",
     "MAXTEXT_REPO_ROOT = os.path.dirname(os.path.dirname(MAXTEXT_PKG_DIR))\n",
     "MAXTEXT_ASSETS_ROOT = os.path.join(MAXTEXT_REPO_ROOT, \"src\", \"maxtext\", \"assets\")\n",
     "\n",

src/maxtext/examples/sft_llama3_demo_gpu.ipynb

@@ -185,12 +185,13 @@
         "\n",
         "import jax\n",
         "import jax.numpy as jnp\n",
-        "import MaxText\n",
+        "import maxtext\n",
         "from maxtext.configs import pyconfig\n",
         "from maxtext.trainers.post_train.sft import train_sft\n",
         "\n",
-        "MAXTEXT_REPO_ROOT = os.path.dirname(MaxText.__file__)\n",
-        "print(f\"MaxText installation path: {MAXTEXT_REPO_ROOT}\")\n",
+        "MAXTEXT_REPO_ROOT = os.path.dirname(os.path.dirname(maxtext.__file__))\n",
+        "print(f\"MaxText installation path: {MAXTEXT_REPO_ROOT}\")\n"
+,
         "\n",
         "print(f\"JAX version: {jax.__version__}\")\n",
         "print(f\"JAX devices: {jax.devices()}\")\n",

src/maxtext/scratch_code/demo_from_config.ipynb

@@ -48,10 +48,10 @@
         }
       ],
       "source": [
-        "import MaxText as mt\n",
+        "import maxtext as mt\n",
         "from maxtext.configs import pyconfig\n",
         "import numpy as np\n",
-        "from MaxText.input_pipeline import _input_pipeline_utils\n",
+        "from maxtext.input_pipeline import input_pipeline_utils\n",
         "import os\n",
         "from maxtext.common import common_types\n",
         "import jax\n",

src/maxtext/utils/layerwise_quantization.py

@@ -19,7 +19,7 @@
 
 Example cmd:
 
-python3 -m MaxText.layerwise_quantization  src/maxtext/configs/base.yml \
+python3 -m maxtext.utils.layerwise_quantization  src/maxtext/configs/base.yml \
   tokenizer_path=${TOKENIZER_PATH?} load_parameters_path=${LOAD_PARAMS_PATH?} \
   model_name=deepseek2-16b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 \
   ici_tensor_parallelism=-1 scan_layers=false weight_dtype=bfloat16 per_device_batch_size=1 \

src/maxtext/utils/muon_utils.py

@@ -21,7 +21,7 @@
 
 This module can also be run as a script to inspect the generated dimension
 numbers for a specific model. Example:
-  python3 -m MaxText.muon_utils qwen3-4b True
+  python3 -m maxtext.utils.muon_utils qwen3-4b True
 """
 
 

tools/data_generation/generate_distillation_data.py

@@ -20,7 +20,7 @@
 This generated dataset can be used to fine-tune a student model.
 
 Example command:
-  python3 -m MaxText.generate_distillation_data \
+  python3 -m tools.data_generation.generate_distillation_data \
     --dataset-path HuggingFaceH4/ultrachat_200k --data-split train_sft --data-columns messages \
     --tokenizer-path deepseek-ai/DeepSeek-V2-Lite-chat \
     --hf-access-token <access token> \
@@ -37,7 +37,7 @@
 `max-target-length` is the max length of prompt tokens and expected completion tokens.
 Set `--remove-local-dataset-files` to remove dataset files created locally after uploading to Hugging Face or GCS.
 `upload-to-hf` will upload the dataset to Hugging Face and `upload-to-gcs` will upload the dataset to GCS.
-For more information, check out `python3 -m MaxText.generate_distillation_data --help`.
+For more information, check out `python3 -m tools.data_generation.generate_distillation_data --help`.
 Note:
 Make sure to run maxengine server in a new terminal before executing this command. Example command to run maxengine server:
   python3 -m maxtext.inference.maxengine.maxengine_server src/maxtext/configs/base.yml \

tools/weight_inspector/weight_inspector.py

@@ -17,7 +17,7 @@
 
 Usage:
 
-python3 -m MaxText.weight_inspector --lhs left_hand.pkl --rhs right_hand.pkl
+python3 -m tools.weight_inspector.weight_inspector --lhs left_hand.pkl --rhs right_hand.pkl
 
 """