feat(scan_layers): dynamically set scan_layers from checkpoint metadata

Author: RexBearIU

docs/guides/checkpointing_solutions/convert_checkpoint.md

@@ -78,7 +78,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 [to the Checkpoints guide](checkpoints) 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).
+- `scan_layers`: Controls whether the output uses a scanned (`scan_layers=true`) or unscanned (`scan_layers=false`) checkpoint format. Refer [to the Checkpoints guide](checkpoints) for more information. **Note:** When resuming or loading a checkpoint in MaxText, this setting will automatically be loaded from the checkpoint metadata, meaning you do not have to manually specify it during model execution. If you do explicitly specify a value for `scan_layers`, it must match the checkpoint's saved configuration, or a `ValueError` mismatch error will be raised.
 - `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.
@@ -298,9 +298,9 @@ Here is an example [PR to add support for gemma3 multi-modal model](https://gith
 
 - **Error:** `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`).
 
-  - **Cause: Configuration mismatch** (e.g., `scan_layers`) between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`).
+  - **Cause: Configuration mismatch** (e.g., `scan_layers`) between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`). (Note: Since MaxText automatically loads `scan_layers` from the checkpoint's saved metadata, you should only encounter this error if you explicitly set a mismatching value on the command line, which raises a `ValueError` mismatch error).
 
-  - **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.
+  - **Solution:** Omit the `scan_layers` parameter from your training or execution command to allow MaxText to automatically resolve it from the checkpoint metadata, or ensure any explicitly specified `scan_layers` parameter matches the format of the loaded checkpoint.
 
 - **Error:** The converted checkpoint loads without errors but produces nonsensical output.
 

docs/reference/core_concepts/checkpoints.md

@@ -74,7 +74,7 @@ In MaxText, the **`scan_layers`** configuration parameter is used to control thi
 - `scan_layers=false` tells MaxText to keep layer parameters unstacked (often required for inference and certain model architectures).
 
 > [!IMPORTANT]
-> **PyTree Structure Compatibility:** Because JAX expects the loaded PyTree structure to exactly match the model's instantiated structure, the value of the `scan_layers` flag during execution (training, SFT, RL, DPO, or decoding) **must** match the format of the checkpoint being loaded. A mismatch will cause PyTree loading or shape/path mismatch errors (which MaxText will intercept to raise a descriptive `ValueError` pointing to the scan_layers setting).
+> **Automatic scan_layers Resolution:** MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify `scan_layers` on the command-line. If you explicitly specify a value for `scan_layers` that conflicts with the checkpoint format, MaxText will raise a descriptive `ValueError` mismatch error to prevent JAX PyTree structure or shape mismatch errors during loading.
 
 ### Takeaways
 

docs/run_maxtext/run_maxtext_localhost.md

@@ -70,7 +70,7 @@ python3 -m maxtext.inference.decode \
 **Note:** Because the model hasn't been properly trained, the output text will be random. To generate meaningful output, you need to load a trained checkpoint using the `load_parameters_path` argument. For instructions on how to convert pre-trained Hugging Face model checkpoints (like Llama or Gemma) to MaxText's Orbax format, please refer to the [Checkpoint Conversion Guide](../guides/checkpointing_solutions/convert_checkpoint.md).
 
 > [!NOTE]
-> **Checkpoints & `scan_layers` compatibility:** When loading an external or converted checkpoint via `load_parameters_path`, the `scan_layers` setting in your command **must** match the setting used to save the checkpoint. If the checkpoint was saved/converted with `scan_layers=False` (common for Hugging Face conversions and inference runs), you must specify `scan_layers=False` in your command. Otherwise, JAX/Orbax will raise PyTree structure mismatch errors.
+> **Checkpoints & `scan_layers` Auto-Resolution:** When resuming or loading an external or converted checkpoint via `load_parameters_path`, MaxText automatically loads `scan_layers` with the checkpoint's saved format (stacked vs unstacked). You do not need to explicitly specify `scan_layers` on your command line when resuming. If you do explicitly specify a `scan_layers` value, it must match the checkpoint's saved configuration, or a `ValueError` mismatch error will be raised.
 
 ### Running models using provided configs
 

docs/tutorials/posttraining/dpo.md

@@ -104,11 +104,11 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/
 ```
 
 > [!IMPORTANT]
-> **Matching the `scan_layers` Parameter:**
-> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`.
+> **Automatic `scan_layers` Resolution:**
+> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line.
 >
-> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.**
-> - If `scan_layers` does not match, MaxText will raise a `ValueError`.
+> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically.
+> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised.
 >   See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details.
 
 ## Running DPO Training

docs/tutorials/posttraining/multimodal.md

@@ -101,10 +101,11 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/
 ```
 
 > [!IMPORTANT]
-> **Matching the `scan_layers` Parameter:**
-> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`.
+> **Automatic `scan_layers` Resolution:**
+> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line.
 >
-> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.**
+> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically.
+> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised.
 
 ## Multimodal Decode
 

docs/tutorials/posttraining/rl_on_multi_host.md

@@ -149,11 +149,11 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/
 ```
 
 > [!IMPORTANT]
-> **Matching the `scan_layers` Parameter:**
-> The `scan_layers` setting during your RL training run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`.
+> **Automatic `scan_layers` Resolution:**
+> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line.
 >
-> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.**
-> - If `scan_layers` does not match, MaxText will raise a `ValueError`.
+> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically.
+> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised.
 >   See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details.
 
 ## Submit your RL workload via Pathways

docs/tutorials/posttraining/sft.md

@@ -89,11 +89,11 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/
 ```
 
 > [!IMPORTANT]
-> **Matching the `scan_layers` Parameter:**
-> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`.
+> **Automatic `scan_layers` Resolution:**
+> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line.
 >
-> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.**
-> - If `scan_layers` does not match, MaxText will raise a `ValueError`.
+> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically.
+> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised.
 >   See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details.
 
 ## Run SFT on Hugging Face Dataset

docs/tutorials/posttraining/sft_on_multi_host.md

@@ -140,11 +140,11 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # gs://my-bucket/my-checkpoint-directory/0/
 ```
 
 > [!IMPORTANT]
-> **Matching the `scan_layers` Parameter:**
-> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`.
+> **Automatic `scan_layers` Resolution:**
+> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line.
 >
-> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.**
-> - If `scan_layers` does not match, MaxText will raise a `ValueError`.
+> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically.
+> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised.
 >   See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details.
 
 ## Submit workload on GKE cluster

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -98,6 +98,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils.globals import HF_IDS
 from maxtext.utils.lora_utils import sync_lora_metadata
+from maxtext.utils.model_creation_utils import verify_and_sync_scan_layers
 
 
 flags.DEFINE_bool(
@@ -443,6 +444,14 @@ def main(argv: Sequence[str]) -> None:
   assert (
       config.load_full_state_path == ""
   ), "This script expects parameters, not a full state. Use generate_param_only_checkpoint first if needed."
+
+  if not config.load_parameters_path and config.lora.lora_restore_path:
+    # Standalone LoRA conversion uses lora_restore_path for metadata
+    temp_config = config.model_copy(update={"load_parameters_path": config.lora.lora_restore_path})
+    temp_config = verify_and_sync_scan_layers(temp_config)
+    config = config.model_copy(update={"scan_layers": temp_config.scan_layers})
+  else:
+    config = verify_and_sync_scan_layers(config)
   max_utils.print_system_information()
   overall_start = time.time()
 

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -1070,6 +1070,7 @@ def _eager_getter(key):
       simulated_cpu_devices_count,
       config.checkpoint_storage_use_ocdbt,
       config.checkpoint_storage_use_zarr3,
+      config=config,
   )
 
   print_ram_usage("Program Ends")