PR #3252: Migrate remaining src/MaxText utils to src/maxtext/utils

Imported from GitHub PR #3252

# Description

* Move the remaining utils from `src/MaxText` to `src/maxtext/utils`

# Tests

Successfully ran vllm_decode:
```
python3 -m maxtext.inference.vllm_decode src/maxtext/configs/base.yml \
    model_name=qwen3-30b-a3b \
    tokenizer_path=Qwen/Qwen3-30B-A3B \
    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
    ici_tensor_parallelism=4 \
    hbm_utilization_vllm=0.5 \
    prompt="Suggest some famous landmarks in London." \
    decode_sampling_temperature=0.0 \
    decode_sampling_nucleus_p=1.0 \
    decode_sampling_top_k=0.0 \
    use_chat_template=True
```

```
python3 -m maxtext.utils.generate_param_only_checkpoint src/maxtext/configs/base.yml \
    base_output_directory=<base_output_directory> \
    load_parameters_path=<load_parameters_path> \
    run_name=<run_name> \
    model_name=gemma-2b \
    force_unroll=true
```

# Checklist

Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label.
- [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

Copybara import of the project:

--
62f4d58 by Branden Vandermoon <bvandermoon@google.com>:

Migrate remaining src/MaxText utils to src/maxtext/utils

Merging this change closes #3252

COPYBARA_INTEGRATE_REVIEW=#3252 from AI-Hypercomputer:bvandermoon-repo-restructure 62f4d58
PiperOrigin-RevId: 875884352

Author: bvandermoon

src/MaxText/pyconfig_deprecated.py

@@ -28,7 +28,7 @@
 
 import omegaconf
 
-from MaxText import accelerator_to_spec_map
+from maxtext.utils import accelerator_to_spec_map
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR
 from maxtext.common.common_types import AttentionType, DecoderBlockType, ShardMode
 from maxtext.utils import gcs_utils

src/maxtext/checkpoint_conversion/standalone_scripts/llama_mistral_mixtral_orbax_to_hf.py

@@ -48,7 +48,7 @@
 from transformers import LlamaForCausalLM, MistralForCausalLM, AutoModelForCausalLM, AutoConfig
 
 from MaxText import pyconfig
-from MaxText.generate_param_only_checkpoint import _read_train_checkpoint
+from maxtext.utils.generate_param_only_checkpoint import _read_train_checkpoint
 from maxtext.checkpoint_conversion.standalone_scripts import llama_or_mistral_ckpt
 from maxtext.common import checkpointing
 from maxtext.utils import max_logging

src/maxtext/configs/types.py

@@ -31,7 +31,7 @@
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
-from MaxText import accelerator_to_spec_map
+from maxtext.utils import accelerator_to_spec_map
 from pydantic.config import ConfigDict
 from pydantic.fields import Field
 from pydantic.functional_validators import field_validator, model_validator

src/maxtext/vllm_decode.py src/maxtext/inference/vllm_decode.pysrc/maxtext/vllm_decode.py renamed to src/maxtext/inference/vllm_decode.py

@@ -14,25 +14,19 @@
 """
 An example script to perform decoding using vLLM via Tunix or via MaxText on vLLM.
 
-Example usage with Tunix:
-  python3 -m maxtext.vllm_decode maxtext/configs/base.yml \
-    model_name=llama3.1-8b tokenizer_path=meta-llama/Llama-3.1-8B-Instruct \
-    tokenizer_type=huggingface hf_access_token=<your_hf_token> \
-    load_parameters_path=<your_checkpoint_path> \
-    per_device_batch_size=1 run_name=vllm_decode_test max_target_length=64 \
-    use_chat_template=False prompt="Suggest some famous landmarks in London." \
-    decode_sampling_temperature=0.0 decode_sampling_nucleus_p=1.0 decode_sampling_top_k=0.0 \
-    --use_tunix \
-  
-Or without Tunix using the MaxText vLLM integration:
-  python3 -m maxtext.vllm_decode maxtext/configs/base.yml \
-    model_name=qwen3-30b-a3b \
-    tokenizer_path=Qwen/Qwen3-30B-A3B \
-    vllm_hf_config_path=src/MaxText/integration/vllm/maxtext_vllm_adapter \
-    load_parameters_path=<your_checkpoint_path> \
-    ici_tensor_parallelism=4 \
-    hbm_utilization_vllm=0.5 \
-    prompt="Suggest some famous landmarks in London."
+Example usage:
+  python3 -m maxtext.inference.vllm_decode src/maxtext/configs/base.yml \
+      model_name=qwen3-30b-a3b \
+      tokenizer_path=Qwen/Qwen3-30B-A3B \
+      load_parameters_path=<your_checkpoint_path> \
+      vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+      ici_tensor_parallelism=4 \
+      hbm_utilization_vllm=0.5 \
+      prompt="Suggest some famous landmarks in London." \
+      decode_sampling_temperature=0.0 \
+      decode_sampling_nucleus_p=1.0 \
+      decode_sampling_top_k=0.0 \
+      use_chat_template=True
 """
 
 import os

src/maxtext/trainers/pre_train/train_compile.py

@@ -32,7 +32,7 @@
 from jax.experimental.serialize_executable import serialize
 from jax.experimental.topologies import get_topology_desc
 from jax.sharding import AxisType, Mesh
-from MaxText import accelerator_to_spec_map
+from maxtext.utils import accelerator_to_spec_map
 from MaxText import pyconfig
 from maxtext.common.common_types import MODEL_MODE_TRAIN, ShardMode
 from maxtext.layers import quantizations

src/MaxText/accelerator_to_spec_map.py …maxtext/utils/accelerator_to_spec_map.pysrc/MaxText/accelerator_to_spec_map.py renamed to src/maxtext/utils/accelerator_to_spec_map.py src/MaxText/accelerator_to_spec_map.py renamed to src/maxtext/utils/accelerator_to_spec_map.py

@@ -27,7 +27,7 @@ export MODEL_BUCKET=gs://maxtext-gemma/gemma3
 
 python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gemma3_chkpt --base_model_path ${CHKPT_BUCKET}/${MODEL_VARIATION} --maxtext_model_path ${MODEL_BUCKET}/${MODEL_VARIATION}/${idx} --model_size ${MODEL_VARIATION}
 
-# Current MaxText.generate_param_only_checkpoint will need to skip on GPU due to cpu process error. reuse the unscanned ckpt generated separately.
+# Current maxtext.utils.generate_param_only_checkpoint will need to skip on GPU due to cpu process error. reuse the unscanned ckpt generated separately.
 
 # # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 # export DATASET_PATH=gs://maxtext-dataset
@@ -38,8 +38,8 @@ python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gemma3_chkpt
 # export RUN_NAME=unscanned_chkpt_${idx}
 # export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/0/items
 # Note that the `CONVERTED_CHECKPOINT` is in a `scanned` format which is great for training but for efficient decoding performance we want the checkpoint in an `unscanned` format.
-# We can do this by running `src/MaxText/generate_param_only_checkpoint.py` on `CONVERTED_CHECKPOINT` with `force_unroll=true`.
-#JAX_PLATFORMS=cpu python3 -m MaxText.generate_param_only_checkpoint "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml base_output_directory=${BASE_OUTPUT_DIRECTORY} load_parameters_path=${CONVERTED_CHECKPOINT} run_name=${RUN_NAME} model_name=${MODEL_NAME} force_unroll=true
+# We can do this by running `src/maxtext/utils/generate_param_only_checkpoint.py` on `CONVERTED_CHECKPOINT` with `force_unroll=true`.
+#JAX_PLATFORMS=cpu python3 -m maxtext.utils.generate_param_only_checkpoint "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml base_output_directory=${BASE_OUTPUT_DIRECTORY} load_parameters_path=${CONVERTED_CHECKPOINT} run_name=${RUN_NAME} model_name=${MODEL_NAME} force_unroll=true
 
 export UNSCANNED_CKPT_PATH=gs://runner-maxtext-logs/unscanned_chkpt_2025-04-16-00-01/checkpoints/0/items
 export NVTE_FUSED_ATTN=1