Integrate SSD TBE stage 1 #2078
Closed
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
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D57452256 |
Summary: # Plan Stage 1 aims to ensure that it can run, and won't break from normal operations (e.g. checkpointing). Checkpointing (i.e. state_dict and load_state_dict) are still work in progress. We also need to guarantee checkpointing for optimizer states. Stage 2: save state_dict (mostly on fbgemm side) * current hope is we can rely on flush to save state dict Stage 3: load_state_dict (need more thoughts) * solution should be similar to that of PS Stage 4: optimizer states checkpointing (torchrec side, should be pretty standard) * should be straightforward * need fbgemm to support split_embedding_weights api # Outstanding issues: * init is not the same as before * SSD TBE doesn't support mixed dim # design doc https://docs.google.com/document/d/1SL1d2Os8KG46ETkCFzrIO0_QMOlFWoTrb8CJiCqGDdk/ # tests should cover * state dict and load state dict (done) * should copy dense parts and not break * deterministics output (done) * numerical equivalence to normal TBE (done) * changing learning rate and warm up policy (done) * work for different sharding types (done) * work with mixed kernel (done) * work with mixed sharding types (done) # OSS NOTE: SSD TBE won't work in an OSS environment, due to some rocksdb problem. # ad hoc * SSD kernel is guarded, user must specify it in constraints to use it # Next steps * add multi-gpu training (todo) * add optimizer checkpointing support (stage 4) * add equivalence of CacheParams for PS and SSD configs * support multi dim via bucketizer * modify tests to check for multi dim, and remove assertion Reviewed By: dstaay-fb Differential Revision: D57452256
|
This pull request was exported from Phabricator. Differential Revision: D57452256 |
henrylhtsang
force-pushed
the
export-D57452256
branch
from
8ac8c0e to
5570843
Compare
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.
Summary:
Plan
Stage 1 aims to ensure that it can run, and won't break from normal operations (e.g. checkpointing).
Checkpointing (i.e. state_dict and load_state_dict) are still work in progress. We also need to guarantee checkpointing for optimizer states.
Stage 2: save state_dict (mostly on fbgemm side)
Stage 3: load_state_dict (need more thoughts)
Stage 4: optimizer states checkpointing (torchrec side, should be pretty standard)
Outstanding issues:
design doc
TODO:
tests should cover
OSS
NOTE: SSD TBE won't work in an OSS environment, due to some rocksdb problem.
ad hoc
Differential Revision: D57452256