Poc elastic training #1310
Draft
Poc elastic training #1310
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
lukebaumann
requested review from
gobbleturk,
khatwanimohit,
bvandermoon,
vipannalla,
RissyRan,
richjames0,
rni418 and
gagika
as code owners
lukebaumann
force-pushed
the
poc-elastic-training
branch
from
506da52 to
0c975ce
Compare
gagika
reviewed
Feb 27, 2025
MaxText/train.py
Outdated
There was a problem hiding this comment.
Could you move it to a separate elastic_train.py file similar to sft_trainer.py: https://github.com/AI-Hypercomputer/maxtext/blob/sft-maxtext/MaxText/sft_trainer.py
MaxText/train.py
Outdated
| @elasticutils.timeit | ||
| def reshard_fn(config: pyconfig.HyperParameters): | ||
| """Reshard function.""" | ||
| while True: |
There was a problem hiding this comment.
Could you add num of reshards instead of infinite loop?
lukebaumann
force-pushed
the
poc-elastic-training
branch
from
0c975ce to
28044f7
Compare
Adding ElasticUtils to config Added elasticutils and gkeutils Added a watchdog/timebomb to each step Completely working test run. Further test runs and optimizations to follow Fixed a bug if DATA_LOSS occurs during a save added timeit to reshard_fn Updated the watchdog to repeatably stack trace every timeout intervals. send a fatal log if the failures > max failures in slice_down() instead of in the training loop in order to fail correctly if there is a reshard/failure loop within the reshard handler Updated elasticutils and added a fake elasticutils Working host memory offloading. Added a max reshard retry count Updated elasticutils to how it will be structured in pathwaysutils Delete checkpoint if we are rewinding behind when it was saved. Fixed a bug in put_array. Added max_reshard_failure_count. Added save nbytes log Adding elastic trainer
lukebaumann
force-pushed
the
poc-elastic-training
branch
from
28044f7 to
9fdab26
Compare
…ynchronized checkpointing might
… the number of values returned from reshard_fn
…ext manager of the old mesh when resharding up
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):