Update checkpoint conversion documentation and fix src/maxtext references in docs

PiperOrigin-RevId: 888829289

Author: SurbhiJainUSC

.coveragerc

@@ -10,7 +10,7 @@ omit =
 [paths]
 source =
     src/MaxText
-    src/MaxText
+    src/maxtext
     */site-packages/MaxText
     */site-packages/maxtext
 

.pre-commit-config.yaml

@@ -42,7 +42,7 @@ repos:
   #        args:
   #          - '--jobs=auto'
   #          - '--keep-going'
-  #          - 'src/MaxText/'
+  #          - 'src/maxtext/'
 
   - repo: https://github.com/google/pyink
     rev: 24.10.1

README.md

@@ -57,7 +57,7 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies
 * \[October 10, 2025\] Post-Training (SFT, RL) via [Tunix](https://github.com/google/tunix) is now available.
 * \[September 26, 2025\] Vocabulary tiling ([PR](https://github.com/AI-Hypercomputer/maxtext/pull/2242)) is now supported in MaxText! Adjust config `num_vocab_tiling` to unlock more efficient memory usage.
 * \[September 24, 2025\] The GPT-OSS family of models (20B, 120B) is now supported.
-* \[September 15, 2025\] MaxText is now available as a [PyPI package](https://pypi.org/project/maxtext). Users can now [install maxtext through pip](https://maxtext.readthedocs.io/en/latest/guides/install_maxtext.html).
+* \[September 15, 2025\] MaxText is now available as a [PyPI package](https://pypi.org/project/maxtext). Users can now [install maxtext through pip](https://maxtext.readthedocs.io/en/latest/install_maxtext.html).
 * \[September 5, 2025\] MaxText has moved to an `src` layout as part of [RESTRUCTURE.md](https://github.com/AI-Hypercomputer/maxtext/blob/aca5b24931ebcbadb55a82e56ebffe8024874028/RESTRUCTURE.md). For existing environments, please run `pip install -e .` from MaxText root.
 * \[August 13, 2025\] The Qwen3 2507 MoE family of models is now supported: MoEs: 235B Thinking & 480B Coder as well as existing dense models: 0.6B, 4B, 8B, 14B, and 32B.
 * \[July 27, 2025\] Updated TFLOPS/s calculation ([PR](https://github.com/AI-Hypercomputer/maxtext/pull/1988)) to account for causal attention, dividing the attention flops in half. Accounted for sliding window and chunked attention reduced attention flops in [PR](https://github.com/AI-Hypercomputer/maxtext/pull/2009) and [PR](https://github.com/AI-Hypercomputer/maxtext/pull/2030). Changes impact large sequence configs, as explained in this [doc](https://maxtext.readthedocs.io/en/latest/reference/performance_metrics.html)

benchmarks/convergence/c4_exp.py

@@ -23,7 +23,6 @@
 
 from benchmarks.benchmark_utils import MaxTextModel, _add_to_model_dictionary
 from benchmarks.convergence.convergence_utils import DatasetHParams, ConvHParams, _setup_model_convergence_
-
 from benchmarks.maxtext_v5p_model_configs import deepseek_v3_ep_256_v5p_512
 
 c4_pretrain_model_dict = {}

benchmarks/disruption_management/monitor.py

@@ -29,7 +29,6 @@
 import time
 
 from benchmarks.disruption_management.disruption_utils import wait_for_pod_to_start
-
 from benchmarks.disruption_management.disruption_handler import DisruptionConfig
 from benchmarks.disruption_management.disruption_handler import TriggerType
 

benchmarks/llama2_v6e-256_benchmarks.py

@@ -17,13 +17,12 @@
 on a specific v6e-256 hardware setup using the XPK runner.
 """
 
-import maxtext_trillium_model_configs as model_configs
+import os
 
-from maxtext_xpk_runner import BenchmarkRunner
-from maxtext_xpk_runner import HWConfig
-from maxtext_xpk_runner import SWconfig
-from maxtext_xpk_runner import xpk_benchmark_runner
-from maxtext_xpk_runner import XpkConfig
+from benchmarks import maxtext_trillium_model_configs as model_configs
+from benchmarks.maxtext_xpk_runner import WorkloadConfig
+from benchmarks.maxtext_xpk_runner import xpk_benchmark_runner
+from benchmarks.maxtext_xpk_runner import XpkClusterConfig
 
 
 DATE = "20241009"
@@ -35,34 +34,37 @@
 DEVICE_TYPE = "v6e-256"
 NUM_SLICES = 1
 BASE_OUTPUT_DIR = "gs://maxtext-experiments-tpem/"
-
-v6e_env_configs = SWconfig(base_docker_image=BASE_DOCKER_IMAGE, libtpu_version=DATE)
-v6e_256_configs = HWConfig(num_slices=NUM_SLICES, device_type=DEVICE_TYPE)
-
-llama2_70b_4096 = BenchmarkRunner(
-    model_name=model_configs.llama2_70b_4096,
-    software_config=v6e_env_configs,
-    hardware_config=v6e_256_configs,
-)
-
-llama2_7b_4096 = BenchmarkRunner(
-    model_name=model_configs.llama2_7b_4096,
-    software_config=v6e_env_configs,
-    hardware_config=v6e_256_configs,
-)
+XPK_PATH = os.path.join("~", "xpk")
+BENCHMARK_STEPS = 20
 
 
 def main() -> None:
-  cluster_config = XpkConfig(
+  cluster_config = XpkClusterConfig(
       cluster_name=CLUSTER_NAME,
       project=PROJECT,
       zone=ZONE,
-      num_slices=NUM_SLICES,
       device_type=DEVICE_TYPE,
-      base_output_directory=BASE_OUTPUT_DIR,
   )
 
-  xpk_benchmark_runner(cluster_config, [llama2_7b_4096, llama2_70b_4096])
+  workload_configs = []
+  for model in [model_configs.llama2_7b_4096, model_configs.llama2_70b_4096]:
+    workload_configs.append(
+        WorkloadConfig(
+            model=model,
+            num_slices=NUM_SLICES,
+            device_type=DEVICE_TYPE,
+            base_output_directory=BASE_OUTPUT_DIR,
+            base_docker_image=BASE_DOCKER_IMAGE,
+            libtpu_type=None,
+            libtpu_nightly_version=DATE,
+            pathways_config=None,
+            xpk_path=XPK_PATH,
+            num_steps=BENCHMARK_STEPS,
+            priority="medium",
+        )
+    )
+
+  xpk_benchmark_runner(cluster_config, workload_configs)
 
 
 if __name__ == "__main__":

benchmarks/maxtext_xpk_runner.py

@@ -35,9 +35,9 @@
 import omegaconf
 
 import benchmarks.maxtext_trillium_model_configs as model_configs
+import benchmarks.xla_flags_library as xla_flags
 from benchmarks.globals import MAXTEXT_PKG_DIR
 from benchmarks.command_utils import run_command_with_updates
-import benchmarks.xla_flags_library as xla_flags
 from benchmarks.disruption_management.disruption_handler import DisruptionConfig
 from benchmarks.disruption_management.disruption_manager import DisruptionManager
 from benchmarks.xpk_configs import XpkClusterConfig

benchmarks/recipes/mcjax_long_running_recipe.py

@@ -27,7 +27,7 @@
 import benchmarks.maxtext_trillium_model_configs as model_configs
 import benchmarks.maxtext_xpk_runner as mxr
 from benchmarks.xpk_configs import XpkClusterConfig
-from . import user_configs
+from benchmarks.recipes import user_configs
 
 # Cluster Params
 CLUSTER = "v6e-256-cluster"

benchmarks/recipes/pw_elastic_training_recipe.py

@@ -25,11 +25,11 @@
 
 parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 sys.path.append(parent_dir)
-from . import args_helper as helper
-from . import user_configs
 
 from benchmarks.disruption_management.disruption_handler import DisruptionMethod
-from .runner_utils import generate_and_run_workloads
+from benchmarks.recipes import args_helper as helper
+from benchmarks.recipes import user_configs
+from benchmarks.recipes.runner_utils import generate_and_run_workloads
 
 user_configs.USER_CONFIG.max_restarts = 10
 COMPARE_WITH_MCJAX = True

benchmarks/recipes/pw_headless_mode.py

@@ -22,8 +22,8 @@
 """
 
 import benchmarks.recipes.args_helper as helper
-from .. import maxtext_xpk_runner as mxr
-from ..recipes.user_configs import USER_CONFIG
+from benchmarks import maxtext_xpk_runner as mxr
+from benchmarks.recipes.user_configs import USER_CONFIG
 
 
 def main() -> int: