feat(scan_layers): dynamically set scan_layers from checkpoint metadata

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 and explicit user configurations. (Note: MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming, so you 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")

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -1190,6 +1190,7 @@ def save_weights_to_checkpoint(
     device_count: int,
     use_ocdbt: bool,
     use_zarr3: bool,
+    config=None,
 ):
   """Saves model weights to a MaxText-compatible checkpoint with optional sharding.
 
@@ -1205,6 +1206,7 @@ def save_weights_to_checkpoint(
       use_ocdbt: If True, enables the Optimized Checkpoint Database with Transactions
           (OCDBT) format for improved metadata handling.
       use_zarr3: If True, uses the Zarr3 storage format for the underlying array data.
+      config: Optional config to save along with checkpoint metadata.
   """
   mem_info = psutil.Process()
   logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
@@ -1241,7 +1243,7 @@ def save_weights_to_checkpoint(
   )
 
   logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
-  if checkpointing.save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state_new):
+  if checkpointing.save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state_new, config=config):
     max_logging.log(f"saved a checkpoint at step {step_number_to_save_new_ckpt}")
   # Upon preemption, exit when and only when all ongoing saves are complete.
   checkpoint_manager.wait_until_finished()

src/maxtext/configs/base.yml

@@ -365,7 +365,8 @@ engram: 'remat'
 
 optimizer_memory_host_offload: false
 parameter_memory_host_offload: false
-scan_layers: true # We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations.
+# Whether to use jax.lax.scan over layers (stacked/unstacked checkpoint). We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations. When resuming from a checkpoint, this flag is auto-determined from metadata.
+scan_layers: true
 param_scan_axis: 1
 
 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores

src/maxtext/configs/pyconfig.py

@@ -498,6 +498,7 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf
   # 6. Handle environment variable overrides
   cli_keys = frozenset(omegaconf.OmegaConf.to_container(cli_cfg, resolve=True).keys())
   kwargs_keys = frozenset(kwargs.keys())
+  env_keys = set()
   for k in tuple(raw_keys_dict.keys()):
     env_key = yaml_key_to_env_key(k)
     if env_key in os.environ:
@@ -513,6 +514,7 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf
             "This is not allowed, unless setting `override_model_config=True`."
         )
 
+      env_keys.add(k)
       new_proposal = os.environ.get(env_key)
       original_value = raw_keys_dict.get(k)
       parser = None
@@ -551,6 +553,8 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf
       compilation_cache.set_cache_dir(os.path.expanduser(pydantic_kwargs["jax_cache_dir"]))
 
   pydantic_config = types.MaxTextConfig(**pydantic_kwargs)
+  explicit_keys = set((cli_keys | kwargs_keys | env_keys) & frozenset(types.MaxTextConfig.model_fields.keys()))
+  object.__setattr__(pydantic_config, "__pydantic_fields_set__", explicit_keys)
   config = HyperParameters(pydantic_config)
 
   if config.log_config:

src/maxtext/configs/types.py

@@ -972,7 +972,13 @@ class HardwareAndMesh(BaseModel):
   )
   shard_mode: ShardMode = Field("auto", description="can be either auto or explicit")
   inhomogeneous_layer_cycle_interval: int = Field(1, description="The interval of repeated inhomogeneous layer patterns.")
-  scan_layers: bool = Field(True, description="Whether to use jax.lax.scan over layers.")
+  scan_layers: bool = Field(
+      True,
+      description=(
+          "Whether to use jax.lax.scan over layers (stacked/unstacked checkpoint). "
+          "When resuming from a checkpoint, this flag is auto-determined from metadata."
+      ),
+  )
   param_scan_axis: int = Field(1, description="Axis to scan over for parameters.")
   context_parallel_load_balance: bool = Field(True, description="Whether to use load balancing for context parallelism.")
   context_parallel_strategy: str = Field(

src/maxtext/inference/vllm_decode.py

@@ -201,9 +201,7 @@ def decode_with_tunix(
     overrides = config.vllm_hf_overrides
     if isinstance(overrides, str) and "MaxTextForCausalLM" in overrides:
       use_no_op_mappings = True
-    elif isinstance(overrides, dict) and "MaxTextForCausalLM" in overrides.get(
-        "architectures", []
-    ):
+    elif isinstance(overrides, dict) and "MaxTextForCausalLM" in overrides.get("architectures", []):
       use_no_op_mappings = True
 
   tunix_model = TunixMaxTextAdapter(
@@ -238,14 +236,8 @@ def decode_with_tunix(
     )
 
   # MaxText uses -1 to mean "disabled"; vLLM requires top_p in (0, 1].
-  top_p = (
-      config.decode_sampling_nucleus_p
-      if config.decode_sampling_nucleus_p > 0
-      else 1.0
-  )
-  top_k = (
-      config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else -1
-  )
+  top_p = config.decode_sampling_nucleus_p if config.decode_sampling_nucleus_p > 0 else 1.0
+  top_k = config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else -1
 
   rollout_vllm_additional_config = {
       "maxtext_config": {
@@ -280,17 +272,9 @@ def decode_with_tunix(
       rollout_vllm_hbm_utilization=config.hbm_utilization_vllm,
       rollout_vllm_init_with_random_weights=True,
       rollout_vllm_tpu_backend_type="jax",
-      tensor_parallel_size=(
-          config.ici_tensor_parallelism
-          if config.ici_tensor_parallelism > 0
-          else 1
-      ),
+      tensor_parallel_size=(config.ici_tensor_parallelism if config.ici_tensor_parallelism > 0 else 1),
       data_parallel_size=jax.device_count()
-      // (
-          config.ici_tensor_parallelism
-          if config.ici_tensor_parallelism > 0
-          else 1
-      ),
+      // (config.ici_tensor_parallelism if config.ici_tensor_parallelism > 0 else 1),
       rollout_vllm_additional_config=rollout_vllm_additional_config,
       rollout_vllm_kwargs={"hf_overrides": config.vllm_hf_overrides},
   )
@@ -330,9 +314,7 @@ def main(argv: Sequence[str]) -> None:
   if FLAGS.use_tunix:
     maxtext_model, mesh = model_creation_utils.from_pretrained(config)
     if config.lora.enable_lora:
-      maxtext_model = lora_utils.apply_lora_to_model(
-          maxtext_model, mesh, config
-      )
+      maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config)
       if config.lora.lora_restore_path:
         lora_utils.restore_lora_from_path(maxtext_model, config)
     decode_with_tunix(config, model=maxtext_model, mesh=mesh)

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -345,9 +345,7 @@ def load_weights(self, rng_key: jax.Array) -> None:
           self.maxtext_config, mesh=self.mesh, model_mode=self.model_mode, rng_key=rng_key
       )
       if self.maxtext_config.lora.enable_lora:
-        model = lora_utils.apply_lora_to_model(
-            model, self.mesh, self.maxtext_config
-        )
+        model = lora_utils.apply_lora_to_model(model, self.mesh, self.maxtext_config)
         if self.maxtext_config.lora.lora_restore_path:
           lora_utils.restore_lora_from_path(model, self.maxtext_config)
       self.model = nnx.data(model)