Skip to content

Add Multi-tiered Checkpointing Support to Pathways - #3952

Merged
copybara-service[bot] merged 1 commit into
mainfrom
sujinesh/mtc_pathways_support
Jun 26, 2026
Merged

Add Multi-tiered Checkpointing Support to Pathways#3952
copybara-service[bot] merged 1 commit into
mainfrom
sujinesh/mtc_pathways_support

Conversation

@SujeethJinesh

@SujeethJinesh SujeethJinesh commented May 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Adds Multi-tiered checkpointing support to MaxText. This will also work with Elastic Training Pause & Resume as well as Elastic Resizing.

This feature is expected to roughly increase the goodput by ~6.5% and is only dependent on Orbax.

Also pins the minimum Orbax version (0.12.1) to accommodate the necessary changes. Regenerated the TPU requirements using instructions from here.

Tests

Tested internally against candidate Orbax.

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 Bot commented May 20, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.61538% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/common/checkpointing.py 60.00% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

@github-actions

github-actions Bot commented Jun 19, 2026

Copy link
Copy Markdown
Contributor

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

@github-actions github-actions Bot added the stale Automatically applied to stale PRs. label Jun 19, 2026
@SujeethJinesh
SujeethJinesh force-pushed the sujinesh/mtc_pathways_support branch 2 times, most recently from 31e7104 to 3651ea8 Compare June 23, 2026 23:44
@SujeethJinesh SujeethJinesh removed the stale Automatically applied to stale PRs. label Jun 23, 2026
@SujeethJinesh
SujeethJinesh force-pushed the sujinesh/mtc_pathways_support branch from 3651ea8 to 5513d51 Compare June 24, 2026 20:28
@SujeethJinesh
SujeethJinesh force-pushed the sujinesh/mtc_pathways_support branch 6 times, most recently from 25934d5 to bbfd1a7 Compare June 25, 2026 04:29
@github-actions

github-actions Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

🤖 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

github-actions Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

🤖 I'm sorry @abhinavclemson, but I was unable to process your request. Please see the logs for more details.

DmitryKakurin
DmitryKakurin requested changes Jun 25, 2026
View reviewed changes
Comment thread src/maxtext/common/checkpointing.py
@SujeethJinesh
SujeethJinesh force-pushed the sujinesh/mtc_pathways_support branch 2 times, most recently from dfd1730 to 948c186 Compare June 25, 2026 21:19
@github-actions

github-actions Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

🤖 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

github-actions Bot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

🤖 I'm sorry @abhinavclemson, but I was unable to process your request. Please see the logs for more details.

@SujeethJinesh
SujeethJinesh force-pushed the sujinesh/mtc_pathways_support branch from 948c186 to 3071a95 Compare June 26, 2026 00:09
SujeethJinesh
SujeethJinesh commented Jun 26, 2026
View reviewed changes
Comment thread tests/unit/train_compile_test.py
def test_engram_integration(self):
"""AOT test for Engram implementation"""
compiled_trainstep_file = "/tmp/test_engram_integration"
transformers.AutoConfig.register("deepseek_v32", DeepseekV32Config)

@SujeethJinesh SujeethJinesh Jun 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI Removing because the transformers upgrade added the deepseek v3.2 config: huggingface/transformers#41251

Post training also already has an upgrade to transformers as well.

@SujeethJinesh
SujeethJinesh force-pushed the sujinesh/mtc_pathways_support branch from 3071a95 to 73cd6ad Compare June 26, 2026 00:14
@copybara-service
copybara-service Bot merged commit 7549b91 into main Jun 26, 2026
42 checks passed
@copybara-service
copybara-service Bot deleted the sujinesh/mtc_pathways_support branch June 26, 2026 22:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@DmitryKakurin DmitryKakurin DmitryKakurin approved these changes
@hengtaoguo hengtaoguo hengtaoguo approved these changes
@SurbhiJainUSC SurbhiJainUSC SurbhiJainUSC approved these changes
@abhinavclemson abhinavclemson abhinavclemson approved these changes
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha 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
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025
@NicoGrande NicoGrande Awaiting requested review from NicoGrande
@jesselu-google jesselu-google Awaiting requested review from jesselu-google
@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
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@jacoguzo jacoguzo Awaiting requested review from jacoguzo
@parambole parambole Awaiting requested review from parambole parambole is a code owner

Assignees

No one assigned

Labels

gemini-review pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@SujeethJinesh @DmitryKakurin @hengtaoguo @SurbhiJainUSC @abhinavclemson