Multi-tier checkpointing + orbax replicator #1332

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Open

ehorning wants to merge 29 commits into apple:main from ehorning:orbax-mtc-testing

Contributor

ehorning commented Aug 12, 2025 •

edited

Loading

Integrate multi-tier checkpointer + orbax replicator into axlearn

Erin Horning added 14 commits

July 21, 2025 22:04


          orbax

6b87f81


          mount ckpt volume and create replicator file

f61131f


          renaming

f6cacc3


          move stuff

264aea6


          use orbax replicator

f62c377


          remove cleanup

79db845


          wait and rename

93438f2


          jax init + moving stuff

e5252e1


          replicator restore tweaks

0cdc871


          update restore objects + mesh shape

fba65dd


          rip out unnecessary process management logic from oecp

2ae0d0f


          unnecessary jax init

a189306


          comment out non-tensor stuff

6e3fa28


          cleanup

b1199ee

DmitryKakurin reviewed

View reviewed changes

Dockerfile Outdated Show resolved Hide resolved

axlearn/cloud/gcp/jobset_utils.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax.py Outdated Show resolved Hide resolved

DmitryKakurin reviewed

View reviewed changes

axlearn/common/checkpointer_orbax.py Outdated Show resolved Hide resolved

axlearn/experiments/text/gpt/common.py Outdated Show resolved Hide resolved

axlearn/experiments/text/gpt/fuji.py Outdated Show resolved Hide resolved


          cleanup

c28f970

DmitryKakurin reviewed

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

DmitryKakurin reviewed

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

Erin Horning added 2 commits

August 13, 2025 19:47


          more cleanup

4d3f30d


          remove special non-tensor checkpointing

19151d9

samos123 reviewed

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

Erin Horning added 2 commits

August 15, 2025 21:41


          remove unnecessary changes


          more cleanup

ceeef4d

DmitryKakurin reviewed

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py Show resolved Hide resolved

axlearn/cloud/gcp/jobset_utils.py Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Show resolved Hide resolved

DmitryKakurin reviewed

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

Erin Horning added 3 commits

August 21, 2025 17:26


          raise errors

12c1e46


          cleanup

485e184


          conditional volume mount

c7662db

ehorning commented

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

DmitryKakurin reviewed

View reviewed changes

axlearn/cloud/gcp/jobset_utils.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Outdated Show resolved Hide resolved

axlearn/common/checkpointer_orbax_emergency_replicator.py Show resolved Hide resolved

Erin Horning and others added 3 commits

August 22, 2025 01:22


          more cleanup

dc60555


          Merge branch 'main' into orbax-mtc-testing

161b179


          logging

21f4453

DmitryKakurin reviewed

View reviewed changes

axlearn/experiments/text/gpt/common.py Outdated Show resolved Hide resolved

ehorning and others added 2 commits

September 10, 2025 15:27


          Merge branch 'main' into orbax-mtc-testing

c249122


          orbax-cp version + dp var

69b12c4

ehorning marked this pull request as ready for review

September 16, 2025 19:39

ehorning requested review from a team as code owners

September 16, 2025 19:39

samos123 reviewed

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py

    
              FLAGS = flags.FLAGS

              flags.DEFINE_integer(

                  "assume_data_parallelism",

Contributor

samos123 Sep 16, 2025

future follow up. I think orbax has a way of figuring this out automatically since it also needs to know this info. Orbax requires you to specify the batch dimension afair so it can know this.

DmitryKakurin Sep 16, 2025

MaxText sets it to the number of slices. However it may not be correct if there is intra-slice DDP, so we plan to make it configurable.


          orbax group install

9882bd4

DmitryKakurin approved these changes

View reviewed changes

axlearn/common/checkpointer_orbax_emergency_replicator.py

    
              FLAGS = flags.FLAGS

              flags.DEFINE_integer(

                  "assume_data_parallelism",

DmitryKakurin Sep 16, 2025

MaxText sets it to the number of slices. However it may not be correct if there is intra-slice DDP, so we plan to make it configurable.


          Merge branch 'main' into orbax-mtc-testing

7dba635

github-actions bot commented Dec 28, 2025

This pull request has been automatically marked as stale because it has been inactive for 60 days. It will be closed in 7 days if no further activity occurs. If you would like to continue working on this, please remove the stale label or leave a comment.

github-actions bot added the stale label

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels