Make the run directory rank-local; fix checkpoints saving and restoring #215
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
- Sharding the of the run directory accross ranks won't work in multi-node training. This change makes the run directory rank-local - Fixed callbacks and loggers to support rank-local run directories. Specifically, wandb and the run directory uploader now run on all ranks, not just rank zero - When using deepspeed with zero-1+, each rank writes to the checkpoint folder. Previously, only rank zero's data was being stored. Now, each rank's data is being stored by the rank-local run directory uploader. The checkpoint loader takes a checkpoint path that is parameterized by the rank, so each node will load only the shards of the checkpoint that is needed.
…omposer into ravi/rank_local_run_directory
This was referenced Jan 11, 2022
jbloxham
suggested changes
Jan 11, 2022
Removed run_directory.get_relative_to_run_directory
jbloxham
approved these changes
Jan 14, 2022
coryMosaicML
pushed a commit
to coryMosaicML/composer
that referenced
this pull request
Feb 23, 2022
…ng (mosaicml#215) Fixed the run directory and deepspeed checkpointing: - Creating a run directory by default, if one was not set via an environment variable - Removed `run_directory.get_relative_to_run_directory(...)`; replaced that with `os.path.join(run_directory.get_run_directory(...))` - Having one run directory per node (i.e. sharing the run directory across local ranks) won't work in multi-node training. This change makes the run directory rank-local (so each rank sees its own run directory) - Added a helper method `run_directory.get_node_run_directory()`. This should not really be used directly, but rather - Switched the run directory uploader and wandb to run on every rank instead of just rank zero. This ensures that all artifacts will be properly stored. - Fixed deepspeed checkpointing: Specifically, when using deepspeed with zero-1+, each rank wrote to the folder specified. Only rank 0's data was being stored, meaning that we were losing rank 1+'s state (i.e. the optimizer state). Deepspeed checkpointing tests were broken. Now, each rank's data is being stored by the rank-local run directory uploader (see above bullet). Meaning, there are N files per checkpoint now. To restore a checkpoint, you need the global rank zero checkpoint (which contains the model) and the rank n checkpoint (which contains the optimizer state). To support this, the checkpoint loader takes a checkpoint path that is parameterized by the rank, so each rank can fetch the files it needs. As an optimization, the LOCAL rank zero on each node is responsible for downloading the GLOBAL rank zero checkpoint (it would be redundant to download this file multiple times per node).
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.
Fixed the run directory and deepspeed checkpointing:
Creating a run directory by default, if one was not set via an environment variable
Removed
run_directory.get_relative_to_run_directory(...); replaced that withos.path.join(run_directory.get_run_directory(...))Having one run directory per node (i.e. sharing the run directory across local ranks) won't work in multi-node training. This change makes the run directory rank-local (so each rank sees its own run directory)
Added a helper method
run_directory.get_node_run_directory(). This should not really be used directly, but ratherSwitched the run directory uploader and wandb to run on every rank instead of just rank zero. This ensures that all artifacts will be properly stored.
Fixed deepspeed checkpointing:
Specifically, when using deepspeed with zero-1+, each rank wrote to the folder specified. Only rank 0's data was being stored, meaning that we were losing rank 1+'s state (i.e. the optimizer state). Deepspeed checkpointing tests were broken. Now, each rank's data is being stored by the rank-local run directory uploader (see above bullet). Meaning, there are N files per checkpoint now.
To restore a checkpoint, you need the global rank zero checkpoint (which contains the model) and the rank n checkpoint (which contains the optimizer state). To support this, the checkpoint loader takes a checkpoint path that is parameterized by the rank, so each rank can fetch the files it needs. As an optimization, the LOCAL rank zero on each node is responsible for downloading the GLOBAL rank zero checkpoint (it would be redundant to download this file multiple times per node).
TODO: