feat(scan_layers): dynamically set scan_layers from checkpoint metadata

[RexBearIU]

Description

This PR is a follow-up to PR #4304 and PR #4340. It introduces dynamic synchronization of scan_layers configuration from checkpoint metadata when a checkpoint path is provided and scan_layers is not explicitly specified by the user on the command-line, in environment variables, or as programmatic keyword arguments.

Key Changes:

• Robust Tracking of Explicit User Configurations (src/maxtext/configs/pyconfig.py):
  • Keeps track of overridden configuration parameters from the environment (env_keys).
  • Sets Pydantic's internal __pydantic_fields_set__ to only the union of explicit CLI, environment, and keyword argument keys. This avoids the default Pydantic behavior where unpacking a dictionary of default configurations marks every single field as explicitly set.
• Helper for Sync and Verification (src/maxtext/utils/model_creation_utils.py):
  • Extracted dynamic configuration resolution and verification into verify_and_sync_scan_layers(config).
  • Inside from_pretrained, resolved the active config via config = verify_and_sync_scan_layers(config) without verbose comments, keeping the code highly readable.
  • If scan_layers is explicitly specified by the user and mismatches checkpoint metadata, a ValueError configuration mismatch error is raised.

Tests

• Added a comprehensive suite of unit tests in TestVerifyAndSyncScanLayers inside tests/unit/model_creation_utils_test.py:
  • test_sync_to_false_when_implicit: Asserts that scan_layers dynamically updates to False from checkpoint metadata when not explicitly specified by the user.
  • test_sync_to_true_when_implicit: Asserts that scan_layers stays/updates to True when checkpoint metadata is True.
  • test_explicit_match_raises_no_error: Asserts that providing matching explicit values does not trigger validation errors.
  • test_explicit_mismatch_raises_value_error: Asserts that an explicit conflict raises the expected ValueError configuration mismatch error.
• Verified all tests in model_creation_utils_test.py pass successfully:

  PYTHONPATH=src python3 -m pytest tests/unit/model_creation_utils_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• 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.

[codecov]

Codecov Report

❌ Patch coverage is 65.51724% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/checkpoint_conversion/to_huggingface.py	0.00%	6 Missing ⚠️
src/maxtext/utils/model_creation_utils.py	86.36%	1 Missing and 2 partials ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[shralex]

Thanks Jacky!

Can you add some comments to wherever scan_layers is described (types.py, base.yml) describing this behavior when resuming from a checkpoint.

Also, can we change docs in appropriate places to explain this new behavior. I remember I added a bunch of warnings related to scan_layers previously to the docs, these might need to be updated since you no longer need to specify scan_layers when resuming from a checkpoint.

[shralex]

Thanks for the changes. I made a few suggestions to update comments. The main missing thing is updates to conversion scripts: to_huggingface should auto-detect the setting as well from the provided checkpoint (and error if its setting conflicts with the checkpoint:

Consider calling the new helper function in to_huggingface.py's main function right after initialising the config:

Initialize maxtext config

config = pyconfig.initialize_pydantic(argv)

Auto-resolve scan_layers from checkpoint metadata if not explicitly provided

config = verify_and_sync_scan_layers(config)

to_maxtext should write the scan_layers metadata into the checkpoint it creates, but in the current implementation, it does not actually do so. I believe we should provide a config when calling save_weights_to_checkpoint,

[xibinliu]

I think we should also cover the following metadata handling when doing checkpoints conversion:

1. to_maxtext.py
   Generate metadata data with scanned_layers and lora parameters and attach the metadata to the MaxText checkpoints.

2. to_hugging_face.py:
   Besides the lora parameters, we should also read scan_layers from the MaxText metadata and use it for conversion.

[xibinliu]

No need to check if the check in the above comment is implemented.

[RexBearIU]

Correct, with the early return in place, that duplicate check was removed. Thanks!

[xibinliu]

Remove this line:

if saved_scan_layers != config.scan_layers: