Update checkpoint conversion documentation

PiperOrigin-RevId: 888829289

Author: SurbhiJainUSC

.coveragerc

@@ -9,8 +9,8 @@ omit =
 
 [paths]
 source =
-    src/MaxText
-    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

benchmarks/api_server/maxtext_generator.py

@@ -34,11 +34,11 @@
 
 from dataclasses import dataclass, field
 
-from maxtext.configs import pyconfig
-from maxtext.inference.maxengine import maxengine
-from maxtext.multimodal import processor as mm_processor
-from maxtext.multimodal import utils as mm_utils
-from maxtext.utils import max_logging, max_utils
+from maxtext.src.maxtext.configs import pyconfig
+from maxtext.src.maxtext.inference.maxengine import maxengine
+from maxtext.src.maxtext.multimodal import processor as mm_processor
+from maxtext.src.maxtext.multimodal import utils as mm_utils
+from maxtext.src.maxtext.utils import max_logging, max_utils
 
 # Set TF log level to avoid verbose startup messages.
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

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

@@ -19,11 +19,11 @@
 
 import maxtext_trillium_model_configs as model_configs
 
-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.maxtext_xpk_runner import BenchmarkRunner
+from benchmarks.maxtext_xpk_runner import HWConfig
+from benchmarks.maxtext_xpk_runner import SWconfig
+from benchmarks.maxtext_xpk_runner import xpk_benchmark_runner
+from benchmarks.maxtext_xpk_runner import XpkConfig
 
 
 DATE = "20241009"

benchmarks/maxtext_xpk_runner.py

@@ -37,7 +37,7 @@
 import benchmarks.maxtext_trillium_model_configs as model_configs
 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.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/mmlu/mmlu_eval.py

@@ -56,10 +56,10 @@
 
 from tqdm import tqdm
 
-from maxtext.configs import pyconfig
-from maxtext.inference.maxengine import maxengine
-from maxtext.utils import max_logging
-from maxtext.utils import max_utils
+from maxtext.src.maxtext.configs import pyconfig
+from maxtext.src.maxtext.inference.maxengine import maxengine
+from maxtext.src.maxtext.utils import max_logging
+from maxtext.src.maxtext.utils import max_utils
 
 ASCII_UPPERCASE_A = ord("A")  # ASCII value for uppercase 'A'
 

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