diff --git a/docs/guides/checkpointing_solutions/convert_checkpoint.md b/docs/guides/checkpointing_solutions/convert_checkpoint.md index e65627e794..77a8b15ebd 100644 --- a/docs/guides/checkpointing_solutions/convert_checkpoint.md +++ b/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. diff --git a/docs/reference/core_concepts/checkpoints.md b/docs/reference/core_concepts/checkpoints.md index c365572688..899a924b90 100644 --- a/docs/reference/core_concepts/checkpoints.md +++ b/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 diff --git a/docs/run_maxtext/run_maxtext_localhost.md b/docs/run_maxtext/run_maxtext_localhost.md index 0689ac956c..3f1ba692c2 100644 --- a/docs/run_maxtext/run_maxtext_localhost.md +++ b/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 diff --git a/docs/tutorials/posttraining/dpo.md b/docs/tutorials/posttraining/dpo.md index 5593e96742..039b6c9b25 100644 --- a/docs/tutorials/posttraining/dpo.md +++ b/docs/tutorials/posttraining/dpo.md @@ -104,11 +104,11 @@ export MAXTEXT_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 diff --git a/docs/tutorials/posttraining/multimodal.md b/docs/tutorials/posttraining/multimodal.md index 0e3fe53a04..90ceff00bc 100644 --- a/docs/tutorials/posttraining/multimodal.md +++ b/docs/tutorials/posttraining/multimodal.md @@ -101,10 +101,11 @@ export MAXTEXT_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 diff --git a/docs/tutorials/posttraining/rl_on_multi_host.md b/docs/tutorials/posttraining/rl_on_multi_host.md index bc53bcf559..9cca5bf3fa 100644 --- a/docs/tutorials/posttraining/rl_on_multi_host.md +++ b/docs/tutorials/posttraining/rl_on_multi_host.md @@ -149,11 +149,11 @@ export MAXTEXT_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 diff --git a/docs/tutorials/posttraining/sft.md b/docs/tutorials/posttraining/sft.md index eba5a8f76c..61bf68af84 100644 --- a/docs/tutorials/posttraining/sft.md +++ b/docs/tutorials/posttraining/sft.md @@ -89,11 +89,11 @@ export MAXTEXT_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 diff --git a/docs/tutorials/posttraining/sft_on_multi_host.md b/docs/tutorials/posttraining/sft_on_multi_host.md index 4b7805d4b1..db0d876e5a 100644 --- a/docs/tutorials/posttraining/sft_on_multi_host.md +++ b/docs/tutorials/posttraining/sft_on_multi_host.md @@ -140,11 +140,11 @@ export MAXTEXT_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 diff --git a/src/maxtext/checkpoint_conversion/to_huggingface.py b/src/maxtext/checkpoint_conversion/to_huggingface.py index c3abdf60f4..f0e94d4829 100644 --- a/src/maxtext/checkpoint_conversion/to_huggingface.py +++ b/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() diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index c2ff040979..ca50d67680 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/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") diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 75ee3dffcc..2a7d84c3a7 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/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() diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..2c33cd7a2c 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py index f4a8e3965f..a23c3e3eb8 100644 --- a/src/maxtext/configs/pyconfig.py +++ b/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: diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a8c1745..56ab4f6b7b 100644 --- a/src/maxtext/configs/types.py +++ b/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( diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py index 5908ab2e2a..5a89f5f7fe 100644 --- a/src/maxtext/inference/vllm_decode.py +++ b/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) diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index 1d441327ac..ef15a697c5 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/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) diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index 9e989acbda..455ffde2d1 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -790,6 +790,44 @@ def create_models_and_meshes(trainer_config, sampler_config, trainer_devices, sa return reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh +def verify_and_sync_scan_layers(config): + """Verify and sync scan_layers based on checkpoint metadata.""" + if not config.load_parameters_path: + return config + + custom_metadata = checkpointing.load_checkpoint_metadata(config.load_parameters_path) + saved_scan_layers = custom_metadata.get("scan_layers") + if not isinstance(saved_scan_layers, bool): + return config + + if saved_scan_layers == config.scan_layers: + return config + + # Extract the Pydantic config (or use direct config if already a Pydantic model) + pydantic_config = getattr(config, "_pydantic_config", config) + model_fields_set = getattr(pydantic_config, "model_fields_set", None) + + # If model metadata tracking isn't supported, fall back to matching check (True) + is_explicit = "scan_layers" in model_fields_set if model_fields_set is not None else True + + if is_explicit: + if saved_scan_layers != config.scan_layers: + raise ValueError( + f"Configuration mismatch: Your run specifies scan_layers={config.scan_layers}, " + f"but the checkpoint was saved with scan_layers={saved_scan_layers}." + ) + else: + max_logging.log(f"Setting scan_layers={saved_scan_layers} loaded from checkpoint metadata.") + new_pydantic_config = pydantic_config.model_copy(update={"scan_layers": saved_scan_layers}) + # Wrap back in HyperParameters if the original config was wrapped + if getattr(config, "_pydantic_config", None) is not None: + config = pyconfig.HyperParameters(new_pydantic_config) + else: + config = new_pydantic_config + + return config + + def from_pretrained( config, mesh=None, @@ -864,15 +902,8 @@ def from_pretrained( } ) config = pyconfig.HyperParameters(new_config) - # Proactive verification of scan_layers from checkpoint metadata - if config.load_parameters_path: - custom_metadata = checkpointing.load_checkpoint_metadata(config.load_parameters_path) - saved_scan_layers = custom_metadata.get("scan_layers") - if isinstance(saved_scan_layers, bool) and saved_scan_layers != config.scan_layers: - raise ValueError( - f"Configuration mismatch: Your run specifies scan_layers={config.scan_layers}, " - f"but the checkpoint was saved with scan_layers={saved_scan_layers}." - ) + + config = verify_and_sync_scan_layers(config) if config.pure_nnx: _create_model, abstract_model = create_nnx_abstract_model( diff --git a/tests/unit/model_creation_utils_test.py b/tests/unit/model_creation_utils_test.py index b1d034e8b7..798e02349f 100644 --- a/tests/unit/model_creation_utils_test.py +++ b/tests/unit/model_creation_utils_test.py @@ -832,5 +832,93 @@ def test_returns_linen_train_state_and_annotations(self): self.assertIsNotNone(state_mesh_annotations) +class TestVerifyAndSyncScanLayers(unittest.TestCase): + """Tests for verify_and_sync_scan_layers().""" + + def setUp(self): + self.mesh = Mesh(np.array(jax.devices()[:1]), axis_names=("x",)) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_sync_to_false_when_implicit(self, mock_load_meta): + """If scan_layers is not explicit, sync scan_layers to False from checkpoint metadata.""" + mock_load_meta.return_value = {"scan_layers": False} + # Create config without scan_layers in kwargs so it's not explicit + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + # Pre-assertions + pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) + self.assertTrue(cfg.scan_layers) # default is True + self.assertNotIn("scan_layers", pydantic_cfg.model_fields_set) + + # Call verify_and_sync_scan_layers + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + + # Post-assertions + self.assertFalse(synced_cfg.scan_layers) + synced_pydantic_cfg = getattr(synced_cfg, "_pydantic_config", synced_cfg) + self.assertFalse(synced_pydantic_cfg.scan_layers) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_sync_to_true_when_implicit(self, mock_load_meta): + """If scan_layers is not explicit, sync scan_layers to True from checkpoint metadata.""" + mock_load_meta.return_value = {"scan_layers": True} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + self.assertTrue(synced_cfg.scan_layers) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_explicit_match_raises_no_error(self, mock_load_meta): + """If scan_layers is explicit and matches checkpoint metadata, no error is raised.""" + mock_load_meta.return_value = {"scan_layers": True} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt", scan_layers=True) + + # Pre-assertions + pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) + self.assertTrue(cfg.scan_layers) + self.assertIn("scan_layers", pydantic_cfg.model_fields_set) + + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + self.assertTrue(synced_cfg.scan_layers) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_explicit_mismatch_raises_value_error(self, mock_load_meta): + """If scan_layers is explicit and mismatches checkpoint metadata, ValueError is raised.""" + mock_load_meta.return_value = {"scan_layers": False} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt", scan_layers=True) + + # Pre-assertions + pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) + self.assertTrue(cfg.scan_layers) + self.assertIn("scan_layers", pydantic_cfg.model_fields_set) + + with self.assertRaises(ValueError) as context: + model_creation_utils.verify_and_sync_scan_layers(cfg) + + self.assertIn("Configuration mismatch", str(context.exception)) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch("maxtext.utils.model_creation_utils.max_logging.log") + def test_sync_log_and_early_return(self, mock_log, mock_load_meta): + """Test that we log only when auto-resolution actually changes scan_layers, and early return otherwise.""" + # Scenario A: saved_scan_layers == config.scan_layers (default True). + # Since they match, it should return early, and NO log should be printed. + mock_load_meta.return_value = {"scan_layers": True} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + self.assertTrue(synced_cfg.scan_layers) + mock_log.assert_not_called() + + # Scenario B: saved_scan_layers (False) != config.scan_layers (default True), and is not explicit (implicit). + # It should log the auto-resolution and update scan_layers to False. + mock_load_meta.return_value = {"scan_layers": False} + cfg2 = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + synced_cfg2 = model_creation_utils.verify_and_sync_scan_layers(cfg2) + self.assertFalse(synced_cfg2.scan_layers) + mock_log.assert_called_once_with("Setting scan_layers=False loaded from checkpoint metadata.") + + if __name__ == "__main__": unittest.main()