Refactor checkpointing to use Orbax v1.

[copybara-service]

Refactor checkpointing to use Orbax v1.

[codecov]

Codecov Report

❌ Patch coverage is 46.97337% with 219 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/grain_utility.py	31.78%	75 Missing and 13 partials ⚠️
src/maxtext/common/checkpointing.py	53.57%	44 Missing and 8 partials ⚠️
src/maxtext/common/emergency_checkpointing.py	41.89%	43 Missing ⚠️
src/maxtext/common/checkpoint_context.py	50.00%	11 Missing and 10 partials ⚠️
...ners/post_train/distillation/distillation_utils.py	64.28%	1 Missing and 4 partials ⚠️
src/maxtext/utils/rampup_batch.py	20.00%	4 Missing ⚠️
...rc/maxtext/utils/generate_param_only_checkpoint.py	33.33%	2 Missing ⚠️
...axtext/checkpoint_conversion/reshard_checkpoint.py	0.00%	1 Missing ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	1 Missing ⚠️
src/maxtext/utils/standalone_checkpointer.py	0.00%	1 Missing ⚠️
... and 1 more

📢 Thoughts on this report? Let us know!