# Description

[copybara-service]

Description

Refactor MaxText's checkpointing system to primarily use the Orbax v1 API, moving away from the legacy v0 implementations.

FIXES: b/532615853

Tests

Added tests that ensure backward compatibility with v0 as we cutover to v1

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 85.71429% with 25 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	85.52%	13 Missing and 9 partials ⚠️
src/maxtext/common/grain_utility.py	83.33%	2 Missing ⚠️
...trainers/diloco/utils/spmd_diloco_checkpointing.py	88.88%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @abhinavclemson, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request presents an exceptionally high-quality refactoring of MaxText's checkpointing system to primarily use the Orbax v1 API, deprecating legacy v0 implementations. The code is clean, idiomatic, and shows meticulous attention to maintaining backward compatibility and preserving standard robust semantics.

🔍 General Feedback

• Graceful Deprecations: The deprecated option enable_orbax_v1 is handled gracefully with clear warnings rather than hard failures, which ensures seamless user transitions.
• Robustness in Mismatch Diagnostic: The _weight_mismatches logic is highly robust, and we've proposed a small enhancement to catch structural mismatches when a dictionary is expected but a tensor is restored.
• Thorough Test Coverage: Unit and integration tests have been extensively migrated and expanded to thoroughly validate the v1 checkpointer logic, including edge cases like existing steps and async saves.

[github-actions]

🟡 Medium - When the model expects a dictionary/subtree of parameters, but the checkpoint contains a tensor/array (or any other non-dict type), `_weight_mismatches` currently recursively processes children of the expected dict with `have = None`, reporting all nested leaves as "missing" instead of reporting a structural mismatch at the parent key itself.

Adding a structural check when the model expects a dictionary but the checkpoint has a non-dict type improves the diagnostics and error messages significantly.

Suggested change

	if isinstance(want, dict):
	out = []
	is_quant = is_quantized_param or any(k in want for k in ("qvalue", "qarray"))
	if isinstance(want, dict):
	if have is not None and not isinstance(have, dict):
	name = "/".join(str(p) for p in path)
	return [(name, f"structural mismatch: model expects a dict but checkpoint provides {type(have).__name__}")]
	out = []
	is_quant = is_quantized_param or any(k in want for k in ("qvalue", "qarray"))

[github-actions]

🟢 Low - Adding a corresponding test case for the inverse structural mismatch (where the model expects a dictionary, but the checkpoint provides a tensor) to ensure complete and robust test coverage for mismatch scenarios.

Suggested change

	self.assertEqual(list(problems.keys()), ["a/k"])
	self.assertIn("structural mismatch", problems["a/k"])
	self.assertIn("dict", problems["a/k"])
	self.assertEqual(list(problems.keys()), ["a/k"])
	self.assertIn("structural mismatch", problems["a/k"])
	self.assertIn("dict", problems["a/k"])
	def test_weight_mismatches_detects_structural_mismatch_inverse(self):
	want = {
	"a": {
	"k": {
	"w": jax.ShapeDtypeStruct((2,), jnp.float32),
	}
	}
	}
	# k is a tensor instead of a dictionary
	have = {"a": {"k": jnp.ones((2,))}}
	problems = dict(checkpointing._weight_mismatches(want, have)) # pylint: disable=protected-access
	self.assertEqual(list(problems.keys()), ["a/k"])
	self.assertIn("structural mismatch", problems["a/k"])
	self.assertIn("model expects a dict", problems["a/k"])