Merge branch 'main' into yixuann-dev-goodput

Author: YixuanWang-99

.github/workflows/run_jupyter_notebooks.yml

@@ -50,6 +50,9 @@ jobs:
     runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
     container:
       image: gcr.io/tpu-prod-env-multipod/${{ inputs.base_image }}
+      env:
+        VLLM_TARGET_DEVICE: "tpu"
+        UV_TORCH_BACKEND: "cpu"
     steps:
       - name: Checkout MaxText
         if: ${{ !inputs.maxtext_installed }}

.github/workflows/run_tests_against_package.yml

@@ -85,6 +85,8 @@ jobs:
         XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
         TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
         TPU_SKIP_MDS_QUERY: ${{ inputs.device_type == 'cpu' && '1' || '' }}
+        VLLM_TARGET_DEVICE: "tpu"
+        UV_TORCH_BACKEND: "cpu"
         MAXTEXT_PACKAGE_EXTRA: >-
           ${{
             !contains(inputs.pytest_marker, 'not post_training') && 'tpu-post-train'

docs/development.md

@@ -7,4 +7,5 @@ hidden:
 ---
 development/update_dependencies.md
 development/contribute_docs.md
+development/hlo_diff_testing.md
 ```

docs/development/update_dependencies.md

@@ -80,7 +80,7 @@ overrides. You will need to do this separately for the TPU and GPU environments.
 
 ### TPU Pre-Training
 
-> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-requirements.txt` in the repository was generated using JAX build commit hash: efd6cf797ee9c4f29c6c6d5e91ae4432209063be. When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
+> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-requirements.txt` in the repository was generated using JAX build commit hash: [efd6cf797ee9c4f29c6c6d5e91ae4432209063be](https://github.com/jax-ml/jax/commit/efd6cf797ee9c4f29c6c6d5e91ae4432209063be). When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
 
 If you have made changes to TPU pre-training dependencies in `src/dependencies/requirements/base_requirements/tpu-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
@@ -89,7 +89,7 @@ bash src/dependencies/scripts/generate_requirements.sh \
 --base-requirements src/dependencies/requirements/base_requirements/tpu-requirements.txt \
 --generated-requirements tpu-requirements.txt \
 --override-requirements src/dependencies/extra_deps/tpu_overrides.txt \
---seed-commit <jax-build-commit-hash>
+--seed-commit efd6cf797ee9c4f29c6c6d5e91ae4432209063be
 
 # Copy generated requirements to src/dependencies/requirements/generated_requirements
 mv generated_artifacts/python3_12/tpu-requirements.txt \
@@ -98,7 +98,7 @@ mv generated_artifacts/python3_12/tpu-requirements.txt \
 
 ### TPU Post-Training
 
-> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt` in the repository was generated using JAX build commit hash: efd6cf797ee9c4f29c6c6d5e91ae4432209063be. When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
+> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt` in the repository was generated using JAX build commit hash: [a0d62932cbb02835296964e707f1cac0cbb939f8](https://github.com/jax-ml/jax/commit/a0d62932cbb02835296964e707f1cac0cbb939f8). When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
 
 If you have made changes to the post-training dependencies in `src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt`, you need to regenerate the pinned post-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
@@ -108,7 +108,7 @@ bash src/dependencies/scripts/generate_requirements.sh \
 --base-requirements src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt \
 --generated-requirements tpu-post-train-requirements.txt \
 --override-requirements src/dependencies/extra_deps/tpu_post_train_overrides.txt \
---seed-commit <jax-build-commit-hash>
+--seed-commit a0d62932cbb02835296964e707f1cac0cbb939f8
 
 # Copy generated requirements to src/dependencies/requirements/generated_requirements
 mv generated_artifacts/python3_12/tpu-post-train-requirements.txt \
@@ -117,15 +117,15 @@ mv generated_artifacts/python3_12/tpu-post-train-requirements.txt \
 
 ### GPU Pre-Training
 
-> **Note:** The current `src/dependencies/requirements/generated_requirements/cuda12-requirements.txt` in the repository was generated using JAX build commit hash: efd6cf797ee9c4f29c6c6d5e91ae4432209063be. When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
+> **Note:** The current `src/dependencies/requirements/generated_requirements/cuda12-requirements.txt` in the repository was generated using JAX build commit hash: [efd6cf797ee9c4f29c6c6d5e91ae4432209063be](https://github.com/jax-ml/jax/commit/efd6cf797ee9c4f29c6c6d5e91ae4432209063be). When regenerating the requirements, either use the same commit hash or update this hash if you use a different one.
 
 If you have made changes to the GPU pre-training dependencies in `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
 
 ```bash
 bash src/dependencies/scripts/generate_requirements.sh \
 --base-requirements src/dependencies/requirements/base_requirements/cuda12-requirements.txt \
 --generated-requirements cuda12-requirements.txt \
---seed-commit <jax-build-commit-hash> \
+--seed-commit efd6cf797ee9c4f29c6c6d5e91ae4432209063be \
 --override-requirements src/dependencies/extra_deps/cuda12_overrides.txt \
 --hardware cuda12
 

docs/guides/checkpointing_solutions/convert_checkpoint.md

@@ -70,7 +70,7 @@ You can find your converted checkpoint files under `${BASE_OUTPUT_DIRECTORY}/0/i
 ### Key Parameters
 
 - `model_name`: The specific model identifier. It must match a supported entry in the MaxText [globals.py](https://github.com/AI-Hypercomputer/maxtext/blob/16b684840db9b96b19e24e84ac49f06af7204ae3/src/maxtext/utils/globals.py#L46C1-L46C7).
-- `scan_layers`: Controls whether the output uses a scanned (`scan_layers=true`) or unscanned (`scan_layers=false`) checkpoint format. Refer [here](../../reference/core_concepts/checkpoints.md) for more information.
+- `scan_layers`: Controls whether the output uses a scanned (`scan_layers=true`) or unscanned (`scan_layers=false`) checkpoint format. Refer [here](../../reference/core_concepts/checkpoints.md) for more information. **IMPORTANT:** This setting *must* match the `scan_layers` value used during model training or loading. A mismatch will cause PyTree loading errors (though MaxText will intercept these and raise a descriptive `ValueError` explaining the mismatch).
 - `use_multimodal`: Indicates if multimodality is used, important for Gemma3.
 - `base_output_directory`: The path where the converted Orbax checkpoint will be stored; it can be Google Cloud Storage (GCS) or local.
 - `hardware=cpu`: The conversion script runs on a CPU machine.
@@ -239,7 +239,10 @@ Here is an example [PR to add support for gemma3 multi-modal model](https://gith
 
 ### Common Errors
 
-- "Type ShapeDtypeStruct is not a valid JAX type": Usually caused by a mismatch in the `scan_layers` flag.
+- "Type ShapeDtypeStruct is not a valid JAX type" or generic **PyTree structure/shape mismatches** (e.g., Orbax reporting `"X/Y paths matched"`, such as `143/145 paths`):
+  This is almost always caused by a mismatch in the `scan_layers` configuration between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`).
+
+  - **Solution:** Ensure the `scan_layers` flag is set to the exact same value (`True` or `False`) in both the conversion command and your training/execution command.
 
 - If the converted checkpoint loads without errors but produces nonsensical output, likely an error in the Q/K/V weight reshaping logic during conversion.