Skip to content

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

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_940558946
Open

Refactor MaxText's checkpointing system to primarily use the Orbax v1 API, moving away from the legacy v0 implementations.#4314
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_940558946

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Jul 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

@codecov

codecov Bot commented Jul 1, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 66.48250% with 182 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/common/grain_utility.py 62.70% 78 Missing and 13 partials ⚠️
src/maxtext/common/checkpointing.py 69.40% 32 Missing and 9 partials ⚠️
src/maxtext/common/emergency_checkpointing.py 48.00% 38 Missing and 1 partial ⚠️
...rc/maxtext/utils/generate_param_only_checkpoint.py 40.00% 3 Missing ⚠️
src/maxtext/utils/elastic_utils.py 33.33% 1 Missing and 1 partial ⚠️
src/maxtext/utils/train_utils.py 66.66% 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/rampup_batch.py 50.00% 1 Missing ⚠️
src/maxtext/utils/standalone_checkpointer.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@copybara-service copybara-service Bot force-pushed the test_940558946 branch 4 times, most recently from e1df263 to d171b74 Compare July 1, 2026 21:56
@copybara-service copybara-service Bot force-pushed the test_940558946 branch 9 times, most recently from 50c7315 to 252cde4 Compare July 8, 2026 09:36
@copybara-service copybara-service Bot changed the title Refactor checkpointing to use Orbax v1. This CL refactors MaxText's checkpointing system to primarily use the Orbax v1 API, moving away from the legacy v0 implementations. Jul 8, 2026
@copybara-service copybara-service Bot force-pushed the test_940558946 branch 6 times, most recently from 54fbad1 to fa4b5ef Compare July 8, 2026 10:20
@copybara-service copybara-service Bot changed the title This CL refactors MaxText's checkpointing system to primarily use the Orbax v1 API, moving away from the legacy v0 implementations. Refactor MaxText's checkpointing system to primarily use the Orbax v1 API, moving away from the legacy v0 implementations. Jul 8, 2026
@angel-core @Google-ML-Automation
Refactor MaxText's checkpointing system to primarily use the Orbax v1…
93b4298 
… API, moving away from the legacy v0 implementations.

PiperOrigin-RevId: 940558946
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jacoguzo jacoguzo Awaiting requested review from jacoguzo
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@parambole parambole Awaiting requested review from parambole parambole is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@xuefgu xuefgu Awaiting requested review from xuefgu xuefgu is a code owner
@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@angel-core