Skip to content

Clarify and improve scan_layers mismatch error handling#4099

Open
shralex wants to merge 1 commit into
mainfrom
shralex_ux3
Open

Clarify and improve scan_layers mismatch error handling#4099
shralex wants to merge 1 commit into
mainfrom
shralex_ux3

Conversation

@shralex
Copy link
Copy Markdown
Collaborator

@shralex shralex commented Jun 6, 2026

When scan_layers is inconsistent between checkpoint creation and training, Orbax throws an error.
This PR documents this and improves error handling across all possible restore paths.

FIXES: b/520780284

Tests

Unit and integration tests were added to verify error handling.

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
Copy link
Copy Markdown

codecov Bot commented Jun 6, 2026

Codecov Report

❌ Patch coverage is 71.92982% with 16 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/common/checkpointing.py 74.07% 11 Missing and 3 partials ⚠️
src/maxtext/utils/model_creation_utils.py 33.33% 1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@shralex shralex force-pushed the shralex_ux3 branch 4 times, most recently from 339e43f to 5ace97a Compare June 6, 2026 19:46
@shralex shralex changed the title docs & feat: clarify and improve scan_layers mismatch error handling … Clarify and improve scan_layers mismatch error handling Jun 6, 2026
@shralex shralex force-pushed the shralex_ux3 branch 2 times, most recently from c832537 to 21c7a8f Compare June 6, 2026 20:43
@shralex shralex force-pushed the shralex_ux3 branch 2 times, most recently from 173c04a to f833914 Compare June 7, 2026 01:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant