[Checkpointing - PR4] Refactored the CheckpointLoader into a load_checkpoint function
#693
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
This PR is the first in a series for cleaning up the checkpoint API. One of the prerequesites is storing the seed on the state. Here, only the rank zero seed is stored on state, since only the rank zero state is persisted in a checkpoint. The trainer uses a distributed reduction to share the seed across states, so the same seed will be restored when resuming from checkpointing, even if a seed was not originally specified. This PR ignores the `seed` parameter passed into the trainer when resuming from a checkpoint. For the time being, if a new seed is desired, the `seed` attribute must be removed from the checkpoint state dict. #497 will introduce a cleaner API for this (edge) use case.
1. RNG serialization / deserialization is moved from `composer.trainer._checkpoint` to `composer.utils.reproducibility`. This change is needed to refactor the checkpoint saver into a public module. 2. Moved helper methods from `composer.trainer._deepspeed` to `composer.core.state` to determine whether the model is deepspeed 3. Added a similar helper for `is_model_ddp`. 3. Refactored how the state_dict was serialized and deserialized to support serialization of `@property`s. Stopped storing leading underscores in the checkpoint, as that is a state implementation detail and not something that should be persisted through the checkpoint.
…oint` function Since the checkpoint loading happens in `Trainer.__init__` (except for the restoration of the rng state), there is no need for a checkpoint loader class. This class is replaced with a function `load_checkpoint`, and all private members are converted into private, module-level helper functions.
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.
Since the checkpoint loading happens in
Trainer.__init__(except for the restoration of the rng state), there is no need for a checkpoint loader class.load_checkpoint, and all private members are converted into private, module-level helper functions.composer.utils.file_retriever, with their own test cases.